chainstream-sdk 2.0.11

//! Authentication middleware for ChainStream API.
//!
//! Supports three modes:
//! - API Key: `X-API-KEY` header (via `ApiKeyAuthMiddleware`)
//! - Wallet signature: `X-Wallet-*` headers for x402 (via `WalletAuthMiddleware`)
//! - Bearer token: `Authorization: Bearer` header (via `BearerAuthMiddleware` in tests)

use async_trait::async_trait;
use rand::Rng;
use std::time::{SystemTime, UNIX_EPOCH};

const SIGNATURE_PREFIX: &str = "chainstream";

/// Trait for wallet signers (EVM or Solana).
#[async_trait]
pub trait WalletSigner: Send + Sync {
    /// Returns "evm" or "solana".
    fn chain(&self) -> &str;
    /// Returns the wallet address (hex for EVM, Base58 for Solana).
    fn address(&self) -> &str;
    /// Sign a message and return the signature string.
    async fn sign_message(
        &self,
        message: &str,
    ) -> Result<String, Box<dyn std::error::Error + Send + Sync>>;
}

/// Build the message that must be signed for wallet auth.
pub fn build_sign_message(chain: &str, address: &str, timestamp: &str, nonce: &str) -> String {
    format!("{SIGNATURE_PREFIX}:{chain}:{address}:{timestamp}:{nonce}")
}

fn generate_nonce() -> String {
    let bytes: [u8; 16] = rand::thread_rng().gen();
    hex::encode(bytes)
}

/// Wallet auth headers for HTTP requests.
pub struct WalletAuthHeaders {
    pub address: String,
    pub chain: String,
    pub signature: String,
    pub timestamp: String,
    pub nonce: String,
}

/// Generate X-Wallet-* headers for a request.
pub async fn create_wallet_auth_headers(
    signer: &dyn WalletSigner,
) -> Result<WalletAuthHeaders, Box<dyn std::error::Error + Send + Sync>> {
    let timestamp = SystemTime::now()
        .duration_since(UNIX_EPOCH)?
        .as_secs()
        .to_string();
    let nonce = generate_nonce();
    let message = build_sign_message(signer.chain(), signer.address(), &timestamp, &nonce);
    let signature = signer.sign_message(&message).await?;

    Ok(WalletAuthHeaders {
        address: signer.address().to_string(),
        chain: signer.chain().to_string(),
        signature,
        timestamp,
        nonce,
    })
}

/// Build a SIWX (EIP-4361) message for authentication.
pub fn build_siwx_message(
    address: &str,
    chain: &str,
    nonce: &str,
    issued_at: &str,
    expires_at: &str,
) -> String {
    let chain_label = if chain == "solana" {
        "Solana"
    } else {
        "Ethereum"
    };
    let chain_id = if chain == "solana" { "mainnet" } else { "8453" };

    format!(
        "{domain} wants you to sign in with your {chain_label} account:\n\
         {address}\n\n\
         Sign in to ChainStream API\n\n\
         URI: https://{domain}\n\
         Version: 1\n\
         Chain ID: {chain_id}\n\
         Nonce: {nonce}\n\
         Issued At: {issued_at}\n\
         Expiration Time: {expires_at}",
        domain = "api.chainstream.io",
        chain_label = chain_label,
        address = address,
        chain_id = chain_id,
        nonce = nonce,
        issued_at = issued_at,
        expires_at = expires_at,
    )
}

/// Middleware for reqwest that adds Authorization: SIWX header.
/// Uses SIWX token (sign once, reuse for 1h) instead of per-request X-Wallet-* headers.
pub struct SiwxAuthMiddleware {
    signer: Box<dyn WalletSigner>,
}

impl SiwxAuthMiddleware {
    pub fn new(signer: Box<dyn WalletSigner>) -> Self {
        Self { signer }
    }
}

#[async_trait]
impl reqwest_middleware::Middleware for SiwxAuthMiddleware {
    async fn handle(
        &self,
        mut req: reqwest::Request,
        extensions: &mut http::Extensions,
        next: reqwest_middleware::Next<'_>,
    ) -> reqwest_middleware::Result<reqwest::Response> {
        let nonce = generate_nonce();
        let now = chrono::Utc::now();
        let expires_at = now + chrono::Duration::hours(1);
        let issued_at_str = now.to_rfc3339();
        let expires_at_str = expires_at.to_rfc3339();

        let message = build_siwx_message(
            self.signer.address(),
            self.signer.chain(),
            &nonce,
            &issued_at_str,
            &expires_at_str,
        );

        let signature = self
            .signer
            .sign_message(&message)
            .await
            .map_err(|e| reqwest_middleware::Error::Middleware(anyhow::anyhow!("{e}")))?;

        use base64::Engine;
        let message_b64 = base64::engine::general_purpose::STANDARD.encode(message.as_bytes());
        let token = format!("{}.{}", message_b64, signature);

        req.headers_mut()
            .insert("Authorization", format!("SIWX {}", token).parse().unwrap());

        next.run(req, extensions).await
    }
}

/// Legacy middleware for reqwest that adds X-Wallet-* headers (deprecated).
pub struct WalletAuthMiddleware {
    signer: Box<dyn WalletSigner>,
}

impl WalletAuthMiddleware {
    pub fn new(signer: Box<dyn WalletSigner>) -> Self {
        Self { signer }
    }
}

#[async_trait]
impl reqwest_middleware::Middleware for WalletAuthMiddleware {
    async fn handle(
        &self,
        mut req: reqwest::Request,
        extensions: &mut http::Extensions,
        next: reqwest_middleware::Next<'_>,
    ) -> reqwest_middleware::Result<reqwest::Response> {
        let headers = create_wallet_auth_headers(self.signer.as_ref())
            .await
            .map_err(|e| reqwest_middleware::Error::Middleware(anyhow::anyhow!("{e}")))?;

        let h = req.headers_mut();
        h.insert("X-Wallet-Address", headers.address.parse().unwrap());
        h.insert("X-Wallet-Chain", headers.chain.parse().unwrap());
        h.insert("X-Wallet-Signature", headers.signature.parse().unwrap());
        h.insert("X-Wallet-Timestamp", headers.timestamp.parse().unwrap());
        h.insert("X-Wallet-Nonce", headers.nonce.parse().unwrap());

        next.run(req, extensions).await
    }
}

/// Middleware for reqwest that adds X-API-KEY header.
pub struct ApiKeyAuthMiddleware {
    api_key: String,
}

impl ApiKeyAuthMiddleware {
    pub fn new(api_key: impl Into<String>) -> Self {
        Self {
            api_key: api_key.into(),
        }
    }
}

#[async_trait]
impl reqwest_middleware::Middleware for ApiKeyAuthMiddleware {
    async fn handle(
        &self,
        mut req: reqwest::Request,
        extensions: &mut http::Extensions,
        next: reqwest_middleware::Next<'_>,
    ) -> reqwest_middleware::Result<reqwest::Response> {
        req.headers_mut()
            .insert("X-API-KEY", self.api_key.parse().unwrap());
        next.run(req, extensions).await
    }
}