cosigner-client 0.4.0

//! `BatchSigner` against the in-process proxy and mock Turnkey (offline, 0 live
//! signatures).
//!
//! The queued path is never exercised through a local backend — `spawn` passes
//! a local signer straight through — so its coverage lives here rather than in
//! any consumer's CI.

use std::sync::Arc;
use std::time::{Duration, Instant};

use arch_program::sanitized::ArchMessage;
use cosigner_client::{ArchSigner, ArchSignerT, BatchSigner, SignError};

mod common;
use common::{proxy_fixture, test_keypair, transfer_message, verifies_under};

/// A `BatchSigner` over the fixture's proxy, tuned for tests.
fn batched(fixture: &common::ProxyFixture) -> BatchSigner {
    BatchSigner::spawn(ArchSigner::remote(
        &fixture.url,
        &fixture.token,
        "arb",
        fixture.role_pubkey,
    ))
}

#[tokio::test]
async fn concurrent_signs_coalesce_into_few_requests() {
    let fixture = proxy_fixture(true).await;
    let signer = Arc::new(batched(&fixture).with_max_batch(32).with_max_rps(50.0));
    let messages: Vec<ArchMessage> = (1..=100)
        .map(|seed| transfer_message(fixture.role_pubkey, seed))
        .collect();

    let mut tasks = Vec::new();
    for message in &messages {
        let signer = Arc::clone(&signer);
        let message = message.clone();
        tasks.push(tokio::spawn(async move {
            signer.sign_message(&message).await.map(|r| r.signature)
        }));
    }

    let mut signatures = Vec::new();
    for task in tasks {
        signatures.push(task.await.expect("task").expect("signed"));
    }
    assert_eq!(signatures.len(), 100);

    // Every message still gets a signature that verifies for it — coalescing
    // must not shuffle results between callers.
    for (signature, message) in signatures.iter().zip(&messages) {
        assert!(verifies_under(
            &message.hash(),
            &fixture.role_pubkey,
            *signature
        ));
    }

    // 100 messages, at most ceil(100/32) = 4 activities.
    let requests = fixture.proxy_requests();
    assert!(
        (1..=4).contains(&requests),
        "100 concurrent signs took {requests} proxy requests"
    );
    assert_eq!(fixture.mock.signs(), 100, "one signature per message");
}

#[tokio::test]
async fn dispatch_rate_stays_within_max_rps() {
    let fixture = proxy_fixture(true).await;
    // max_batch = 1 forces one activity per message, so the elapsed time is
    // governed by the limiter rather than by batching.
    let signer = Arc::new(batched(&fixture).with_max_batch(1).with_max_rps(20.0));

    let started = Instant::now();
    let mut tasks = Vec::new();
    for seed in 1..=6 {
        let signer = Arc::clone(&signer);
        let message = transfer_message(fixture.role_pubkey, seed);
        tasks.push(tokio::spawn(
            async move { signer.sign_message(&message).await },
        ));
    }
    for task in tasks {
        task.await.expect("task").expect("signed");
    }
    let elapsed = started.elapsed();

    // Six activities at 20/s need five 50ms gaps after the first.
    assert!(
        elapsed >= Duration::from_millis(250),
        "6 activities at 20 rps finished in {elapsed:?}, faster than the budget allows"
    );
    assert_eq!(fixture.proxy_requests(), 6);
}

#[tokio::test]
async fn retry_attempts_consume_rate_permits() {
    // The inner signer is built with retries cleared, so the dispatcher owns
    // retry and each attempt spends a permit. Were it otherwise, the retried
    // attempts below would burst past the configured rate.
    let fixture = proxy_fixture(true).await;
    let signer = batched(&fixture)
        .with_max_batch(1)
        .with_max_rps(20.0)
        .with_retries(2)
        .with_deadline(Duration::from_secs(30));

    // Two 500s: the first attempt and its first retry fail, the second retry
    // succeeds — three attempts, three permits.
    fixture.mock.fail_next(500, 2);
    let started = Instant::now();
    let message = transfer_message(fixture.role_pubkey, 3);
    signer.sign_message(&message).await.expect("recovers");
    let elapsed = started.elapsed();

    assert!(
        elapsed >= Duration::from_millis(100),
        "3 attempts at 20 rps finished in {elapsed:?}, so retries skipped the limiter"
    );
    assert_eq!(fixture.proxy_requests(), 3, "one request per attempt");
}

#[tokio::test]
async fn full_queue_rejects_instead_of_blocking() {
    let fixture = proxy_fixture(true).await;
    // One slot, and a rate slow enough that the dispatcher is still waiting for
    // its permit while the rest of the requests arrive.
    let signer = Arc::new(
        batched(&fixture)
            .with_queue_depth(1)
            .with_max_batch(1)
            .with_max_rps(0.5),
    );

    let mut tasks = Vec::new();
    for seed in 1..=20 {
        let signer = Arc::clone(&signer);
        let message = transfer_message(fixture.role_pubkey, seed);
        tasks.push(tokio::spawn(
            async move { signer.sign_message(&message).await },
        ));
    }

    let mut rejections = 0;
    for task in tasks {
        if let Err(SignError::Signing(detail)) = task.await.expect("task") {
            assert_eq!(detail, "batch queue full");
            rejections += 1;
        }
    }
    assert!(
        rejections > 0,
        "a full queue must reject rather than grow without bound"
    );
}

#[tokio::test]
async fn expired_job_is_dropped_not_signed() {
    let fixture = proxy_fixture(true).await;
    // The deadline expires long before the limiter admits the second activity,
    // so the queued requests behind the first are dropped unsigned.
    let signer = Arc::new(
        batched(&fixture)
            .with_max_batch(1)
            .with_max_rps(0.4)
            .with_queue_depth(64)
            .with_deadline(Duration::from_millis(120)),
    );

    let mut tasks = Vec::new();
    for seed in 1..=5 {
        let signer = Arc::clone(&signer);
        let message = transfer_message(fixture.role_pubkey, seed);
        tasks.push(tokio::spawn(
            async move { signer.sign_message(&message).await },
        ));
    }

    let mut expired = 0;
    for task in tasks {
        if let Err(SignError::Signing(detail)) = task.await.expect("task") {
            assert_eq!(detail, "batch deadline exceeded");
            expired += 1;
        }
    }
    assert!(expired >= 3, "expected several expiries, saw {expired}");
    // A dropped request must not have cost a signature.
    assert!(
        fixture.mock.signs() <= 5 - expired as u64,
        "an expired request was signed anyway"
    );
}

#[tokio::test]
async fn single_sign_adds_no_latency_when_idle() {
    let fixture = proxy_fixture(true).await;
    let signer = batched(&fixture).with_max_rps(8.0);

    let started = Instant::now();
    let message = transfer_message(fixture.role_pubkey, 7);
    signer.sign_message(&message).await.expect("signed");
    let elapsed = started.elapsed();

    // The first permit is free and the drain finds nothing, so a lone request
    // must not wait for any linger window.
    assert!(
        elapsed < Duration::from_millis(100),
        "a lone request waited {elapsed:?}"
    );
    assert_eq!(fixture.proxy_requests(), 1);
}

#[tokio::test]
async fn local_signer_is_not_batched() {
    let (keypair, pubkey) = test_keypair(21);
    let signer = BatchSigner::spawn(ArchSigner::local(keypair))
        // Accepted and ignored: there is no budget to configure.
        .with_queue_depth(1)
        .with_max_rps(0.001);

    assert!(!signer.is_batching());
    assert_eq!(signer.pubkey(), pubkey);

    // No queue means no rate limit and no queue-full failure mode: many
    // concurrent signs all succeed immediately.
    let started = Instant::now();
    let signer = Arc::new(signer);
    let mut tasks = Vec::new();
    for seed in 1..=20 {
        let signer = Arc::clone(&signer);
        let message = transfer_message(pubkey, seed);
        tasks.push(tokio::spawn(async move {
            signer.sign_message(&message).await.map(|r| r.signature)
        }));
    }
    for task in tasks {
        let signature = task.await.expect("task").expect("signed");
        assert_eq!(signature.len(), 64);
    }
    assert!(
        started.elapsed() < Duration::from_secs(1),
        "a local signer must not inherit the rate budget"
    );
}

#[tokio::test]
async fn batched_signature_verifies_like_a_local_one() {
    // BIP340 signing is randomized, so byte equality is meaningless; what must
    // hold is that both signatures pass the check the validator runs.
    let fixture = proxy_fixture(true).await;
    let signer = batched(&fixture);
    let message = transfer_message(fixture.role_pubkey, 9);

    let response = signer.sign_message(&message).await.expect("signed");
    assert!(verifies_under(
        &message.hash(),
        &fixture.role_pubkey,
        response.signature
    ));
    assert_eq!(
        response.arch_account_pubkey,
        fixture.role_pubkey.serialize()
    );
    // Batch metadata still reaches the caller.
    assert!(response.turnkey_activity_id.is_some());
    assert!(response.digest_hex.is_some());
}

#[tokio::test]
async fn sign_messages_through_the_queue_keeps_input_order() {
    let fixture = proxy_fixture(true).await;
    let signer = batched(&fixture).with_max_batch(8);
    let messages: Vec<ArchMessage> = (1..=5)
        .map(|seed| transfer_message(fixture.role_pubkey, seed))
        .collect();

    let results = signer.sign_messages(&messages).await.expect("queued batch");
    assert_eq!(results.len(), 5);
    for (result, message) in results.iter().zip(&messages) {
        let response = result.as_ref().expect("signed");
        assert!(verifies_under(
            &message.hash(),
            &fixture.role_pubkey,
            response.signature
        ));
    }
    // Enqueued as individual jobs, then coalesced by the dispatcher.
    assert_eq!(fixture.proxy_requests(), 1);
}

#[tokio::test]
async fn transaction_assembly_works_through_the_queue() {
    // sign_transaction is a provided method over sign_message, so it comes
    // along for free — this pins that it actually does.
    let fixture = proxy_fixture(true).await;
    let signer = batched(&fixture);
    let message = transfer_message(fixture.role_pubkey, 4);

    let tx = signer
        .sign_transaction(message.clone())
        .await
        .expect("assembled");
    assert_eq!(tx.signatures.len(), 1);
    assert!(verifies_under(
        &message.hash(),
        &fixture.role_pubkey,
        tx.signatures[0].0
    ));
}