hypercall_api/

middleware.rs

use axum::{
    extract::{Request, State},
    http::{Method, StatusCode},
    middleware::Next,
    response::{IntoResponse, Response},
};
use sonic_rs::json;
use std::sync::Arc;
use tokio::sync::Mutex;

use crate::sonic_json::SonicJson;

use crate::caches::rate_limit::{RateLimitAction, RateLimitCache, RateLimitError, RateLimitInfo};
use crate::observability_boundary::AuthFailureRecorder;
use crate::request_auth::{authorize_signer, RequestAuthError};
use crate::runtime_status::{ReadinessGate, StandbyReplayProgress};
use crate::signed_actions::recover_signed_action;
use hypercall_types::{WalletAddress, API_ROUTE_PROFILE_IMAGE};
use std::str::FromStr;

const SIGNED_JSON_BODY_LIMIT_BYTES: usize = 1_048_576;
const PROFILE_IMAGE_BODY_LIMIT_BYTES: usize = 7_000_000;

#[derive(Clone)]
pub struct SignatureMiddlewareState {
    pub agent_auth: Arc<dyn hypercall_runtime_api::AgentAuthProvider>,
    pub auth_failure_recorder: Arc<dyn AuthFailureRecorder>,
    pub signing_chain_id: u64,
}

#[cfg(feature = "otel-tracing")]
use opentelemetry::propagation::TextMapPropagator;
#[cfg(feature = "otel-tracing")]
use opentelemetry_sdk::propagation::TraceContextPropagator;
#[cfg(feature = "otel-tracing")]
use tracing_opentelemetry::OpenTelemetrySpanExt;

/// Helper function to create a JSON error response
fn json_error_response(status: StatusCode, error: &str, message: &str) -> Response {
    let body = json!({
        "error": error,
        "message": message
    });
    (status, SonicJson(body)).into_response()
}

/// A header extractor that implements OpenTelemetry's Extractor trait for axum headers.
#[cfg(feature = "otel-tracing")]
struct HeaderExtractor<'a>(&'a axum::http::HeaderMap);

#[cfg(feature = "otel-tracing")]
impl opentelemetry::propagation::Extractor for HeaderExtractor<'_> {
    fn get(&self, key: &str) -> Option<&str> {
        self.0.get(key).and_then(|v| v.to_str().ok())
    }

    fn keys(&self) -> Vec<&str> {
        self.0.keys().map(|k| k.as_str()).collect()
    }
}

/// Middleware that extracts W3C Trace Context from incoming HTTP headers.
///
/// When the `otel-tracing` feature is enabled, this extracts `traceparent` and `tracestate`
/// headers and sets the OpenTelemetry context on the current tracing span, enabling
/// distributed tracing across service boundaries.
///
/// This middleware should be applied early in the middleware stack (after TraceLayer)
/// so that the extracted context is available for all subsequent handlers.
#[cfg(feature = "otel-tracing")]
pub async fn trace_context_middleware(req: Request, next: Next) -> Response {
    let propagator = TraceContextPropagator::new();
    let parent_cx = propagator.extract(&HeaderExtractor(req.headers()));

    // Set the extracted context as the parent for the current span
    tracing::Span::current().set_parent(parent_cx);

    next.run(req).await
}

/// Middleware that blocks all requests (except /health, /ready, /metrics) until the service is ready.
///
/// This ensures clients don't receive stale data while the portfolio is being restored
/// from a snapshot or catching up from event streams.
///
/// Uses the ReadinessRegistry to check all component readiness and provides detailed
/// component breakdown in the 503 response.
#[derive(Clone)]
pub struct ReadinessMiddlewareState {
    pub readiness: Arc<dyn ReadinessGate>,
    pub standby_progress: Option<Arc<dyn StandbyReplayProgress>>,
    pub standby_promote: Option<Arc<Mutex<Option<tokio::sync::oneshot::Sender<()>>>>>,
}

impl ReadinessMiddlewareState {
    pub fn primary(readiness: Arc<dyn ReadinessGate>) -> Self {
        Self {
            readiness,
            standby_progress: None,
            standby_promote: None,
        }
    }

    pub fn new(
        readiness: Arc<dyn ReadinessGate>,
        standby_progress: Option<Arc<dyn StandbyReplayProgress>>,
        standby_promote: Option<Arc<Mutex<Option<tokio::sync::oneshot::Sender<()>>>>>,
    ) -> Self {
        Self {
            readiness,
            standby_progress,
            standby_promote,
        }
    }
}

pub async fn readiness_middleware(
    State(state): State<ReadinessMiddlewareState>,
    req: Request,
    next: Next,
) -> Response {
    let path = req.uri().path();

    if is_readiness_bypass_path(path) {
        return next.run(req).await;
    }

    // Block all other requests until ready
    if !state.readiness.all_ready() {
        if is_standby_monitoring_get(req.method(), path)
            && standby_is_caught_up_and_unpromoted(&state).await
        {
            return next.run(req).await;
        }

        let components = state.readiness.reports();
        return (
            StatusCode::SERVICE_UNAVAILABLE,
            SonicJson(json!({
                "error": "service_not_ready",
                "message": "Service is starting up. Please retry shortly.",
                "components": components
            })),
        )
            .into_response();
    }

    next.run(req).await
}

fn is_readiness_bypass_path(path: &str) -> bool {
    matches!(
        path,
        "/health"
            | "/ready"
            | "/metrics"
            | "/standby-ready"
            | "/admin/promote"
            | "/admin/drain"
            | "/admin/undrain"
            | "/drain-status"
            | "/recovery-safety"
            | "/recovery-safety/alert"
            | "/monitoring/recovery-safety"
    )
}

fn is_standby_monitoring_get(method: &Method, path: &str) -> bool {
    if method != Method::GET {
        return false;
    }

    matches!(
        path,
        "/monitoring/integrity"
            | "/monitoring/orderbooks"
            | "/monitoring/accounts"
            | "/monitoring/positions"
            | "/monitoring/directive-outbox"
            | "/monitoring/deposits"
            | "/monitoring/engine-state-digest"
            | "/recovery-safety"
            | "/recovery-safety/alert"
            | "/monitoring/recovery-safety"
            | "/monitoring/trading-halts"
            | "/monitoring/vol-oracles"
            | "/monitoring/price-oracles"
            | "/monitoring/vol-surface"
    )
}

async fn standby_is_caught_up_and_unpromoted(state: &ReadinessMiddlewareState) -> bool {
    let Some(progress) = &state.standby_progress else {
        return false;
    };
    if !progress.is_caught_up() {
        return false;
    }

    match &state.standby_promote {
        Some(promote) => promote.lock().await.is_some(),
        None => false,
    }
}

#[derive(Clone)]
pub struct AuthenticatedUser {
    pub wallet_address: WalletAddress,
}

#[derive(Clone)]
pub struct AccountContext {
    pub account_contract_address: WalletAddress,
    pub api_wallet_address: WalletAddress,
}

#[derive(Clone, Debug)]
pub struct SignerContext {
    pub wallet_address: WalletAddress, // The account being acted upon
    pub signer_address: WalletAddress, // The agent that signed (may equal wallet)
}

/// Middleware for signature-based authentication with agent authorization
pub async fn signature_and_agent_middleware(
    State(state): State<SignatureMiddlewareState>,
    req: Request,
    next: Next,
) -> Response {
    match signature_and_agent_middleware_inner(state, req, next).await {
        Ok(response) => response,
        Err(response) => response,
    }
}

async fn signature_and_agent_middleware_inner(
    state: SignatureMiddlewareState,
    req: Request,
    next: Next,
) -> Result<Response, Response> {
    // Extract the request body to get signature, nonce, and wallet
    let (mut parts, body) = req.into_parts();
    let body_limit = match (parts.uri.path(), &parts.method) {
        (API_ROUTE_PROFILE_IMAGE, &Method::POST) => PROFILE_IMAGE_BODY_LIMIT_BYTES,
        _ => SIGNED_JSON_BODY_LIMIT_BYTES,
    };

    let body_bytes = match axum::body::to_bytes(body, body_limit).await {
        Ok(bytes) => bytes,
        Err(_) => {
            return Err(json_error_response(
                StatusCode::PAYLOAD_TOO_LARGE,
                "payload_too_large",
                "Request body exceeds the allowed size",
            ));
        }
    };

    // Parse the JSON body
    let request_data: sonic_rs::Value = match sonic_rs::from_slice(&body_bytes) {
        Ok(json) => json,
        Err(e) => {
            tracing::error!("Failed to parse request body: {}", e);
            return Err(json_error_response(
                StatusCode::BAD_REQUEST,
                "invalid_json",
                &format!("Failed to parse JSON: {}", e),
            ));
        }
    };

    let recovered = recover_signed_action(
        parts.uri.path(),
        &parts.method,
        &request_data,
        state.signing_chain_id,
    )
    .map_err(|e| {
        if let Some(reason) = e.auth_failure_reason {
            state.auth_failure_recorder.record_auth_failure(reason);
        }
        tracing::error!(
            error = e.error,
            message = %e.message,
            "Failed to recover signed action"
        );
        json_error_response(e.status, e.error, &e.message)
    })?;

    let signer_ctx = authorize_signer(
        state.agent_auth.as_ref(),
        recovered.wallet,
        recovered.signer,
    )
    .map_err(|e| match e {
        RequestAuthError::Unauthorized { wallet, signer } => {
            tracing::warn!(
                "Signer {} is not authorized to act for wallet {}",
                signer,
                wallet
            );
            json_error_response(StatusCode::UNAUTHORIZED, "unauthorized", &e.to_string())
        }
        RequestAuthError::Signature(_) => json_error_response(
            StatusCode::BAD_REQUEST,
            "signature_verification_failed",
            &e.to_string(),
        ),
    })?;

    // Inject SignerContext into request extensions
    parts.extensions.insert(signer_ctx);

    // Reconstruct the request with the original body
    let req = Request::from_parts(parts, axum::body::Body::from(body_bytes));

    Ok(next.run(req).await)
}

/// Create a 429 Too Many Requests response with rate limit headers.
fn rate_limit_exceeded_response(err: &RateLimitError) -> Response {
    match err {
        RateLimitError::Exceeded {
            retry_after_secs,
            limit,
            ..
        } => {
            let body = json!({
                "error": "rate_limit_exceeded",
                "message": format!("{}", err),
                "retry_after_secs": retry_after_secs,
                "limit": limit,
            });

            let mut response = (StatusCode::TOO_MANY_REQUESTS, SonicJson(body)).into_response();

            // Add Retry-After header
            response
                .headers_mut()
                .insert("Retry-After", retry_after_secs.to_string().parse().unwrap());
            response
                .headers_mut()
                .insert("X-RateLimit-Limit", limit.to_string().parse().unwrap());
            response
                .headers_mut()
                .insert("X-RateLimit-Remaining", "0".parse().unwrap());

            response
        }
        RateLimitError::ServiceUnavailable { message } => {
            let body = json!({
                "error": "service_unavailable",
                "message": message,
            });

            (StatusCode::SERVICE_UNAVAILABLE, SonicJson(body)).into_response()
        }
    }
}

/// Add rate limit headers to a response.
fn add_rate_limit_headers(response: &mut Response, info: &RateLimitInfo) {
    response
        .headers_mut()
        .insert("X-RateLimit-Limit", info.limit.to_string().parse().unwrap());
    response.headers_mut().insert(
        "X-RateLimit-Remaining",
        info.remaining.to_string().parse().unwrap(),
    );
    response.headers_mut().insert(
        "X-RateLimit-Reset",
        info.reset_at.to_string().parse().unwrap(),
    );
}

/// Extract wallet address from SignerContext extension.
fn extract_wallet(req: &Request) -> Option<WalletAddress> {
    req.extensions()
        .get::<SignerContext>()
        .map(|ctx| ctx.wallet_address)
}

/// State required for rate limiting middleware.
#[derive(Clone)]
pub struct RateLimitState {
    pub rate_limiter: Arc<RateLimitCache>,
}

/// Middleware to check order placement rate limits.
///
/// Must be applied AFTER signature middleware (needs SignerContext).
pub async fn order_rate_limit_middleware(
    State(state): State<RateLimitState>,
    req: Request,
    next: Next,
) -> Response {
    // Extract wallet from SignerContext (set by signature middleware)
    let wallet = req
        .extensions()
        .get::<SignerContext>()
        .map(|ctx| ctx.wallet_address);

    if let Some(wallet) = wallet {
        // Check order placement rate limit
        match state
            .rate_limiter
            .check_and_increment(&wallet, RateLimitAction::OrderPlacement)
            .await
        {
            Ok(info) => {
                let mut response = next.run(req).await;
                // Add rate limit headers to successful response
                response
                    .headers_mut()
                    .insert("X-RateLimit-Limit", info.limit.to_string().parse().unwrap());
                response.headers_mut().insert(
                    "X-RateLimit-Remaining",
                    info.remaining.to_string().parse().unwrap(),
                );
                response.headers_mut().insert(
                    "X-RateLimit-Reset",
                    info.reset_at.to_string().parse().unwrap(),
                );
                response
            }
            Err(err) => rate_limit_exceeded_response(&err),
        }
    } else {
        // No wallet context - allow through (signature middleware will handle auth)
        next.run(req).await
    }
}

/// Middleware to check order cancellation rate limits.
///
/// Must be applied AFTER signature middleware (needs SignerContext).
pub async fn cancel_rate_limit_middleware(
    State(state): State<RateLimitState>,
    req: Request,
    next: Next,
) -> Response {
    // Extract wallet from SignerContext
    let wallet = req
        .extensions()
        .get::<SignerContext>()
        .map(|ctx| ctx.wallet_address);

    if let Some(wallet) = wallet {
        match state
            .rate_limiter
            .check_and_increment(&wallet, RateLimitAction::OrderCancellation)
            .await
        {
            Ok(info) => {
                let mut response = next.run(req).await;
                response
                    .headers_mut()
                    .insert("X-RateLimit-Limit", info.limit.to_string().parse().unwrap());
                response.headers_mut().insert(
                    "X-RateLimit-Remaining",
                    info.remaining.to_string().parse().unwrap(),
                );
                response.headers_mut().insert(
                    "X-RateLimit-Reset",
                    info.reset_at.to_string().parse().unwrap(),
                );
                response
            }
            Err(err) => rate_limit_exceeded_response(&err),
        }
    } else {
        next.run(req).await
    }
}

/// Middleware to check general API rate limits.
///
/// This can be applied to any endpoint. For authenticated endpoints,
/// it uses the wallet from SignerContext. For public endpoints,
/// it attempts to extract wallet from query parameters.
pub async fn api_rate_limit_middleware(
    State(state): State<RateLimitState>,
    req: Request,
    next: Next,
) -> Response {
    // Try to get wallet from SignerContext first (authenticated endpoints)
    let wallet = req
        .extensions()
        .get::<SignerContext>()
        .map(|ctx| ctx.wallet_address);

    // If no SignerContext, try to extract from query params (for public endpoints)
    let wallet = wallet.or_else(|| {
        req.uri()
            .query()
            .and_then(|q| {
                // Parse query string to find "wallet=" parameter
                q.split('&')
                    .find(|p| p.starts_with("wallet="))
                    .map(|p| &p[7..]) // Skip "wallet="
            })
            .and_then(|w| WalletAddress::from_str(w).ok())
    });

    if let Some(wallet) = wallet {
        match state
            .rate_limiter
            .check_and_increment(&wallet, RateLimitAction::ApiRequest)
            .await
        {
            Ok(info) => {
                let mut response = next.run(req).await;
                response
                    .headers_mut()
                    .insert("X-RateLimit-Limit", info.limit.to_string().parse().unwrap());
                response.headers_mut().insert(
                    "X-RateLimit-Remaining",
                    info.remaining.to_string().parse().unwrap(),
                );
                response.headers_mut().insert(
                    "X-RateLimit-Reset",
                    info.reset_at.to_string().parse().unwrap(),
                );
                response
            }
            Err(err) => rate_limit_exceeded_response(&err),
        }
    } else {
        // No wallet identified - allow through without rate limiting
        // (could add IP-based rate limiting here in the future)
        next.run(req).await
    }
}

/// Combined middleware for write routes that applies appropriate rate limit
/// based on the request path and method.
///
/// Must be applied AFTER signature middleware (needs SignerContext).
///
/// Rate limit mapping:
/// - POST /order, PUT /order -> OrderPlacement
/// - DELETE /order, DELETE /order_cloid -> OrderCancellation
/// - All other write routes -> ApiRequest
pub async fn write_route_rate_limit_middleware(
    State(state): State<RateLimitState>,
    req: Request,
    next: Next,
) -> Response {
    use axum::http::Method;
    use hypercall_types::{API_ROUTE_ORDER, API_ROUTE_ORDER_CLOID, API_ROUTE_RFQ_REQUEST};

    // Determine which rate limit action applies based on path and method
    let action = match (req.method(), req.uri().path()) {
        (&Method::POST, API_ROUTE_ORDER) | (&Method::PUT, API_ROUTE_ORDER) => {
            RateLimitAction::OrderPlacement
        }
        (&Method::DELETE, API_ROUTE_ORDER) | (&Method::DELETE, API_ROUTE_ORDER_CLOID) => {
            RateLimitAction::OrderCancellation
        }
        // RFQ submit fans out to every connected QP — throttle per wallet
        (&Method::POST, API_ROUTE_RFQ_REQUEST) => RateLimitAction::RfqSubmit,
        // MMP and admin endpoints use general API rate limit
        _ => RateLimitAction::ApiRequest,
    };

    // Extract wallet from SignerContext (set by signature middleware)
    let wallet = extract_wallet(&req);

    if let Some(wallet) = wallet {
        match state
            .rate_limiter
            .check_and_increment(&wallet, action)
            .await
        {
            Ok(info) => {
                let mut response = next.run(req).await;
                add_rate_limit_headers(&mut response, &info);
                response
            }
            Err(err) => rate_limit_exceeded_response(&err),
        }
    } else {
        // No wallet context yet - allow through (signature middleware will handle auth)
        next.run(req).await
    }
}

// Extension trait to extract authenticated user from request
pub mod auth_ext {
    use super::{AccountContext, AuthenticatedUser, SignerContext};
    use crate::error::ApiError;
    use axum::{async_trait, extract::FromRequestParts, http::request::Parts};

    #[async_trait]
    impl<S> FromRequestParts<S> for AuthenticatedUser
    where
        S: Send + Sync,
    {
        type Rejection = ApiError;

        async fn from_request_parts(
            parts: &mut Parts,
            _state: &S,
        ) -> Result<Self, Self::Rejection> {
            parts
                .extensions
                .get::<AuthenticatedUser>()
                .cloned()
                .ok_or_else(|| ApiError::unauthorized("Authentication required"))
        }
    }

    #[async_trait]
    impl<S> FromRequestParts<S> for AccountContext
    where
        S: Send + Sync,
    {
        type Rejection = ApiError;

        async fn from_request_parts(
            parts: &mut Parts,
            _state: &S,
        ) -> Result<Self, Self::Rejection> {
            parts
                .extensions
                .get::<AccountContext>()
                .cloned()
                .ok_or_else(|| ApiError::unauthorized("Account context required"))
        }
    }

    #[async_trait]
    impl<S> FromRequestParts<S> for SignerContext
    where
        S: Send + Sync,
    {
        type Rejection = ApiError;

        async fn from_request_parts(
            parts: &mut Parts,
            _state: &S,
        ) -> Result<Self, Self::Rejection> {
            parts
                .extensions
                .get::<SignerContext>()
                .cloned()
                .ok_or_else(|| {
                    ApiError::unauthorized("Signer context required - valid signature is required")
                })
        }
    }
}

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

    struct TestReadinessGate {
        ready: bool,
    }

    impl ReadinessGate for TestReadinessGate {
        fn reports(&self) -> Vec<crate::models::ReadinessComponentReport> {
            vec![crate::models::ReadinessComponentReport {
                name: "engine".to_string(),
                ready: self.ready,
                detail: Some(if self.ready { "Ready" } else { "CatchingUp" }.to_string()),
            }]
        }

        fn all_ready(&self) -> bool {
            self.ready
        }
    }

    #[derive(Clone, Default)]
    struct TestStandbyReplayProgress {
        caught_up: bool,
    }

    impl StandbyReplayProgress for TestStandbyReplayProgress {
        fn commands_replayed(&self) -> u64 {
            0
        }

        fn is_caught_up(&self) -> bool {
            self.caught_up
        }

        fn replay_cursor_seq(&self) -> Option<i64> {
            None
        }

        fn latest_stream_seq(&self) -> Option<u64> {
            None
        }

        fn stream_lag(&self) -> Option<u64> {
            None
        }

        fn last_replayed_seq(&self) -> Option<i64> {
            None
        }

        fn last_replay_unix_ms(&self) -> Option<u64> {
            None
        }
    }

    fn not_ready_registry() -> Arc<dyn ReadinessGate> {
        Arc::new(TestReadinessGate { ready: false })
    }

    #[test]
    fn recovery_safety_bypasses_readiness_gate() {
        assert!(is_readiness_bypass_path("/monitoring/recovery-safety"));
        assert!(is_readiness_bypass_path("/recovery-safety"));
        assert!(is_readiness_bypass_path("/recovery-safety/alert"));
    }

    #[test]
    fn standby_monitoring_bypass_only_allows_read_only_state_routes() {
        assert!(is_standby_monitoring_get(
            &Method::GET,
            "/monitoring/integrity"
        ));
        assert!(is_standby_monitoring_get(
            &Method::GET,
            "/monitoring/orderbooks"
        ));
        assert!(is_standby_monitoring_get(
            &Method::GET,
            "/monitoring/accounts"
        ));
        assert!(is_standby_monitoring_get(
            &Method::GET,
            "/monitoring/positions"
        ));
        assert!(is_standby_monitoring_get(
            &Method::GET,
            "/monitoring/directive-outbox"
        ));
        assert!(is_standby_monitoring_get(
            &Method::GET,
            "/monitoring/engine-state-digest"
        ));
        assert!(is_standby_monitoring_get(
            &Method::GET,
            "/monitoring/recovery-safety"
        ));
        assert!(is_standby_monitoring_get(&Method::GET, "/recovery-safety"));
        assert!(is_standby_monitoring_get(
            &Method::GET,
            "/recovery-safety/alert"
        ));
        assert!(!is_standby_monitoring_get(
            &Method::POST,
            "/monitoring/integrity"
        ));
        assert!(!is_standby_monitoring_get(
            &Method::GET,
            "/monitoring/reload_instruments_cache"
        ));
        assert!(!is_standby_monitoring_get(
            &Method::GET,
            hypercall_types::API_ROUTE_ORDERS
        ));
    }

    #[tokio::test]
    async fn caught_up_unpromoted_standby_can_bypass_monitoring_readiness() {
        let progress = TestStandbyReplayProgress { caught_up: true };
        let (promote_tx, _promote_rx) = tokio::sync::oneshot::channel();
        let state = ReadinessMiddlewareState::new(
            not_ready_registry(),
            Some(Arc::new(progress)),
            Some(Arc::new(Mutex::new(Some(promote_tx)))),
        );

        assert!(standby_is_caught_up_and_unpromoted(&state).await);
    }

    #[tokio::test]
    async fn standby_bypass_is_disabled_before_catchup_and_after_promote() {
        let progress = TestStandbyReplayProgress { caught_up: false };
        let (promote_tx, _promote_rx) = tokio::sync::oneshot::channel();
        let state = ReadinessMiddlewareState::new(
            not_ready_registry(),
            Some(Arc::new(progress.clone())),
            Some(Arc::new(Mutex::new(Some(promote_tx)))),
        );
        assert!(!standby_is_caught_up_and_unpromoted(&state).await);

        let promoted_state = ReadinessMiddlewareState::new(
            not_ready_registry(),
            Some(Arc::new(TestStandbyReplayProgress { caught_up: true })),
            None,
        );
        assert!(!standby_is_caught_up_and_unpromoted(&promoted_state).await);
    }

    #[test]
    fn primary_middleware_state_never_has_standby_bypass_inputs() {
        let state = ReadinessMiddlewareState::primary(not_ready_registry());

        assert!(state.standby_progress.is_none());
        assert!(state.standby_promote.is_none());
    }
}