baiyue-firstleg/scripts/order-submit-dependency-fault-proxy.py

"""TEST-only scoped synchronous dependency proxy. Default/expired mode is PASS.

Configuration and control token are read from a mode-0600 file, never argv values.
The runtime owns lifecycle; this file never invokes Docker, SQL or a shell.
"""
import argparse
import hmac
import http.client
import json
import os
import re
import threading
import time
import urllib.parse
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

HOP = {'connection', 'keep-alive', 'proxy-authenticate', 'proxy-authorization',
       'te', 'trailer', 'transfer-encoding', 'upgrade', 'host', 'content-length'}


def forward_headers(headers, length):
    connection = {k.lower(): v for k, v in headers.items()}.get('connection', '')
    excluded = HOP | {x.strip().lower() for x in connection.split(',')}
    result = {k: v for k, v in headers.items() if k.lower() not in excluded}
    result['Content-Length'] = str(length)
    return result


def transport(upstream, method, path, headers, body):
    url = urllib.parse.urlsplit(upstream)
    cls = http.client.HTTPSConnection if url.scheme == 'https' else http.client.HTTPConnection
    conn = cls(url.hostname, url.port, timeout=8)
    try:
        conn.request(method, path, body, forward_headers(headers, len(body)))
        response = conn.getresponse()
        return response.status, dict(response.getheaders()), response.read()
    finally:
        conn.close()


class Router:
    def __init__(self, config, forward=transport, clock=time.monotonic, record=None):
        self.config, self.forward, self.clock = config, forward, clock
        self.mode, self.deadline, self.generation = 'PASS', 0, 0
        self.lock = threading.Lock()
        self.record = record or (lambda row: None)
        self.counts = {'forwarded': 0, 'faulted': 0}
        customer = config['customerId']
        if not re.fullmatch(r'[1-9][0-9]{0,19}', customer):
            raise ValueError('INVALID_CUSTOMER')
        warehouse = config.get('receivingWarehouseId', '')
        if not re.fullmatch(r'[1-9][0-9]{0,19}', warehouse):
            raise ValueError('INVALID_RECEIVING_WAREHOUSE')
        if not config.get('requestIds') or any(not re.fullmatch(r'[A-Za-z0-9_.:-]{1,128}', value) for value in config['requestIds']):
            raise ValueError('BOUNDED_REQUEST_IDS_REQUIRED')
        self.paths = {'IAM': '/internal/iam/tenants/' + customer + '/order-eligibility',
                      'FINANCE': '/internal/finance/v1/wallet/accounts/' + customer,
                      'WAREHOUSE': '/internal/warehouses/0/' + warehouse + '/active'}

    def status(self):
        with self.lock:
            if self.mode != 'PASS' and self.clock() >= self.deadline:
                self.mode, self.deadline = 'PASS', 0
            return {'mode': self.mode, 'generation': self.generation,
                    'remainingSeconds': max(0, self.deadline - self.clock()), **self.counts}

    def set_mode(self, token, mode, ttl=0):
        if not hmac.compare_digest(str(token), self.config['controlToken']):
            raise PermissionError('CONTROL_AUTH_REQUIRED')
        with self.lock:
            # Even an authenticated malformed command disarms first.
            self.mode, self.deadline = 'PASS', 0
            if mode not in ('PASS', 'IAM', 'FINANCE', 'WAREHOUSE') or (mode != 'PASS' and
                    (type(ttl) is not int or not 1 <= ttl <= 60)):
                raise ValueError('INVALID_MODE_OR_TTL')
            self.mode = mode
            self.deadline = self.clock() + ttl if mode != 'PASS' else 0
            self.generation += 1
        return self.status()

    def dispatch(self, method, path, headers, body, peer):
        prefix, separator, rest = path.lstrip('/').partition('/')
        dependency = {'iam': 'IAM', 'finance': 'FINANCE', 'warehouse': 'WAREHOUSE'}.get(prefix)
        if not dependency or not separator:
            return 404, {}, b'Unknown dependency route'
        actual_path = '/' + rest
        normalized = {k.lower(): v for k, v in headers.items()}
        token = normalized.get('x-internal-service-token', '')
        authenticated = bool(token) and hmac.compare_digest(token, self.config['internalToken'])
        target = method == 'GET' and actual_path == self.paths[dependency]
        request_id = normalized.get('x-request-id', '')
        selected = target and peer in self.config['orderIps'] and authenticated and request_id in self.config['requestIds']
        mode = self.status()['mode']
        fault = selected and mode == dependency
        # Only bounded correlation IDs and this customer's selected paths enter evidence.
        safe_id = request_id if re.fullmatch(r'[A-Za-z0-9_.:-]{1,128}', request_id) else None
        if fault:
            result = (503, {'Content-Type': 'text/plain'}, b'TEST dependency temporarily unavailable')
        else:
            result = self.forward(self.config['upstreams'][dependency], method, actual_path, headers, body)
        with self.lock:
            self.counts['faulted' if fault else 'forwarded'] += 1
        if selected:
            self.record({'requestId': safe_id, 'method': method, 'path': actual_path,
                         'dependency': dependency, 'matched': fault, 'responseStatus': result[0],
                         'forwarded': not fault, 'internalTokenPresent': True,
                         'internalTokenValidated': True, 'mode': mode})
        return result


def main(path):
    config_path = Path(path)
    if config_path.is_symlink() or config_path.stat().st_mode & 0o077:
        raise RuntimeError('PRIVATE_CONFIG_REQUIRED')
    config = json.loads(config_path.read_text())
    evidence_path = config_path.parent / 'proxy-evidence.jsonl'
    evidence_lock = threading.Lock()

    def record(row):
        with evidence_lock:
            fd = os.open(evidence_path, os.O_WRONLY | os.O_APPEND | os.O_CREAT, 0o600)
            with os.fdopen(fd, 'w') as stream:
                stream.write(json.dumps(row) + '\n'); stream.flush(); os.fsync(stream.fileno())

    router = Router(config, record=record)

    class Handler(BaseHTTPRequestHandler):
        def log_message(self, *_args):
            pass

        def dispatch(self):
            try:
                length = int(self.headers.get('Content-Length', '0'))
                if length < 0 or length > 2 * 1024 * 1024:
                    raise ValueError('BODY_LIMIT')
                body = self.rfile.read(length)
                if self.server.control:
                    token = self.headers.get('X-Order-Fault-Control', '')
                    if not hmac.compare_digest(token, config['controlToken']):
                        raise PermissionError('CONTROL_AUTH_REQUIRED')
                    if self.command == 'GET' and self.path == '/status':
                        value = router.status()
                    elif self.command == 'POST' and self.path == '/mode':
                        value = json.loads(body)
                        value = router.set_mode(token, value.get('mode'), value.get('ttl', 0))
                    else:
                        raise ValueError('UNKNOWN_CONTROL')
                    status, headers, response = 200, {'Content-Type': 'application/json'}, json.dumps(value).encode()
                else:
                    status, headers, response = router.dispatch(self.command, self.path, dict(self.headers), body, self.client_address[0])
            except PermissionError:
                status, headers, response = 403, {}, b'Forbidden'
            except (ValueError, KeyError):
                if self.server.control and hmac.compare_digest(self.headers.get('X-Order-Fault-Control', ''), config['controlToken']):
                    router.set_mode(config['controlToken'], 'PASS')
                status, headers, response = 400, {}, b'Invalid request'
            except Exception:
                status, headers, response = 502, {}, b'Upstream transport failed'
            self.send_response(status)
            for name, value in headers.items():
                if name.lower() not in HOP:
                    self.send_header(name, value)
            self.send_header('Content-Length', str(len(response))); self.end_headers()
            if self.command != 'HEAD':
                self.wfile.write(response)

        do_GET = do_POST = do_PUT = do_PATCH = do_DELETE = do_HEAD = do_OPTIONS = dispatch

    business = ThreadingHTTPServer(('0.0.0.0', config['proxyPort']), Handler); business.control = False
    control = ThreadingHTTPServer(('127.0.0.1', config['controlPort']), Handler); control.control = True
    (config_path.parent / 'proxy.pid').write_text(str(os.getpid()))
    os.chmod(config_path.parent / 'proxy.pid', 0o600)
    threading.Thread(target=business.serve_forever, daemon=True).start()
    control.serve_forever()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(); parser.add_argument('--config', required=True)
    main(parser.parse_args().config)