129 lines
4.0 KiB
Python
129 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
k_server — protected resource backend.
|
|
|
|
Exposes a monotonic counter behind a shared proxy token. Only k_proxy
|
|
is expected to reach this service; k_client should have no direct path.
|
|
All state is process-local and resets on restart.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import ssl
|
|
import threading
|
|
import time
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from typing import Any
|
|
from urllib.parse import urlparse
|
|
|
|
|
|
class ServerState:
|
|
# All state is process-local; a restart resets the counter to zero.
|
|
def __init__(self, proxy_token: str):
|
|
self.proxy_token = proxy_token
|
|
self.counter = 0
|
|
self.lock = threading.Lock()
|
|
|
|
def next_counter(self) -> int:
|
|
with self.lock:
|
|
self.counter += 1
|
|
return self.counter
|
|
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
state: ServerState
|
|
protocol_version = "HTTP/1.1"
|
|
|
|
def _json(self, status: int, payload: dict[str, Any]) -> None:
|
|
body = json.dumps(payload).encode("utf-8")
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def _discard_request_body(self) -> None:
|
|
# HTTP/1.1 keep-alive: the connection is reused, so the body must be fully
|
|
# consumed before we send the response, even for endpoints that ignore it.
|
|
length = int(self.headers.get("Content-Length", "0"))
|
|
if length > 0:
|
|
self.rfile.read(length)
|
|
|
|
def _is_proxy_authorized(self) -> bool:
|
|
return self.headers.get("X-Proxy-Token") == self.state.proxy_token
|
|
|
|
def do_GET(self) -> None: # noqa: N802
|
|
path = urlparse(self.path).path
|
|
if path == "/health":
|
|
self._json(
|
|
200,
|
|
{
|
|
"ok": True,
|
|
"service": "k_server",
|
|
"time": int(time.time()),
|
|
},
|
|
)
|
|
return
|
|
self.send_error(404)
|
|
|
|
def do_POST(self) -> None: # noqa: N802
|
|
path = urlparse(self.path).path
|
|
if path != "/resource/counter":
|
|
self.send_error(404)
|
|
return
|
|
self._discard_request_body()
|
|
if not self._is_proxy_authorized():
|
|
self._json(401, {"ok": False, "error": "unauthorized proxy"})
|
|
return
|
|
|
|
value = self.state.next_counter()
|
|
self._json(
|
|
200,
|
|
{
|
|
"ok": True,
|
|
"resource": "counter",
|
|
"value": value,
|
|
"time": int(time.time()),
|
|
},
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Run k_server counter service")
|
|
parser.add_argument("--host", default="127.0.0.1")
|
|
parser.add_argument("--port", type=int, default=8780)
|
|
parser.add_argument("--tls-certfile", help="PEM certificate chain for HTTPS listener")
|
|
parser.add_argument("--tls-keyfile", help="PEM private key for HTTPS listener")
|
|
parser.add_argument(
|
|
"--proxy-token",
|
|
default="dev-proxy-token",
|
|
help="Shared token expected in X-Proxy-Token from k_proxy",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
if bool(args.tls_certfile) != bool(args.tls_keyfile):
|
|
raise SystemExit("Both --tls-certfile and --tls-keyfile are required to enable HTTPS")
|
|
|
|
state = ServerState(proxy_token=args.proxy_token)
|
|
Handler.state = state
|
|
server = ThreadingHTTPServer((args.host, args.port), Handler)
|
|
scheme = "http"
|
|
if args.tls_certfile:
|
|
context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
|
context.load_cert_chain(certfile=args.tls_certfile, keyfile=args.tls_keyfile)
|
|
server.socket = context.wrap_socket(server.socket, server_side=True)
|
|
scheme = "https"
|
|
|
|
print(f"k_server listening on {scheme}://{args.host}:{args.port}")
|
|
server.serve_forever()
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|