vti_common/idempotency/
store.rs

1//! Persistent idempotency cache: `(principal, key) → CacheEntry`.
2
3use std::net::IpAddr;
4
5use axum::extract::ConnectInfo;
6use axum::http::request::Parts;
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Serialize};
9use sha2::{Digest, Sha256};
10
11use super::class::IdempotencyClass;
12use crate::error::AppError;
13use crate::store::KeyspaceHandle;
14
15/// Identifier scoping the idempotency cache. **Never plaintext on
16/// disk** — the principal bytes are hashed and the hash becomes part
17/// of the storage key. Different principals therefore inhabit
18/// disjoint namespaces.
19#[derive(Debug, Clone, PartialEq, Eq, Hash)]
20pub enum Principal {
21    /// Authenticated request — principal is the bearer credential
22    /// itself (hashed at storage time, never persisted in the clear).
23    /// Different tokens are different principals; token rotation
24    /// resets the cache namespace (conservatively).
25    AuthToken(Vec<u8>),
26    /// Unauthenticated request scoped to the source IP. Phase-0
27    /// unauth surfaces are `/v1/join-requests` and `/v1/install/*`;
28    /// the IP-scoping prevents one IP's idempotent retry returning
29    /// another IP's cached response.
30    Ip(IpAddr),
31    /// Fallback when neither Authorization nor `ConnectInfo` is
32    /// available (e.g. unit tests). Cache is effectively shared
33    /// across anonymous callers — acceptable because no Phase-0
34    /// production path lacks both signals.
35    Anonymous,
36}
37
38impl Principal {
39    /// 32-byte hash of the principal — the actual cache namespace.
40    /// Stable across calls; equal `Principal`s hash to equal bytes.
41    pub fn hash(&self) -> [u8; 32] {
42        let mut hasher = Sha256::new();
43        match self {
44            Principal::AuthToken(bytes) => {
45                hasher.update(b"auth-token:");
46                hasher.update(bytes);
47            }
48            Principal::Ip(ip) => {
49                hasher.update(b"ip:");
50                hasher.update(ip.to_string().as_bytes());
51            }
52            Principal::Anonymous => {
53                hasher.update(b"anonymous");
54            }
55        }
56        hasher.finalize().into()
57    }
58}
59
60/// Derive a [`Principal`] from request parts.
61///
62/// Prefers the Authorization header (hashed) when present, falls
63/// back to `ConnectInfo<SocketAddr>` (which Axum populates from
64/// `into_make_service_with_connect_info`), and finally to
65/// [`Principal::Anonymous`].
66///
67/// Public so a service can inspect / log the principal without
68/// re-implementing the precedence.
69pub fn principal_from_request(parts: &Parts) -> Principal {
70    if let Some(auth) = parts.headers.get(axum::http::header::AUTHORIZATION) {
71        return Principal::AuthToken(auth.as_bytes().to_vec());
72    }
73    if let Some(ConnectInfo(addr)) = parts.extensions.get::<ConnectInfo<std::net::SocketAddr>>() {
74        return Principal::Ip(addr.ip());
75    }
76    Principal::Anonymous
77}
78
79// ---------------------------------------------------------------------------
80// Cache entry
81// ---------------------------------------------------------------------------
82
83/// Persisted cache record. The response is held in full so a retry
84/// reproduces every header + body byte the original delivered.
85#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
86pub struct CacheEntry {
87    pub idempotency_key: String,
88    /// SHA-256 over the request body. Differing hashes for the same
89    /// `(principal, key)` cause [`AppError::IdempotencyKeyConflict`].
90    pub request_hash: [u8; 32],
91    pub response_status: u16,
92    pub response_headers: Vec<(String, String)>,
93    pub response_body: Vec<u8>,
94    pub class: IdempotencyClass,
95    pub created_at: DateTime<Utc>,
96    pub expires_at: DateTime<Utc>,
97}
98
99impl CacheEntry {
100    pub fn is_expired(&self, now: DateTime<Utc>) -> bool {
101        self.expires_at <= now
102    }
103}
104
105// ---------------------------------------------------------------------------
106// IdempotencyStore
107// ---------------------------------------------------------------------------
108
109/// Wraps an `idempotency` keyspace. Cheap to clone — the underlying
110/// keyspace handle is `Arc`-shared.
111#[derive(Clone)]
112pub struct IdempotencyStore {
113    ks: KeyspaceHandle,
114}
115
116impl IdempotencyStore {
117    pub fn new(ks: KeyspaceHandle) -> Self {
118        Self { ks }
119    }
120
121    /// Look up an existing entry. **Expired entries are treated as
122    /// absent** so a long-stale cached response is never served, even
123    /// if a background sweeper hasn't yet reclaimed the disk space.
124    pub async fn get(
125        &self,
126        principal_hash: &[u8; 32],
127        key: &str,
128    ) -> Result<Option<CacheEntry>, AppError> {
129        let storage_key = storage_key(principal_hash, key);
130        let entry: Option<CacheEntry> = self.ks.get(storage_key).await?;
131        let now = Utc::now();
132        Ok(entry.filter(|e| !e.is_expired(now)))
133    }
134
135    /// Insert or replace a cache entry. Caller is responsible for
136    /// setting `expires_at = created_at + class.ttl_seconds()`.
137    pub async fn put(&self, principal_hash: &[u8; 32], entry: &CacheEntry) -> Result<(), AppError> {
138        let storage_key = storage_key(principal_hash, &entry.idempotency_key);
139        self.ks.insert(storage_key, entry).await
140    }
141}
142
143fn storage_key(principal_hash: &[u8; 32], key: &str) -> Vec<u8> {
144    // Hex-encode the principal hash so the resulting fjall key stays
145    // ASCII for grepping during debugging. Newlines / NUL bytes are
146    // rejected upstream by the idempotency middleware's header
147    // validation, so the unencoded `key` part is safe to embed
148    // directly.
149    let mut out = Vec::with_capacity(64 + key.len() + 5);
150    out.extend_from_slice(b"idem:");
151    out.extend_from_slice(hex::encode(principal_hash).as_bytes());
152    out.push(b':');
153    out.extend_from_slice(key.as_bytes());
154    out
155}
156
157// ---------------------------------------------------------------------------
158// Tests
159// ---------------------------------------------------------------------------
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164    use crate::config::StoreConfig;
165    use crate::store::Store;
166    use chrono::Duration;
167
168    fn temp_store() -> (IdempotencyStore, tempfile::TempDir) {
169        let dir = tempfile::tempdir().expect("tempdir");
170        let cfg = StoreConfig {
171            data_dir: dir.path().to_path_buf(),
172        };
173        let store = Store::open(&cfg).expect("store");
174        let ks = store.keyspace("idempotency-test").expect("ks");
175        (IdempotencyStore::new(ks), dir)
176    }
177
178    fn sample_entry() -> CacheEntry {
179        let now = Utc::now();
180        CacheEntry {
181            idempotency_key: "key-1".into(),
182            request_hash: [0xAB; 32],
183            response_status: 201,
184            response_headers: vec![("content-type".into(), "application/json".into())],
185            response_body: br#"{"ok":true}"#.to_vec(),
186            class: IdempotencyClass::NonDestructive,
187            created_at: now,
188            expires_at: now
189                + Duration::seconds(IdempotencyClass::NonDestructive.ttl_seconds() as i64),
190        }
191    }
192
193    #[test]
194    fn principal_hash_is_stable_and_distinct_across_kinds() {
195        let a = Principal::AuthToken(b"Bearer abc".to_vec());
196        let a_again = Principal::AuthToken(b"Bearer abc".to_vec());
197        let b = Principal::AuthToken(b"Bearer xyz".to_vec());
198        let ip = Principal::Ip(IpAddr::V4("127.0.0.1".parse().unwrap()));
199        let anon = Principal::Anonymous;
200
201        assert_eq!(a.hash(), a_again.hash());
202        assert_ne!(a.hash(), b.hash());
203        assert_ne!(a.hash(), ip.hash());
204        assert_ne!(a.hash(), anon.hash());
205        assert_ne!(ip.hash(), anon.hash());
206    }
207
208    #[tokio::test]
209    async fn put_then_get_returns_entry() {
210        let (store, _dir) = temp_store();
211        let principal = Principal::AuthToken(b"Bearer t".to_vec()).hash();
212        let entry = sample_entry();
213
214        store.put(&principal, &entry).await.unwrap();
215        let got = store.get(&principal, &entry.idempotency_key).await.unwrap();
216        assert_eq!(got.as_ref(), Some(&entry));
217    }
218
219    #[tokio::test]
220    async fn entries_are_scoped_by_principal() {
221        let (store, _dir) = temp_store();
222        let a = Principal::AuthToken(b"alice".to_vec()).hash();
223        let b = Principal::AuthToken(b"bob".to_vec()).hash();
224        let entry = sample_entry();
225
226        store.put(&a, &entry).await.unwrap();
227        let got_a = store.get(&a, &entry.idempotency_key).await.unwrap();
228        let got_b = store.get(&b, &entry.idempotency_key).await.unwrap();
229        assert!(got_a.is_some());
230        assert!(got_b.is_none(), "principal scoping leaked");
231    }
232
233    #[tokio::test]
234    async fn expired_entries_are_filtered_at_read_time() {
235        let (store, _dir) = temp_store();
236        let principal = Principal::AuthToken(b"Bearer t".to_vec()).hash();
237        let mut entry = sample_entry();
238        entry.expires_at = Utc::now() - Duration::seconds(1);
239        store.put(&principal, &entry).await.unwrap();
240
241        let got = store.get(&principal, &entry.idempotency_key).await.unwrap();
242        assert!(got.is_none(), "stale entry served");
243    }
244
245    #[tokio::test]
246    async fn put_overwrites_existing_entry_under_same_key() {
247        let (store, _dir) = temp_store();
248        let principal = Principal::AuthToken(b"Bearer t".to_vec()).hash();
249        let first = sample_entry();
250        store.put(&principal, &first).await.unwrap();
251
252        let mut second = first.clone();
253        second.response_status = 204;
254        second.response_body = b"updated".to_vec();
255        store.put(&principal, &second).await.unwrap();
256
257        let got = store.get(&principal, &first.idempotency_key).await.unwrap();
258        assert_eq!(got.unwrap().response_status, 204);
259    }
260}
vti_common/idempotency/store.rs

vti_common/idempotency/
store.rs
