libchat/core/conversations/src/conversation/privatev1.rs

use blake2::{
    Blake2b, Blake2bMac, Digest,
    digest::{FixedOutput, consts::U18},
};
use chat_proto::logoschat::{
    convos::private_v1::{PrivateV1Frame, private_v1_frame::FrameType},
    encryption::{Doubleratchet, EncryptedPayload, encrypted_payload::Encryption},
};
use crypto::{PrivateKey, PublicKey, SymmetricKey32};
use double_ratchets::{Header, InstallationKeyPair, RatchetState, restore_ratchet_state};
use prost::{Message, bytes::Bytes};
use std::fmt::Debug;
use std::sync::{Arc, Mutex};
use storage::{ConversationKind, ConversationMeta, ConversationStore};

use crate::{
    context::{ConversationIdOwned, PRIVATE_V1_INBOX_ADDRESS},
    conversation::{ChatError, ConversationId, Convo, Id},
    errors::EncryptionError,
    event::Event,
    proto,
    types::AddressedEncryptedPayload,
    utils::timestamp_millis,
};
use double_ratchets::{to_ratchet_record, to_skipped_key_records};
use storage::RatchetStore;

// Represents the potential participant roles in this Conversation
enum Role {
    Initiator,
    Responder,
}

impl Role {
    const fn as_str(&self) -> &'static str {
        match self {
            Self::Initiator => "I",
            Self::Responder => "R",
        }
    }
}

struct BaseConvoId([u8; 18]);

impl BaseConvoId {
    fn new(key: &SymmetricKey32) -> Self {
        let base = Blake2bMac::<U18>::new_with_salt_and_personal(key.as_bytes(), b"", b"L-PV1-CID")
            .expect("fixed inputs should never fail");
        Self(base.finalize_fixed().into())
    }

    fn id_for_participant(&self, role: Role) -> String {
        let hash = Blake2b::<U18>::new()
            .chain_update(self.0)
            .chain_update(role.as_str())
            .finalize();
        hex::encode(hash)
    }
}

pub struct PrivateV1Convo<S: ConversationStore + RatchetStore> {
    local_convo_id: String,
    remote_convo_id: String,
    dr_state: RatchetState,
    store: Arc<Mutex<S>>,
}

impl<S: ConversationStore + RatchetStore> PrivateV1Convo<S> {
    /// Reconstructs a PrivateV1Convo from persisted metadata and ratchet state.
    pub fn new(
        store: Arc<Mutex<S>>,
        local_convo_id: String,
        remote_convo_id: String,
    ) -> Result<Self, ChatError> {
        let dr_record = store.lock().unwrap().load_ratchet_state(&local_convo_id)?;
        let skipped_keys = store.lock().unwrap().load_skipped_keys(&local_convo_id)?;
        let dr_state: RatchetState = restore_ratchet_state(dr_record, skipped_keys);

        Ok(Self {
            local_convo_id,
            remote_convo_id,
            dr_state,
            store,
        })
    }

    pub fn new_initiator(
        store: Arc<Mutex<S>>,
        seed_key: SymmetricKey32,
        remote: PublicKey,
    ) -> Self {
        let base_convo_id = BaseConvoId::new(&seed_key);
        let local_convo_id = base_convo_id.id_for_participant(Role::Initiator);
        let remote_convo_id = base_convo_id.id_for_participant(Role::Responder);

        // TODO: Danger - Fix double-ratchets types to Accept SymmetricKey32
        // perhaps update the  DH to work with cryptocrate.
        // init_sender doesn't take ownership of the key so a reference can be used.
        let shared_secret: [u8; 32] = seed_key.DANGER_to_bytes();
        let dr_state = RatchetState::init_sender(shared_secret, *remote);

        Self {
            local_convo_id,
            remote_convo_id,
            dr_state,
            store,
        }
    }

    pub fn new_responder(
        store: Arc<Mutex<S>>,
        seed_key: SymmetricKey32,
        dh_self: &PrivateKey,
    ) -> Self {
        let base_convo_id = BaseConvoId::new(&seed_key);
        let local_convo_id = base_convo_id.id_for_participant(Role::Responder);
        let remote_convo_id = base_convo_id.id_for_participant(Role::Initiator);

        // TODO: (P3) Rename; This accepts a Ephemeral key in most cases
        let dh_self_installation_keypair =
            InstallationKeyPair::from_secret_bytes(dh_self.DANGER_to_bytes());
        // TODO: Danger - Fix double-ratchets types to Accept SymmetricKey32
        let dr_state =
            RatchetState::init_receiver(seed_key.DANGER_to_bytes(), dh_self_installation_keypair);

        Self {
            local_convo_id,
            remote_convo_id,
            dr_state,
            store,
        }
    }

    fn encrypt(&mut self, frame: PrivateV1Frame) -> EncryptedPayload {
        let encoded_bytes = frame.encode_to_vec();
        let (cipher_text, header) = self.dr_state.encrypt_message(&encoded_bytes);

        EncryptedPayload {
            encryption: Some(Encryption::Doubleratchet(Doubleratchet {
                dh: Bytes::from(Vec::from(header.dh_pub.to_bytes())),
                msg_num: header.msg_num,
                prev_chain_len: header.prev_chain_len,
                ciphertext: Bytes::from(cipher_text),
                aux: "".into(),
            })),
        }
    }

    fn decrypt(&mut self, payload: EncryptedPayload) -> Result<PrivateV1Frame, EncryptionError> {
        // Validate and extract the encryption header or return errors
        let dr_header = if let Some(enc) = payload.encryption {
            if let proto::Encryption::Doubleratchet(dr) = enc {
                dr
            } else {
                return Err(EncryptionError::Decryption(
                    "incorrect encryption type".into(),
                ));
            }
        } else {
            return Err(EncryptionError::Decryption("missing payload".into()));
        };

        // Turn the bytes into a PublicKey
        let byte_arr: [u8; 32] = dr_header
            .dh
            .to_vec()
            .try_into()
            .map_err(|_| EncryptionError::Decryption("invalid public key length".into()))?;
        let dh_pub = PublicKey::from(byte_arr);

        // Build the Header that DR impl expects
        let header = Header {
            dh_pub: *dh_pub,
            msg_num: dr_header.msg_num,
            prev_chain_len: dr_header.prev_chain_len,
        };

        // Decrypt into Frame
        let content_bytes = self
            .dr_state
            .decrypt_message(&dr_header.ciphertext, header)
            .map_err(|e| EncryptionError::Decryption(e.to_string()))?;
        Ok(PrivateV1Frame::decode(content_bytes.as_slice()).unwrap())
    }

    // Handler for application content
    fn handle_content(&self, data: Vec<u8>) -> Vec<Event> {
        vec![Event::MessageReceived {
            conversation_id: Arc::from(self.id()),
            data,
        }]
    }

    /// Persists a conversation's metadata and ratchet state to DB.
    pub fn persist(&mut self) -> Result<ConversationIdOwned, ChatError> {
        let convo_info = ConversationMeta {
            local_convo_id: self.id().to_string(),
            remote_convo_id: self.remote_id(),
            kind: self.convo_type(),
        };
        self.store.lock().unwrap().save_conversation(&convo_info)?;
        self.save_ratchet_state(&mut *self.store.lock().unwrap())?;
        Ok(Arc::from(self.id()))
    }

    pub fn save_ratchet_state<T: RatchetStore>(&self, storage: &mut T) -> Result<(), ChatError> {
        let record = to_ratchet_record(&self.dr_state);
        let skipped_keys = to_skipped_key_records(&self.dr_state.skipped_keys());
        storage.save_ratchet_state(&self.local_convo_id, &record, &skipped_keys)?;
        Ok(())
    }
}

impl<S: ConversationStore + RatchetStore> Id for PrivateV1Convo<S> {
    fn id(&self) -> ConversationId<'_> {
        &self.local_convo_id
    }
}

impl<S: ConversationStore + RatchetStore> Convo for PrivateV1Convo<S> {
    fn send_message(
        &mut self,
        content: &[u8],
    ) -> Result<Vec<AddressedEncryptedPayload>, ChatError> {
        let frame = PrivateV1Frame {
            conversation_id: self.id().into(),
            sender: "delete".into(),
            timestamp: timestamp_millis(),
            frame_type: Some(FrameType::Content(content.to_vec().into())),
        };

        let data = self.encrypt(frame);

        self.save_ratchet_state::<S>(&mut *self.store.lock().unwrap())?;

        Ok(vec![AddressedEncryptedPayload {
            delivery_address: PRIVATE_V1_INBOX_ADDRESS.into(),
            data,
        }])
    }

    fn handle_frame(&mut self, encoded_payload: EncryptedPayload) -> Result<Vec<Event>, ChatError> {
        // Extract expected frame
        let frame = self
            .decrypt(encoded_payload)
            .map_err(|_| ChatError::Protocol("decryption".into()))?;

        let Some(frame_type) = frame.frame_type else {
            return Err(ChatError::ProtocolExpectation("None", "Some".into()));
        };

        self.save_ratchet_state(&mut *self.store.lock().unwrap())?;

        // Handle FrameTypes
        let output = match frame_type {
            FrameType::Content(bytes) => self.handle_content(bytes.into()),
            FrameType::Placeholder(_) => Vec::new(),
        };

        Ok(output)
    }

    fn remote_id(&self) -> String {
        self.remote_convo_id.clone()
    }

    fn convo_type(&self) -> ConversationKind {
        ConversationKind::PrivateV1
    }
}

impl<S: ConversationStore + RatchetStore> Debug for PrivateV1Convo<S> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PrivateV1Convo")
            .field("dr_state", &"******")
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use chat_sqlite::{ChatStorage, StorageConfig};
    use crypto::PrivateKey;

    use super::*;

    #[test]
    fn test_encrypt_roundtrip() {
        let saro = PrivateKey::random();
        let raya = PrivateKey::random();

        let saro_storage = Arc::new(Mutex::new(
            ChatStorage::new(StorageConfig::InMemory).unwrap(),
        ));

        let raya_storage = Arc::new(Mutex::new(
            ChatStorage::new(StorageConfig::InMemory).unwrap(),
        ));

        let pub_raya = PublicKey::from(&raya);

        let seed_key = saro.diffie_hellman(&pub_raya).DANGER_to_bytes();
        let seed_key_saro = SymmetricKey32::from(seed_key);
        let seed_key_raya = SymmetricKey32::from(seed_key);
        let send_content_bytes = vec![0, 2, 4, 6, 8];
        let mut sr_convo = PrivateV1Convo::new_initiator(saro_storage, seed_key_saro, pub_raya);
        let mut rs_convo = PrivateV1Convo::new_responder(raya_storage, seed_key_raya, &raya);

        let send_frame = PrivateV1Frame {
            conversation_id: "_".into(),
            sender: Bytes::new(),
            timestamp: timestamp_millis(),
            frame_type: Some(FrameType::Content(Bytes::from(send_content_bytes.clone()))),
        };
        let payload = sr_convo.encrypt(send_frame.clone());
        let recv_frame = rs_convo.decrypt(payload).unwrap();

        assert!(
            recv_frame == send_frame,
            "{:?}. {:?}",
            recv_frame,
            send_content_bytes
        );
    }
}