sdk_common/src/prd.rs

use std::collections::HashSet;
use std::str::FromStr;

use anyhow::{Error, Result};
use serde::{Deserialize, Serialize};

use serde_json::{Map, Value};
use sp_client::bitcoin::hashes::{sha256t_hash_newtype, Hash, HashEngine};
use sp_client::bitcoin::hex::FromHex;
use sp_client::bitcoin::secp256k1::SecretKey;
use sp_client::bitcoin::{OutPoint, Psbt, XOnlyPublicKey};
use sp_client::silentpayments::utils::SilentPaymentAddress;
use sp_client::spclient::SpWallet;
use tsify::Tsify;

use crate::pcd::{AnkPcdHash, Member, Pcd};
use crate::signature::{AnkHash, AnkMessageHash, Proof};

#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
#[allow(non_camel_case_types)]
pub enum PrdType {
    #[default]
    None,
    Message,
    Update, // Update an existing process
    List,   // request a list of items
    Response,
    Confirm,
    TxProposal, // Send a psbt asking for recipient signature, used for login not sure about other use cases
}

sha256t_hash_newtype! {
    pub struct AnkPrdTag = hash_str("4nk/Prd");

    #[hash_newtype(forward)]
    pub struct AnkPrdHash(_);
}

impl AnkPrdHash {
    pub fn from_value(value: &Value) -> Self {
        let mut eng = AnkPrdHash::engine();
        eng.input(value.to_string().as_bytes());
        AnkPrdHash::from_engine(eng)
    }

    pub fn from_map(map: &Map<String, Value>) -> Self {
        let value = Value::Object(map.clone());
        let mut eng = AnkPrdHash::engine();
        eng.input(value.to_string().as_bytes());
        AnkPrdHash::from_engine(eng)
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
#[allow(non_camel_case_types)]
pub struct Prd {
    pub prd_type: PrdType,
    pub root_commitment: String,
    pub sender: String,
    pub keys: Map<String, Value>, // key is a key in pcd, value is the key to decrypt it
    pub validation_tokens: Vec<Proof>,
    pub payload: String,      // Payload depends on the actual type
    pub proof: Option<Proof>, // This must be None up to the creation of the network message
}

impl Prd {
    pub fn new_update(
        root_commitment: OutPoint,
        sender: String, // Should take Member as argument
        encrypted_pcd: Map<String, Value>,
        keys: Map<String, Value>,
    ) -> Self {
        Self {
            prd_type: PrdType::Update,
            root_commitment: root_commitment.to_string(),
            sender,
            validation_tokens: vec![],
            keys,
            payload: Value::Object(encrypted_pcd).to_string(),
            proof: None,
        }
    }

    pub fn new_tx_proposal(root_commitment: OutPoint, sender: Member, psbt: Psbt) -> Self {
        Self {
            prd_type: PrdType::TxProposal,
            root_commitment: root_commitment.to_string(),
            sender: serde_json::to_string(&sender).unwrap(),
            validation_tokens: vec![],
            keys: Map::new(),
            payload: serde_json::to_string(&psbt).unwrap(),
            proof: None,
        }
    }

    pub fn new_response(
        root_commitment: OutPoint,
        sender: String,
        validation_tokens: Vec<Proof>,
        pcd_commitment: AnkPcdHash,
    ) -> Self {
        Self {
            prd_type: PrdType::Response,
            root_commitment: root_commitment.to_string(),
            sender,
            validation_tokens: validation_tokens,
            keys: Map::new(),
            payload: pcd_commitment.to_string(),
            proof: None,
        }
    }

    pub fn new_confirm(
        root_commitment: OutPoint,
        sender: Member,
        pcd_commitment: AnkPcdHash,
    ) -> Self {
        Self {
            prd_type: PrdType::Confirm,
            root_commitment: root_commitment.to_string(),
            sender: serde_json::to_string(&sender).unwrap(),
            validation_tokens: vec![],
            keys: Map::new(),
            payload: pcd_commitment.to_string(),
            proof: None,
        }
    }

    fn _extract_from_message(plain: &[u8], commitment: Option<&AnkPrdHash>) -> Result<Self> {
        let prd: Prd = serde_json::from_slice(plain)?;
        if let Some(commitment) = commitment {
            // check that the hash of the prd is consistent with what's commited in the op_return
            if prd.create_commitment() != *commitment {
                return Err(anyhow::Error::msg(
                    "Received prd is not what was commited in the transaction",
                ));
            }
        }
        // check that the proof is consistent
        let sender: Member = serde_json::from_str(&prd.sender)?;
        if let Some(proof) = prd.proof {
            // take the spending keys in sender
            let addresses = sender.get_addresses();
            let mut spend_keys: Vec<XOnlyPublicKey> = vec![];
            for address in addresses {
                spend_keys.push(
                    <SilentPaymentAddress>::try_from(address)?
                        .get_spend_key()
                        .x_only_public_key()
                        .0,
                );
            }
            // The key in proof must be one of the sender keys
            let proof_key = proof.get_key();
            let mut known_key = false;
            for key in spend_keys {
                if key == proof_key {
                    known_key = true;
                    break;
                }
            }
            if !known_key {
                return Err(anyhow::Error::msg("Proof signed with an unknown key"));
            }
            proof.verify()?;
        }
        // check that the commitment outpoint is valid, just in case
        OutPoint::from_str(&prd.root_commitment)?;
        Ok(prd)
    }

    pub fn extract_from_message(plain: &[u8]) -> Result<Self> {
        Self::_extract_from_message(plain, None)
    }

    pub fn extract_from_message_with_commitment(
        plain: &[u8],
        commitment: &AnkPrdHash,
    ) -> Result<Self> {
        Self::_extract_from_message(plain, Some(commitment))
    }

    pub fn filter_keys(&mut self, to_keep: HashSet<String>) {
        let current_keys = self.keys.clone();
        let filtered_keys: Map<String, Value> = current_keys
            .into_iter()
            .filter(|(field, _)| to_keep.contains(field))
            .collect();
        self.keys = filtered_keys;
    }

    /// We commit to everything except the keys and the proof
    /// Because 1) we need one commitment to common data for all recipients of the transaction
    /// 2) we already commit to the keys in the sender proof anyway
    pub fn create_commitment(&self) -> AnkPrdHash {
        let mut to_commit = self.clone();
        to_commit.keys = Map::new();
        to_commit.proof = None;

        if to_commit.payload.len() != 64 && Vec::from_hex(&to_commit.payload).is_err() {
            to_commit.payload = Value::from_str(&to_commit.payload)
                .unwrap()
                .tagged_hash()
                .to_string();
        }

        AnkPrdHash::from_value(&to_commit.to_value())
    }

    /// Generate the signed proof and serialize to send over the network
    pub fn to_network_msg(&self, sp_wallet: &SpWallet) -> Result<String> {
        let spend_sk: SecretKey = sp_wallet.get_client().get_spend_key().try_into()?;
        let to_sign = self.clone(); // we sign the whole prd, incl the keys, for each recipient

        let message_hash =
            AnkHash::Message(AnkMessageHash::from_message(to_sign.to_string().as_bytes()));

        let proof = Proof::new(message_hash, spend_sk);

        let mut res = self.clone();
        res.proof = Some(proof);

        Ok(res.to_string())
    }

    pub fn to_string(&self) -> String {
        serde_json::to_string(self).unwrap()
    }

    pub fn to_value(&self) -> Value {
        Value::from_str(&self.to_string()).unwrap()
    }
}