61 lines
4.1 KiB
Python
61 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Loopback-first compute-pool entry with authenticated worker heartbeats."""
|
|
from __future__ import annotations
|
|
import json, os, secrets, ssl, threading
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
|
|
STATE=Path(os.environ.get("GH_POOL_STATE","/var/lib/guanghu-compute-pool/state.json"))
|
|
TOKEN=Path(os.environ.get("GH_POOL_TOKEN","/etc/guanghu-compute-pool/token"))
|
|
TLS_CERT=Path(os.environ.get("GH_POOL_TLS_CERT","/etc/guanghu-compute-pool/tls/server.crt"))
|
|
TLS_KEY=Path(os.environ.get("GH_POOL_TLS_KEY","/etc/guanghu-compute-pool/tls/server.key"))
|
|
TLS_CA=Path(os.environ.get("GH_POOL_TLS_CA","/etc/guanghu-compute-pool/tls/ca.crt"))
|
|
NODE=os.environ.get("GH_POOL_NODE","QY-LH-MAIN-PROD-01")
|
|
LOCK=threading.Lock()
|
|
|
|
def read_state():
|
|
if STATE.exists(): return json.loads(STATE.read_text())
|
|
return {"schema":"guanghu.compute-pool-entry-state/v1","node_id":NODE,"role":"LIGHTWEIGHT_ENTERPRISE_COORDINATOR","state":"READY_LOOPBACK_ONLY","workers":{},"leases":{},"authority_granted":False}
|
|
|
|
def write_state(value):
|
|
STATE.parent.mkdir(parents=True, exist_ok=True); tmp=STATE.with_suffix('.tmp'); tmp.write_text(json.dumps(value,ensure_ascii=False,indent=2)+'\n'); tmp.replace(STATE)
|
|
|
|
def normalized_state():
|
|
value=read_state(); value.setdefault('workers',{}); value.setdefault('leases',{}); value.setdefault('registered_workers',len(value['workers'])); return value
|
|
|
|
def authorized(handler):
|
|
try:
|
|
subject=handler.connection.getpeercert().get('subject',())
|
|
return any(key=='commonName' and value for group in subject for key,value in group)
|
|
except Exception: return False
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def send_json(self, code, value):
|
|
body=json.dumps(value,ensure_ascii=False).encode(); self.send_response(code); self.send_header('Content-Type','application/json'); self.send_header('Content-Length',str(len(body))); self.end_headers(); self.wfile.write(body)
|
|
def do_GET(self):
|
|
if self.path=='/health': self.send_json(200,{'ok':True,'service':'guanghu-compute-pool-entry','node_id':NODE,'authority_granted':False}); return
|
|
if not authorized(self): self.send_json(401,{'error':'MTLS_CLIENT_CERT_REQUIRED'}); return
|
|
if self.path=='/v1/pool/status': self.send_json(200,read_state()); return
|
|
self.send_json(404,{'error':'NOT_FOUND'})
|
|
def do_POST(self):
|
|
if not authorized(self): self.send_json(401,{'error':'MTLS_CLIENT_CERT_REQUIRED'}); return
|
|
try: payload=json.loads(self.rfile.read(int(self.headers.get('Content-Length','0')) or 0) or b'{}')
|
|
except Exception: self.send_json(400,{'error':'JSON_INVALID'}); return
|
|
with LOCK:
|
|
state=normalized_state()
|
|
if self.path=='/v1/worker/heartbeat':
|
|
node=payload.get('node_id');
|
|
if not node or node=='JD-FD-PRIMARY': self.send_json(403,{'error':'MOTHER_NODE_FORBIDDEN'}); return
|
|
peer=next((value for group in self.connection.getpeercert().get('subject',()) for key,value in group if key=='commonName'),None)
|
|
if node != peer: self.send_json(403,{'error':'WORKER_CERT_ID_MISMATCH'}); return
|
|
state['workers'][node]={**payload,'last_seen':payload.get('observed_at'),'authority_granted':False}; state['registered_workers']=len(state['workers']); write_state(state); self.send_json(200,{'outcome':'HEARTBEAT_ACCEPTED','node_id':node,'authority_granted':False}); return
|
|
if self.path=='/v1/lease/request':
|
|
self.send_json(200,{'outcome':'LEASE_QUEUED','state':'PENDING_DETERMINISTIC_SELECTION','authority_granted':False}); return
|
|
self.send_json(404,{'error':'NOT_FOUND'})
|
|
def log_message(self,*args): pass
|
|
|
|
if __name__=='__main__':
|
|
bind=os.environ.get('GH_POOL_BIND','127.0.0.1'); port=int(os.environ.get('GH_POOL_PORT','8787')); write_state(read_state())
|
|
server=ThreadingHTTPServer((bind,port),Handler)
|
|
context=ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER); context.load_cert_chain(TLS_CERT,TLS_KEY); context.load_verify_locations(TLS_CA); context.verify_mode=ssl.CERT_OPTIONAL
|
|
server.socket=context.wrap_socket(server.socket,server_side=True); server.serve_forever()
|