enigma_protocol/
session.rs

1use std::collections::HashMap;
2
3use enigma_packet::{Message, MessageMeta};
4use uuid::Uuid;
5
6use crate::crypto::{conversation_ad_bytes, SessionRatchet};
7use crate::error::{EnigmaProtocolError, Result};
8use crate::types::SessionBootstrap;
9
10pub struct Session {
11    local_id: String,
12    remote_id: String,
13    ratchet: SessionRatchet,
14    attachments: HashMap<Uuid, AttachmentTracker>,
15}
16
17struct AttachmentTracker {
18    meta: enigma_packet::AttachmentMeta,
19    received_chunks: u32,
20    received_size: u64,
21}
22
23pub enum AttachmentUpdate {
24    Init {
25        attachment_id: Uuid,
26        total_chunks: u32,
27    },
28    Chunk {
29        attachment_id: Uuid,
30        received_chunks: u32,
31        total_chunks: u32,
32    },
33    End {
34        attachment_id: Uuid,
35        total_size: u64,
36    },
37}
38
39impl Session {
40    pub fn new(local: &str, remote: &str, bootstrap: &SessionBootstrap) -> Result<Self> {
41        let ratchet = SessionRatchet::from_bootstrap(bootstrap, local, remote)?;
42        Ok(Self {
43            local_id: local.to_string(),
44            remote_id: remote.to_string(),
45            ratchet,
46            attachments: HashMap::new(),
47        })
48    }
49
50    pub fn conversation_id(&self) -> Uuid {
51        self.ratchet.conversation_id()
52    }
53
54    pub fn local_id(&self) -> &str {
55        &self.local_id
56    }
57
58    pub fn remote_id(&self) -> &str {
59        &self.remote_id
60    }
61
62    pub fn encrypt_message(&mut self, message: &Message) -> Result<Vec<u8>> {
63        let encoded = enigma_packet::encode_message(message)?;
64        let key = self.ratchet.next_send_key()?;
65        let box_ = enigma_aead::AeadBox::new(key);
66        let ad = conversation_ad_bytes(&self.conversation_id(), &self.local_id, &self.remote_id);
67        let ciphertext = box_.encrypt(&encoded, &ad)?;
68        Ok(ciphertext)
69    }
70
71    pub fn decrypt_packet(&mut self, packet: &[u8]) -> Result<Message> {
72        let key = self.ratchet.next_recv_key()?;
73        let box_ = enigma_aead::AeadBox::new(key);
74        let ad = conversation_ad_bytes(&self.conversation_id(), &self.remote_id, &self.local_id);
75        let plaintext = box_.decrypt(packet, &ad)?;
76        let message = enigma_packet::decode_message(&plaintext)?;
77        Ok(message)
78    }
79
80    pub fn handle_attachment_message(
81        &mut self,
82        message: &Message,
83    ) -> Result<Option<AttachmentUpdate>> {
84        match &message.meta {
85            MessageMeta::AttachmentInit(meta) => {
86                let tracker = AttachmentTracker {
87                    meta: meta.clone(),
88                    received_chunks: 0,
89                    received_size: 0,
90                };
91                self.attachments.insert(meta.attachment_id, tracker);
92                Ok(Some(AttachmentUpdate::Init {
93                    attachment_id: meta.attachment_id,
94                    total_chunks: meta.chunk_count,
95                }))
96            }
97            MessageMeta::AttachmentChunk(chunk) => {
98                let tracker = self
99                    .attachments
100                    .get_mut(&chunk.attachment_id)
101                    .ok_or(EnigmaProtocolError::Attachment)?;
102                if chunk.index != tracker.received_chunks {
103                    return Err(EnigmaProtocolError::Attachment);
104                }
105                tracker.received_chunks = tracker
106                    .received_chunks
107                    .checked_add(1)
108                    .ok_or(EnigmaProtocolError::Attachment)?;
109                tracker.received_size = tracker
110                    .received_size
111                    .checked_add(chunk.chunk_size as u64)
112                    .ok_or(EnigmaProtocolError::Attachment)?;
113                Ok(Some(AttachmentUpdate::Chunk {
114                    attachment_id: chunk.attachment_id,
115                    received_chunks: tracker.received_chunks,
116                    total_chunks: tracker.meta.chunk_count,
117                }))
118            }
119            MessageMeta::AttachmentEnd {
120                attachment_id,
121                total_size,
122                chunk_count,
123                ..
124            } => {
125                let tracker = self
126                    .attachments
127                    .remove(attachment_id)
128                    .ok_or(EnigmaProtocolError::Attachment)?;
129                if &tracker.meta.chunk_count != chunk_count {
130                    return Err(EnigmaProtocolError::Attachment);
131                }
132                if &tracker.meta.total_size != total_size {
133                    return Err(EnigmaProtocolError::Attachment);
134                }
135                Ok(Some(AttachmentUpdate::End {
136                    attachment_id: *attachment_id,
137                    total_size: *total_size,
138                }))
139            }
140            _ => Ok(None),
141        }
142    }
143}
enigma_protocol/session.rs

enigma_protocol/
session.rs
