newton-aggregator 0.4.12

//! Production JSON-RPC implementation of [`StateCommitOperatorClient`].
//!
//! Wraps `newt_getStateCommitProposal` (Prepare phase) and `newt_signStateCommit`
//! (Commit phase) HTTP calls per operator. Wire schema lives in
//! [`newton_prover_core::rpc::state_commit`]; the operator-side handlers live in
//! `crates/operator/src/state_commit_rpc.rs`.
//!
//! ## Design choices
//!
//! - **Raw `reqwest::Client` + hand-rolled JSON-RPC envelope.** Mirrors the
//!   established workspace pattern in `crates/gateway/src/task/operator.rs`.
//!   No `jsonrpsee::http_client` dependency.
//! - **HTTP/2 prior knowledge, 5s per-call timeout, TCP keepalive 60s.**
//!   Same defaults as `gateway::task::operator::OperatorClient`. The 5s
//!   timeout is well below the 60s `STATE_COMMIT_RECEIPT_TIMEOUT` so we
//!   surface a fast `Timeout` error to the orchestrator.
//! - **Error mapping** (in priority order):
//!   - `tokio::time::timeout` elapsed → [`OperatorClientError::Timeout`]
//!   - `reqwest::Error::is_timeout()` true → [`OperatorClientError::Timeout`]
//!   - Other `reqwest::Error` (connect, TLS, body decode, non-2xx HTTP) →
//!     [`OperatorClientError::Transport`]
//!   - JSON-RPC `error` field present:
//!     - On `sign_state_commit`, code `-32616` (operator's `ContractError`)
//!       AND message contains `"digest"` or `"newStateRoot"` →
//!       [`OperatorClientError::DigestDisagreement`]. This pattern-matches
//!       the operator's `handle_sign_state_commit` rejection. TODO: harden
//!       by introducing a dedicated `DigestMismatch` variant in
//!       `OperatorError` and a stable RPC code.
//!     - All other JSON-RPC errors → [`OperatorClientError::Malformed`]
//!   - Body parse failure / missing `result` field → [`OperatorClientError::Malformed`]
//! - **Socket discovery** is delegated to the caller — the orchestrator owns
//!   the socket map. The client takes a snapshot via constructor and a setter
//!   for live updates.
//!
//! See `docs/PRIVATE_DATA_STORAGE.md` §6 (Commit Protocol).

use std::{collections::HashMap, sync::Arc, time::Duration};

use alloy::{
    primitives::{Address, B256},
    signers::local::PrivateKeySigner,
};
use async_trait::async_trait;
use eigensdk::{crypto_bls::BlsG1Point, types::operator::OperatorId};
use newton_prover_chainio::operator_rpc_auth::{sign_authenticated, DEFAULT_EXPIRY_SECS};
use newton_prover_core::{
    rpc::state_commit::{
        GetStateCommitProposalRequest, GetStateCommitProposalResponse, SignStateCommitRequest, SignStateCommitResponse,
        StateCommitWire,
    },
    state_commit_registry::IStateRootCommittable::StateCommit,
};
use reqwest::{Client, ClientBuilder};
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::{debug, instrument, warn};

use super::operator_client::{OperatorClientError, OperatorProposal, StateCommitOperatorClient};

/// JSON-RPC error code returned by the operator for `OperatorError::ContractError`.
///
/// Matches `crates/operator/src/error.rs:181`. Used together with a message
/// substring match to recognize the operator's digest/root-mismatch refusal in
/// the Commit phase.
const OPERATOR_CONTRACT_ERROR_CODE: i64 = -32616;

/// Per-call HTTP timeout. Well below the 60s `STATE_COMMIT_RECEIPT_TIMEOUT` and
/// the 120s commit cadence so transient operators surface fast.
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);

/// Production `reqwest`-backed implementation of [`StateCommitOperatorClient`].
///
/// Holds a single shared HTTP client (TCP connection pool reused across
/// operators) plus a per-operator socket map kept in an `RwLock` so the
/// orchestrator can refresh it on registry rotation without rebuilding the
/// client.
///
/// Carries the aggregator's task-generator signer alongside the per-chain
/// task manager address so that `newt_getStateCommitProposal` and
/// `newt_signStateCommit` requests are wrapped in an `Authenticated<T>`
/// EIP-712 envelope before they go on the wire — operator-side dispatch
/// rejects unsigned or tampered envelopes via `recover_operator_rpc_signer`
/// + on-chain `isTaskGenerator(signer)`.
pub struct HttpStateCommitOperatorClient {
    http: Client,
    sockets: Arc<RwLock<HashMap<OperatorId, String>>>,
    timeout: Duration,
    /// Aggregator's task-generator signer used to wrap each outgoing
    /// state-commit RPC body in an EIP-712 `Authenticated<T>` envelope. The
    /// operator side recovers this signer and gates dispatch on
    /// `isTaskGenerator(signer)`.
    signer: PrivateKeySigner,
    /// Chain ID baked into both the `OperatorRpcCall.chainId` field and the
    /// EIP-712 domain so an envelope minted for one chain cannot replay
    /// against an operator handling another.
    chain_id: u64,
    /// Per-chain task manager address used as the EIP-712 verifying contract.
    /// Pinning to the task manager (not a generic constant) means an envelope
    /// minted against one chain's task manager can never verify under another
    /// chain's manager even if the chain id were spoofed.
    task_manager: Address,
}

impl std::fmt::Debug for HttpStateCommitOperatorClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("HttpStateCommitOperatorClient")
            .field("timeout_ms", &self.timeout.as_millis())
            .field("chain_id", &self.chain_id)
            .field("task_manager", &self.task_manager)
            .finish_non_exhaustive()
    }
}

impl HttpStateCommitOperatorClient {
    /// Construct a new client with the workspace-standard pool / keepalive
    /// settings and a 5-second per-call timeout.
    pub fn new(
        sockets: HashMap<OperatorId, String>,
        signer: PrivateKeySigner,
        chain_id: u64,
        task_manager: Address,
    ) -> Result<Self, reqwest::Error> {
        Self::with_timeout(sockets, DEFAULT_TIMEOUT, signer, chain_id, task_manager)
    }

    /// Construct with a caller-specified per-call timeout. Tests use this to
    /// drive the `Timeout` mapping path with sub-second deadlines.
    pub fn with_timeout(
        sockets: HashMap<OperatorId, String>,
        timeout: Duration,
        signer: PrivateKeySigner,
        chain_id: u64,
        task_manager: Address,
    ) -> Result<Self, reqwest::Error> {
        // reqwest's own timeout acts as a safety net set 1s above the
        // tokio::time::timeout wrapper around `send_fut` — the wrapper owns
        // structured error mapping (`OperatorClientError::Timeout`), and
        // matching reqwest's deadline to the same value would race the two,
        // randomly producing `Timeout` vs `Transport(timeout)` for the same
        // condition.
        let http = ClientBuilder::new()
            .pool_max_idle_per_host(8)
            .pool_idle_timeout(Duration::from_secs(90))
            .timeout(timeout + Duration::from_secs(1))
            .tcp_keepalive(Some(Duration::from_secs(60)))
            .tcp_nodelay(true)
            .build()?;

        Ok(Self {
            http,
            sockets: Arc::new(RwLock::new(sockets)),
            timeout,
            signer,
            chain_id,
            task_manager,
        })
    }

    /// Replace the per-operator socket map. Called by the orchestrator after
    /// a registry-driven operator-set rotation.
    pub async fn set_sockets(&self, sockets: HashMap<OperatorId, String>) {
        *self.sockets.write().await = sockets;
    }

    /// Resolve an operator id to a fully-qualified URL (`http://` prefixed).
    async fn url_for(&self, operator_id: &OperatorId) -> Result<String, OperatorClientError> {
        let sockets = self.sockets.read().await;
        let raw = sockets.get(operator_id).ok_or_else(|| OperatorClientError::Malformed {
            operator_id: hex::encode(operator_id),
            reason: format!("no socket registered for operator {}", hex::encode(operator_id)),
        })?;
        Ok(format_socket_as_http_url(raw))
    }

    /// Submit a single JSON-RPC call and return the parsed response envelope.
    ///
    /// Maps every failure mode onto [`OperatorClientError`] per the rules
    /// documented at the module level.
    async fn rpc_call<P, R>(
        &self,
        operator_id: &OperatorId,
        url: &str,
        method: &str,
        params: P,
    ) -> Result<R, OperatorClientError>
    where
        P: Serialize,
        R: for<'de> Deserialize<'de>,
    {
        let envelope = JsonRpcRequest {
            jsonrpc: "2.0",
            method,
            params,
            id: 1,
        };

        let send_fut = self.http.post(url).json(&envelope).send();
        let response = match tokio::time::timeout(self.timeout, send_fut).await {
            Ok(Ok(r)) => r,
            Ok(Err(e)) if e.is_timeout() => {
                return Err(OperatorClientError::Timeout {
                    operator_id: hex::encode(operator_id),
                    timeout_ms: self.timeout.as_millis() as u64,
                });
            }
            Ok(Err(e)) => {
                return Err(OperatorClientError::Transport {
                    operator_id: hex::encode(operator_id),
                    source: Box::new(e),
                });
            }
            Err(_elapsed) => {
                return Err(OperatorClientError::Timeout {
                    operator_id: hex::encode(operator_id),
                    timeout_ms: self.timeout.as_millis() as u64,
                });
            }
        };

        // HTTP-layer status check: non-2xx is transport, even with a JSON body.
        let status = response.status();
        if !status.is_success() {
            let body = response.text().await.unwrap_or_default();
            return Err(OperatorClientError::Transport {
                operator_id: hex::encode(operator_id),
                source: format!("HTTP {status}: {body}").into(),
            });
        }

        let envelope: JsonRpcResponse<R> = response.json().await.map_err(|e| {
            // A body that fails to deserialize either as `result` or `error` is
            // a schema violation — non-transient.
            OperatorClientError::Malformed {
                operator_id: hex::encode(operator_id),
                reason: format!("response body parse: {e}"),
            }
        })?;

        if let Some(err) = envelope.error {
            // The caller will branch on method to decide whether `-32616` +
            // digest/root substring counts as DigestDisagreement.
            return Err(OperatorClientError::Malformed {
                operator_id: hex::encode(operator_id),
                reason: format!("rpc error code={} message={}", err.code, err.message),
            });
        }

        envelope.result.ok_or_else(|| OperatorClientError::Malformed {
            operator_id: hex::encode(operator_id),
            reason: "rpc envelope missing both `result` and `error`".to_string(),
        })
    }
}

/// Lift a raw socket (`host:port` or `http://...` or `https://...`) into a
/// fully-qualified URL, defaulting to `http://`. Mirrors the gateway client
/// behaviour at `crates/gateway/src/task/operator.rs`.
fn format_socket_as_http_url(socket: &str) -> String {
    if socket.starts_with("http://") || socket.starts_with("https://") {
        socket.to_string()
    } else {
        format!("http://{socket}")
    }
}

/// Detect the operator's "I disagree" signal from the Commit-phase RPC error.
///
/// Matches code `-32616` (operator's `ContractError`) AND a message substring
/// that distinguishes digest/root mismatch from unrelated `ContractError`s
/// (e.g., a `pcr0_provider` failure that bubbled up as a `ContractError`).
///
/// TODO(post-MVP): replace with a dedicated `DigestMismatch` variant in
/// `OperatorError` and a stable RPC code.
fn is_digest_disagreement_error(reason: &str) -> bool {
    let code_match = reason.contains(&format!("code={OPERATOR_CONTRACT_ERROR_CODE}"));
    let msg_match =
        reason.contains("digest") || reason.contains("newStateRoot") || reason.contains("diverges from commit");
    code_match && msg_match
}

#[async_trait]
impl StateCommitOperatorClient for HttpStateCommitOperatorClient {
    #[instrument(skip(self), fields(operator_id = %hex::encode(operator_id), sequence_no))]
    async fn get_state_commit_proposal(
        &self,
        operator_id: &OperatorId,
        sequence_no: u64,
    ) -> Result<OperatorProposal, OperatorClientError> {
        let url = self.url_for(operator_id).await?;
        let req = GetStateCommitProposalRequest { sequence_no };

        // Wrap the request body in an EIP-712 `Authenticated<T>` envelope.
        // Signing is local-only (no I/O) and only fails on bincode size caps
        // for non-trivial inner bodies — `Malformed` is the closest fit since
        // the failure is on our side, not the operator's.
        let envelope = sign_authenticated(
            &self.signer,
            "newt_getStateCommitProposal",
            self.chain_id,
            self.task_manager,
            DEFAULT_EXPIRY_SECS,
            req,
        )
        .map_err(|e| OperatorClientError::Malformed {
            operator_id: hex::encode(operator_id),
            reason: format!("authenticate envelope: {e}"),
        })?;

        let resp: GetStateCommitProposalResponse = self
            .rpc_call(operator_id, &url, "newt_getStateCommitProposal", envelope)
            .await?;

        debug!(
            operator_id = %hex::encode(operator_id),
            sequence_no,
            new_state_root = %resp.new_state_root,
            "state-commit prepare: received proposal"
        );

        Ok(OperatorProposal {
            new_state_root: resp.new_state_root,
            da_cert_hash: resp.da_cert_hash,
            pcr0_commitment: resp.pcr0_commitment,
        })
    }

    #[instrument(
        skip(self, digest, commit),
        fields(
            operator_id = %hex::encode(operator_id),
            digest = %digest,
            sequence_no = commit.sequenceNo,
        )
    )]
    async fn sign_state_commit(
        &self,
        operator_id: &OperatorId,
        digest: B256,
        commit: &StateCommit,
        reference_timestamp: u32,
    ) -> Result<BlsG1Point, OperatorClientError> {
        let url = self.url_for(operator_id).await?;
        let sequence_no = commit.sequenceNo;

        let req = SignStateCommitRequest {
            digest,
            commit: StateCommitWire::from(commit),
            reference_timestamp,
        };

        // Wrap the request body in an EIP-712 `Authenticated<T>` envelope.
        // Signing is local-only (no I/O); a failure here is a local bug, not
        // an operator-side issue, and is reported via `Malformed`.
        let envelope = sign_authenticated(
            &self.signer,
            "newt_signStateCommit",
            self.chain_id,
            self.task_manager,
            DEFAULT_EXPIRY_SECS,
            req,
        )
        .map_err(|e| OperatorClientError::Malformed {
            operator_id: hex::encode(operator_id),
            reason: format!("authenticate envelope: {e}"),
        })?;

        let resp: SignStateCommitResponse = match self
            .rpc_call::<_, SignStateCommitResponse>(operator_id, &url, "newt_signStateCommit", envelope)
            .await
        {
            Ok(r) => r,
            Err(OperatorClientError::Malformed {
                operator_id: id,
                reason,
            }) if is_digest_disagreement_error(&reason) => {
                warn!(
                    operator_id = %id,
                    sequence_no,
                    reason = %reason,
                    "state-commit sign: operator refused (digest disagreement)"
                );
                return Err(OperatorClientError::DigestDisagreement {
                    operator_id: id,
                    sequence_no,
                });
            }
            Err(e) => return Err(e),
        };

        // Decode 2 × 32-byte big-endian U256 coordinates into a BlsG1Point.
        if resp.signature_bytes.len() != 64 {
            return Err(OperatorClientError::Malformed {
                operator_id: hex::encode(operator_id),
                reason: format!("expected 64-byte G1 encoding, got {}", resp.signature_bytes.len()),
            });
        }

        let g1 = decode_g1_point(&resp.signature_bytes).map_err(|reason| OperatorClientError::Malformed {
            operator_id: hex::encode(operator_id),
            reason,
        })?;

        Ok(g1)
    }
}

/// JSON-RPC 2.0 request envelope.
#[derive(Debug, Serialize)]
struct JsonRpcRequest<'a, P> {
    jsonrpc: &'a str,
    method: &'a str,
    params: P,
    id: u64,
}

/// JSON-RPC 2.0 response envelope.
#[derive(Debug, Deserialize)]
struct JsonRpcResponse<R> {
    #[allow(dead_code)]
    jsonrpc: Option<String>,
    result: Option<R>,
    error: Option<JsonRpcError>,
    #[allow(dead_code)]
    id: Option<u64>,
}

/// JSON-RPC 2.0 error object.
#[derive(Debug, Deserialize)]
struct JsonRpcError {
    code: i64,
    message: String,
}

/// Decode 64 raw bytes (2 × 32-byte big-endian coordinates: X, Y) into a
/// `BlsG1Point`. Mirrors the encoding in
/// `crates/operator/src/state_commit_rpc.rs::handle_sign_state_commit`.
fn decode_g1_point(bytes: &[u8]) -> Result<BlsG1Point, String> {
    use alloy::primitives::U256;
    use ark_bn254::{Fq, G1Affine};
    use ark_ff::PrimeField;

    if bytes.len() != 64 {
        return Err(format!("g1 point: want 64 bytes, got {}", bytes.len()));
    }
    let x = U256::from_be_slice(&bytes[0..32]);
    let y = U256::from_be_slice(&bytes[32..64]);

    let to_fq = |u: U256| -> Fq { Fq::from_be_bytes_mod_order(&u.to_be_bytes::<32>()) };
    let affine = G1Affine::new_unchecked(to_fq(x), to_fq(y));

    // Sanity: identity is acceptable (matches the test fake), so we do not
    // require subgroup membership here. The aggregator's BLS verification
    // step is the authoritative check.
    Ok(BlsG1Point::new(affine))
}

#[cfg(test)]
mod tests {
    use super::*;

    fn op_id(byte: u8) -> OperatorId {
        OperatorId::repeat_byte(byte)
    }

    /// Build a test client with a random signer, chain id 1, and zero task
    /// manager address. The state-commit envelope is honored end-to-end by
    /// the operator side; these tests exercise socket lookup + timeout +
    /// transport mappings only, so the auth fields are irrelevant — but the
    /// constructor signature requires them.
    fn test_client(sockets: HashMap<OperatorId, String>) -> HttpStateCommitOperatorClient {
        HttpStateCommitOperatorClient::new(sockets, PrivateKeySigner::random(), 1u64, Address::ZERO)
            .expect("build client")
    }

    fn test_client_with_timeout(
        sockets: HashMap<OperatorId, String>,
        timeout: Duration,
    ) -> HttpStateCommitOperatorClient {
        HttpStateCommitOperatorClient::with_timeout(sockets, timeout, PrivateKeySigner::random(), 1u64, Address::ZERO)
            .expect("build client")
    }

    #[test]
    fn format_socket_adds_http_prefix() {
        assert_eq!(format_socket_as_http_url("127.0.0.1:9000"), "http://127.0.0.1:9000");
        assert_eq!(format_socket_as_http_url("http://op:9000"), "http://op:9000");
        assert_eq!(format_socket_as_http_url("https://op:9000"), "https://op:9000");
    }

    #[test]
    fn digest_disagreement_pattern_matches() {
        assert!(is_digest_disagreement_error(
            "rpc error code=-32616 message=state-commit sign: digest mismatch — expected 0xaa got 0xbb"
        ));
        assert!(is_digest_disagreement_error(
            "rpc error code=-32616 message=state-commit sign: local root 0xaa diverges from commit.newStateRoot 0xbb — stale view, refusing to sign"
        ));
        // Different code — not a disagreement signal.
        assert!(!is_digest_disagreement_error(
            "rpc error code=-32603 message=internal error"
        ));
        // Same code but unrelated message (e.g., a pcr0_provider failure).
        assert!(!is_digest_disagreement_error(
            "rpc error code=-32616 message=pcr0_provider: enclave unreachable"
        ));
    }

    #[tokio::test]
    async fn url_for_unknown_operator_is_malformed() {
        let client = test_client(HashMap::new());
        let err = client.url_for(&op_id(0x01)).await.expect_err("unknown operator");
        assert!(matches!(
            err,
            OperatorClientError::Malformed { ref operator_id, ref reason }
                if operator_id == &hex::encode(op_id(0x01))
                    && reason.contains("no socket registered")
        ));
    }

    #[tokio::test]
    async fn set_sockets_replaces_map() {
        let client = test_client(HashMap::new());
        let mut next = HashMap::new();
        next.insert(op_id(0x42), "127.0.0.1:9000".to_string());
        client.set_sockets(next).await;

        let url = client.url_for(&op_id(0x42)).await.expect("registered");
        assert_eq!(url, "http://127.0.0.1:9000");
    }

    #[tokio::test]
    async fn timeout_classifies_as_timeout_error() {
        // Bind a TCP listener that accepts connections but never responds —
        // forces the per-call timeout to elapse.
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind");
        let addr = listener.local_addr().expect("addr");
        let _accept_loop = tokio::spawn(async move {
            // Accept and hold; never write a response.
            while let Ok((_socket, _)) = listener.accept().await {
                tokio::time::sleep(Duration::from_secs(60)).await;
            }
        });

        let mut sockets = HashMap::new();
        sockets.insert(op_id(0xab), addr.to_string());
        let client = test_client_with_timeout(sockets, Duration::from_millis(150));

        let err = client
            .get_state_commit_proposal(&op_id(0xab), 1)
            .await
            .expect_err("expect timeout");
        assert!(
            matches!(err, OperatorClientError::Timeout { ref operator_id, timeout_ms }
                if operator_id == &hex::encode(op_id(0xab)) && timeout_ms == 150),
            "got {err:?}"
        );
    }

    #[tokio::test]
    async fn connection_refused_classifies_as_transport() {
        // Reserve a port then drop the listener so the address is free.
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind");
        let addr = listener.local_addr().expect("addr");
        drop(listener);

        let mut sockets = HashMap::new();
        sockets.insert(op_id(0x07), addr.to_string());
        let client = test_client_with_timeout(sockets, Duration::from_secs(2));

        let err = client
            .get_state_commit_proposal(&op_id(0x07), 1)
            .await
            .expect_err("expect transport error");
        assert!(
            matches!(err, OperatorClientError::Transport { ref operator_id, .. }
                if operator_id == &hex::encode(op_id(0x07))),
            "got {err:?}"
        );
    }
}