tinylayer

//! Minimal fail-stop BIP340 signer for a Bitcoin statechain.

// Reject unsafe Rust anywhere in this crate.
#![forbid(unsafe_code)]

// The enclave is deliberately just an in-memory map of coin states.
use std::collections::HashMap;

// Ordinary libsecp256k1 key generation and BIP340 Schnorr signing.
use secp256k1::{Keypair, Message, Secp256k1, SecretKey, XOnlyPublicKey, rand, schnorr::Signature};
// JSON request and response encoding.
use serde::{Deserialize, Serialize};
// BIP340-style tagged SHA-256 authorization commitments.
use sha2::{Digest, Sha256};
// Stable, readable protocol errors.
use thiserror::Error;

// Changing protocol semantics requires a new version and hash domains.
pub const PROTOCOL_VERSION: u32 = 1;
// Fixed-width aliases keep the wire protocol small and unambiguous.
pub type CoinId = [u8; 32];
pub type Capability = [u8; 32];
pub type HandoffToken = [u8; 32];
// A new coin begins before the enclave has generated its first handoff.
pub const INITIAL_HANDOFF: HandoffToken = [0; 32];
// The production alias caps stored coins while tests can use smaller limits.
pub type Enclave = Signer<21_000>;

// A client chooses a random coin ID and commits to its first capability.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
// Reject misspelled or future fields instead of silently ignoring them.
#[serde(deny_unknown_fields)]
pub struct RegisterRequest {
    pub coin_id: CoinId,
    // The enclave does not need the capability preimage until signing.
    pub initial_capability_hash: [u8; 32],
}

// Equality lets the enclave recognize one exact request retry.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct SignRequest {
    // Select the independently generated per-coin enclave key.
    pub coin_id: CoinId,
    // Prove knowledge of the current owner's secret capability.
    pub current_capability: Capability,
    // Prove receipt of the handoff generated by the previous transition.
    pub current_handoff: HandoffToken,
    // Commit to the receiver's capability without learning its preimage.
    pub next_capability_hash: [u8; 32],
    // The enclave signs a digest; Bitcoin policy stays in the client.
    pub sighash: [u8; 32],
}

// The response is fixed-size and Copy, which makes cached retries simple.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct SignResponse {
    // A complete ordinary BIP340 signature, not a partial signature.
    pub signature: Signature,
    // Fresh entropy required alongside the receiver's capability.
    pub next_handoff: HandoffToken,
}

// Public state lets clients independently verify the live transition.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct CoinStatus {
    pub coin_id: CoinId,
    // The private half never leaves enclave memory.
    pub signing_pubkey: XOnlyPublicKey,
    // This commits to capability and handoff but reveals neither preimage.
    pub authorization: [u8; 32],
    // Clients compare this with the complete signed recovery history.
    pub signature_count: u64,
}

// Every protocol call is one tagged JSON request.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case", tag = "method", content = "params")]
pub enum Request {
    Register(RegisterRequest),
    Status { coin_id: CoinId },
    Sign(SignRequest),
}

// Registration and status both return CoinStatus; signing returns a signature.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case", tag = "method", content = "result")]
pub enum Response {
    Status(CoinStatus),
    Signature(SignResponse),
}

// There are only five state-machine failures.
#[derive(Clone, Debug, Eq, Error, PartialEq)]
pub enum Error {
    #[error("coin is not registered")]
    UnknownCoin,
    #[error("enclave coin capacity is exhausted")]
    CapacityReached,
    #[error("current capability or handoff is stale")]
    Unauthorized,
    #[error("next capability is unchanged")]
    UnchangedCapability,
    #[error("signature count is exhausted")]
    SignatureCountOverflow,
}

// Domain separation prevents this hash from being confused with another hash.
pub fn capability_hash(capability: &Capability) -> [u8; 32] {
    tagged_hash(b"Tinylayer/Capability/v1", &[capability])
}

// Authorization binds one coin to both pieces of current ownership state.
pub fn authorization(
    coin_id: &CoinId,
    capability_hash: &[u8; 32],
    handoff: &HandoffToken,
) -> [u8; 32] {
    tagged_hash(
        b"Tinylayer/Authorization/v1",
        &[coin_id, capability_hash, handoff],
    )
}

// This is the same tag-hash/tag-hash prefix construction used by BIP340.
fn tagged_hash(tag: &[u8], parts: &[&[u8]]) -> [u8; 32] {
    let tag_hash = Sha256::digest(tag);
    let mut hash = Sha256::new();
    hash.update(tag_hash);
    hash.update(tag_hash);
    // Field order is part of the protocol.
    parts.iter().for_each(|part| hash.update(part));
    hash.finalize().into()
}

// All persistent in-process protocol state lives inside this bounded signer.
pub struct Signer<const LIMIT: usize> {
    coins: HashMap<CoinId, Coin>,
}

// The complete private state for one funded coin.
struct Coin {
    signing_key: SecretKey,
    authorization: [u8; 32],
    signature_count: u64,
    // Cache only the latest transition, enough to retry one lost response while
    // this enclave process remains alive.
    last: Option<(SignRequest, SignResponse)>,
}

impl<const LIMIT: usize> Signer<LIMIT> {
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Self {
            coins: HashMap::new(),
        }
    }

    // Keep HTTP outside the state machine: dispatch a typed request here.
    pub fn handle(&mut self, request: Request) -> Result<Response, Error> {
        match request {
            Request::Register(request) => self.register(request).map(Response::Status),
            Request::Status { coin_id } => self.status(coin_id).map(Response::Status),
            Request::Sign(request) => self.sign(request).map(Response::Signature),
        }
    }

    pub fn register(&mut self, request: RegisterRequest) -> Result<CoinStatus, Error> {
        // First registration wins. A duplicate can inspect but cannot replace it.
        if self.coins.contains_key(&request.coin_id) {
            return self.status(request.coin_id);
        }
        if self.coins.len() >= LIMIT {
            return Err(Error::CapacityReached);
        }
        let coin = Coin {
            // The client must verify this commitment before funding Bitcoin.
            authorization: authorization(
                &request.coin_id,
                &request.initial_capability_hash,
                &INITIAL_HANDOFF,
            ),
            // Generate an independent enclave key for this coin.
            signing_key: SecretKey::new(&mut rand::thread_rng()),
            signature_count: 0,
            last: None,
        };
        // Derive public status before moving the private record into the map.
        let status = coin.status(request.coin_id);
        self.coins.insert(request.coin_id, coin);
        Ok(status)
    }

    pub fn status(&self, coin_id: CoinId) -> Result<CoinStatus, Error> {
        self.coins
            .get(&coin_id)
            // Expose the coin ID, public key, authorization commitment, and count.
            .map(|coin| coin.status(coin_id))
            .ok_or(Error::UnknownCoin)
    }

    pub fn sign(&mut self, request: SignRequest) -> Result<SignResponse, Error> {
        let coin = self
            .coins
            .get_mut(&request.coin_id)
            .ok_or(Error::UnknownCoin)?;

        // The previous success already rotated authorization. Return its exact
        // response before checking current authorization so a lost response can
        // be retried without issuing another signature or handoff.
        if let Some((last, response)) = &coin.last
            && last == &request
        {
            return Ok(*response);
        }

        // Rebuild the public commitment from the two secret current preimages.
        let current_capability_hash = capability_hash(&request.current_capability);
        let expected = authorization(
            &request.coin_id,
            &current_capability_hash,
            &request.current_handoff,
        );
        if expected != coin.authorization {
            return Err(Error::Unauthorized);
        }
        // Every successful transition must rotate to a different capability.
        if request.next_capability_hash == current_capability_hash {
            return Err(Error::UnchangedCapability);
        }

        // Never let the public transition counter wrap around.
        let next_count = coin
            .signature_count
            .checked_add(1)
            .ok_or(Error::SignatureCountOverflow)?;

        // Produce one standard BIP340 signature over the client's digest.
        let secp = Secp256k1::new();
        let keypair = Keypair::from_secret_key(&secp, &coin.signing_key);
        let response = SignResponse {
            signature: secp
                .sign_schnorr_no_aux_rand(&Message::from_digest(request.sighash), &keypair),
            // The enclave contributes fresh state unknown before this success.
            next_handoff: rand::random(),
        };

        // The receiver knows the next capability; the sender learns this
        // handoff from the response. The reference wallet transfers it only
        // after durable commit.
        let next_authorization = authorization(
            &request.coin_id,
            &request.next_capability_hash,
            &response.next_handoff,
        );

        // Apply the transition and retain exactly what an identical retry needs.
        (coin.authorization, coin.signature_count, coin.last) =
            (next_authorization, next_count, Some((request, response)));
        Ok(response)
    }
}

impl Coin {
    // Convert private state into the only state clients are allowed to inspect.
    fn status(&self, coin_id: CoinId) -> CoinStatus {
        CoinStatus {
            coin_id,
            signing_pubkey: self.signing_key.x_only_public_key(&Secp256k1::new()).0,
            authorization: self.authorization,
            signature_count: self.signature_count,
        }
    }
}

// The state machine can be built without the HTTP workload.
#[cfg(feature = "workload")]
pub mod workload {
    use std::sync::Arc;

    use axum::{
        Json, Router,
        extract::{DefaultBodyLimit, State},
        http::StatusCode,
        routing::{get, post},
    };
    // One mutex serializes requests; the short alias keeps the bootstrap small.
    pub(crate) use tokio::{net::TcpListener as Tcp, sync::Mutex};

    use crate::{Enclave, Request, Response};

    pub fn router(enclave: Enclave) -> Router {
        Router::new()
            // Health says the process is reachable, not that a coin exists.
            .route("/health", get(|| async { "ok" }))
            // All typed protocol requests use one small versioned endpoint.
            .route("/v1", post(handle))
            // Bound JSON parser allocation and accidental oversized requests.
            .layer(DefaultBodyLimit::max(4 * 1024))
            .with_state(Arc::new(Mutex::new(enclave)))
    }

    async fn handle(
        State(enclave): State<Arc<Mutex<Enclave>>>,
        Json(request): Json<Request>,
    ) -> Result<Json<Response>, (StatusCode, String)> {
        // No await occurs inside handle(request), so the complete operation is
        // one serialized critical section with no visible intermediate state.
        let result = enclave.lock().await.handle(request);
        result
            .map(Json)
            // A tiny prototype uses one status code for every policy conflict.
            .map_err(|error| (StatusCode::CONFLICT, error.to_string()))
    }
}

// The same file is also the workload executable when this feature is enabled.
#[cfg(feature = "workload")]
#[tokio::main]
pub async fn main() {
    // Enclavia routes the production connection to this workload listener.
    let listener = workload::Tcp::bind("0.0.0.0:8080").await.unwrap();
    // No persistence is configured; state starts empty on every process boot.
    axum::serve(listener, workload::router(Enclave::new()))
        .await
        .expect("serve workload");
}