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>
This commit is contained in:
kaffa
2026-02-01 13:38:00 +00:00
parent f17b02fddf
commit 4c4ec24848

View File

@@ -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}"