securedrop_protocol_minimal/
storage.rs

1use alloc::vec::Vec;
2use hashbrown::HashMap;
3use rand_core::{CryptoRng, RngCore};
4use uuid::Uuid;
5
6use crate::message::MessagePublicKey;
7use crate::primitives::ristretto255::DHPublicKey;
8use crate::sign::{JournalistLongTermKey, NewsroomOnJournalist, Signature, VerifyingKey};
9use crate::{Enrollment, Envelope, SignedKeyBundlePublic, SignedLongtermKeyBundle};
10
11pub type ServerMessageStore = HashMap<Uuid, Envelope>;
12
13#[derive(Default)]
14pub struct ServerStorage {
15    /// Journalists with their long/medium term keys, self-signature, newsroom signature.
16    journalists: HashMap<
17        Uuid,
18        (
19            VerifyingKey,
20            SignedLongtermKeyBundle,
21            Signature<NewsroomOnJournalist>,
22        ),
23    >,
24
25    /// Journalists ephemeral keystore
26    /// Maps journalist ID to a vector of ephemeral key sets
27    /// Each journalist maintains a pool of ephemeral keys that are randomly selected and removed when fetched
28    /// TODO recheck lifetime
29    ephemeral_keys: HashMap<Uuid, Vec<SignedKeyBundlePublic>>,
30
31    /// Store of messages
32    messages: HashMap<Uuid, Envelope>,
33}
34
35impl ServerStorage {
36    /// Create a new ServerStorage instance
37    pub fn new() -> Self {
38        Self::default()
39    }
40
41    /// Add ephemeral keys for a journalist
42    pub fn add_ephemeral_keys(&mut self, journalist_id: Uuid, keys: Vec<SignedKeyBundlePublic>) {
43        let journalist_keys = self
44            .ephemeral_keys
45            .entry(journalist_id)
46            // avoid `or_insert_with(Vec::new)` because hax doesn't accept FnMut/FnOnce closure
47            .or_insert(Vec::new());
48        journalist_keys.extend(keys);
49    }
50
51    /// Get a random ephemeral key set for a journalist and remove it from the pool
52    /// Returns None if no keys are available for this journalist
53    ///
54    /// Note: This method deletes the ephemeral key from storage.
55    /// The returned key is permanently removed from the journalist's ephemeral key pool.
56    pub fn pop_random_ephemeral_keys<R: RngCore + CryptoRng>(
57        &mut self,
58        journalist_id: Uuid,
59        rng: &mut R,
60    ) -> Option<SignedKeyBundlePublic> {
61        if let Some(keys) = self.ephemeral_keys.get_mut(&journalist_id) {
62            if keys.is_empty() {
63                return None;
64            }
65
66            // Select a "random" index (note: Modulo bias, Toy purposes only!)
67            let index = rng.next_u32() as usize % keys.len();
68
69            // Remove and return the selected key set
70            Some(keys.remove(index))
71        } else {
72            None
73        }
74    }
75
76    /// Get random ephemeral keys for all journalists
77    /// Returns a vector of (journalist_id, ephemeral_keys) pairs
78    /// Only includes journalists that have available keys
79    ///
80    /// Note: This method deletes the ephemeral keys from storage.
81    /// Each call removes the returned keys from the journalist's ephemeral key pool.
82    pub fn get_all_ephemeral_keys<R: RngCore + CryptoRng>(
83        &mut self,
84        rng: &mut R,
85    ) -> Vec<(Uuid, SignedKeyBundlePublic)> {
86        let mut result = Vec::new();
87        let journalist_ids: Vec<Uuid> = self.ephemeral_keys.keys().copied().collect();
88
89        for journalist_id in journalist_ids {
90            if let Some(keys) = self.pop_random_ephemeral_keys(journalist_id, rng) {
91                result.push((journalist_id, keys));
92            }
93        }
94
95        result
96    }
97
98    /// Check how many ephemeral keys are available for a journalist
99    pub fn ephemeral_keys_count(&self, journalist_id: Uuid) -> usize {
100        self.ephemeral_keys
101            .get(&journalist_id)
102            .map_or(0, |keys| keys.len())
103    }
104
105    /// Check if a journalist has any ephemeral keys available
106    pub fn has_ephemeral_keys(&self, journalist_id: Uuid) -> bool {
107        self.ephemeral_keys_count(journalist_id) > 0
108    }
109
110    /// Get all journalists
111    pub fn get_journalists(
112        &self,
113    ) -> &HashMap<
114        Uuid,
115        (
116            VerifyingKey,
117            SignedLongtermKeyBundle,
118            Signature<NewsroomOnJournalist>,
119        ),
120    > {
121        &self.journalists
122    }
123
124    /// Add a journalist to storage and return the generated UUID
125    pub fn add_journalist(
126        &mut self,
127        journalist: Enrollment,
128        newsroom_signature: Signature<NewsroomOnJournalist>,
129    ) -> Uuid {
130        let journalist_id = Uuid::new_v4();
131
132        // match hashmap above
133        let bundle = journalist.bundle;
134        let values = (journalist.verification_key, bundle, newsroom_signature);
135
136        self.journalists.insert(journalist_id, values);
137        journalist_id
138    }
139
140    /// Find a journalist by their verifying key
141    /// Returns the journalist ID if found
142    ///
143    /// TODO: Remove?
144    pub fn find_journalist_by_verifying_key(&self, verifying_key: &VerifyingKey) -> Option<Uuid> {
145        for (journalist_id, (stored_vk, _, _)) in &self.journalists {
146            if stored_vk.into_bytes() == verifying_key.into_bytes() {
147                return Some(*journalist_id);
148            }
149        }
150        None
151    }
152
153    pub(crate) fn deterministic_uuid<R: RngCore + CryptoRng>(&mut self, rng: &mut R) -> Uuid {
154        let mut bytes = [0u8; 16];
155        rng.fill_bytes(&mut bytes);
156
157        uuid::Builder::from_random_bytes(bytes).into_uuid()
158    }
159
160    /// Get all messages
161    pub fn get_messages(&self) -> &HashMap<Uuid, Envelope> {
162        &self.messages
163    }
164
165    /// Add a message to storage
166    pub fn add_message(&mut self, message_id: Uuid, message: Envelope) {
167        self.messages.insert(message_id, message);
168    }
169}
securedrop_protocol_minimal/storage.rs

securedrop_protocol_minimal/
storage.rs
