store sessions in the DB

packages/tracker/meshnet_tracker/server.py

@@ -21,6 +21,7 @@ HTTP API contract:
   Response 400/404/503: {"error": str}
 """
 
+import http.cookies
 import http.server
 import hashlib
 import itertools
@@ -47,14 +48,15 @@ from .wallet_proof import binding_message, verify_wallet_signature
 from .billing import DEFAULT_BILLING_DB_PATH, BillingLedger
 from .calibration import DEFAULT_CALIBRATION_DB_PATH, ToplocCalibrationStore
 from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH, HfPricingLog, refresh_preset_price
-from .gossip import NodeGossip
-from .logging_setup import tracker_logger
-from .model_files import files_for_layer_range, snapshot_dir_for_repo
+from .gossip import NodeGossip
+from .logging_setup import tracker_logger
+from .model_files import files_for_layer_range, snapshot_dir_for_repo
 from .raft import RaftNode
 
 
 _CONSOLE_LIMIT = 300
 _PROXY_PROGRESS_LOG_INTERVAL = 5.0
+_SESSION_COOKIE_NAME = "meshnet_session"
 
 
 def _preset_price_keys(name: str, preset: dict) -> set[str]:
@@ -1935,6 +1937,38 @@ def _api_key_from_headers(headers) -> str | None:
     return auth.strip() or None
 
 
+def _session_token_from_headers(headers) -> str | None:
+    token = _api_key_from_headers(headers)
+    if token:
+        return token
+    cookie_header = headers.get("Cookie")
+    if not cookie_header:
+        return None
+    cookie = http.cookies.SimpleCookie()
+    try:
+        cookie.load(cookie_header)
+    except http.cookies.CookieError:
+        return None
+    morsel = cookie.get(_SESSION_COOKIE_NAME)
+    if morsel is None:
+        return None
+    return morsel.value.strip() or None
+
+
+def _session_cookie_header(token: str | None) -> str:
+    cookie = http.cookies.SimpleCookie()
+    cookie[_SESSION_COOKIE_NAME] = token or ""
+    morsel = cookie[_SESSION_COOKIE_NAME]
+    morsel["path"] = "/"
+    morsel["httponly"] = True
+    morsel["samesite"] = "Lax"
+    if token:
+        morsel["max-age"] = str(int(7 * 86400))
+    else:
+        morsel["max-age"] = "0"
+    return morsel.OutputString()
+
+
 def _usage_total_tokens(payload: dict) -> int | None:
     usage = payload.get("usage")
     if not isinstance(usage, dict):
@@ -2094,7 +2128,7 @@ def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -
     return None
 
 
-def _tracker_log(
+def _tracker_log(
     server: "_TrackerHTTPServer",
     level: str,
     message: str,
@@ -2102,18 +2136,18 @@ def _tracker_log(
     stdout: bool = True,
     update_console_key: str | None = None,
     **fields: Any,
-) -> None:
-    log_level = {
-        "debug": 10,
-        "info": 20,
-        "warn": 30,
-        "warning": 30,
-        "error": 40,
-    }.get(level.lower(), 20)
-    event = {
-        "ts": time.time(),
-        "level": level,
-        "message": message,
+) -> None:
+    log_level = {
+        "debug": 10,
+        "info": 20,
+        "warn": 30,
+        "warning": 30,
+        "error": 40,
+    }.get(level.lower(), 20)
+    event = {
+        "ts": time.time(),
+        "level": level,
+        "message": message,
         "fields": {
             key: value
             for key, value in fields.items()
@@ -2134,13 +2168,13 @@ def _tracker_log(
                     break
             if not updated:
                 server.console_events.append(event)
-        else:
-            server.console_events.append(event)
-    extras = " ".join(f"{key}={value}" for key, value in event["fields"].items())
-    suffix = f"  {extras}" if extras else ""
-    tracker_logger().log(log_level, f"{message}{suffix}")
-    if stdout:
-        print(f"[tracker] {level}: {message}{suffix}", flush=True)
+        else:
+            server.console_events.append(event)
+    extras = " ".join(f"{key}={value}" for key, value in event["fields"].items())
+    suffix = f"  {extras}" if extras else ""
+    tracker_logger().log(log_level, f"{message}{suffix}")
+    if stdout:
+        print(f"[tracker] {level}: {message}{suffix}", flush=True)
 
 
 @dataclass
@@ -2344,11 +2378,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
     def log_message(self, fmt, *args):  # noqa: suppress request logs in tests
         pass
 
-    def _send_json(self, status: int, data: dict) -> None:
+    def _send_json(self, status: int, data: dict, headers: dict[str, str] | None = None) -> None:
         body = json.dumps(data).encode()
         self.send_response(status)
         self.send_header("Content-Type", "application/json")
         self.send_header("Content-Length", str(len(body)))
+        for name, value in (headers or {}).items():
+            self.send_header(name, value)
         self.end_headers()
         try:
             self.wfile.write(body)
@@ -2365,7 +2401,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
         inference and wallet binding only, never operator endpoints.
         """
         server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
-        token = _api_key_from_headers(self.headers)
+        token = _session_token_from_headers(self.headers)
         if not token:
             return None, None
         if is_validator_token(token, server.validator_service_token):
@@ -4209,7 +4245,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
         server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
         if server.accounts is None:
             return None
-        return server.accounts.session_account(_api_key_from_headers(self.headers))
+        return server.accounts.session_account(_session_token_from_headers(self.headers))
 
     def _require_accounts(self) -> "AccountStore | None":
         server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
@@ -4249,7 +4285,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             f"[tracker] account registered: {account.get('email') or account.get('wallet')} "
             f"role={account['role']}", flush=True,
         )
-        self._send_json(200, {"account": account, "session_token": token, "api_key": api_key})
+        self._send_json(
+            200,
+            {"account": account, "session_token": token, "api_key": api_key},
+            headers={"Set-Cookie": _session_cookie_header(token)},
+        )
 
     def _handle_auth_login(self):
         accounts = self._require_accounts()
@@ -4264,14 +4304,18 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             self._send_json(401, {"error": "invalid credentials"})
             return
         token = accounts.create_session(account["account_id"])
-        self._send_json(200, {"account": account, "session_token": token})
+        self._send_json(
+            200,
+            {"account": account, "session_token": token},
+            headers={"Set-Cookie": _session_cookie_header(token)},
+        )
 
     def _handle_auth_logout(self):
         accounts = self._require_accounts()
         if accounts is None:
             return
-        accounts.destroy_session(_api_key_from_headers(self.headers))
-        self._send_json(200, {"ok": True})
+        accounts.destroy_session(_session_token_from_headers(self.headers))
+        self._send_json(200, {"ok": True}, headers={"Set-Cookie": _session_cookie_header(None)})
 
     def _handle_account_me(self):
         """Balance, usage, and API keys for the logged-in account."""