Add capture worker v0.4 with durable BugAnalysis API handoff

deploy/capture-worker/buganalysis_handoff.py

@@ -0,0 +1,162 @@
+"""Durable, streaming BugAnalysis intake client. Python stdlib only."""
+from contextlib import contextmanager
+import hashlib
+import http.client
+import json
+import os
+from pathlib import Path
+import re
+import secrets
+import sqlite3
+import time
+from urllib.parse import urlsplit
+
+
+class HandoffError(Exception):
+    pass
+
+
+def settings(raw,base):
+    if raw and type(raw.get('enabled',False)) is not bool:raise HandoffError('HANDOFF_CONFIG')
+    if not raw or not raw.get('enabled',False):return None
+    value=dict(raw)
+    allowed={'enabled','url','project_id','profile_id','token_env','outbox','timeout_seconds','retry_seconds','producer_id'}
+    if set(value)-allowed:raise HandoffError('HANDOFF_CONFIG')
+    value.setdefault('profile_id','vcore');value.setdefault('token_env','VBA_INTAKE_TOKEN')
+    value.setdefault('timeout_seconds',120);value.setdefault('retry_seconds',60)
+    for field in ('project_id','profile_id','producer_id'):
+        if not re.fullmatch(r'[A-Za-z0-9][A-Za-z0-9_.-]{0,63}',str(value.get(field,''))):raise HandoffError('HANDOFF_CONFIG')
+    if not re.fullmatch(r'[A-Za-z_][A-Za-z0-9_]*',value['token_env']):raise HandoffError('HANDOFF_CONFIG')
+    parsed=urlsplit(value.get('url',''))
+    if (parsed.scheme not in {'https','http'} or not parsed.hostname or parsed.username or parsed.password
+            or parsed.query or parsed.fragment or parsed.path not in {'','/'}
+            or (parsed.scheme=='http' and parsed.hostname not in {'127.0.0.1','localhost','::1'})):
+        raise HandoffError('HANDOFF_URL')
+    if any(type(value[f]) is not int or not 1<=value[f]<=3600 for f in ('timeout_seconds','retry_seconds')):
+        raise HandoffError('HANDOFF_CONFIG')
+    value['outbox']=str((base/Path(value.get('outbox','buganalysis-outbox.sqlite3'))).resolve())
+    return value
+
+
+def digest(path):
+    h=hashlib.sha256()
+    with Path(path).open('rb') as stream:
+        for chunk in iter(lambda:stream.read(1024*1024),b''):h.update(chunk)
+    return h.hexdigest()
+
+
+class Outbox:
+    def __init__(self,config,*,clock=time.time):
+        self.config=config;self.clock=clock;self.path=Path(config['outbox'])
+        self.path.parent.mkdir(parents=True,exist_ok=True)
+        if not self.path.exists():
+            try:
+                fd=os.open(self.path,os.O_CREAT|os.O_EXCL|os.O_WRONLY,0o600);os.close(fd)
+            except FileExistsError:pass
+        with self.db() as db:
+            db.execute('CREATE TABLE IF NOT EXISTS deliveries (job TEXT PRIMARY KEY, idempotency TEXT NOT NULL UNIQUE, payload TEXT NOT NULL,status TEXT NOT NULL, attempts INTEGER NOT NULL DEFAULT 0, next_try REAL NOT NULL DEFAULT 0, code TEXT, receipt TEXT, lease_until REAL NOT NULL DEFAULT 0)')
+
+    @contextmanager
+    def db(self):
+        conn=sqlite3.connect(self.path,timeout=10);conn.row_factory=sqlite3.Row
+        try:
+            with conn:yield conn
+        finally:conn.close()
+
+    def stage(self,run_id,bug_id,diag,src,revision):
+        files=[]
+        for role,path in (('wowdiag',diag),('wowsrc',src)):
+            path=Path(path).resolve()
+            files.append({'role':role,'path':str(path),'size':path.stat().st_size,'sha256':digest(path)})
+        # The adapters import actual observations/title from the unchanged archives.
+        # Do not invent an observation merely to populate envelope metadata.
+        metadata={'kind':'new_case','project_id':self.config['project_id'],'profile_id':self.config['profile_id'],
+            'case_id':f'BUG-{bug_id:06d}','report':{'title':f'BUG-{bug_id:06d}'},'versions':{'capture_server':revision}}
+        payload=json.dumps({'url':self.config['url'],'metadata':metadata,'files':files},sort_keys=True)
+        job=f"{self.config['producer_id']}:{run_id}"
+        with self.db() as db:
+            db.execute('BEGIN IMMEDIATE')
+            row=db.execute('SELECT payload FROM deliveries WHERE job=?',(job,)).fetchone()
+            if row:
+                if row['payload']!=payload:raise HandoffError('HANDOFF_JOB_CONTENT_CHANGED')
+                return
+            db.execute('INSERT INTO deliveries(job,idempotency,payload,status) VALUES (?,?,?,?)',
+                (job,'capture-'+secrets.token_hex(24),payload,'pending'))
+
+    def status(self):
+        with self.db() as db:
+            return [dict(r) for r in db.execute('SELECT job,status,attempts,next_try,code,receipt FROM deliveries ORDER BY job')]
+
+    def flush(self,*,send=None,force=False,maximum=5):
+        send=send or upload
+        for _ in range(maximum):
+            now=self.clock()
+            with self.db() as db:
+                db.execute('BEGIN IMMEDIATE')
+                row=db.execute("SELECT * FROM deliveries WHERE status='pending' AND lease_until<=? AND (next_try<=? OR ?) ORDER BY next_try,job LIMIT 1",(now,now,int(force))).fetchone()
+                if not row:return
+                # Lease prevents concurrent --retry-handoffs and worker submissions.
+                db.execute('UPDATE deliveries SET lease_until=?,attempts=attempts+1 WHERE job=?',
+                    (now+self.config['timeout_seconds']+60,row['job']))
+            status='pending';receipt=None
+            try:
+                payload=json.loads(row['payload'])
+                for file in payload['files']:
+                    path=Path(file['path'])
+                    if not path.is_file() or path.stat().st_size!=file['size'] or digest(path)!=file['sha256']:
+                        raise HandoffError('HANDOFF_ARCHIVE_CHANGED')
+                code,receipt=send(self.config,payload,row['idempotency'])
+                if code=='OK':status='sent'
+                elif code.startswith('HTTP_') and code not in {'HTTP_401','HTTP_403','HTTP_408','HTTP_425','HTTP_429','HTTP_500','HTTP_502','HTTP_503','HTTP_504','HTTP_507'}:status='attention'
+            except HandoffError as error:
+                code=str(error)
+                if code=='HANDOFF_ARCHIVE_CHANGED':status='attention'
+            except Exception:
+                # Never persist arbitrary exception text or HTTP bodies (may contain secrets).
+                code='HANDOFF_TRANSPORT'
+            with self.db() as db:
+                db.execute('UPDATE deliveries SET status=?,next_try=?,code=?,receipt=?,lease_until=0 WHERE job=?',
+                    (status,self.clock()+self.config['retry_seconds'],code,json.dumps(receipt) if receipt else None,row['job']))
+            print(json.dumps({'handoff':row['job'],'status':status,'code':code}),flush=True)
+            if status=='pending':return # avoid retrying this request repeatedly in the same pass
+
+
+def upload(config,payload,idempotency):
+    token=os.environ.get(config['token_env'],'')
+    if not re.fullmatch(r'[A-Za-z0-9_-]{32,256}',token):raise HandoffError('HANDOFF_CREDENTIALS')
+    target=urlsplit(payload['url']);boundary='vba-'+secrets.token_hex(24)
+    prefix=(f'--{boundary}\r\nContent-Disposition: form-data; name="metadata"\r\n\r\n'+json.dumps(payload['metadata'])+'\r\n').encode()
+    parts=[]
+    for file in payload['files']:
+        role=file['role']
+        header=f'--{boundary}\r\nContent-Disposition: form-data; name="{role}"; filename="{role}.tar.gz"\r\nContent-Type: application/gzip\r\n\r\n'.encode()
+        parts.append((header,file))
+    ending=f'--{boundary}--\r\n'.encode()
+    length=len(prefix)+len(ending)+sum(len(header)+file['size']+2 for header,file in parts)
+    connection=(http.client.HTTPSConnection if target.scheme=='https' else http.client.HTTPConnection)(target.hostname,target.port,timeout=config['timeout_seconds'])
+    started=time.monotonic()
+    def transmit(chunk):
+        remaining=config['timeout_seconds']-(time.monotonic()-started)
+        if remaining<=0:raise HandoffError('HANDOFF_TIMEOUT')
+        if connection.sock:connection.sock.settimeout(remaining)
+        connection.send(chunk)
+    try:
+        connection.request('POST','/v1/submissions',headers={'Authorization':'Bearer '+token,
+            'Idempotency-Key':idempotency,'Content-Type':'multipart/form-data; boundary='+boundary,
+            'Content-Length':str(length)},body=None)
+        transmit(prefix)
+        for header,file in parts:
+            transmit(header)
+            with Path(file['path']).open('rb') as stream:
+                for chunk in iter(lambda:stream.read(1024*1024),b''):transmit(chunk)
+            transmit(b'\r\n')
+        transmit(ending)
+        response=connection.getresponse();body=response.read(65537)
+        if response.status not in {200,201}:return f'HTTP_{response.status}',None
+        if len(body)>65536:raise HandoffError('HANDOFF_RECEIPT_INVALID')
+        try:data=json.loads(body)
+        except ValueError:raise HandoffError('HANDOFF_RECEIPT_INVALID') from None
+        if (not isinstance(data,dict) or not isinstance(data.get('submission_id'),str)
+                or data.get('case_id')!=payload['metadata']['case_id']):raise HandoffError('HANDOFF_RECEIPT_INVALID')
+        return 'OK',{field:data[field] for field in ('submission_id','case_id','status') if field in data}
+    finally:connection.close()