PyMkDB 0.1.0__py3-none-any.whl

src/server/coms/socket.py

@@ -0,0 +1,461 @@
+"""
+SocketServer — persistent TCP server for MkDB data-plane operations.
+
+Message framing: see socket_protocol.py (4-byte big-endian length prefix + JSON body)
+
+Client → server message format:
+{
+  "type":      "request",
+  "id":        "<uuid4 correlation ID>",
+  "action":    "read" | "write" | "delete" | "query" | "subscribe" | "ping",
+  "store":     "<store_name>",
+  "record_id": "<id>",              # for read / write / delete
+  "delta":     {<flat-path>: <v>},  # for write
+  "filter":    {<query filter>},    # for query
+}
+
+Server → client response:
+{
+  "type":   "response",
+  "id":     "<same correlation ID>",
+  "status": "ok" | "error",
+  "data":   <result> | None,
+  "error":  <message> | None,
+}
+
+Subscribe message (client → server):
+  {"type": "subscribe", "store": "<name>"}
+  → server adds client to subscription set for that store
+
+Broadcast event (server → subscribed clients):
+  {"type": "update", "store": "<name>", "record_id": "<id>", "op": "write"|"delete"}
+"""
+
+import logging
+import socket
+import threading
+import time
+import uuid
+from dataclasses import dataclass
+from typing import Optional
+
+from src.server.coms.socket_protocol import read_frame, write_frame
+from src.server.coms.actions import execute as _execute
+
+logger = logging.getLogger(__name__)
+
+# Default write password enforced when no users are configured in the DB.
+_DEFAULT_WRITE_PASSWORD = "mk_db"
+
+
+@dataclass
+class ClientSession:
+    addr:          tuple
+    conn:          socket.socket
+    connected_at:  float
+    last_pong:     float
+    authenticated: bool  = False
+    username:      str   = ""      # empty = unauthenticated / open-mode
+    can_read:      bool  = True    # True for R and RW sessions
+    can_write:     bool  = False   # True for W and RW sessions after auth
+
+
+class SocketServer:
+    def __init__(self, host: str, port: int, database, heartbeat_interval: float = 5.0,
+                 max_clients: int = 100, recv_timeout: float = 30.0):
+        self.host               = host
+        self.port               = port
+        self.database           = database
+        self.heartbeat_interval = heartbeat_interval
+        self.max_clients        = max_clients
+        self.recv_timeout       = recv_timeout
+
+        self._server_sock: Optional[socket.socket] = None
+        self._clients:     dict  = {}    # addr_str -> ClientSession
+        self._subscriptions: dict = {}   # store_name -> set of addr_str keys
+        self._running:     bool  = False
+        self._lock:        threading.Lock = threading.Lock()
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def start(self) -> None:
+        self._server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self._server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        self._server_sock.bind((self.host, self.port))
+        self._server_sock.listen(self.max_clients)
+        self._running = True
+        logger.info("SocketServer listening on %s:%d", self.host, self.port)
+        print(f"Socket server listening on {self.host}:{self.port}")
+
+        accept_thread = threading.Thread(
+            target=self._accept_loop, daemon=True, name="SocketServer-accept"
+        )
+        accept_thread.start()
+
+    def stop(self) -> None:
+        self._running = False
+        if self._server_sock:
+            try:
+                self._server_sock.close()
+            except Exception:
+                pass
+        # Close all client connections
+        with self._lock:
+            for session in list(self._clients.values()):
+                try:
+                    write_frame(session.conn, {"type": "disconnect"})
+                    session.conn.close()
+                except Exception:
+                    pass
+            self._clients.clear()
+
+    # ------------------------------------------------------------------
+    # Accept loop
+    # ------------------------------------------------------------------
+
+    def _accept_loop(self) -> None:
+        while self._running:
+            try:
+                conn, addr = self._server_sock.accept()
+            except Exception:
+                break
+            with self._lock:
+                if len(self._clients) >= self.max_clients:
+                    try:
+                        write_frame(conn, {"type": "error", "error": "Server full"})
+                        conn.close()
+                    except Exception:
+                        pass
+                    continue
+            conn.settimeout(self.recv_timeout)
+            thread = threading.Thread(
+                target=self._client_loop, args=(conn, addr),
+                daemon=True, name=f"SocketClient-{addr}"
+            )
+            thread.start()
+
+    # ------------------------------------------------------------------
+    # Client loop
+    # ------------------------------------------------------------------
+
+    # HTTP method prefixes — first 4 bytes of any HTTP/1.x request
+    _HTTP_PREFIXES = {b"GET ", b"POST", b"PUT ", b"DELE", b"HEAD", b"PATC", b"OPTI"}
+
+    def _client_loop(self, conn: socket.socket, addr: tuple) -> None:
+        addr_str = f"{addr[0]}:{addr[1]}"
+        now = time.time()
+        session = ClientSession(addr=addr, conn=conn, connected_at=now, last_pong=now)
+
+        with self._lock:
+            self._clients[addr_str] = session
+
+        logger.info("SocketServer: client connected %s", addr_str)
+
+        # ── HTTP misrouting detection ─────────────────────────────────────────
+        # Peek at the first 4 bytes without consuming them; if they look like
+        # an HTTP method verb the client connected to the wrong port.
+        # Use a short timeout: HTTP clients write immediately on connect, while
+        # proper MkDB clients wait for the server's permissions frame first.
+        try:
+            conn.settimeout(0.15)
+            first4 = conn.recv(4, socket.MSG_PEEK)
+        except (socket.timeout, OSError):
+            first4 = b""
+        finally:
+            conn.settimeout(self.recv_timeout)
+
+        if first4 and first4 in self._HTTP_PREFIXES:
+            db = self.database
+            http_port = None
+            if db is not None:
+                http_port = getattr(
+                    getattr(db.config, "servers", None),
+                    "http_server", None,
+                )
+                http_port = getattr(http_port, "port", None)
+
+            hint = (
+                f"Use the HTTP data-plane port ({http_port})."
+                if http_port else "Check your client configuration."
+            )
+            msg = (
+                f"This port speaks the MkDB binary socket protocol, not HTTP. {hint}"
+            )
+            http_response = (
+                "HTTP/1.1 400 Bad Request\r\n"
+                "Content-Type: application/json\r\n"
+                f"Content-Length: {len(msg) + 12}\r\n"
+                "Connection: close\r\n\r\n"
+                f'{{"error":"{msg}"}}'
+            )
+            try:
+                conn.sendall(http_response.encode("utf-8"))
+            except Exception:
+                pass
+            self._unregister(addr_str)
+            try:
+                conn.close()
+            except Exception:
+                pass
+            logger.warning("SocketServer: rejected HTTP client at %s — wrong port", addr_str)
+            return
+
+        # ── Handshake ─────────────────────────────────────────────────────────
+        # 1. Server sends PERMISSIONS message describing what needs auth.
+        # 2. Client replies with their desired access level: "R" or "RW".
+        # 3. If auth is required for that level, server challenges with AUTH.
+        # 4. Client sends {type: "auth", username: ..., password: ...}.
+        db = self.database
+        has_users     = db is not None and bool(getattr(db.config, "users", {}))
+        protect_reads = (
+            db is not None
+            and has_users
+            and any(getattr(sc, "protect_reads", False) for sc in db.config.stores.values())
+        )
+
+        try:
+            write_frame(conn, {
+                "type":           "permissions",
+                "read_protected":  has_users and protect_reads,
+                # Writes are always protected — real users or the default password
+                "write_protected": True,
+            })
+
+            # Wait for the client's access-level declaration: "R", "W", or "RW"
+            access_msg = read_frame(conn)
+            if not access_msg:
+                return
+            client_access = str(access_msg.get("access", "R")).upper()
+            if client_access not in ("R", "W", "RW"):
+                write_frame(conn, {"type": "error", "error": "access must be 'R', 'W', or 'RW'"})
+                return
+
+            # Writes always need auth (real users, or default password if none configured).
+            # Reads need auth only when protect_reads is True.
+            need_auth = client_access in ("W", "RW") or (protect_reads and client_access == "R")
+
+            if need_auth:
+                write_frame(conn, {"type": "auth_required"})
+                auth_msg = read_frame(conn)
+                if not auth_msg or auth_msg.get("type") != "auth":
+                    write_frame(conn, {"type": "error", "error": "Expected auth message"})
+                    return
+                username = str(auth_msg.get("username", ""))
+                password = str(auth_msg.get("password", ""))
+
+                if has_users:
+                    from src.server.control.server import verify_password
+                    user = db.config.users.get(username)
+                    if user is None or not verify_password(password, user.password_hash):
+                        write_frame(conn, {"type": "error", "error": "Invalid credentials"})
+                        return
+                    session.username = username
+                else:
+                    # No users configured — enforce the default write password
+                    if password != _DEFAULT_WRITE_PASSWORD:
+                        write_frame(conn, {"type": "error", "error": "Invalid credentials"})
+                        return
+
+                session.authenticated = True
+                session.can_read      = client_access in ("R", "RW")
+                session.can_write     = client_access in ("W", "RW")
+                write_frame(conn, {
+                    "type":      "auth_ok",
+                    "username":  username,
+                    "can_read":  session.can_read,
+                    "can_write": session.can_write,
+                })
+                logger.info("SocketServer: %s authenticated as '%s' (access=%s)",
+                            addr_str, username, client_access)
+            else:
+                # Read-only, no auth required
+                session.authenticated = True
+                session.can_read      = True
+                session.can_write     = False
+                write_frame(conn, {"type": "ready", "can_read": True, "can_write": False})
+
+        except (ConnectionError, OSError) as exc:
+            logger.warning("SocketServer: handshake error with %s: %s", addr_str, exc)
+            self._unregister(addr_str)
+            try:
+                conn.close()
+            except Exception:
+                pass
+            return
+
+        # ── Start heartbeat ───────────────────────────────────────────────────
+        hb_thread = threading.Thread(
+            target=self._heartbeat_loop, args=(conn, addr_str, session),
+            daemon=True, name=f"SocketHB-{addr_str}"
+        )
+        hb_thread.start()
+
+        try:
+            while self._running:
+                try:
+                    msg = read_frame(conn)
+                except (ConnectionError, OSError, ValueError) as exc:
+                    if isinstance(exc, ValueError):
+                        logger.warning("SocketServer: oversized frame from %s: %s", addr_str, exc)
+                        try:
+                            write_frame(conn, {"type": "error", "error": str(exc)})
+                        except Exception:
+                            pass
+                    break
+
+                if not msg:
+                    continue
+
+                msg_type = msg.get("type")
+
+                # Handle pong
+                if msg_type == "pong":
+                    session.last_pong = time.time()
+                    continue
+
+                # Handle subscribe
+                if msg_type == "subscribe":
+                    store_name = msg.get("store", "")
+                    with self._lock:
+                        self._subscriptions.setdefault(store_name, set()).add(addr_str)
+                    try:
+                        write_frame(conn, {"type": "subscribed", "store": store_name})
+                    except Exception:
+                        break
+                    continue
+
+                # Handle requests
+                if msg_type == "request":
+                    response = self._handle_request(msg, session)
+                    try:
+                        write_frame(conn, response)
+                    except Exception:
+                        break
+
+        finally:
+            self._unregister(addr_str)
+            try:
+                conn.close()
+            except Exception:
+                pass
+            logger.info("SocketServer: client disconnected %s", addr_str)
+
+    # ------------------------------------------------------------------
+    # Heartbeat
+    # ------------------------------------------------------------------
+
+    def _heartbeat_loop(self, conn: socket.socket, addr_str: str, session: ClientSession) -> None:
+        while self._running and addr_str in self._clients:
+            time.sleep(self.heartbeat_interval)
+            if addr_str not in self._clients:
+                break
+            try:
+                write_frame(conn, {"type": "ping"})
+            except Exception:
+                break
+            # Check pong timeout
+            if time.time() - session.last_pong > 2 * self.heartbeat_interval:
+                logger.warning("SocketServer: client %s timed out (no pong)", addr_str)
+                self._unregister(addr_str)
+                try:
+                    conn.close()
+                except Exception:
+                    pass
+                break
+
+    # ------------------------------------------------------------------
+    # Broadcast (pub-sub)
+    # ------------------------------------------------------------------
+
+    def broadcast(self, store_name: str, event: dict) -> None:
+        """Fan-out an event to all subscribers for a store."""
+        with self._lock:
+            subscribers = set(self._subscriptions.get(store_name, set()))
+        dead = set()
+        for addr_str in subscribers:
+            session = self._clients.get(addr_str)
+            if session is None:
+                dead.add(addr_str)
+                continue
+            try:
+                write_frame(session.conn, event)
+            except Exception:
+                dead.add(addr_str)
+        # Clean up dead subscribers
+        if dead:
+            with self._lock:
+                subs = self._subscriptions.get(store_name, set())
+                subs -= dead
+
+    # ------------------------------------------------------------------
+    # Request dispatch
+    # ------------------------------------------------------------------
+
+    def _handle_request(self, msg: dict, session: ClientSession) -> dict:
+        correlation_id = msg.get("id", str(uuid.uuid4()))
+        action     = msg.get("action")
+        store_name = msg.get("store", "")
+
+        def ok(data=None):
+            return {"type": "response", "id": correlation_id, "status": "ok", "data": data, "error": None}
+
+        def err(message: str):
+            return {"type": "response", "id": correlation_id, "status": "error", "data": None, "error": message}
+
+        if not action:
+            return err("Missing 'action' field")
+
+        db = self.database
+        if db is None:
+            return err("Database not available")
+
+        # ── Permission checks ────────────────────────────────────────────────
+        is_write = action in ("write", "delete")
+
+        # Session-level capability gates (set during handshake)
+        if is_write and not session.can_write:
+            return err("Write access denied — reconnect with 'W' or 'RW' access")
+        if action in ("read", "query") and not session.can_read:
+            return err("Read access denied — reconnect with 'R' or 'RW' access")
+
+        has_users = bool(getattr(db.config, "users", {}))
+        if has_users and session.username:
+            user = db.config.users.get(session.username)
+            if user is None:
+                return err("Authenticated user no longer exists")
+            if store_name:
+                perm = user.stores.get(store_name)
+                if perm is None:
+                    return err(f"No access to store '{store_name}'")
+                if is_write and not perm.write:
+                    return err("Write access denied")
+        elif has_users and is_write:
+            return err("Authentication required for write operations")
+        elif has_users and not is_write and action in ("read", "query"):
+            store_cfg = db.config.stores.get(store_name)
+            if getattr(store_cfg, "protect_reads", False) and not session.username:
+                return err(f"Authentication required to read from store '{store_name}'")
+
+        # ── Execute ──────────────────────────────────────────────────────────
+        client_key = session.username or session.addr[0]
+        params = {
+            "record_id": msg.get("record_id", ""),
+            "delta":     msg.get("delta", {}),
+            "filter":    msg.get("filter", {}),
+            "hydrate":   bool(msg.get("hydrate", False)),
+        }
+        r = _execute(db, action, store_name, params, client_key, "socket",
+                     on_broadcast=self.broadcast)
+        return ok(r.data) if r.ok else err(r.error)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _unregister(self, addr_str: str) -> None:
+        with self._lock:
+            self._clients.pop(addr_str, None)
+            for subs in self._subscriptions.values():
+                subs.discard(addr_str)

src/server/coms/socket_protocol.py

@@ -0,0 +1,54 @@
+"""
+Socket message framing helpers.
+
+Wire format: [4-byte big-endian uint32 length][UTF-8 JSON payload bytes]
+"""
+
+import json
+import socket
+import struct
+
+_MAX_FRAME_SIZE = 10 * 1024 * 1024  # 10 MB hard limit
+
+
+def encode_message(payload: dict) -> bytes:
+    """Serialize dict to JSON, prefix with 4-byte big-endian length."""
+    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
+    return struct.pack(">I", len(body)) + body
+
+
+def decode_message(data: bytes) -> dict:
+    """Strip 4-byte length prefix, parse JSON payload."""
+    return json.loads(data[4:].decode("utf-8"))
+
+
+def read_frame(sock: socket.socket) -> dict:
+    """
+    Read exactly 4 bytes for the length, then exactly N bytes for the payload.
+    Raises ConnectionError on premature close.
+    Raises ValueError if the frame exceeds 10 MB.
+    """
+    length_bytes = _recv_exact(sock, 4)
+    length = struct.unpack(">I", length_bytes)[0]
+    if length == 0:
+        return {}
+    if length > _MAX_FRAME_SIZE:
+        raise ValueError(f"Frame too large: {length} bytes (max {_MAX_FRAME_SIZE})")
+    payload_bytes = _recv_exact(sock, length)
+    return json.loads(payload_bytes.decode("utf-8"))
+
+
+def write_frame(sock: socket.socket, payload: dict) -> None:
+    """Encode and send one frame."""
+    sock.sendall(encode_message(payload))
+
+
+def _recv_exact(sock: socket.socket, n: int) -> bytes:
+    """Receive exactly n bytes from sock; raises ConnectionError if connection closes early."""
+    buf = b""
+    while len(buf) < n:
+        chunk = sock.recv(n - len(buf))
+        if not chunk:
+            raise ConnectionError("Socket closed before all bytes received")
+        buf += chunk
+    return buf