vti_common/idempotency/
middleware.rs

1//! Axum middleware that implements the `Idempotency-Key` cache.
2//!
3//! Wired in via [`axum::middleware::from_fn_with_state`] using
4//! [`IdempotencyLayerState`] for the captured store + class. The
5//! middleware buffers both the request body (for hashing) and the
6//! response body (for caching), so consumers attaching it to a route
7//! MUST ensure the route's expected body fits inside
8//! [`MAX_BODY_BYTES`] (1 MB — matches the workspace's global cap,
9//! spec §14.4).
10//!
11//! # Flow
12//!
13//! 1. Request arrives; pull `Idempotency-Key` header. **Absent →
14//!    pass through.** Idempotency keys are optional per the spec
15//!    (only retries need them).
16//! 2. Derive [`Principal`] from request parts (Auth-token bytes
17//!    hashed, else IP, else anonymous).
18//! 3. Buffer the request body and hash it.
19//! 4. Look up `(principal, key)` in the store:
20//!    - Hit, request hash matches → **replay** the cached response.
21//!    - Hit, request hash differs → 422 [`AppError::IdempotencyKeyConflict`].
22//!    - Miss → forward to inner service.
23//! 5. After the inner service responds, buffer + store the response,
24//!    then return it unchanged. Storage failures are logged but **do
25//!    not** fail the request — a cache miss is always survivable.
26
27use axum::body::{Body, Bytes};
28use axum::extract::{Request, State};
29use axum::http::{HeaderMap, HeaderName, HeaderValue, Response, StatusCode};
30use axum::middleware::Next;
31use chrono::{Duration, Utc};
32use http_body_util::BodyExt;
33use sha2::{Digest, Sha256};
34use tracing::{debug, warn};
35
36use super::class::IdempotencyClass;
37use super::store::{CacheEntry, IdempotencyStore, principal_from_request};
38use crate::error::AppError;
39
40/// Canonical `Idempotency-Key` HTTP header name.
41pub const IDEMPOTENCY_HEADER: &str = "Idempotency-Key";
42
43/// Hard cap on request + response body sizes the middleware will
44/// buffer. Matches the workspace's global 1 MB body cap. Routes whose
45/// bodies legitimately exceed this (e.g. website upload) should not
46/// attach the idempotency layer.
47pub const MAX_BODY_BYTES: usize = 1024 * 1024;
48
49/// State passed to [`idempotency_middleware`] via
50/// [`axum::middleware::from_fn_with_state`]. Cheap to clone — the
51/// underlying [`IdempotencyStore`] is `Arc`-shared.
52#[derive(Clone)]
53pub struct IdempotencyLayerState {
54    pub store: IdempotencyStore,
55    pub class: IdempotencyClass,
56}
57
58/// Axum-compatible middleware function. Wire this in via
59/// [`axum::middleware::from_fn_with_state`] (the workspace's
60/// canonical state-capturing-middleware pattern):
61///
62/// ```ignore
63/// use axum::middleware::from_fn_with_state;
64/// use vti_common::idempotency::{
65///     idempotency_middleware, IdempotencyLayerState, IdempotencyClass,
66/// };
67///
68/// let state = IdempotencyLayerState {
69///     store: store.clone(),
70///     class: IdempotencyClass::Destructive,
71/// };
72/// router.route_with_task(
73///     "/v1/members/{did}",
74///     delete(remove_handler)
75///         .layer(from_fn_with_state(state, idempotency_middleware)),
76///     task,
77/// )
78/// ```
79pub async fn idempotency_middleware(
80    State(state): State<IdempotencyLayerState>,
81    request: Request,
82    next: Next,
83) -> Result<Response<Body>, AppError> {
84    run(state.store, state.class, request, next).await
85}
86
87async fn run(
88    store: IdempotencyStore,
89    class: IdempotencyClass,
90    request: Request,
91    next: Next,
92) -> Result<Response<Body>, AppError> {
93    // 1. Header present?
94    let idempotency_key = match request
95        .headers()
96        .get(IDEMPOTENCY_HEADER)
97        .and_then(|v| v.to_str().ok())
98        .map(str::to_owned)
99    {
100        Some(k) if !k.is_empty() => k,
101        // No key, or empty → pass straight through. Optional header.
102        _ => return Ok(next.run(request).await),
103    };
104
105    // Reject control characters in the key — they would corrupt the
106    // storage key (`idem:<hash>:<key>`) and could mask cache entries.
107    if idempotency_key
108        .chars()
109        .any(|c| c.is_control() || c == ':' || c == '\n' || c == '\r')
110    {
111        return Err(AppError::Validation(format!(
112            "Idempotency-Key contains a disallowed character: {:?}",
113            idempotency_key
114        )));
115    }
116
117    // 2. Derive principal.
118    let (parts, body) = request.into_parts();
119    let principal = principal_from_request(&parts);
120    let principal_hash = principal.hash();
121
122    // 3. Buffer + hash the request body.
123    let body_bytes = collect_body(body).await?;
124    let request_hash = sha256(&body_bytes);
125
126    // 4. Cache lookup.
127    if let Some(existing) = store.get(&principal_hash, &idempotency_key).await? {
128        if existing.request_hash == request_hash {
129            debug!(
130                idempotency_key = %idempotency_key,
131                "serving cached idempotent response"
132            );
133            return rebuild_response(&existing);
134        }
135        warn!(
136            idempotency_key = %idempotency_key,
137            "Idempotency-Key conflict — same key, different request body"
138        );
139        return Err(AppError::IdempotencyKeyConflict);
140    }
141
142    // 5. Forward to handler.
143    let request = Request::from_parts(parts, Body::from(body_bytes));
144    let response = next.run(request).await;
145
146    // Cache the response. Storage failures are non-fatal — they just
147    // mean a retry will hit the handler again rather than serve from
148    // cache. We log and proceed.
149    let (resp_parts, resp_body) = response.into_parts();
150    let resp_bytes = match collect_body(resp_body).await {
151        Ok(b) => b,
152        Err(e) => {
153            warn!(error = %e, "could not buffer response for idempotency cache; skipping cache");
154            // Re-emit a response with an empty body — the original is
155            // gone. The caller still sees the status + headers.
156            return Ok(Response::from_parts(resp_parts, Body::empty()));
157        }
158    };
159
160    let now = Utc::now();
161    let entry = CacheEntry {
162        idempotency_key: idempotency_key.clone(),
163        request_hash,
164        response_status: resp_parts.status.as_u16(),
165        response_headers: capture_headers(&resp_parts.headers),
166        response_body: resp_bytes.to_vec(),
167        class,
168        created_at: now,
169        expires_at: now + Duration::seconds(class.ttl_seconds() as i64),
170    };
171    if let Err(e) = store.put(&principal_hash, &entry).await {
172        warn!(error = %e, "idempotency cache write failed; response will not be replayed on retry");
173    }
174
175    let mut rebuilt = Response::new(Body::from(resp_bytes));
176    *rebuilt.status_mut() = resp_parts.status;
177    *rebuilt.headers_mut() = resp_parts.headers;
178    Ok(rebuilt)
179}
180
181async fn collect_body(body: Body) -> Result<Bytes, AppError> {
182    let bytes = http_body_util::Limited::new(body, MAX_BODY_BYTES)
183        .collect()
184        .await
185        .map_err(|e| AppError::Validation(format!("body exceeds idempotency buffer limit: {e}")))?
186        .to_bytes();
187    Ok(bytes)
188}
189
190fn sha256(bytes: &[u8]) -> [u8; 32] {
191    let mut hasher = Sha256::new();
192    hasher.update(bytes);
193    hasher.finalize().into()
194}
195
196fn capture_headers(headers: &HeaderMap) -> Vec<(String, String)> {
197    headers
198        .iter()
199        .filter_map(|(k, v)| match v.to_str() {
200            Ok(value) => Some((k.as_str().to_string(), value.to_string())),
201            Err(_) => {
202                // Non-UTF-8 header values (rare) can't round-trip through
203                // our cache. Drop them with a warning rather than fail
204                // the response.
205                warn!(header = %k, "skipping non-UTF-8 response header in idempotency cache");
206                None
207            }
208        })
209        .collect()
210}
211
212fn rebuild_response(entry: &CacheEntry) -> Result<Response<Body>, AppError> {
213    let mut response = Response::new(Body::from(entry.response_body.clone()));
214    *response.status_mut() = StatusCode::from_u16(entry.response_status)
215        .map_err(|e| AppError::Internal(format!("invalid cached status: {e}")))?;
216    for (name, value) in &entry.response_headers {
217        let name = HeaderName::from_bytes(name.as_bytes())
218            .map_err(|e| AppError::Internal(format!("invalid cached header name: {e}")))?;
219        let value = HeaderValue::from_str(value)
220            .map_err(|e| AppError::Internal(format!("invalid cached header value: {e}")))?;
221        response.headers_mut().insert(name, value);
222    }
223    Ok(response)
224}
225
226// ---------------------------------------------------------------------------
227// Tests
228// ---------------------------------------------------------------------------
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233    use crate::config::StoreConfig;
234    use crate::store::Store;
235    use axum::Router;
236    use axum::http::Request as HttpRequest;
237    use axum::routing::post;
238    use std::sync::Arc;
239    use std::sync::atomic::{AtomicU32, Ordering};
240    use tower::ServiceExt;
241
242    fn temp_store() -> (IdempotencyStore, tempfile::TempDir) {
243        let dir = tempfile::tempdir().expect("tempdir");
244        let cfg = StoreConfig {
245            data_dir: dir.path().to_path_buf(),
246        };
247        let store = Store::open(&cfg).expect("store");
248        let ks = store.keyspace("idempotency-mw-test").expect("ks");
249        (IdempotencyStore::new(ks), dir)
250    }
251
252    /// Test handler that counts invocations and echoes the request
253    /// body in a JSON response.
254    fn counted_router(
255        store: IdempotencyStore,
256        class: IdempotencyClass,
257        counter: Arc<AtomicU32>,
258    ) -> Router {
259        let state = IdempotencyLayerState { store, class };
260        Router::new().route(
261            "/",
262            post({
263                let counter = counter.clone();
264                move |body: Bytes| {
265                    let counter = counter.clone();
266                    async move {
267                        counter.fetch_add(1, Ordering::SeqCst);
268                        Response::builder()
269                            .status(StatusCode::CREATED)
270                            .header("content-type", "application/json")
271                            .body(Body::from(format!(
272                                r#"{{"echo":"{}"}}"#,
273                                String::from_utf8_lossy(&body)
274                            )))
275                            .unwrap()
276                    }
277                }
278            })
279            .layer(axum::middleware::from_fn_with_state(
280                state,
281                idempotency_middleware,
282            )),
283        )
284    }
285
286    fn req_with_key(key: Option<&str>, body: &str) -> HttpRequest<Body> {
287        let mut builder = HttpRequest::builder().method("POST").uri("/");
288        if let Some(k) = key {
289            builder = builder.header(IDEMPOTENCY_HEADER, k);
290        }
291        builder
292            .header("authorization", "Bearer test-token")
293            .body(Body::from(body.to_string()))
294            .unwrap()
295    }
296
297    async fn body_text(resp: Response<Body>) -> (StatusCode, String) {
298        let status = resp.status();
299        let bytes = resp.into_body().collect().await.unwrap().to_bytes();
300        (status, String::from_utf8(bytes.to_vec()).unwrap())
301    }
302
303    #[tokio::test]
304    async fn no_header_passes_through_and_does_not_cache() {
305        let (store, _dir) = temp_store();
306        let counter = Arc::new(AtomicU32::new(0));
307        let app = counted_router(
308            store.clone(),
309            IdempotencyClass::NonDestructive,
310            counter.clone(),
311        );
312
313        // Two requests, neither carries a key: handler called twice.
314        for body in ["one", "two"] {
315            let resp = app.clone().oneshot(req_with_key(None, body)).await.unwrap();
316            assert_eq!(resp.status(), StatusCode::CREATED);
317        }
318        assert_eq!(counter.load(Ordering::SeqCst), 2);
319    }
320
321    #[tokio::test]
322    async fn same_key_same_body_replays_cached_response() {
323        let (store, _dir) = temp_store();
324        let counter = Arc::new(AtomicU32::new(0));
325        let app = counted_router(
326            store.clone(),
327            IdempotencyClass::NonDestructive,
328            counter.clone(),
329        );
330
331        let first = app
332            .clone()
333            .oneshot(req_with_key(Some("abc-123"), "payload"))
334            .await
335            .unwrap();
336        let (status1, body1) = body_text(first).await;
337        assert_eq!(status1, StatusCode::CREATED);
338        assert!(body1.contains("payload"));
339        assert_eq!(counter.load(Ordering::SeqCst), 1);
340
341        let second = app
342            .clone()
343            .oneshot(req_with_key(Some("abc-123"), "payload"))
344            .await
345            .unwrap();
346        let (status2, body2) = body_text(second).await;
347        assert_eq!(status2, StatusCode::CREATED);
348        assert_eq!(body2, body1, "cached body should match");
349        assert_eq!(
350            counter.load(Ordering::SeqCst),
351            1,
352            "handler should not be re-invoked"
353        );
354    }
355
356    #[tokio::test]
357    async fn same_key_different_body_returns_422() {
358        let (store, _dir) = temp_store();
359        let counter = Arc::new(AtomicU32::new(0));
360        let app = counted_router(
361            store.clone(),
362            IdempotencyClass::NonDestructive,
363            counter.clone(),
364        );
365
366        let _first = app
367            .clone()
368            .oneshot(req_with_key(Some("key-x"), "body-A"))
369            .await
370            .unwrap();
371        let second = app
372            .clone()
373            .oneshot(req_with_key(Some("key-x"), "body-B"))
374            .await
375            .unwrap();
376        assert_eq!(second.status(), StatusCode::UNPROCESSABLE_ENTITY);
377
378        let (_, body) = body_text(second).await;
379        let v: serde_json::Value = serde_json::from_str(&body).unwrap();
380        assert_eq!(v["error"], "IdempotencyKeyConflict");
381        assert_eq!(
382            counter.load(Ordering::SeqCst),
383            1,
384            "second handler call must be rejected pre-handler"
385        );
386    }
387
388    #[tokio::test]
389    async fn destructive_class_caches_for_60s_only() {
390        let (store, _dir) = temp_store();
391        let counter = Arc::new(AtomicU32::new(0));
392        let app = counted_router(
393            store.clone(),
394            IdempotencyClass::Destructive,
395            counter.clone(),
396        );
397
398        let resp = app
399            .clone()
400            .oneshot(req_with_key(Some("d-1"), "del"))
401            .await
402            .unwrap();
403        assert_eq!(resp.status(), StatusCode::CREATED);
404
405        // Cache entry has the short TTL.
406        let principal =
407            super::super::store::Principal::AuthToken(b"Bearer test-token".to_vec()).hash();
408        let entry = store.get(&principal, "d-1").await.unwrap().unwrap();
409        assert_eq!(entry.class, IdempotencyClass::Destructive);
410        let ttl = (entry.expires_at - entry.created_at).num_seconds();
411        assert!((59..=61).contains(&ttl), "ttl was {ttl}s");
412    }
413
414    #[tokio::test]
415    async fn different_principals_have_separate_caches() {
416        let (store, _dir) = temp_store();
417        let counter = Arc::new(AtomicU32::new(0));
418        let app = counted_router(
419            store.clone(),
420            IdempotencyClass::NonDestructive,
421            counter.clone(),
422        );
423
424        // Alice
425        let a = HttpRequest::builder()
426            .method("POST")
427            .uri("/")
428            .header(IDEMPOTENCY_HEADER, "shared-key")
429            .header("authorization", "Bearer alice")
430            .body(Body::from("a"))
431            .unwrap();
432        let _ = app.clone().oneshot(a).await.unwrap();
433
434        // Bob using the same idempotency-key but a different bearer
435        let b = HttpRequest::builder()
436            .method("POST")
437            .uri("/")
438            .header(IDEMPOTENCY_HEADER, "shared-key")
439            .header("authorization", "Bearer bob")
440            .body(Body::from("b"))
441            .unwrap();
442        let resp = app.clone().oneshot(b).await.unwrap();
443        assert_eq!(resp.status(), StatusCode::CREATED);
444
445        // Both reached the handler — same key, different principals
446        assert_eq!(counter.load(Ordering::SeqCst), 2);
447    }
448
449    #[tokio::test]
450    async fn idempotency_key_with_control_chars_rejected() {
451        let (store, _dir) = temp_store();
452        let counter = Arc::new(AtomicU32::new(0));
453        let app = counted_router(
454            store.clone(),
455            IdempotencyClass::NonDestructive,
456            counter.clone(),
457        );
458
459        let bad = HttpRequest::builder()
460            .method("POST")
461            .uri("/")
462            .header(IDEMPOTENCY_HEADER, "evil:injected")
463            .header("authorization", "Bearer test")
464            .body(Body::from("x"))
465            .unwrap();
466        let resp = app.clone().oneshot(bad).await.unwrap();
467        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
468        assert_eq!(counter.load(Ordering::SeqCst), 0);
469    }
470}
vti_common/idempotency/middleware.rs

vti_common/idempotency/
middleware.rs
