zoe_wire_protocol/message/
store_key.rs

1use std::cmp::Ordering;
2
3use serde::{Deserialize, Serialize};
4
5// PQXDH types are available through the main crate re-exports
6
7/// PQXDH-based inbox protocols
8///
9/// Reserved ranges:
10/// - 10000-10999: Core PQXDH protocols (group invites, direct messages, etc.)
11/// - 10000-10999: Ephemeral group invites (randomized within this range)
12/// - 11000-11999: PQXDH RPC services
13/// - 15001-19999: Available for custom/experimental PQXDH protocols
14#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
15#[serde(from = "u32", into = "u32")]
16#[repr(u32)]
17pub enum PqxdhInboxProtocol {
18    // Reserved for Core protocols (10000-10999)
19    // RPC services (11000-11999)
20    EchoService = 1000, // 11000 - Echo/ping RPC service
21
22    // WhatsApp Bot Protocol 11202
23    WhatsAppBot = 1202,
24
25    /// 12000-12999 - Ephemeral Inboxes, using a randomized value in this range
26    /// expected to only live for a short time and a specific purpose or connection
27    /// like sharing a group invite
28    Ephemeral(u32),
29    // Custom protocols (anything else)
30    CustomProtocol(u32), // 15001+ - Custom PQXDH protocols
31}
32
33impl PartialOrd for PqxdhInboxProtocol {
34    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
35        Some(self.cmp(other))
36    }
37}
38
39impl Ord for PqxdhInboxProtocol {
40    fn cmp(&self, other: &Self) -> Ordering {
41        u32::from(self).cmp(&u32::from(other))
42    }
43}
44
45impl PqxdhInboxProtocol {
46    pub fn as_u32(&self) -> u32 {
47        u32::from(self)
48    }
49    pub fn into_bytes(&self) -> Vec<u8> {
50        self.as_u32().to_le_bytes().to_vec()
51    }
52
53    pub fn from_bytes(bytes: &[u8]) -> Result<Self, String> {
54        let value = u32::from_le_bytes(bytes.try_into().map_err(|_| "Invalid bytes")?);
55        Ok(Self::from(value))
56    }
57}
58
59#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
60#[serde(from = "u32", into = "u32")]
61pub enum StoreKey {
62    PublicUserInfo,
63    MlsKeyPackage,
64    PqxdhInbox(PqxdhInboxProtocol), // PQXDH inbox system
65    CustomKey(u32),                 // yet to be known variant
66}
67
68impl From<u32> for PqxdhInboxProtocol {
69    fn from(value: u32) -> Self {
70        match value {
71            // RPC services (1000-1999 maps to 11000-11999)
72            1000 => PqxdhInboxProtocol::EchoService,
73
74            // Bot services (1200-1299)
75            1202 => PqxdhInboxProtocol::WhatsAppBot,
76
77            // Ephemeral protocols (0- 999 maps to 12000-12999)
78            2000..=2999 => PqxdhInboxProtocol::Ephemeral(value - 2000),
79
80            // Everything else is custom
81            a => PqxdhInboxProtocol::CustomProtocol(a),
82        }
83    }
84}
85
86impl std::fmt::Display for PqxdhInboxProtocol {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        match self {
89            PqxdhInboxProtocol::Ephemeral(id) => write!(f, "Ephemeral({id})"),
90            PqxdhInboxProtocol::EchoService => write!(f, "EchoService"),
91            PqxdhInboxProtocol::WhatsAppBot => write!(f, "WhatsAppBot"),
92            PqxdhInboxProtocol::CustomProtocol(a) => write!(f, "CustomProtocol({a})"),
93        }
94    }
95}
96
97impl From<&PqxdhInboxProtocol> for u32 {
98    fn from(val: &PqxdhInboxProtocol) -> Self {
99        match val {
100            // Core protocols (0-999 maps to 10000-10999)
101            PqxdhInboxProtocol::Ephemeral(id) => *id,
102
103            // RPC services
104            PqxdhInboxProtocol::EchoService => 1000,
105
106            // Bot services
107            PqxdhInboxProtocol::WhatsAppBot => 1202,
108
109            // Custom protocols
110            PqxdhInboxProtocol::CustomProtocol(a) => *a,
111        }
112    }
113}
114
115impl From<PqxdhInboxProtocol> for u32 {
116    fn from(val: PqxdhInboxProtocol) -> Self {
117        u32::from(&val)
118    }
119}
120
121impl From<u32> for StoreKey {
122    fn from(value: u32) -> Self {
123        match value {
124            0 => StoreKey::PublicUserInfo,
125            100 => StoreKey::MlsKeyPackage,
126            10000..=19999 => {
127                let protocol_id = value - 10000;
128                StoreKey::PqxdhInbox(PqxdhInboxProtocol::from(protocol_id))
129            }
130            a => StoreKey::CustomKey(a),
131        }
132    }
133}
134
135impl From<StoreKey> for u32 {
136    fn from(val: StoreKey) -> Self {
137        match val {
138            StoreKey::PublicUserInfo => 0,
139            StoreKey::MlsKeyPackage => 100,
140            StoreKey::PqxdhInbox(protocol) => 10000 + u32::from(&protocol),
141            StoreKey::CustomKey(a) => a,
142        }
143    }
144}
145
146impl From<&StoreKey> for u32 {
147    fn from(val: &StoreKey) -> Self {
148        match val {
149            StoreKey::PublicUserInfo => 0,
150            StoreKey::MlsKeyPackage => 100,
151            StoreKey::PqxdhInbox(protocol) => 10000 + u32::from(protocol),
152            StoreKey::CustomKey(a) => *a,
153        }
154    }
155}
156
157#[cfg(test)]
158mod test {
159
160    use super::*;
161    use postcard;
162
163    #[test]
164    fn test_store_key_serialization() {
165        let keys = vec![
166            StoreKey::PublicUserInfo,
167            StoreKey::MlsKeyPackage,
168            StoreKey::PqxdhInbox(PqxdhInboxProtocol::EchoService),
169            StoreKey::PqxdhInbox(PqxdhInboxProtocol::CustomProtocol(5001)),
170            StoreKey::CustomKey(10),
171        ];
172        // Serialize and deserialize
173        let serialized = postcard::to_stdvec(&keys).unwrap();
174        let deserialized: Vec<StoreKey> = postcard::from_bytes(&serialized).unwrap();
175
176        assert_eq!(keys, deserialized);
177    }
178
179    #[test]
180    fn test_store_key_regular_is_flat() {
181        let key = StoreKey::PublicUserInfo;
182        // Serialize and deserialize
183        let serialized = postcard::to_stdvec(&key).unwrap();
184        assert_eq!(serialized.len(), 1);
185        let deserialized: StoreKey = postcard::from_bytes(&serialized).unwrap();
186
187        assert_eq!(key, deserialized);
188    }
189
190    #[test]
191    fn test_store_key_custom_type_is_flat() {
192        let key = StoreKey::CustomKey(42);
193        // Serialize and deserialize
194        let serialized = postcard::to_stdvec(&key).unwrap();
195        assert_eq!(serialized.len(), 1); // same length
196        let deserialized: StoreKey = postcard::from_bytes(&serialized).unwrap();
197
198        assert_eq!(key, deserialized);
199    }
200
201    #[test]
202    fn test_store_key_long_custom_type_is_flat() {
203        let key = StoreKey::CustomKey(1024);
204        // Serialize and deserialize
205        let serialized = postcard::to_stdvec(&key).unwrap();
206        assert_eq!(serialized.len(), 2); // varint
207        let deserialized: StoreKey = postcard::from_bytes(&serialized).unwrap();
208
209        assert_eq!(key, deserialized);
210    }
211
212    #[test]
213    fn test_pqxdh_inbox_protocol_mapping() {
214        // Test EchoService mapping
215        let echo_protocol = PqxdhInboxProtocol::EchoService;
216        let echo_u32: u32 = echo_protocol.clone().into();
217        assert_eq!(echo_u32, 1000);
218        let echo_back = PqxdhInboxProtocol::from(echo_u32);
219        assert_eq!(echo_protocol, echo_back);
220
221        // Test CustomProtocol mapping
222        let custom_protocol = PqxdhInboxProtocol::CustomProtocol(5001);
223        let custom_u32: u32 = custom_protocol.clone().into();
224        assert_eq!(custom_u32, 5001);
225        let custom_back = PqxdhInboxProtocol::from(custom_u32);
226        assert_eq!(custom_protocol, custom_back);
227    }
228
229    #[test]
230    fn test_pqxdh_store_key_mapping() {
231        // Test EchoService StoreKey mapping
232        let echo_store_key = StoreKey::PqxdhInbox(PqxdhInboxProtocol::EchoService);
233        let echo_u32: u32 = echo_store_key.clone().into();
234        assert_eq!(echo_u32, 11000); // 10000 + 1000
235        let echo_back = StoreKey::from(echo_u32);
236        assert_eq!(echo_store_key, echo_back);
237
238        // Test CustomProtocol StoreKey mapping
239        let custom_store_key = StoreKey::PqxdhInbox(PqxdhInboxProtocol::CustomProtocol(5001));
240        let custom_u32: u32 = custom_store_key.clone().into();
241        assert_eq!(custom_u32, 15001); // 10000 + 5001
242        let custom_back = StoreKey::from(custom_u32);
243        assert_eq!(custom_store_key, custom_back);
244    }
245
246    #[test]
247    fn test_pqxdh_inbox_protocol_serialization() {
248        let protocols = vec![
249            PqxdhInboxProtocol::EchoService,
250            PqxdhInboxProtocol::CustomProtocol(5001),
251            PqxdhInboxProtocol::CustomProtocol(9999),
252        ];
253
254        // Serialize and deserialize
255        let serialized = postcard::to_stdvec(&protocols).unwrap();
256        let deserialized: Vec<PqxdhInboxProtocol> = postcard::from_bytes(&serialized).unwrap();
257
258        assert_eq!(protocols, deserialized);
259    }
260
261    #[test]
262    fn test_pqxdh_store_key_ranges() {
263        // Test that PQXDH inbox range is correctly handled
264
265        // Test lower bound (10000 = EchoService)
266        let lower_bound = StoreKey::from(11000u32);
267        assert_eq!(
268            lower_bound,
269            StoreKey::PqxdhInbox(PqxdhInboxProtocol::EchoService)
270        );
271
272        // Test custom protocol in range
273        let custom_in_range = StoreKey::from(15001u32);
274        assert_eq!(
275            custom_in_range,
276            StoreKey::PqxdhInbox(PqxdhInboxProtocol::CustomProtocol(5001))
277        );
278
279        // Test upper bound (19999)
280        let upper_bound = StoreKey::from(19999u32);
281        assert_eq!(
282            upper_bound,
283            StoreKey::PqxdhInbox(PqxdhInboxProtocol::CustomProtocol(9999))
284        );
285
286        // Test outside range becomes CustomKey
287        let outside_range = StoreKey::from(20000u32);
288        assert_eq!(outside_range, StoreKey::CustomKey(20000));
289    }
290
291    #[test]
292    fn test_store_key_reference_conversion() {
293        let echo_store_key = StoreKey::PqxdhInbox(PqxdhInboxProtocol::EchoService);
294        let echo_u32: u32 = (&echo_store_key).into();
295        assert_eq!(echo_u32, 11000);
296
297        let custom_store_key = StoreKey::PqxdhInbox(PqxdhInboxProtocol::CustomProtocol(5001));
298        let custom_u32: u32 = (&custom_store_key).into();
299        assert_eq!(custom_u32, 15001);
300    }
301}