mixnet/core/
surb_keystore.rs

1// Copyright 2022 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21//! Keystore for SURB payload encryption keys.
22
23use super::{
24	fragment::MessageId,
25	sphinx::{SurbId, SurbPayloadEncryptionKeys, SURB_ID_SIZE},
26};
27use hashlink::{linked_hash_map, LinkedHashMap};
28use log::debug;
29use rand::{CryptoRng, Rng};
30
31struct Value {
32	keys: SurbPayloadEncryptionKeys,
33	message_id: MessageId,
34}
35
36pub struct Entry<'a>(linked_hash_map::OccupiedEntry<'a, SurbId, Value>);
37
38impl<'a> Entry<'a> {
39	pub fn keys(&self) -> &SurbPayloadEncryptionKeys {
40		&self.0.get().keys
41	}
42
43	pub fn message_id(&self) -> &MessageId {
44		&self.0.get().message_id
45	}
46
47	pub fn remove(self) {
48		self.0.remove();
49	}
50}
51
52pub struct SurbKeystore {
53	/// Maximum number of SURBs to keep keys for.
54	capacity: usize,
55	/// In creation order: oldest SURBs at the front, newest SURBs at the back.
56	surbs: LinkedHashMap<SurbId, Value>,
57}
58
59impl SurbKeystore {
60	pub fn new(capacity: usize) -> Self {
61		debug_assert_ne!(capacity, 0);
62		Self { capacity, surbs: LinkedHashMap::with_capacity(capacity) }
63	}
64
65	/// Create an entry for a new SURB. Returns the randomly generated ID and a mutable reference
66	/// to the keys, which should be filled in by the caller.
67	pub fn insert(
68		&mut self,
69		rng: &mut (impl Rng + CryptoRng),
70		message_id: &MessageId,
71		log_target: &str,
72	) -> (SurbId, &mut SurbPayloadEncryptionKeys) {
73		// Discard the oldest SURB if we're already at capacity
74		debug_assert!(self.surbs.len() <= self.capacity);
75		if self.surbs.len() == self.capacity {
76			debug!(target: log_target, "Too many entries in SURB keystore; evicting oldest");
77			self.surbs.pop_front();
78		}
79
80		let mut id = [0; SURB_ID_SIZE];
81		rng.fill_bytes(&mut id);
82		match self.surbs.entry(id) {
83			linked_hash_map::Entry::Occupied(_) => panic!(
84				"Randomly generated SURB ID matches an existing SURB ID; something wrong with RNG?"
85			),
86			linked_hash_map::Entry::Vacant(entry) => {
87				let value = entry.insert(Value {
88					keys: SurbPayloadEncryptionKeys::new(),
89					message_id: *message_id,
90				});
91				(id, &mut value.keys)
92			},
93		}
94	}
95
96	/// Returns the entry for a SURB, or [`None`] if the ID is not recognised.
97	pub fn entry(&mut self, id: &SurbId) -> Option<Entry> {
98		match self.surbs.entry(*id) {
99			linked_hash_map::Entry::Occupied(entry) => Some(Entry(entry)),
100			linked_hash_map::Entry::Vacant(_) => None,
101		}
102	}
103}