cosigner-client 0.1.0

//! `cosigner-client` — the signer abstraction the bot repos lack.
//!
//! One trait, two implementations, two transaction helpers:
//!
//! ```text
//! ArchSigner                        LocalSigner   (testnet + ephemeral keys;
//!   ├─ pubkey()                                    byte-compatible with
//!   └─ sign_message(&ArchMessage)                  arch_sdk::sign_message_bip322)
//!                                   RemoteSigner  (POSTs to arch-cosigner;
//!                                                  role, token, intent labels,
//!                                                  response verification)
//!
//! sign_transaction(signer, msg)                 → RuntimeTransaction
//! sign_transaction_mixed(signer, msg, &[kp])    → remote-signs the signer's
//!     slot, locally signs each ephemeral cosigner, places every signature by
//!     its pubkey's position in message.account_keys — the same positional
//!     logic as arch_sdk::build_and_sign_transaction.
//! ```
//!
//! Migration contract: swapping `LocalSigner` in for today's direct
//! `build_and_sign_transaction` / `sign_message_bip322` calls is a
//! zero-behavior-change refactor (property-tested in
//! `tests/signer_equivalence.rs`); flipping the same bot to mainnet is then
//! only an env change (`from_env`).

#![warn(missing_docs)]

use std::sync::Arc;
use std::time::Duration;

use arch_program::pubkey::Pubkey;
use arch_program::sanitized::ArchMessage;
use arch_sdk::{RuntimeTransaction, Signature};
use async_trait::async_trait;
use base64::Engine;
use bitcoin::key::UntweakedKeypair;
use bitcoin::secp256k1::XOnlyPublicKey;

/// Typed error surface. The variants encode the *required bot reaction*:
/// `Halted`/`Unreachable` ⇒ halt-and-alert (no local fallback key,
/// no crash-loop); `Denied` ⇒ never retried, page ops; `Transient` ⇒ bounded
/// retry is acceptable (the SDK already did one round); the rest are caller
/// bugs or configuration errors.
#[derive(Debug, thiserror::Error)]
pub enum SignerError {
    /// Proxy answered 503 — the halt machine is engaged (or the operator
    /// stopped signing by shutting the service down).
    /// **Halt and alert.**
    #[error("cosigner halted: {0}")]
    Halted(String),
    /// Proxy unreachable (connection refused / timeout). Same reaction as
    /// `Halted`: halt-and-alert.
    #[error("cosigner unreachable: {0}")]
    Unreachable(String),
    /// Turnkey policy denied the signature. **Never retry — page.** With
    /// correct config this indicates role/key confusion or an attack.
    #[error("denied by turnkey policy: {0}")]
    Denied(String),
    /// 401 bad_token / 403 role_mismatch — token or role wiring is wrong.
    #[error("unauthorized at proxy: {0}")]
    Unauthorized(String),
    /// 400 — the message bytes did not parse as an ArchMessage. Builder bug.
    #[error("proxy rejected message as malformed: {0}")]
    MalformedMessage(String),
    /// 502 — upstream Turnkey trouble after the proxy's own retries.
    #[error("transient cosigner failure: {0}")]
    Transient(String),
    /// The returned signature failed local verification — treat as an
    /// incident (corrupted or malicious proxy), not a retry.
    #[error("response verification failed: {0}")]
    Verification(String),
    /// A required signer in `account_keys` has no matching key or signer.
    #[error("no signer available for required key {0}")]
    MissingSigner(String),
    /// Local signing failure (arch_sdk).
    #[error("local signing failed: {0}")]
    Signing(String),
    /// from_env: configuration incomplete/invalid.
    #[error("signer configuration: {0}")]
    Config(String),
}

/// The signer abstraction: a BIP322 signature over `message.hash()`,
/// byte-compatible with `arch_sdk::sign_message_bip322` semantics.
#[async_trait]
pub trait ArchSigner: Send + Sync {
    /// The Arch account key this signer signs for.
    fn pubkey(&self) -> Pubkey;
    /// Produce a 64-byte BIP340 `r‖s` over the message's BIP322 digest.
    async fn sign_message(&self, message: &ArchMessage) -> Result<[u8; 64], SignerError>;
}

// ---------------------------------------------------------------------------
// LocalSigner
// ---------------------------------------------------------------------------

/// Wraps a local keypair — testnet flows and valueless ephemeral keys.
/// Delegates to `arch_sdk::sign_message_bip322` so behavior is identical to
/// today's direct signing (equivalence-tested against `arch_sdk`).
pub struct LocalSigner {
    keypair: UntweakedKeypair,
    pubkey: Pubkey,
    network: bitcoin::Network,
}

impl LocalSigner {
    /// Wrap an in-memory keypair for `network`.
    pub fn new(keypair: UntweakedKeypair, network: bitcoin::Network) -> Self {
        let pubkey = Pubkey::from_slice(&XOnlyPublicKey::from_keypair(&keypair).0.serialize());
        Self {
            keypair,
            pubkey,
            network,
        }
    }

    /// Load from a key file in `arch_sdk::with_secret_key_file` format.
    /// Fails if the file does not exist (unlike arch_sdk, which would
    /// silently generate a fresh key).
    pub fn from_key_file(path: &str, network: bitcoin::Network) -> Result<Self, SignerError> {
        if !std::path::Path::new(path).exists() {
            return Err(SignerError::Config(format!("key file {path} not found")));
        }
        let (keypair, _pubkey) = arch_sdk::with_secret_key_file(path)
            .map_err(|e| SignerError::Config(format!("loading {path}: {e}")))?;
        Ok(Self::new(keypair, network))
    }
}

#[async_trait]
impl ArchSigner for LocalSigner {
    fn pubkey(&self) -> Pubkey {
        self.pubkey
    }

    async fn sign_message(&self, message: &ArchMessage) -> Result<[u8; 64], SignerError> {
        arch_sdk::sign_message_bip322(&self.keypair, &message.hash(), self.network)
            .map_err(|e| SignerError::Signing(e.to_string()))
    }
}

// ---------------------------------------------------------------------------
// RemoteSigner
// ---------------------------------------------------------------------------

/// The full, verified `/v1/sign` response — what [`RemoteSigner::sign_detailed`]
/// returns to callers that need more than the bare signature.
#[derive(Debug, Clone)]
pub struct SignResponse {
    /// 64-byte BIP340 `r‖s`, already BIP322-verified for the submitted message.
    pub signature: [u8; 64],
    /// The role's x-only Arch account key the proxy signed with — checked
    /// against the signer's configured pubkey before this struct is built.
    pub arch_account_pubkey: [u8; 32],
    /// The 32-byte digest the proxy reports having sent to Turnkey, hex.
    pub digest_hex: String,
    /// Turnkey activity id for reconciliation against the proxy's audit log.
    pub turnkey_activity_id: String,
}

/// POSTs to `arch-cosigner`'s `/v1/sign` and cryptographically verifies every
/// response (BIP322 for the exact submitted message under the configured role
/// pubkey) before handing the signature to the caller.
#[derive(Clone)]
pub struct RemoteSigner {
    http: reqwest::Client,
    url: Arc<str>,
    token: Arc<str>,
    role: Arc<str>,
    intent: Arc<str>,
    pubkey: Pubkey,
    network: bitcoin::Network,
    /// Bounded retries on `Transient` errors only.
    retries: u32,
    retry_backoff: Duration,
}

impl RemoteSigner {
    /// `pubkey` is the role's Arch account key from the bot's config — the
    /// same value used for on-chain preflights (`state.operator ==
    /// signer.pubkey()`); every response is verified against it.
    pub fn new(
        url: &str,
        token: &str,
        role: &str,
        pubkey: Pubkey,
        network: bitcoin::Network,
    ) -> Self {
        Self {
            // client timeout must exceed the proxy's worst case
            // ((1 + retries) × turnkey_timeout + backoff ≈ 31 s on defaults)
            http: reqwest::Client::builder()
                .timeout(Duration::from_secs(35))
                .build()
                .expect("client construction with static config cannot fail"),
            url: url.trim_end_matches('/').into(),
            token: token.into(),
            role: role.into(),
            intent: "unlabeled".into(),
            pubkey,
            network,
            retries: 2,
            retry_backoff: Duration::from_millis(250),
        }
    }

    /// Cheap per-intent handle: `signer.with_intent("sweep")`. Labels are
    /// audit-only today but a future validation engine will enforce them —
    /// label honestly from day one.
    pub fn with_intent(&self, intent: &str) -> Self {
        let mut s = self.clone();
        s.intent = intent.into();
        s
    }

    /// Bounded retry count for `Transient` errors (default 2).
    pub fn with_retries(mut self, retries: u32) -> Self {
        self.retries = retries;
        self
    }

    /// One `POST /v1/sign` round trip returning the full typed response.
    /// The signature is verified exactly as in [`ArchSigner::sign_message`]
    /// (pubkey match + BIP322 for the exact submitted message) before it is
    /// returned; no retries are performed at this level.
    pub async fn sign_detailed(&self, message: &ArchMessage) -> Result<SignResponse, SignerError> {
        let b64 = base64::engine::general_purpose::STANDARD.encode(message.serialize());
        self.sign_once(&b64, message).await
    }

    async fn sign_once(
        &self,
        message_b64: &str,
        message: &ArchMessage,
    ) -> Result<SignResponse, SignerError> {
        let resp = self
            .http
            .post(format!("{}/v1/sign", self.url))
            .bearer_auth(self.token.as_ref())
            .json(&serde_json::json!({
                "role": self.role.as_ref(),
                "intent_type": self.intent.as_ref(),
                "unsigned_message_b64": message_b64,
            }))
            .send()
            .await
            .map_err(|e| {
                if e.is_timeout() || e.is_connect() {
                    SignerError::Unreachable(e.to_string())
                } else {
                    SignerError::Transient(e.to_string())
                }
            })?;

        let status = resp.status().as_u16();
        let body: serde_json::Value = resp.json().await.unwrap_or_default();
        let err_str = || {
            body["error"]
                .as_str()
                .or(body["halted"].as_str())
                .unwrap_or("<no detail>")
                .to_string()
        };
        match status {
            200 => {}
            503 => return Err(SignerError::Halted(err_str())),
            400 => return Err(SignerError::MalformedMessage(err_str())),
            401 => return Err(SignerError::Unauthorized(err_str())),
            403 if body["error"] == "role_mismatch" => {
                return Err(SignerError::Unauthorized(err_str()))
            }
            403 => return Err(SignerError::Denied(err_str())),
            502 => return Err(SignerError::Transient(err_str())),
            other => return Err(SignerError::Transient(format!("http {other}: {body}"))),
        }

        let sig: [u8; 64] = hex::decode(body["signature_hex"].as_str().unwrap_or_default())
            .ok()
            .and_then(|v| v.try_into().ok())
            .ok_or_else(|| SignerError::Verification("signature_hex not 64 bytes".into()))?;

        // the proxy must be signing with OUR role's key…
        if body["arch_account_pubkey"].as_str()
            != Some(hex::encode(self.pubkey.serialize()).as_str())
        {
            return Err(SignerError::Verification(format!(
                "proxy signed with {} but this signer is configured for {}",
                body["arch_account_pubkey"],
                hex::encode(self.pubkey.serialize())
            )));
        }
        // …and the signature must verify for the EXACT message we submitted
        // (recomputes the whole BIP322 pipeline locally — microseconds).
        arch_digest::verify_message_signature(
            message,
            &self.pubkey.serialize(),
            &sig,
            self.network,
        )
        .map_err(|e| SignerError::Verification(e.to_string()))?;

        Ok(SignResponse {
            signature: sig,
            // The equality check above proved the proxy signed with the
            // configured key.
            arch_account_pubkey: self.pubkey.serialize(),
            digest_hex: body["digest_hex"].as_str().unwrap_or_default().to_string(),
            turnkey_activity_id: body["turnkey_activity_id"]
                .as_str()
                .unwrap_or_default()
                .to_string(),
        })
    }
}

#[async_trait]
impl ArchSigner for RemoteSigner {
    fn pubkey(&self) -> Pubkey {
        self.pubkey
    }

    async fn sign_message(&self, message: &ArchMessage) -> Result<[u8; 64], SignerError> {
        let b64 = base64::engine::general_purpose::STANDARD.encode(message.serialize());
        let mut last = None;
        for attempt in 0..=self.retries {
            if attempt > 0 {
                tokio::time::sleep(self.retry_backoff * 2u32.pow(attempt - 1)).await;
            }
            match self.sign_once(&b64, message).await {
                Ok(resp) => return Ok(resp.signature),
                // Only Transient is retryable; everything else has a
                // different mandated reaction (see the SignerError docs).
                Err(SignerError::Transient(e)) => last = Some(SignerError::Transient(e)),
                Err(other) => return Err(other),
            }
        }
        Err(last.expect("the loop always runs at least one attempt"))
    }
}

// ---------------------------------------------------------------------------
// env resolution
// ---------------------------------------------------------------------------

/// Resolve a signer from the environment:
///
/// ```text
/// COSIGNER_URL + COSIGNER_TOKEN + COSIGNER_ROLE + COSIGNER_PUBKEY → RemoteSigner   (mainnet)
///   (+ optional COSIGNER_NETWORK, default "bitcoin")
/// ARCH_KEY_PATH → LocalSigner   (testnet; + optional ARCH_NETWORK, default "testnet")
/// ```
///
/// Remote wins if both are configured. Bots keep ONE code path; the
/// deployment environment decides local vs remote.
pub fn from_env() -> Result<Arc<dyn ArchSigner>, SignerError> {
    let get = |k: &str| std::env::var(k).ok().filter(|v| !v.is_empty());

    if let (Some(url), Some(token), Some(role)) = (
        get("COSIGNER_URL"),
        get("COSIGNER_TOKEN"),
        get("COSIGNER_ROLE"),
    ) {
        let pubkey_hex = get("COSIGNER_PUBKEY").ok_or_else(|| {
            SignerError::Config(
                "COSIGNER_PUBKEY (role's 32-byte Arch pubkey, hex) is required with COSIGNER_URL"
                    .into(),
            )
        })?;
        let bytes: [u8; 32] = hex::decode(&pubkey_hex)
            .map_err(|e| SignerError::Config(format!("COSIGNER_PUBKEY: {e}")))?
            .try_into()
            .map_err(|_| SignerError::Config("COSIGNER_PUBKEY must be 32 bytes".into()))?;
        let network = parse_network(&get("COSIGNER_NETWORK").unwrap_or_else(|| "bitcoin".into()))?;
        return Ok(Arc::new(RemoteSigner::new(
            &url,
            &token,
            &role,
            Pubkey::from_slice(&bytes),
            network,
        )));
    }

    if let Some(path) = get("ARCH_KEY_PATH") {
        let network = parse_network(&get("ARCH_NETWORK").unwrap_or_else(|| "testnet".into()))?;
        return Ok(Arc::new(LocalSigner::from_key_file(&path, network)?));
    }

    Err(SignerError::Config(
        "set COSIGNER_URL/COSIGNER_TOKEN/COSIGNER_ROLE/COSIGNER_PUBKEY (remote) \
         or ARCH_KEY_PATH (local)"
            .into(),
    ))
}

/// Parse a network name (`bitcoin`/`mainnet`, `testnet`, `signet`,
/// `regtest`) as accepted by the proxy config and the `*_NETWORK` env vars.
pub fn parse_network(s: &str) -> Result<bitcoin::Network, SignerError> {
    match s {
        "bitcoin" | "mainnet" => Ok(bitcoin::Network::Bitcoin),
        "testnet" => Ok(bitcoin::Network::Testnet),
        "signet" => Ok(bitcoin::Network::Signet),
        "regtest" => Ok(bitcoin::Network::Regtest),
        other => Err(SignerError::Config(format!("unknown network {other:?}"))),
    }
}

// ---------------------------------------------------------------------------
// transaction helpers
// ---------------------------------------------------------------------------

/// Single-signer convenience — replaces
/// `build_and_sign_transaction(msg, vec![kp], network)`.
pub async fn sign_transaction(
    signer: &dyn ArchSigner,
    message: ArchMessage,
) -> Result<RuntimeTransaction, SignerError> {
    sign_transaction_mixed(signer, message, &[]).await
}

/// Mixed-signer — replaces
/// `build_and_sign_transaction(msg, vec![kp, ephemeral…], network)`.
///
/// Remote-signs (or local-signs, whatever `signer` is) the signer's slot and
/// locally signs each extra `Keypair`, placing every signature by its
/// pubkey's position within the first `num_required_signatures` entries of
/// `message.account_keys` — the positional logic bots hand-roll today.
///
/// Ephemeral keys (position mints, IDL buffers) are always `local_cosigners`:
/// valueless one-shots whose secrets exist only in process memory.
pub async fn sign_transaction_mixed(
    signer: &dyn ArchSigner,
    message: ArchMessage,
    local_cosigners: &[UntweakedKeypair],
) -> Result<RuntimeTransaction, SignerError> {
    let digest = message.hash();
    let required = message.header.num_required_signatures as usize;
    let mut signatures = Vec::with_capacity(required);

    for key in message.account_keys.iter().take(required) {
        if *key == signer.pubkey() {
            signatures.push(Signature(signer.sign_message(&message).await?));
        } else if let Some(kp) = local_cosigners
            .iter()
            .find(|kp| XOnlyPublicKey::from_keypair(kp).0.serialize() == key.serialize())
        {
            // Note: BIP322 P2TR signatures are network-independent (the
            // script_pubkey carries no network; pinned by arch-digest's
            // golden vectors), so Bitcoin is safe for ephemeral cosigners
            // regardless of deployment network.
            let sig = arch_sdk::sign_message_bip322(kp, &digest, bitcoin::Network::Bitcoin)
                .map_err(|e| SignerError::Signing(e.to_string()))?;
            signatures.push(Signature(sig));
        } else {
            return Err(SignerError::MissingSigner(hex::encode(key.serialize())));
        }
    }

    Ok(RuntimeTransaction {
        version: 0,
        signatures,
        message,
    })
}