cycles-quartz/enclaves/quartz/src/mtcs_server.rs

use std::{
    collections::BTreeMap,
    sync::{Arc, Mutex},
};

use cosmwasm_std::HexBinary;
use cw_tee_mtcs::{
    msg::execute::SubmitSetoffsMsg,
    state::{RawCipherText, RawHash},
};
use ecies::{decrypt, encrypt};
use k256::ecdsa::{SigningKey, VerifyingKey};
use mtcs::{
    algo::mcmf::primal_dual::PrimalDual, impls::complex_id::ComplexIdMtcs,
    obligation::SimpleObligation, prelude::DefaultMtcs, setoff::SimpleSetoff, Mtcs,
};
use tonic::{Request, Response, Result as TonicResult, Status};

use crate::{
    attestor::Attestor,
    proto::{clearing_server::Clearing, RunClearingRequest, RunClearingResponse},
};

#[derive(Clone, Debug)]
pub struct MtcsService<A> {
    sk: Arc<Mutex<Option<SigningKey>>>,
    _attestor: A,
}

impl<A> MtcsService<A>
where
    A: Attestor,
{
    pub fn new(sk: Arc<Mutex<Option<SigningKey>>>, _attestor: A) -> Self {
        Self { sk, _attestor }
    }
}

#[tonic::async_trait]
impl<A> Clearing for MtcsService<A>
where
    A: Attestor + Send + Sync + 'static,
{
    async fn run(
        &self,
        request: Request<RunClearingRequest>,
    ) -> TonicResult<Response<RunClearingResponse>> {
        let message = request.into_inner().message;
        let obligations_enc: BTreeMap<RawHash, RawCipherText> =
            serde_json::from_str(&message).map_err(|e| Status::invalid_argument(e.to_string()))?;

        let sk = self.sk.lock().unwrap();
        let (digests, obligations): (Vec<RawHash>, Vec<SimpleObligation<HexBinary, i64>>) =
            obligations_enc
                .into_iter()
                .map(|(digest, ciphertext)| {
                    let o = decrypt(&sk.as_ref().unwrap().to_bytes(), &ciphertext).unwrap();
                    (digest, serde_json::from_slice(&o).unwrap())
                })
                .unzip();

        let mut mtcs = ComplexIdMtcs::wrapping(DefaultMtcs::new(PrimalDual::default()));
        let setoffs: Vec<SimpleSetoff<HexBinary, i64>> = mtcs.run(obligations).unwrap();

        let setoffs_enc: BTreeMap<RawHash, RawCipherText> = setoffs
            .into_iter()
            .zip(digests)
            .flat_map(|(so, digest)| {
                let debtor_pk = VerifyingKey::from_sec1_bytes(&so.debtor).unwrap();
                let creditor_pk = VerifyingKey::from_sec1_bytes(&so.creditor).unwrap();

                let so_ser = serde_json::to_string(&so).expect("infallible serializer");
                let so_debtor = encrypt(&debtor_pk.to_sec1_bytes(), so_ser.as_bytes()).unwrap();
                let so_creditor = encrypt(&creditor_pk.to_sec1_bytes(), so_ser.as_bytes()).unwrap();

                [
                    (digest.clone(), so_debtor.into()),
                    (digest, so_creditor.into()),
                ]
            })
            .collect();

        let message = serde_json::to_string(&SubmitSetoffsMsg { setoffs_enc }).unwrap();

        Ok(Response::new(RunClearingResponse { message }))
    }
}