sdk_common/src/network.rs

use std::collections::HashMap;

use aes_gcm::aead::{Aead, Payload};
use aes_gcm::{Aes256Gcm, KeyInit};
use anyhow::{Error, Result};
use js_sys::Date;
use rand::{thread_rng, RngCore};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use sp_client::bitcoin::consensus::serialize;
use sp_client::bitcoin::hex::{DisplayHex, FromHex};
use sp_client::bitcoin::{OutPoint, Transaction};
use tsify::Tsify;

use crate::crypto::AAD;
use crate::error::AnkError;
use crate::pcd::{Member, RoleDefinition};
use crate::signature::Proof;

#[derive(Debug, Serialize, Deserialize, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
pub enum AnkFlag {
    NewTx,
    Faucet,
    Cipher,
    Commit,
    Unknown,
}

impl From<&str> for AnkFlag {
    fn from(value: &str) -> Self {
        match value {
            "NewTx" => Self::NewTx,
            "Faucet" => Self::Faucet,
            "Cipher" => Self::Cipher,
            "Commit" => Self::Commit,
            _ => Self::Unknown,
        }
    }
}

impl From<String> for AnkFlag {
    fn from(value: String) -> Self {
        (&value[..]).into()
    }
}

impl AnkFlag {
    pub fn new_from_byte(byte: u8) -> Self {
        match byte {
            0 => Self::NewTx,
            1 => Self::Faucet,
            2 => Self::Cipher,
            3 => Self::Commit,
            _ => Self::Unknown,
        }
    }

    pub fn as_str(&self) -> &str {
        match self {
            Self::NewTx => "NewTx",
            Self::Faucet => "Faucet",
            Self::Cipher => "Cipher",
            Self::Commit => "Commit",
            Self::Unknown => "Unknown",
        }
    }
}

/// Message sent to the server to commit some state in a transaction
/// Client must first send a commit message with empty validation_tokens
/// Relay will ignore a commit message for an update he's not aware of that also bears validation_tokens
#[derive(Debug, PartialEq, Clone, Serialize, Deserialize, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
pub struct CommitMessage {
    pub init_tx: String, // Can be tx or txid of the first transaction of the chain, which is maybe not ideal
    pub pcd_commitment: Value, // map of field <=> hash of the clear value
    pub roles: HashMap<String, RoleDefinition>, // Can be hashed and compared with the value above
    pub validation_tokens: Vec<Proof>,
    pub error: Option<AnkError>,
}

impl CommitMessage {
    /// Create a new commitment message for the first transaction of the chain
    /// init_tx must be the hex string of the transaction
    /// validation_tokens must be empty
    pub fn new_first_commitment(
        transaction: Transaction,
        pcd_commitment: Value,
        roles: HashMap<String, RoleDefinition>,
    ) -> Self {
        Self {
            init_tx: serialize(&transaction).to_lower_hex_string(),
            pcd_commitment,
            roles,
            validation_tokens: vec![],
            error: None,
        }
    }

    /// Create a new commitment message for an update transaction
    /// init_tx must be the hex string of the txid of the first commitment transaction
    /// validation_tokens must be empty
    pub fn new_update_commitment(
        init_tx: OutPoint,
        pcd_commitment: Value,
        roles: HashMap<String, RoleDefinition>,
    ) -> Self {
        Self {
            init_tx: init_tx.to_string(),
            pcd_commitment,
            roles,
            validation_tokens: vec![],
            error: None,
        }
    }

    /// Set the validation tokens for a pending commitment
    pub fn set_validation_tokens(&mut self, validation_tokens: Vec<Proof>) {
        self.validation_tokens = validation_tokens;
    }
}

#[derive(Debug, Serialize, Deserialize, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
pub struct FaucetMessage {
    pub sp_address: String,
    pub commitment: String,
    pub error: Option<AnkError>,
}

impl FaucetMessage {
    pub fn new(sp_address: String) -> Self {
        let mut buf = [0u8; 32];
        thread_rng().fill_bytes(&mut buf);
        Self {
            sp_address,
            commitment: buf.to_lower_hex_string(),
            error: None,
        }
    }

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

#[derive(Debug, PartialEq, Serialize, Deserialize, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
pub struct NewTxMessage {
    pub transaction: String,
    pub tweak_data: Option<String>,
    pub error: Option<AnkError>,
}

impl NewTxMessage {
    pub fn new(transaction: String, tweak_data: Option<String>) -> Self {
        Self {
            transaction,
            tweak_data,
            error: None,
        }
    }

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

#[derive(Debug, Serialize, Deserialize)]
pub struct Envelope {
    pub flag: AnkFlag,
    pub content: String,
}

impl Envelope {
    pub fn new(flag: AnkFlag, raw: &str) -> Self {
        Self {
            flag,
            content: raw.into(),
        }
    }

    pub fn from_string(json: &str) -> Result<Self> {
        let res: Self = serde_json::from_str(json)?;
        Ok(res)
    }

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

#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Tsify, Clone)]
pub enum CachedMessageStatus {
    #[default]
    NoStatus,
    CipherWaitingTx,
    TxWaitingPrd,
}

/// Unique struct for both 4nk messages and notification/key exchange, both rust and ts
#[derive(Debug, Default, Serialize, Deserialize, Tsify, Clone, PartialEq)]
#[tsify(into_wasm_abi, from_wasm_abi)]
#[allow(non_camel_case_types)]
pub struct CachedMessage {
    pub id: u32,
    pub status: CachedMessageStatus,
    pub transaction: Option<String>,
    pub commitment: Option<String>,  // content of the op_return
    pub sender: Option<Member>,      // Never None when message sent
    pub shared_secrets: Vec<String>, // Max 2 secrets in case we send to both address of the recipient
    pub cipher: Vec<String>, // Max 2 ciphers in case we send to both address of the recipient
    pub timestamp: u64,
}

impl CachedMessage {
    pub fn new() -> Self {
        let mut new = Self::default();
        let mut buf = [0u8; 4];
        thread_rng().fill_bytes(&mut buf);
        new.id = u32::from_be_bytes(buf);
        new.timestamp = Date::now().floor() as u64;
        new
    }

    pub fn from_string(json: &str) -> Result<Self> {
        let res = serde_json::from_str(json)?;
        Ok(res)
    }

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

    pub fn try_decrypt_message(&self, cipher: &[u8]) -> Result<Vec<u8>> {
        if self.shared_secrets.is_empty() {
            return Err(Error::msg(
                "Can't try decrypt this message, there's no shared secret",
            ));
        }
        for shared_secret in &self.shared_secrets {
            let mut key = [0u8; 32];
            let mut nonce = [0u8; 12];
            key.copy_from_slice(&Vec::from_hex(shared_secret)?);
            nonce.copy_from_slice(&cipher[..12]);

            let engine = Aes256Gcm::new(&key.into());
            let payload = Payload {
                msg: &cipher[12..],
                aad: AAD,
            };
            match engine.decrypt(&nonce.into(), payload) {
                Ok(plain) => return Ok(plain),
                Err(_) => continue,
            }
        }

        Err(Error::msg("Failed to decrypt message"))
    }

    pub fn try_decrypt_with_shared_secret(&self, shared_secret: [u8; 32]) -> Result<Vec<u8>> {
        if self.cipher.is_empty() || !self.shared_secrets.is_empty() {
            return Err(Error::msg(
                "Can't try decrypt this message, ciphertext is none or shared_secret already found",
            ));
        }
        for prd_cipher in &self.cipher {
            let cipher = Vec::from_hex(prd_cipher)?;
            let mut nonce = [0u8; 12];
            nonce.copy_from_slice(&cipher[..12]);

            let engine = Aes256Gcm::new(&shared_secret.into());
            let payload = Payload {
                msg: &cipher[12..],
                aad: AAD,
            };
            match engine.decrypt(&nonce.into(), payload) {
                Ok(plain) => return Ok(plain),
                Err(_) => continue,
            }
        }

        Err(Error::msg("Failed to decrypt message"))
    }
}