Complete statechain signer source

//! Rollback-protected BIP448 statechain linearizer.

#![forbid(unsafe_code)]

use std::collections::HashMap;

use secp256k1::{Keypair, Message, Secp256k1, SecretKey, XOnlyPublicKey, rand, schnorr::Signature};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use thiserror::Error;

pub const PROTOCOL_VERSION: u32 = 2;
pub const INITIAL_STATE_NUMBER: u64 = 1;
pub type CoinId = [u8; 32];
pub type Capability = [u8; 32];
pub type TemplateHash = [u8; 32];
pub type Enclave = Signer<21_000>;

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct ActivationRequest {
    pub coin_id: CoinId,
    pub initial_capability_hash: [u8; 32],
    pub initial_template_hash: TemplateHash,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct SignRequest {
    pub coin_id: CoinId,
    pub current_capability: Capability,
    pub next_capability_hash: [u8; 32],
    pub next_state_number: u64,
    pub template_hash: TemplateHash,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct SignResponse {
    pub state_number: u64,
    pub template_hash: TemplateHash,
    pub signature: Signature,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct ServiceInfo {
    pub protocol_version: u32,
    pub signing_pubkey: XOnlyPublicKey,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct CoinStatus {
    pub coin_id: CoinId,
    pub signing_pubkey: XOnlyPublicKey,
    pub capability_hash: [u8; 32],
    pub state_number: u64,
    pub template_hash: TemplateHash,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case", tag = "method", content = "params")]
pub enum Request {
    Info,
    Activate(ActivationRequest),
    Status { coin_id: CoinId },
    Sign(SignRequest),
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case", tag = "method", content = "result")]
pub enum Response {
    Info(ServiceInfo),
    Status(CoinStatus),
    Signature(SignResponse),
}

#[derive(Clone, Debug, Eq, Error, PartialEq)]
pub enum Error {
    #[error("coin is not activated")]
    UnknownCoin,
    #[error("activation conflicts with existing coin state")]
    ActivationConflict,
    #[error("enclave coin capacity is exhausted")]
    CapacityReached,
    #[error("current capability is stale")]
    Unauthorized,
    #[error("next capability is unchanged")]
    UnchangedCapability,
    #[error("state number must increase by exactly one")]
    InvalidStateNumber,
    #[error("next template is unchanged")]
    UnchangedTemplate,
    #[error("state number is exhausted")]
    StateNumberOverflow,
}

pub fn capability_hash(capability: &Capability) -> [u8; 32] {
    tagged_hash(b"Tinylayer/Capability/v2", &[capability])
}

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);
    parts.iter().for_each(|part| hash.update(part));
    hash.finalize().into()
}

pub struct Signer<const LIMIT: usize> {
    signing_key: SecretKey,
    coins: HashMap<CoinId, Coin>,
}

struct Coin {
    capability_hash: [u8; 32],
    state_number: u64,
    template_hash: TemplateHash,
    activation: ActivationRequest,
    last: Option<(SignRequest, SignResponse)>,
}

impl<const LIMIT: usize> Signer<LIMIT> {
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Self::with_signing_key(SecretKey::new(&mut rand::thread_rng()))
    }

    pub fn with_signing_key(signing_key: SecretKey) -> Self {
        Self {
            signing_key,
            coins: HashMap::new(),
        }
    }

    pub fn handle(&mut self, request: Request) -> Result<Response, Error> {
        match request {
            Request::Info => Ok(Response::Info(self.info())),
            Request::Activate(request) => self.activate(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 info(&self) -> ServiceInfo {
        ServiceInfo {
            protocol_version: PROTOCOL_VERSION,
            signing_pubkey: self.signing_pubkey(),
        }
    }

    pub fn activate(&mut self, request: ActivationRequest) -> Result<CoinStatus, Error> {
        if let Some(coin) = self.coins.get(&request.coin_id) {
            if coin.activation == request {
                return Ok(coin.status(request.coin_id, self.signing_pubkey()));
            }
            return Err(Error::ActivationConflict);
        }
        if self.coins.len() >= LIMIT {
            return Err(Error::CapacityReached);
        }
        let request_coin_id = request.coin_id;
        let coin = Coin {
            capability_hash: request.initial_capability_hash,
            state_number: INITIAL_STATE_NUMBER,
            template_hash: request.initial_template_hash,
            activation: request,
            last: None,
        };
        let status = coin.status(request_coin_id, self.signing_pubkey());
        self.coins.insert(request_coin_id, coin);
        Ok(status)
    }

    pub fn status(&self, coin_id: CoinId) -> Result<CoinStatus, Error> {
        self.coins
            .get(&coin_id)
            .map(|coin| coin.status(coin_id, self.signing_pubkey()))
            .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)?;
        if let Some((last, response)) = &coin.last
            && last == &request
        {
            return Ok(*response);
        }
        let current_capability_hash = capability_hash(&request.current_capability);
        if current_capability_hash != coin.capability_hash {
            return Err(Error::Unauthorized);
        }
        if request.next_capability_hash == current_capability_hash {
            return Err(Error::UnchangedCapability);
        }
        let next_state_number = coin
            .state_number
            .checked_add(1)
            .ok_or(Error::StateNumberOverflow)?;
        if request.next_state_number != next_state_number {
            return Err(Error::InvalidStateNumber);
        }
        if request.template_hash == coin.template_hash {
            return Err(Error::UnchangedTemplate);
        }

        let keypair = Keypair::from_secret_key(&Secp256k1::new(), &self.signing_key);
        let response = SignResponse {
            state_number: request.next_state_number,
            template_hash: request.template_hash,
            signature: Secp256k1::new()
                .sign_schnorr_no_aux_rand(&Message::from_digest(request.template_hash), &keypair),
        };
        (
            coin.capability_hash,
            coin.state_number,
            coin.template_hash,
            coin.last,
        ) = (
            request.next_capability_hash,
            next_state_number,
            request.template_hash,
            Some((request, response)),
        );
        Ok(response)
    }

    pub fn coin_count(&self) -> usize {
        self.coins.len()
    }

    fn signing_pubkey(&self) -> XOnlyPublicKey {
        self.signing_key.x_only_public_key(&Secp256k1::new()).0
    }
}

impl Coin {
    fn status(&self, coin_id: CoinId, signing_pubkey: XOnlyPublicKey) -> CoinStatus {
        CoinStatus {
            coin_id,
            signing_pubkey,
            capability_hash: self.capability_hash,
            state_number: self.state_number,
            template_hash: self.template_hash,
        }
    }
}

#[cfg(feature = "workload")]
pub mod workload {
    use std::sync::Arc;

    use axum::{
        Json, Router,
        extract::{DefaultBodyLimit, State},
        http::StatusCode,
        routing::{get, post},
    };
    use tokio::sync::Mutex;

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

    pub fn router(enclave: Enclave) -> Router {
        Router::new()
            .route("/health", get(|| async { "ok" }))
            .route("/v2", post(handle))
            .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)> {
        enclave
            .lock()
            .await
            .handle(request)
            .map(Json)
            .map_err(|error| (StatusCode::CONFLICT, error.to_string()))
    }
}