Async Routing for MPC: Procedural Workflow Guide

import os
import uuid
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import serialization
from pika import BlockingConnection, ConnectionParameters, PlainCredentials, BasicProperties
from pika.exceptions import AMQPConnectionError

class MPCQueueProvisioner:
    def __init__(self, broker_url: str, session_id: str, hmac_key: bytes):
        self.session_id = session_id
        self.exchange = f"mpc.{session_id}.exchange"
        self.queue = f"mpc.{session_id}.shares"
        self.dlq = f"mpc.{session_id}.dlq"
        # Shared HMAC key so producers and consumers can verify integrity.
        # Distribute out-of-band over the same secure channel that publishes
        # session_params.
        self.hmac_key = hmac_key
        self.conn_params = ConnectionParameters(
            host=broker_url,
            credentials=PlainCredentials(os.getenv("MQ_USER"), os.getenv("MQ_PASS")),
            ssl_options=None,  # TLS 1.3 enforced at broker level
            heartbeat=600
        )

    def initialize_session(self) -> dict:
        """Generate session crypto params and provision routing topology."""
        private_key = ec.generate_private_key(ec.SECP256R1())
        public_key = private_key.public_key()
        
        session_params = {
            "session_id": self.session_id,
            "curve": "secp256r1",
            "threshold": 2,
            "total_shares": 3,
            "nonce_range": (0, 2**128)
        }
        
        try:
            conn = BlockingConnection(self.conn_params)
            channel = conn.channel()
            channel.exchange_declare(exchange=self.exchange, exchange_type="direct", durable=True)
            channel.queue_declare(queue=self.queue, durable=True, arguments={"x-dead-letter-exchange": self.dlq})
            channel.queue_declare(queue=self.dlq, durable=True)
            channel.queue_bind(queue=self.queue, exchange=self.exchange, routing_key="shares")
            conn.close()
        except AMQPConnectionError as e:
            raise RuntimeError(f"Broker provisioning failed: {e}") from e
            
        return session_params

import hmac
import hashlib
import json
import struct
from typing import List, Tuple

class MPCShareProducer:
    def __init__(self, session_params: dict, provisioner: MPCQueueProvisioner):
        self.session_id = session_params["session_id"]
        self.exchange = f"mpc.{self.session_id}.exchange"
        # Reuse the provisioner's shared HMAC key so consumers can verify.
        self.provisioner = provisioner
        self.hmac_key = provisioner.hmac_key

    def _compute_hmac(self, payload: bytes) -> bytes:
        return hmac.new(self.hmac_key, payload, hashlib.sha256).digest()

    def _pack_session_id(self) -> bytes:
        # Pad/truncate to a fixed 16-byte field so the wire frame is stable.
        return self.session_id.encode("utf-8")[:16].ljust(16, b"\x00")

    def publish_coordinate_shares(self, coordinates: List[Tuple[float, float]], shares: List[bytes]) -> None:
        """Split coordinates, attach HMAC, and dispatch to async queue."""
        conn = BlockingConnection(self.provisioner.conn_params)
        channel = conn.channel()
        session_bytes = self._pack_session_id()

        for idx, share in enumerate(shares):
            # Frame: [session_id(16)] [share_index(4)] [hmac(32)] [payload]
            # Use little-endian, no-alignment-padding format ("<") so the
            # 52-byte header size is portable across architectures.
            frame = struct.pack("<16sI32s", session_bytes, idx, self._compute_hmac(share)) + share

            channel.basic_publish(
                exchange=self.exchange,
                routing_key="shares",
                body=frame,
                properties=BasicProperties(
                    delivery_mode=2,  # Persistent
                    content_type="application/octet-stream",
                    message_id=str(uuid.uuid4())
                )
            )
        conn.close()

class MPCShareConsumer:
    def __init__(self, provisioner: MPCQueueProvisioner, expected_shares: int):
        self.session_id = provisioner.session_id
        self.expected = expected_shares
        self.buffer: dict[int, bytes] = {}
        # Reuse the same provisioner instance so the HMAC key matches the
        # producer's key (consumers MUST NOT create their own provisioner).
        self.provisioner = provisioner

    def on_message(self, ch, method, properties, body: bytes):
        """Callback for async share ingestion with HMAC verification."""
        if len(body) < 52:
            ch.basic_nack(delivery_tag=method.delivery_tag, requeue=False)
            return

        session_bytes, idx, received_hmac = struct.unpack("<16sI32s", body[:52])
        # Strip the null padding applied by the producer before comparing.
        if session_bytes.rstrip(b"\x00").decode("utf-8") != self.session_id:
            ch.basic_nack(delivery_tag=method.delivery_tag, requeue=False)
            return

        payload = body[52:]
        expected_hmac = hmac.new(self.provisioner.hmac_key, payload, hashlib.sha256).digest()
        if not hmac.compare_digest(received_hmac, expected_hmac):
            ch.basic_nack(delivery_tag=method.delivery_tag, requeue=False)
            return

        self.buffer[idx] = payload
        ch.basic_ack(delivery_tag=method.delivery_tag)

        if len(self.buffer) >= self.expected:
            self._trigger_round_completion()

    def _trigger_round_completion(self) -> None:
        """Reconstruct spatial payload and dispatch to MPC compute kernel."""
        reconstructed = b"".join(self.buffer[i] for i in sorted(self.buffer))
        # Pass to MPC computation engine (e.g., MPyC, ABY, or custom spatial kernel)
        print(f"[MPC Router] Round {self.session_id} complete. Payload size: {len(reconstructed)} bytes")