feat(pool): add mTLS worker transport and heartbeat agent

This commit is contained in:
冰朔 2026-09-11 14:58:19 +08:00
commit c40c3fdb61
3 changed files with 52 additions and 7 deletions

View file

@ -1,12 +1,15 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
"""Loopback-first compute-pool entry with authenticated worker heartbeats.""" """Loopback-first compute-pool entry with authenticated worker heartbeats."""
from __future__ import annotations from __future__ import annotations
import json, os, secrets, threading import json, os, secrets, ssl, threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path from pathlib import Path
STATE=Path(os.environ.get("GH_POOL_STATE","/var/lib/guanghu-compute-pool/state.json")) 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")) 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") NODE=os.environ.get("GH_POOL_NODE","QY-LH-MAIN-PROD-01")
LOCK=threading.Lock() LOCK=threading.Lock()
@ -18,20 +21,21 @@ 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) 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 authorized(handler): def authorized(handler):
if not TOKEN.exists(): return False try:
expected=TOKEN.read_text().strip(); actual=handler.headers.get('Authorization','') subject=handler.connection.getpeercert().get('subject',())
return secrets.compare_digest(actual, 'Bearer '+expected) if expected else False return any(key=='commonName' and value for group in subject for key,value in group)
except Exception: return False
class Handler(BaseHTTPRequestHandler): class Handler(BaseHTTPRequestHandler):
def send_json(self, code, value): 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) 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): 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 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':'AUTH_REQUIRED'}); 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 if self.path=='/v1/pool/status': self.send_json(200,read_state()); return
self.send_json(404,{'error':'NOT_FOUND'}) self.send_json(404,{'error':'NOT_FOUND'})
def do_POST(self): def do_POST(self):
if not authorized(self): self.send_json(401,{'error':'AUTH_REQUIRED'}); return 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'{}') 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 except Exception: self.send_json(400,{'error':'JSON_INVALID'}); return
with LOCK: with LOCK:
@ -39,6 +43,8 @@ class Handler(BaseHTTPRequestHandler):
if self.path=='/v1/worker/heartbeat': if self.path=='/v1/worker/heartbeat':
node=payload.get('node_id'); node=payload.get('node_id');
if not node or node=='JD-FD-PRIMARY': self.send_json(403,{'error':'MOTHER_NODE_FORBIDDEN'}); return 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}; write_state(state); self.send_json(200,{'outcome':'HEARTBEAT_ACCEPTED','node_id':node,'authority_granted':False}); return state['workers'][node]={**payload,'last_seen':payload.get('observed_at'),'authority_granted':False}; write_state(state); self.send_json(200,{'outcome':'HEARTBEAT_ACCEPTED','node_id':node,'authority_granted':False}); return
if self.path=='/v1/lease/request': if self.path=='/v1/lease/request':
self.send_json(200,{'outcome':'LEASE_QUEUED','state':'PENDING_DETERMINISTIC_SELECTION','authority_granted':False}); return self.send_json(200,{'outcome':'LEASE_QUEUED','state':'PENDING_DETERMINISTIC_SELECTION','authority_granted':False}); return
@ -46,4 +52,7 @@ class Handler(BaseHTTPRequestHandler):
def log_message(self,*args): pass def log_message(self,*args): pass
if __name__=='__main__': 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()); ThreadingHTTPServer((bind,port),Handler).serve_forever() 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()

View file

@ -11,4 +11,7 @@ class EntryServiceTests(unittest.TestCase):
def test_state_has_authority_false(self): def test_state_has_authority_false(self):
value=M.read_state(); self.assertFalse(value['authority_granted']) value=M.read_state(); self.assertFalse(value['authority_granted'])
def test_worker_certificate_identity_is_required_by_contract(self):
self.assertIn('MTLS_CLIENT_CERT_REQUIRED', open(ROOT/'server-tools/persona-compute-pool/entry_service.py').read())
if __name__=='__main__': unittest.main() if __name__=='__main__': unittest.main()

View file

@ -0,0 +1,33 @@
#!/usr/bin/env python3
"""Opt-in worker heartbeat agent; it never executes arbitrary scheduler commands."""
from __future__ import annotations
import json, os, ssl, time, urllib.request, shutil
from datetime import datetime, timezone
NODE=os.environ["GH_WORKER_NODE_ID"]
CONTROLLER=os.environ.get("GH_POOL_CONTROLLER","https://124.222.54.198:8787")
CA=os.environ.get("GH_WORKER_CA", "/etc/guanghu-compute-pool/tls/ca.crt")
CERT=os.environ.get("GH_WORKER_CERT", f"/etc/guanghu-compute-pool/tls/{NODE}.crt")
KEY=os.environ.get("GH_WORKER_KEY", f"/etc/guanghu-compute-pool/tls/{NODE}.key")
INTERVAL=float(os.environ.get("GH_WORKER_INTERVAL","30"))
def metrics():
mem={}
try:
for line in open('/proc/meminfo'):
key,value,*_=line.split(); mem[key]=int(value)*1024
except OSError: pass
usage=shutil.disk_usage('/')
return {'cpu_cores':os.cpu_count() or 1,'memory_total_bytes':mem.get('MemTotal'),'memory_available_bytes':mem.get('MemAvailable'),'disk_free_bytes':usage.free,'load_1m':os.getloadavg()[0] if hasattr(os,'getloadavg') else None}
def heartbeat():
payload={'node_id':NODE,'observed_at':datetime.now(timezone.utc).isoformat(),'metrics':metrics(),'state':'ACTIVE_WORKER','authority_granted':False}
request=urllib.request.Request(CONTROLLER+'/v1/worker/heartbeat',data=json.dumps(payload).encode(),headers={'Content-Type':'application/json'},method='POST')
context=ssl.create_default_context(cafile=CA); context.load_cert_chain(CERT,KEY); context.check_hostname=False; context.verify_mode=ssl.CERT_REQUIRED
with urllib.request.urlopen(request,context=context,timeout=10) as response: return json.loads(response.read())
if __name__=='__main__':
while True:
try: heartbeat()
except Exception: pass
time.sleep(INTERVAL)