secret_manager/
secret_rotation.rs1#[cfg(feature = "arc-swap")]
2use arc_swap::ArcSwap;
3#[cfg(feature = "parking-lot")]
4use parking_lot::RwLock as ParkingRwLock;
5#[cfg(not(any(feature = "arc-swap", feature = "parking-lot")))]
6use std::sync::RwLock as StdRwLock;
7
8#[cfg(feature = "arc-swap")]
9use std::sync::Arc;
10
11#[derive(Clone)]
16struct SecretInner<const V: usize, const S: usize> {
17 keys: [Option<[u8; S]>; V],
19 current_version: u8,
20}
21
22pub trait SecretGroup<const V: usize = 256, const S: usize = 32>: Send + Sync {
28 fn current(&self) -> (u8, [u8; S]);
30
31 fn resolve(&self, version: u8) -> Option<[u8; S]>;
33}
34
35pub struct InMemorySecretGroup<const V: usize = 256, const S: usize = 32> {
37 #[cfg(feature = "arc-swap")]
38 inner: ArcSwap<SecretInner<V, S>>,
39
40 #[cfg(all(feature = "parking-lot", not(feature = "arc-swap")))]
41 inner: ParkingRwLock<SecretInner<V, S>>,
42
43 #[cfg(not(any(feature = "arc-swap", feature = "parking-lot")))]
44 inner: StdRwLock<SecretInner<V, S>>,
45}
46
47impl<const V: usize, const S: usize> InMemorySecretGroup<V, S> {
48 pub fn new(version: u8, initial_key: [u8; S]) -> Self {
50 assert!(
51 (version as usize) < V,
52 "version {} out of range for ring buffer of size {V}",
53 version
54 );
55 let mut keys: [Option<[u8; S]>; V] = std::array::from_fn(|_| None);
56 keys[version as usize] = Some(initial_key);
57
58 let inner_val = SecretInner {
59 keys,
60 current_version: version,
61 };
62
63 Self {
64 #[cfg(feature = "arc-swap")]
65 inner: ArcSwap::from_pointee(inner_val),
66
67 #[cfg(all(feature = "parking-lot", not(feature = "arc-swap")))]
68 inner: ParkingRwLock::new(inner_val),
69
70 #[cfg(not(any(feature = "arc-swap", feature = "parking-lot")))]
71 inner: StdRwLock::new(inner_val),
72 }
73 }
74
75 pub fn store_key(&self, version: u8, key: [u8; S]) {
77 assert!(
78 (version as usize) < V,
79 "version {} out of range for ring buffer of size {V}",
80 version
81 );
82
83 #[cfg(feature = "arc-swap")]
84 {
85 let mut inner = (**self.inner.load()).clone();
86 inner.keys[version as usize] = Some(key);
87 self.inner.store(Arc::new(inner));
88 }
89
90 #[cfg(all(feature = "parking-lot", not(feature = "arc-swap")))]
91 {
92 let mut inner = self.inner.write();
93 inner.keys[version as usize] = Some(key);
94 }
95
96 #[cfg(not(any(feature = "arc-swap", feature = "parking-lot")))]
97 {
98 let mut inner = self.inner.write().expect("lock poisoned");
99 inner.keys[version as usize] = Some(key);
100 }
101 }
102
103 pub fn promote(&self, version: u8) {
105 assert!(
106 (version as usize) < V,
107 "version {} out of range for ring buffer of size {V}",
108 version
109 );
110
111 #[cfg(feature = "arc-swap")]
112 {
113 let mut inner = (**self.inner.load()).clone();
114 if inner.keys[version as usize].is_none() {
115 panic!("cannot promote version {version} before it is stored");
116 }
117 inner.current_version = version;
118 self.inner.store(Arc::new(inner));
119 }
120
121 #[cfg(all(feature = "parking-lot", not(feature = "arc-swap")))]
122 {
123 let mut inner = self.inner.write();
124 if inner.keys[version as usize].is_none() {
125 panic!("cannot promote version {version} before it is stored");
126 }
127 inner.current_version = version;
128 }
129
130 #[cfg(not(any(feature = "arc-swap", feature = "parking-lot")))]
131 {
132 let mut inner = self.inner.write().expect("lock poisoned");
133 if inner.keys[version as usize].is_none() {
134 panic!("cannot promote version {version} before it is stored");
135 }
136 inner.current_version = version;
137 }
138 }
139
140 pub fn apply(&self, version: u8, key: [u8; S]) {
142 assert!(
143 (version as usize) < V,
144 "version {} out of range for ring buffer of size {V}",
145 version
146 );
147
148 #[cfg(feature = "arc-swap")]
149 {
150 let mut inner = (**self.inner.load()).clone();
151 inner.keys[version as usize] = Some(key);
152 inner.current_version = version;
153 self.inner.store(Arc::new(inner));
154 }
155
156 #[cfg(all(feature = "parking-lot", not(feature = "arc-swap")))]
157 {
158 let mut inner = self.inner.write();
159 inner.keys[version as usize] = Some(key);
160 inner.current_version = version;
161 }
162
163 #[cfg(not(any(feature = "arc-swap", feature = "parking-lot")))]
164 {
165 let mut inner = self.inner.write().expect("lock poisoned");
166 inner.keys[version as usize] = Some(key);
167 inner.current_version = version;
168 }
169 }
170}
171
172impl<const V: usize, const S: usize> SecretGroup<V, S> for InMemorySecretGroup<V, S> {
173 fn current(&self) -> (u8, [u8; S]) {
174 #[cfg(feature = "arc-swap")]
175 let (v, keys) = {
176 let inner = self.inner.load();
177 (inner.current_version, inner.keys)
178 };
179
180 #[cfg(all(feature = "parking-lot", not(feature = "arc-swap")))]
181 let (v, keys) = {
182 let inner = self.inner.read();
183 (inner.current_version, inner.keys)
184 };
185
186 #[cfg(not(any(feature = "arc-swap", feature = "parking-lot")))]
187 let (v, keys) = {
188 let inner = self.inner.read().expect("lock poisoned");
189 (inner.current_version, inner.keys)
190 };
191
192 let key = keys[v as usize].expect("current_version slot must always be populated");
193 (v, key)
194 }
195
196 fn resolve(&self, version: u8) -> Option<[u8; S]> {
197 #[cfg(feature = "arc-swap")]
198 return self.inner.load().keys[version as usize];
199
200 #[cfg(all(feature = "parking-lot", not(feature = "arc-swap")))]
201 return self.inner.read().keys[version as usize];
202
203 #[cfg(not(any(feature = "arc-swap", feature = "parking-lot")))]
204 return self.inner.read().expect("lock poisoned").keys[version as usize];
205 }
206}
207
208#[cfg(test)]
213mod tests {
214 use super::*;
215 use std::sync::Arc;
216
217 const KEY_A: [u8; 32] = [1u8; 32];
218 const KEY_B: [u8; 32] = [2u8; 32];
219
220 #[test]
221 fn new_returns_initial_key_as_current() {
222 let sg = InMemorySecretGroup::<256, 32>::new(0, KEY_A);
223 let (v, k) = sg.current();
224 assert_eq!(v, 0);
225 assert_eq!(k, KEY_A);
226 }
227
228 #[test]
229 fn resolve_returns_none_for_unpopulated_slot() {
230 let sg = InMemorySecretGroup::<256, 32>::new(0, KEY_A);
231 assert!(sg.resolve(1).is_none());
232 assert!(sg.resolve(255).is_none());
233 }
234
235 #[test]
236 fn resolve_returns_some_for_populated_slot() {
237 let sg = InMemorySecretGroup::<256, 32>::new(0, KEY_A);
238 assert_eq!(sg.resolve(0), Some(KEY_A));
239 }
240
241 #[test]
242 fn apply_updates_current_and_ring() {
243 let sg = InMemorySecretGroup::<256, 32>::new(0, KEY_A);
244 sg.apply(1, KEY_B);
245 let (v, k) = sg.current();
246 assert_eq!(v, 1);
247 assert_eq!(k, KEY_B);
248 assert_eq!(sg.resolve(0), Some(KEY_A));
249 assert_eq!(sg.resolve(1), Some(KEY_B));
250 }
251
252 #[tokio::test]
253 async fn concurrent_reads_during_apply_are_safe() {
254 let sg = Arc::new(InMemorySecretGroup::<256, 32>::new(0, KEY_A));
255 let sg2 = sg.clone();
256
257 let reader = tokio::spawn(async move {
258 for _ in 0..1000 {
259 let _ = sg2.current();
260 let _ = sg2.resolve(0);
261 let _ = sg2.resolve(1);
262 tokio::task::yield_now().await;
263 }
264 });
265
266 for i in 0u8..10 {
267 sg.apply(i, KEY_B);
268 }
269
270 reader.await.expect("reader must not panic");
271 }
272}