cycles-quartz/apps/mtcs/enclave/src/mtcs_server.rs

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

use cosmwasm_std::{Addr, HexBinary, Uint128};
use cw_tee_mtcs::{
    msg::execute::SubmitSetoffsMsg,
    state::{LiquiditySource, LiquiditySourceType, RawHash, SettleOff, Transfer},
};
use ecies::decrypt;
use k256::ecdsa::SigningKey;
use mtcs::{
    algo::mcmf::primal_dual::PrimalDual, impls::complex_id::ComplexIdMtcs,
    obligation::SimpleObligation, prelude::DefaultMtcs, setoff::SimpleSetoff, Mtcs,
};
use quartz_common::{
    contract::{msg::execute::attested::RawAttested, state::Config},
    enclave::{attestor::Attestor, server::IntoServer},
};
use tonic::{Request, Response, Result as TonicResult, Status};
use uuid::Uuid;

use crate::{
    proto::{
        clearing_server::{Clearing, ClearingServer},
        RunClearingRequest, RunClearingResponse,
    },
    types::{ContractObligation, RunClearingMessage},
};

pub type RawCipherText = HexBinary;

impl<A: Attestor> IntoServer for MtcsService<A> {
    type Server = ClearingServer<MtcsService<A>>;

    fn into_server(self) -> Self::Server {
        ClearingServer::new(self)
    }
}

#[derive(Clone, Debug)]
pub struct MtcsService<A> {
    config: Config, // TODO: this config is not used anywhere
    sk: Arc<Mutex<Option<SigningKey>>>,
    attestor: A,
}

impl<A> MtcsService<A>
where
    A: Attestor,
{
    pub fn new(config: Config, sk: Arc<Mutex<Option<SigningKey>>>, attestor: A) -> Self {
        Self {
            config,
            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>> {
        // Light client check
        let message: RunClearingMessage = {
            let message = request.into_inner().message;
            serde_json::from_str(&message).map_err(|e| Status::invalid_argument(e.to_string()))?
        };

        // let (proof_value, message) = message
        //     .verify(self.config.light_client_opts())
        //     .map_err(Status::failed_precondition)?;

        // let proof_value_matches_msg =
        //     serde_json::to_string(&message.intents).is_ok_and(|s| s.as_bytes() == proof_value);
        // if !proof_value_matches_msg {
        //     return Err(Status::failed_precondition("proof verification"));
        // }
        // TODO: ensure no duplicates somewhere else!
        let liquidity_sources: Vec<LiquiditySource> =
            message.liquidity_sources.into_iter().collect();
        let digests_ciphertexts: BTreeMap<HexBinary, HexBinary> = message.intents;
        let (digests, ciphertexts): (Vec<_>, Vec<_>) = digests_ciphertexts.into_iter().unzip();

        let sk = self.sk.lock().unwrap();
        let obligations: Vec<SimpleObligation<LiquiditySource, i64, Uuid>> = ciphertexts
            .into_iter()
            .map(|ciphertext| decrypt_obligation(sk.as_ref().unwrap(), &ciphertext))
            .collect();

        let mut mtcs = ComplexIdMtcs::wrapping(DefaultMtcs::new(PrimalDual::default()));

        let setoffs: Vec<SimpleSetoff<LiquiditySource, i64, Uuid>> = mtcs.run(obligations).unwrap();
        let setoffs_enc: BTreeMap<RawHash, SettleOff> = setoffs
            .into_iter()
            .map(|so| into_settle_offs(so, &liquidity_sources))
            .zip(digests)
            .map(|(settle_off, digest)| (digest, settle_off))
            .collect();

        let msg = SubmitSetoffsMsg { setoffs_enc };
        println!("setoff_msg: {:?}", msg);
        let attestation = self
            .attestor
            .quote(msg.clone())
            .map_err(|e| Status::internal(e.to_string()))?;

        let attested_msg = RawAttested { msg, attestation };
        let message = serde_json::to_string(&attested_msg).unwrap();
        Ok(Response::new(RunClearingResponse { message }))
    }
}

// TODO Switch from Vec<_> to Vec<LiquiditySource>
fn into_settle_offs(
    so: SimpleSetoff<LiquiditySource, i64, Uuid>,
    liquidity_sources: &Vec<LiquiditySource>,
) -> SettleOff {
    println!("\nsetoff: {:?}", so);
    println!("\nliq sources: {:?}", liquidity_sources);

    // TODO: temporary patch, fix issue with liquidity sources becoming type External so that .contains() can be called directly
    let liquidity_sources_addrs = liquidity_sources
        .iter()
        .map(|lqs| lqs.address.clone())
        .collect::<Vec<Addr>>();

    // In tenders and acceptances, the creditor's balance decreases
    if liquidity_sources_addrs.contains(&so.debtor.address)
        || liquidity_sources_addrs.contains(&so.creditor.address)
    {
        SettleOff::Transfer(Transfer {
            payer: so.creditor.address.clone(),
            payee: so.debtor.address.clone(),
            // TODO: Include denominations
            amount: ("peppicoin".to_owned(), Uint128::from(so.set_off as u128)),
        })
    } else {
        // TODO: Tracked by issue #22

        // A no-op for the time being.
        SettleOff::SetOff(vec![])
    }
}

// fn wasm_address(pk: VerifyingKey) -> String {
//     let tm_pk = TmAccountId::from(pk);
//     AccountId::new("wasm", tm_pk.as_bytes())
//         .unwrap()
//         .to_string()
// }

// fn encrypt_setoff(
//     so: SimpleSetoff<HexBinary, i64>,
//     debtor_pk: VerifyingKey,
//     creditor_pk: VerifyingKey,
// ) -> Vec<RawCipherText> {
//     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();

//     vec![so_debtor.into(), so_creditor.into()]
// }

fn decrypt_obligation(
    sk: &SigningKey,
    ciphertext: &RawCipherText,
) -> SimpleObligation<LiquiditySource, i64, Uuid> {
    let o: ContractObligation = {
        let o = decrypt(&sk.to_bytes(), ciphertext).unwrap();
        serde_json::from_slice(&o).unwrap()
    };

    SimpleObligation::new(
        None,
        LiquiditySource {
            address: o.debtor,
            source_type: LiquiditySourceType::External,
        },
        LiquiditySource {
            address: o.creditor,
            source_type: LiquiditySourceType::External,
        },
        i64::try_from(o.amount).unwrap(),
    )
    .unwrap()
}