Skip to main content

secret_manager/
secret_rotation.rs

1#[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// ---------------------------------------------------------------------------
12// Internal ring-buffer state
13// ---------------------------------------------------------------------------
14
15#[derive(Clone)]
16struct SecretInner<const V: usize, const S: usize> {
17    /// Indexed by version (u8 cast to usize). Slot is `None` until first written.
18    keys: [Option<[u8; S]>; V],
19    current_version: u8,
20}
21
22// ---------------------------------------------------------------------------
23// Public thread-safe wrapper
24// ---------------------------------------------------------------------------
25
26/// A trait for versioned secret key groups.
27pub trait SecretGroup<const V: usize = 256, const S: usize = 32>: Send + Sync {
28    /// Return `(current_version, key_bytes)`.
29    fn current(&self) -> (u8, [u8; S]);
30
31    /// Look up a key by version. Returns `None` for slots that have never been written.
32    fn resolve(&self, version: u8) -> Option<[u8; S]>;
33}
34
35/// An in-memory ring buffer of versioned secret keys, safe for concurrent use.
36pub 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    /// Create a new `InMemorySecretGroup` pre-populated with one key at `version`.
49    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    /// Install a key at `version` without making it the `current` signing key.
76    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    /// Advance the `current_version` to `version` and notify subscribers.
104    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    /// Combined operation: store the key and immediately promote it to current.
141    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// ---------------------------------------------------------------------------
209// Tests
210// ---------------------------------------------------------------------------
211
212#[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}