From 4c4ec24848e24f9eafe2f7601b3ea2b6ef766ae3 Mon Sep 17 00:00:00 2001 From: kaffa Date: Sun, 1 Feb 2026 13:38:00 +0000 Subject: [PATCH] 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 --- mcp/server.py | 147 +++++++++++++++++++++----------------------------- 1 file changed, 62 insertions(+), 85 deletions(-) diff --git a/mcp/server.py b/mcp/server.py index 86a3d70..fe129af 100644 --- a/mcp/server.py +++ b/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: - """Save entries to domains.map file atomically. - - Uses temp file + rename for atomic write to prevent race conditions. +def atomic_write_file(file_path: str, content: str) -> None: + """Write content to file atomically using temp file + rename. 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("# Format: domain backend_name\n") - f.write("# Wildcard: .domain.com matches *.domain.com\n\n") - 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 + f.write(content) + os.rename(temp_path, file_path) + temp_path = None # Rename succeeded 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) - 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 + atomic_write_file(SERVERS_FILE, json.dumps(config, indent=2)) 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) - haproxy_cmd(f"add map {MAP_FILE_CONTAINER} {domain} {pool}") - haproxy_cmd(f"add map {MAP_FILE_CONTAINER} .{domain} {pool}") - - # Read current map entries and save to file (persistence) + # Save to disk first (atomic write for persistence) + # If HAProxy update fails after this, state will be correct on restart 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) - haproxy_cmd(f"del map {MAP_FILE_CONTAINER} {domain}") - haproxy_cmd(f"del map {MAP_FILE_CONTAINER} .{domain}") - - # Remove entries from map file (persistence) + # Save to disk first (atomic write for persistence) + # If HAProxy update fails after this, state will be correct on restart 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) - 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 + atomic_write_file(STATE_FILE, state) return "Server state saved" except HaproxyError as e: return f"Error: {e}"