state-hub/api/services/memo_dispatch.py

"""Private, explicitly queued State Hub memo dispatch. No API route or human act."""
from contextlib import contextmanager
import argparse
import fcntl
import json
import os
from pathlib import Path
import signal
import stat
import threading
import urllib.error
import urllib.request

from api.services.informed_decision_contract import MAX_ENVELOPE_BYTES, validate_envelope
from api.services.memo_forwarding import MemoForwarder

ORIGIN = 'http://informed-decision-workload.informed-decision.svc.cluster.local:8081'
TOKEN = '/var/run/secrets/state-hub/informed-decision/intake-token'
MAX_RESPONSE = 65536


class Refused(RuntimeError):
    pass


def read_json(path, limit):
    with Path(path).open('rb') as stream:
        raw = stream.read(limit + 1)
    if len(raw) > limit:
        raise Refused('input_too_large')
    return json.loads(raw)


def token_file(path):
    # Reopen on every attempt: projected files rotate by symlink replacement.
    with Path(path).open('rb') as stream:
        raw = stream.read(16385)
    value = raw.decode('ascii').strip()
    if not value or len(raw) > 16384 or any(c.isspace() for c in value):
        raise Refused('token_unavailable')
    return value


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, *args, **kwargs):
        return None


def transport(method, url, body, headers):
    if method != 'POST' or url != ORIGIN + '/v1/memos':
        raise Refused('route_refused')
    raw = json.dumps(body, allow_nan=False, separators=(',', ':')).encode()
    if len(raw) > MAX_ENVELOPE_BYTES:
        raise Refused('input_too_large')
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect())
    request = urllib.request.Request(url, data=raw, headers=headers, method=method)
    try:
        try:
            response = opener.open(request, timeout=5)
        except urllib.error.HTTPError as error:
            response = error
        with response:
            status = response.code
            raw = response.read(MAX_RESPONSE + 1)
            if len(raw) > MAX_RESPONSE:
                raise OSError('transport_unavailable')
            try:
                body = json.loads(raw)
            except (ValueError, UnicodeError):
                body = {}
            return status, body
    except Exception:
        # Never propagate URLs, response bodies, credentials or transport errors.
        raise OSError('transport_unavailable') from None


class Dispatcher:
    def __init__(self, directory, *, send=transport, read_token=lambda: token_file(TOKEN), clock=None):
        self.directory = Path(directory)
        if not self.directory.is_absolute():
            raise Refused('absolute_queue_path_required')
        self.directory.mkdir(mode=0o700, parents=False, exist_ok=True)
        info = self.directory.lstat()
        if not stat.S_ISDIR(info.st_mode) or info.st_uid != os.getuid() or stat.S_IMODE(info.st_mode) != 0o700:
            raise Refused('private_queue_required')
        # SQLite files inherit 0600; reject existing unsafe custody, including links.
        self.path = self.directory / 'forward.sqlite'
        for name in ('forward.sqlite', 'forward.sqlite-wal', 'forward.sqlite-shm', 'dispatch.lock'):
            path = self.directory / name
            if path.exists() or path.is_symlink():
                info = path.lstat()
                if not stat.S_ISREG(info.st_mode) or info.st_uid != os.getuid() or stat.S_IMODE(info.st_mode) != 0o600:
                    raise Refused('private_queue_required')
        self.send, self.read_token, self.clock = send, read_token, clock
        with self.lock():
            previous = os.umask(0o077)
            try:
                self.ledger = MemoForwarder(self.path, transport=send,
                    intake_base_url=ORIGIN, hub_base_url='http://state-hub.state-hub.svc.cluster.local:8000',
                    token_reader=read_token, **({'clock': clock} if clock else {}))
            finally:
                os.umask(previous)

    @contextmanager
    def lock(self):
        fd = os.open(self.directory / 'dispatch.lock', os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600)
        try:
            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
            previous = os.umask(0o077)
            try:
                yield
            finally:
                os.umask(previous)
        except BlockingIOError:
            raise Refused('queue_busy') from None
        finally:
            os.close(fd)

    def enqueue(self, envelope):
        validate_envelope(envelope)
        if envelope['memo_class'] != 'decision-record' or envelope['review_group'] != 'net-kingdom-admins':
            raise Refused('routing_refused')
        with self.lock():
            return self.ledger.submit(envelope, dispatch=False)

    def once(self):
        # Hold across remote attempt + local receipt to serialize other processes.
        # If killed after POST, restart replays the same envelope/key. Intake owns
        # native create uncertainty and refuses automatic native-approval repost.
        with self.lock():
            return self.ledger.run_due()

    def counts(self):
        with self.lock(), self.ledger._connect() as db:
            return dict(db.execute('SELECT state,COUNT(*) FROM pose_ledger GROUP BY state'))


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--queue', required=True, type=Path)
    commands = parser.add_subparsers(dest='command', required=True)
    commands.add_parser('enqueue').add_argument('--envelope', required=True, type=Path)
    for name in ('once', 'run', 'status'):
        commands.add_parser(name)
    args = parser.parse_args(argv)
    try:
        dispatcher = Dispatcher(args.queue)
        if args.command == 'enqueue':
            result = dispatcher.enqueue(read_json(args.envelope, MAX_ENVELOPE_BYTES))
            print(json.dumps({'status': result.status, 'detail': result.detail}))
        elif args.command == 'once':
            result = dispatcher.once()
            print(json.dumps({'status': result.status if result else 'idle'}))
        elif args.command == 'status':
            print(json.dumps(dispatcher.counts()))
        else:
            stopped = threading.Event()
            for sig in (signal.SIGTERM, signal.SIGINT):
                signal.signal(sig, lambda *_: stopped.set())
            while not stopped.is_set():
                try:
                    dispatcher.once()
                except Refused as exc:
                    if str(exc) != 'queue_busy':
                        raise
                stopped.wait(2)
        return 0
    except Exception:
        print('{"error":"dispatch_refused"}')
        return 1


if __name__ == '__main__':
    raise SystemExit(main())