noq_proto/
cid_generator.rs1use std::hash::Hasher;
2
3use rand::Rng;
4use rand::RngExt;
5
6use crate::Duration;
7use crate::MAX_CID_SIZE;
8use crate::shared::ConnectionId;
9
10pub trait ConnectionIdGenerator: Send + Sync {
12 fn generate_cid(&mut self) -> ConnectionId;
20
21 fn validate(&self, _cid: ConnectionId) -> Result<(), InvalidCid> {
25 Ok(())
26 }
27
28 fn cid_len(&self) -> usize;
30 fn cid_lifetime(&self) -> Option<Duration>;
35}
36
37#[derive(Debug, Copy, Clone)]
39pub struct InvalidCid;
40
41#[derive(Debug, Clone, Copy)]
46pub struct RandomConnectionIdGenerator {
47 cid_len: usize,
48 lifetime: Option<Duration>,
49}
50
51impl Default for RandomConnectionIdGenerator {
52 fn default() -> Self {
53 Self {
54 cid_len: 8,
55 lifetime: None,
56 }
57 }
58}
59
60impl RandomConnectionIdGenerator {
61 pub fn new(cid_len: usize) -> Self {
65 debug_assert!(cid_len <= MAX_CID_SIZE);
66 Self {
67 cid_len,
68 ..Self::default()
69 }
70 }
71
72 pub fn set_lifetime(&mut self, d: Duration) -> &mut Self {
74 self.lifetime = Some(d);
75 self
76 }
77}
78
79impl ConnectionIdGenerator for RandomConnectionIdGenerator {
80 fn generate_cid(&mut self) -> ConnectionId {
81 let mut bytes_arr = [0; MAX_CID_SIZE];
82 rand::rng().fill_bytes(&mut bytes_arr[..self.cid_len]);
83
84 ConnectionId::new(&bytes_arr[..self.cid_len])
85 }
86
87 fn cid_len(&self) -> usize {
89 self.cid_len
90 }
91
92 fn cid_lifetime(&self) -> Option<Duration> {
93 self.lifetime
94 }
95}
96
97pub struct HashedConnectionIdGenerator {
103 key: u64,
104 lifetime: Option<Duration>,
105}
106
107impl HashedConnectionIdGenerator {
108 pub fn new() -> Self {
110 Self::from_key(rand::rng().random())
111 }
112
113 pub fn from_key(key: u64) -> Self {
118 Self {
119 key,
120 lifetime: None,
121 }
122 }
123
124 pub fn set_lifetime(&mut self, d: Duration) -> &mut Self {
126 self.lifetime = Some(d);
127 self
128 }
129}
130
131impl Default for HashedConnectionIdGenerator {
132 fn default() -> Self {
133 Self::new()
134 }
135}
136
137impl ConnectionIdGenerator for HashedConnectionIdGenerator {
138 fn generate_cid(&mut self) -> ConnectionId {
139 let mut bytes_arr = [0; NONCE_LEN + SIGNATURE_LEN];
140 rand::rng().fill_bytes(&mut bytes_arr[..NONCE_LEN]);
141 let mut hasher = rustc_hash::FxHasher::default();
142 hasher.write_u64(self.key);
143 hasher.write(&bytes_arr[..NONCE_LEN]);
144 bytes_arr[NONCE_LEN..].copy_from_slice(&hasher.finish().to_le_bytes()[..SIGNATURE_LEN]);
145 ConnectionId::new(&bytes_arr)
146 }
147
148 fn validate(&self, cid: ConnectionId) -> Result<(), InvalidCid> {
149 let (nonce, signature) = cid.split_at(NONCE_LEN);
150 let mut hasher = rustc_hash::FxHasher::default();
151 hasher.write_u64(self.key);
152 hasher.write(nonce);
153 let expected = hasher.finish().to_le_bytes();
154 match expected[..SIGNATURE_LEN] == signature[..] {
155 true => Ok(()),
156 false => Err(InvalidCid),
157 }
158 }
159
160 fn cid_len(&self) -> usize {
161 NONCE_LEN + SIGNATURE_LEN
162 }
163
164 fn cid_lifetime(&self) -> Option<Duration> {
165 self.lifetime
166 }
167}
168
169const NONCE_LEN: usize = 3; const SIGNATURE_LEN: usize = 8 - NONCE_LEN; #[cfg(test)]
173mod tests {
174 use super::*;
175
176 #[test]
177 fn validate_keyed_cid() {
178 let mut generator = HashedConnectionIdGenerator::new();
179 let cid = generator.generate_cid();
180 generator.validate(cid).unwrap();
181 }
182}