Skip to main content

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