refactor: Improve stability and add IPv6 support

1. Fix order: Save to disk FIRST, then update HAProxy
   - Prevents inconsistency if HAProxy update succeeds but disk write fails
   - Data is preserved correctly on restart

2. Add IPv6 support
   - Use Python ipaddress module instead of regex
   - Now accepts both IPv4 and IPv6 addresses

3. Extract atomic_write_file() helper
   - Eliminates duplicated code in save_map_file, save_servers_config, haproxy_save_state
   - Single source of truth for atomic file operations

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

mcp/server.py

@@ -27,6 +27,7 @@ import os
 import tempfile
 import time
 import fcntl
+import ipaddress
 from typing import Any, Dict, Generator, List, Optional, Set, Tuple
 from mcp.server.fastmcp import FastMCP
 
@@ -191,7 +192,7 @@ def validate_domain(domain: str) -> bool:
 
 
 def validate_ip(ip: str, allow_empty: bool = False) -> bool:
-    """Validate IPv4 address format.
+    """Validate IPv4 or IPv6 address format.
 
     Args:
         ip: The IP address to validate
@@ -202,9 +203,11 @@ def validate_ip(ip: str, allow_empty: bool = False) -> bool:
     """
     if not ip:
         return allow_empty
-    if not IP_PATTERN.match(ip):
+    try:
+        ipaddress.ip_address(ip)
+        return True
+    except ValueError:
         return False
-    return all(0 <= int(octet) <= 255 for octet in ip.split('.'))
 
 
 def validate_backend_name(name: str) -> bool:
@@ -367,33 +370,28 @@ def get_legacy_backend_name(domain: str) -> str:
     return f"{domain_to_backend(domain)}_backend"
 
 
-def save_map_file(entries: List[Tuple[str, str]]) -> None:
+def atomic_write_file(file_path: str, content: str) -> None:
-    """Save entries to domains.map file atomically.
+    """Write content to file atomically using temp file + rename.
-
-    Uses temp file + rename for atomic write to prevent race conditions.
 
     Args:
-        entries: List of (domain, backend) tuples to write
+        file_path: Target file path
+        content: Content to write
 
     Raises:
-        IOError: If the file cannot be written
+        IOError: If write fails
     """
-    dir_path = os.path.dirname(MAP_FILE)
+    dir_path = os.path.dirname(file_path)
     fd = None
     temp_path = None
     try:
-        fd, temp_path = tempfile.mkstemp(dir=dir_path, prefix='.domains.map.')
+        fd, temp_path = tempfile.mkstemp(dir=dir_path, prefix='.tmp.')
         with os.fdopen(fd, 'w', encoding='utf-8') as f:
             fd = None  # fd is now owned by the file object
-            f.write("# Domain to Backend mapping\n")
+            f.write(content)
-            f.write("# Format: domain backend_name\n")
+        os.rename(temp_path, file_path)
-            f.write("# Wildcard: .domain.com matches *.domain.com\n\n")
+        temp_path = None  # Rename succeeded
-            for domain, backend in entries:
-                f.write(f"{domain} {backend}\n")
-        os.rename(temp_path, MAP_FILE)
-        temp_path = None  # Rename succeeded, don't unlink
     except OSError as e:
-        raise IOError(f"Failed to save map file: {e}") from e
+        raise IOError(f"Failed to write {file_path}: {e}") from e
     finally:
         if fd is not None:
             try:
@@ -407,6 +405,27 @@ def save_map_file(entries: List[Tuple[str, str]]) -> None:
                 pass
 
 
+def save_map_file(entries: List[Tuple[str, str]]) -> None:
+    """Save entries to domains.map file atomically.
+
+    Uses temp file + rename for atomic write to prevent race conditions.
+
+    Args:
+        entries: List of (domain, backend) tuples to write
+
+    Raises:
+        IOError: If the file cannot be written
+    """
+    lines = [
+        "# Domain to Backend mapping\n",
+        "# Format: domain backend_name\n",
+        "# Wildcard: .domain.com matches *.domain.com\n\n",
+    ]
+    for domain, backend in entries:
+        lines.append(f"{domain} {backend}\n")
+    atomic_write_file(MAP_FILE, "".join(lines))
+
+
 def load_servers_config() -> Dict[str, Any]:
     """Load servers configuration from JSON file with file locking.
 
@@ -441,29 +460,7 @@ def save_servers_config(config: Dict[str, Any]) -> None:
     Args:
         config: Dictionary with server configurations
     """
-    dir_path = os.path.dirname(SERVERS_FILE)
+    atomic_write_file(SERVERS_FILE, json.dumps(config, indent=2))
-    fd = None
-    temp_path = None
-    try:
-        fd, temp_path = tempfile.mkstemp(dir=dir_path, prefix='.servers.json.')
-        with os.fdopen(fd, 'w', encoding='utf-8') as f:
-            fd = None  # fd is now owned by the file object
-            json.dump(config, f, indent=2)
-        os.rename(temp_path, SERVERS_FILE)
-        temp_path = None  # Rename succeeded, don't unlink
-    except OSError as e:
-        raise IOError(f"Failed to save servers config: {e}") from e
-    finally:
-        if fd is not None:
-            try:
-                os.close(fd)
-            except OSError:
-                pass
-        if temp_path is not None:
-            try:
-                os.unlink(temp_path)
-            except OSError:
-                pass
 
 
 def add_server_to_config(domain: str, slot: int, ip: str, http_port: int) -> None:
@@ -685,34 +682,35 @@ def haproxy_add_domain(domain: str, ip: str = "", http_port: int = 80) -> str:
         return f"Error: No available pools (all {POOL_COUNT} pools are in use)"
 
     try:
-        # Update HAProxy map via Runtime API first (immediate effect)
+        # Save to disk first (atomic write for persistence)
-        haproxy_cmd(f"add map {MAP_FILE_CONTAINER} {domain} {pool}")
+        # If HAProxy update fails after this, state will be correct on restart
-        haproxy_cmd(f"add map {MAP_FILE_CONTAINER} .{domain} {pool}")
-
-        # Read current map entries and save to file (persistence)
         entries = get_map_contents()
         entries.append((domain, pool))
         entries.append((f".{domain}", pool))
         save_map_file(entries)
 
+        # Update HAProxy map via Runtime API (immediate effect)
+        haproxy_cmd(f"add map {MAP_FILE_CONTAINER} {domain} {pool}")
+        haproxy_cmd(f"add map {MAP_FILE_CONTAINER} .{domain} {pool}")
+
         # If IP provided, add server to slot 1
         if ip:
+            # Save server config to disk first
+            add_server_to_config(domain, 1, ip, http_port)
+
             for suffix, port in get_server_suffixes(http_port):
                 server = f"{pool}{suffix}_1"
                 haproxy_cmd(f"set server {pool}/{server} addr {ip} port {port}")
                 haproxy_cmd(f"set server {pool}/{server} state ready")
 
-            # Save to persistent config
-            add_server_to_config(domain, 1, ip, http_port)
-
             return f"Domain {domain} added to {pool} with server {ip}:{http_port}"
 
         return f"Domain {domain} added to {pool} (no servers configured)"
 
-    except HaproxyError as e:
-        return f"Error: {e}"
     except IOError as e:
         return f"Error: Failed to update map file: {e}"
+    except HaproxyError as e:
+        return f"Error: {e}"
 
 
 @mcp.tool()
@@ -738,15 +736,19 @@ def haproxy_remove_domain(domain: str) -> str:
         return f"Error: Cannot remove legacy domain {domain} (uses static backend {backend})"
 
     try:
-        # Clear map entries via Runtime API first (immediate effect)
+        # Save to disk first (atomic write for persistence)
-        haproxy_cmd(f"del map {MAP_FILE_CONTAINER} {domain}")
+        # If HAProxy update fails after this, state will be correct on restart
-        haproxy_cmd(f"del map {MAP_FILE_CONTAINER} .{domain}")
-
-        # Remove entries from map file (persistence)
         entries = get_map_contents()
         new_entries = [(d, b) for d, b in entries if d != domain and d != f".{domain}"]
         save_map_file(new_entries)
 
+        # Remove from persistent server config
+        remove_domain_from_config(domain)
+
+        # Clear map entries via Runtime API (immediate effect)
+        haproxy_cmd(f"del map {MAP_FILE_CONTAINER} {domain}")
+        haproxy_cmd(f"del map {MAP_FILE_CONTAINER} .{domain}")
+
         # Disable all servers in the pool (reset to 0.0.0.0:0)
         for slot in range(1, MAX_SLOTS + 1):
             server = f"{backend}_{slot}"
@@ -756,15 +758,12 @@ def haproxy_remove_domain(domain: str) -> str:
             except HaproxyError:
                 pass  # Ignore errors for individual servers
 
-        # Remove from persistent config
-        remove_domain_from_config(domain)
-
         return f"Domain {domain} removed from {backend}"
 
-    except HaproxyError as e:
-        return f"Error: {e}"
     except IOError as e:
         return f"Error: Failed to update map file: {e}"
+    except HaproxyError as e:
+        return f"Error: {e}"
 
 
 @mcp.tool()
@@ -1317,29 +1316,7 @@ def haproxy_save_state() -> str:
     """
     try:
         state = haproxy_cmd("show servers state")
-        dir_path = os.path.dirname(STATE_FILE)
+        atomic_write_file(STATE_FILE, state)
-        fd = None
-        temp_path = None
-        try:
-            fd, temp_path = tempfile.mkstemp(dir=dir_path, prefix='.servers.state.')
-            with os.fdopen(fd, 'w', encoding='utf-8') as f:
-                fd = None  # fd is now owned by the file object
-                f.write(state)
-            os.rename(temp_path, STATE_FILE)
-            temp_path = None  # Rename succeeded, don't unlink
-        except OSError as e:
-            raise IOError(f"Failed to save state: {e}") from e
-        finally:
-            if fd is not None:
-                try:
-                    os.close(fd)
-                except OSError:
-                    pass
-            if temp_path is not None:
-                try:
-                    os.unlink(temp_path)
-                except OSError:
-                    pass
         return "Server state saved"
     except HaproxyError as e:
         return f"Error: {e}"