1use std::future::Future;
45use std::pin::Pin;
46use std::sync::Arc;
47use std::task::{Context, Poll};
48use std::time::Duration;
49
50use axum::body::{Body, Bytes};
51use axum::extract::{FromRequestParts, OptionalFromRequestParts};
52use axum::http::{HeaderMap, Method, Request, Response, StatusCode};
53use futures::StreamExt as _;
54use sha2::Digest as _;
55use tower::{Layer, Service};
56use uuid::Uuid;
57
58use super::config::SubmitTokenConfig;
59use crate::idempotency::{IdempotencyRecord, IdempotencyStore};
60
61const SUBMIT_TOKEN_REPLAYED: &str = "x-submit-token-replayed";
63
64const MAX_CACHEABLE_RESPONSE_BODY: usize = 10 * 1024 * 1024; const MAX_SCAN_BYTES_CAP: usize = 2 * 1024 * 1024; #[derive(Clone, Debug)]
78pub struct SubmitFormField(pub String);
79
80#[derive(Clone, Debug)]
87pub struct SubmitToken(String);
88
89impl SubmitToken {
90 #[must_use]
92 pub fn token(&self) -> &str {
93 &self.0
94 }
95
96 #[cfg(test)]
97 pub(crate) const fn new(token: String) -> Self {
98 Self(token)
99 }
100}
101
102impl std::fmt::Display for SubmitToken {
103 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104 f.write_str(&self.0)
105 }
106}
107
108impl<S> FromRequestParts<S> for SubmitToken
109where
110 S: Send + Sync,
111{
112 type Rejection = (StatusCode, &'static str);
113
114 async fn from_request_parts(
115 parts: &mut axum::http::request::Parts,
116 _state: &S,
117 ) -> Result<Self, Self::Rejection> {
118 parts.extensions.get::<Self>().cloned().ok_or((
119 StatusCode::INTERNAL_SERVER_ERROR,
120 "Submit token not found in request extensions. Is SubmitTokenLayer enabled?",
121 ))
122 }
123}
124
125impl<S> OptionalFromRequestParts<S> for SubmitToken
126where
127 S: Send + Sync,
128{
129 type Rejection = std::convert::Infallible;
130
131 async fn from_request_parts(
132 parts: &mut axum::http::request::Parts,
133 _state: &S,
134 ) -> Result<Option<Self>, Self::Rejection> {
135 Ok(parts.extensions.get::<Self>().cloned())
136 }
137}
138
139impl<S> FromRequestParts<S> for SubmitFormField
140where
141 S: Send + Sync,
142{
143 type Rejection = (StatusCode, &'static str);
144
145 async fn from_request_parts(
146 parts: &mut axum::http::request::Parts,
147 _state: &S,
148 ) -> Result<Self, Self::Rejection> {
149 parts.extensions.get::<Self>().cloned().ok_or((
150 StatusCode::INTERNAL_SERVER_ERROR,
151 "Submit form field not found in request extensions. Is SubmitTokenLayer enabled?",
152 ))
153 }
154}
155
156impl<S> OptionalFromRequestParts<S> for SubmitFormField
157where
158 S: Send + Sync,
159{
160 type Rejection = std::convert::Infallible;
161
162 async fn from_request_parts(
163 parts: &mut axum::http::request::Parts,
164 _state: &S,
165 ) -> Result<Option<Self>, Self::Rejection> {
166 Ok(parts.extensions.get::<Self>().cloned())
167 }
168}
169
170const fn is_mutating_method(method: &Method) -> bool {
171 matches!(
172 *method,
173 Method::POST | Method::PUT | Method::PATCH | Method::DELETE
174 )
175}
176
177async fn release_lock(store: &Arc<dyn IdempotencyStore>, key: &str, owner: &str) {
182 if let Err(error) = store.unlock(key, owner).await {
183 tracing::warn!(error = %error, "Submit-token unlock failed; the lock expires by its TTL");
184 }
185}
186
187fn storage_key(token: &str) -> String {
191 let mut hasher = sha2::Sha256::new();
192 hasher.update(b"autumn.submit_token:v1:");
193 hasher.update(token.as_bytes());
194 format!("submit:{}", hex::encode(hasher.finalize()))
195}
196
197fn scan_for_token(
200 bytes: &[u8],
201 is_urlencoded: bool,
202 boundary: Option<&str>,
203 field: &str,
204) -> Option<String> {
205 if is_urlencoded {
206 url::form_urlencoded::parse(bytes)
207 .find(|(key, _)| key == field)
208 .map(|(_, value)| value.into_owned())
209 } else if let Some(boundary) = boundary {
210 super::multipart_scan::scan_multipart_field(bytes, boundary, field).map(str::to_owned)
211 } else {
212 None
213 }
214}
215
216enum CollectedBody {
218 Full(Bytes),
220 Oversized { prefix: Bytes, body: Body },
226 Errored(axum::Error),
232}
233
234async fn collect_body(body: Body, limit: usize) -> CollectedBody {
236 let mut buf = Vec::<u8>::new();
237 let mut stream = body.into_data_stream();
238 loop {
239 match stream.next().await {
240 None => break,
242 Some(Err(err)) => return CollectedBody::Errored(err),
247 Some(Ok(chunk)) => {
248 let remaining = limit.saturating_sub(buf.len());
249 if chunk.len() > remaining {
250 let mut prefix_buf = buf.clone();
257 prefix_buf.extend_from_slice(&chunk[..remaining]);
258 let prefix = Bytes::from(prefix_buf);
259 let mut leading = Vec::with_capacity(2);
265 if !buf.is_empty() {
266 leading.push(Ok::<Bytes, axum::Error>(Bytes::from(buf)));
267 }
268 leading.push(Ok::<Bytes, axum::Error>(chunk));
269 let body = Body::from_stream(futures::stream::iter(leading).chain(stream));
270 return CollectedBody::Oversized { prefix, body };
271 }
272 buf.extend_from_slice(&chunk);
273 }
274 }
275 }
276 CollectedBody::Full(Bytes::from(buf))
277}
278
279async fn extract_submitted_token(
286 req: Request<Body>,
287 field: &str,
288 max_scan_bytes: usize,
289) -> Result<(Option<String>, Request<Body>), axum::Error> {
290 let (parts, body) = req.into_parts();
291
292 let content_type = parts
293 .headers
294 .get(axum::http::header::CONTENT_TYPE)
295 .and_then(|v| v.to_str().ok())
296 .unwrap_or_default();
297 let is_urlencoded =
298 crate::media_type::is_media_type(content_type, "application/x-www-form-urlencoded");
299 let boundary = multer::parse_boundary(content_type).ok();
310 if !is_urlencoded && boundary.is_none() {
313 return Ok((None, Request::from_parts(parts, body)));
314 }
315
316 match collect_body(body, max_scan_bytes).await {
317 CollectedBody::Full(bytes) => {
318 let token = scan_for_token(&bytes, is_urlencoded, boundary.as_deref(), field);
319 Ok((token, Request::from_parts(parts, Body::from(bytes))))
320 }
321 CollectedBody::Oversized { prefix, body } => {
322 let token = scan_for_token(&prefix, is_urlencoded, boundary.as_deref(), field);
323 Ok((token, Request::from_parts(parts, body)))
324 }
325 CollectedBody::Errored(err) => Err(err),
326 }
327}
328
329fn replay_headers(headers: &HeaderMap) -> Vec<(String, Vec<u8>)> {
332 const SKIP: &[&str] = &[
333 "connection",
334 "transfer-encoding",
335 "keep-alive",
336 "upgrade",
337 "proxy-authenticate",
338 "proxy-authorization",
339 "te",
340 "trailer",
341 "set-cookie",
342 SUBMIT_TOKEN_REPLAYED,
343 ];
344 headers
345 .iter()
346 .filter(|(name, _)| !SKIP.contains(&name.as_str()))
347 .map(|(name, value)| (name.to_string(), value.as_bytes().to_vec()))
348 .collect()
349}
350
351fn replay_response(record: &IdempotencyRecord) -> Response<Body> {
352 let mut builder = Response::builder().status(record.status);
353 for (name, value) in &record.headers {
354 builder = builder.header(name.as_str(), value.as_slice());
355 }
356 builder
357 .header(SUBMIT_TOKEN_REPLAYED, "true")
358 .body(Body::from(record.body.clone()))
359 .unwrap_or_else(|_| {
360 let mut resp = Response::new(Body::empty());
361 *resp.status_mut() =
362 StatusCode::from_u16(record.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
363 resp
364 })
365}
366
367fn in_flight_conflict_response() -> Response<Body> {
368 Response::builder()
369 .status(StatusCode::CONFLICT)
370 .header("retry-after", "1")
371 .body(Body::from(
372 "this form submission is already being processed; retry after 1 second",
373 ))
374 .unwrap_or_else(|_| Response::new(Body::empty()))
375}
376
377fn body_read_error_response() -> Response<Body> {
381 Response::builder()
382 .status(StatusCode::BAD_REQUEST)
383 .body(Body::from("could not read the request body"))
384 .unwrap_or_else(|_| Response::new(Body::empty()))
385}
386
387fn response_read_error_response() -> Response<Body> {
393 Response::builder()
394 .status(StatusCode::INTERNAL_SERVER_ERROR)
395 .body(Body::from("could not read the response body"))
396 .unwrap_or_else(|_| Response::new(Body::empty()))
397}
398
399struct SubmitTokenSettings {
402 store: Arc<dyn IdempotencyStore>,
403 field_name: String,
404 ttl: Duration,
405 in_flight_ttl: Duration,
406 exempt_paths: Vec<String>,
407 max_scan_bytes: usize,
408}
409
410impl Clone for SubmitTokenSettings {
412 fn clone(&self) -> Self {
413 Self {
414 store: Arc::clone(&self.store),
415 field_name: self.field_name.clone(),
416 ttl: self.ttl,
417 in_flight_ttl: self.in_flight_ttl,
418 exempt_paths: self.exempt_paths.clone(),
419 max_scan_bytes: self.max_scan_bytes,
420 }
421 }
422}
423
424#[derive(Clone)]
432pub struct SubmitTokenLayer {
433 settings: Arc<SubmitTokenSettings>,
434}
435
436impl std::fmt::Debug for SubmitTokenLayer {
437 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
438 f.debug_struct("SubmitTokenLayer")
439 .field("field_name", &self.settings.field_name)
440 .field("ttl", &self.settings.ttl)
441 .finish_non_exhaustive()
442 }
443}
444
445impl SubmitTokenLayer {
446 #[must_use]
449 pub fn new(store: Arc<dyn IdempotencyStore>, config: &SubmitTokenConfig) -> Self {
450 let ttl = Duration::from_secs(config.ttl_secs);
451 let in_flight_ttl = Duration::from_secs(config.in_flight_ttl_secs);
457 Self {
458 settings: Arc::new(SubmitTokenSettings {
459 store,
460 field_name: config.field_name.clone(),
461 ttl,
462 in_flight_ttl,
463 exempt_paths: config.exempt_paths.clone(),
464 max_scan_bytes: MAX_SCAN_BYTES_CAP,
465 }),
466 }
467 }
468
469 #[must_use]
472 pub fn with_max_scan_bytes(mut self, n: usize) -> Self {
473 Arc::make_mut(&mut self.settings).max_scan_bytes = n.min(MAX_SCAN_BYTES_CAP);
474 self
475 }
476
477 #[must_use]
479 pub fn with_exempt_path(mut self, path: impl Into<String>) -> Self {
480 Arc::make_mut(&mut self.settings)
481 .exempt_paths
482 .push(path.into());
483 self
484 }
485}
486
487impl<S> Layer<S> for SubmitTokenLayer {
488 type Service = SubmitTokenService<S>;
489
490 fn layer(&self, inner: S) -> Self::Service {
491 SubmitTokenService {
492 inner,
493 settings: Arc::clone(&self.settings),
494 }
495 }
496}
497
498#[derive(Clone)]
500pub struct SubmitTokenService<S> {
501 inner: S,
502 settings: Arc<SubmitTokenSettings>,
503}
504
505impl<S> Service<Request<Body>> for SubmitTokenService<S>
506where
507 S: Service<Request<Body>, Response = Response<Body>, Error = std::convert::Infallible>
508 + Clone
509 + Send
510 + 'static,
511 S::Future: Send + 'static,
512{
513 type Response = Response<Body>;
514 type Error = std::convert::Infallible;
515 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
516
517 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
518 self.inner.poll_ready(cx)
519 }
520
521 fn call(&mut self, mut req: Request<Body>) -> Self::Future {
522 let minted = Uuid::new_v4().to_string();
525 req.extensions_mut().insert(SubmitToken(minted));
526 req.extensions_mut()
527 .insert(SubmitFormField(self.settings.field_name.clone()));
528
529 let is_guarded = is_mutating_method(req.method())
534 && (self.settings.exempt_paths.is_empty() || {
535 let clean = crate::security::path::clean_path(req.uri().path());
536 !crate::security::path::is_exempt_path(clean.as_str(), &self.settings.exempt_paths)
537 });
538
539 if !is_guarded {
547 return Box::pin(self.inner.call(req));
548 }
549
550 let settings = Arc::clone(&self.settings);
551 let clone = self.inner.clone();
552 let mut inner = std::mem::replace(&mut self.inner, clone);
553
554 Box::pin(async move {
555 let (submitted, req) =
556 match extract_submitted_token(req, &settings.field_name, settings.max_scan_bytes)
557 .await
558 {
559 Ok(pair) => pair,
560 Err(error) => {
561 tracing::warn!(
566 error = %error,
567 "Submit-token body scan failed on a read error; rejecting the request"
568 );
569 return Ok(body_read_error_response());
570 }
571 };
572
573 let Some(token) = submitted.filter(|t| !t.is_empty()) else {
574 return inner.call(req).await;
576 };
577
578 let key = storage_key(&token);
579
580 match settings.store.get(&key).await {
587 Ok(Some(entry)) => return Ok(replay_response(&entry.record)),
588 Ok(None) => {}
589 Err(error) => {
590 tracing::error!(
591 error = %error,
592 "Submit-token consumed-token lookup failed; failing closed"
593 );
594 return Ok(crate::idempotency::persistence_failed_response());
595 }
596 }
597
598 let owner = Uuid::new_v4().to_string();
601 let lock_taken_at = crate::time::ambient_instant();
602 let acquired = settings
603 .store
604 .try_lock(&key, &owner, settings.in_flight_ttl)
605 .await
606 .unwrap_or_else(|error| {
607 tracing::warn!(
609 error = %error,
610 "Submit-token lock unavailable; failing closed"
611 );
612 false
613 });
614 if !acquired {
615 return Ok(match settings.store.get(&key).await {
619 Ok(Some(entry)) => replay_response(&entry.record),
620 Ok(None) | Err(_) => in_flight_conflict_response(),
621 });
622 }
623
624 match settings.store.get(&key).await {
632 Ok(Some(entry)) => {
633 release_lock(&settings.store, &key, &owner).await;
634 return Ok(replay_response(&entry.record));
635 }
636 Ok(None) => {}
637 Err(error) => {
638 tracing::error!(
639 error = %error,
640 "Submit-token consumed-token lookup failed after lock acquisition; failing closed"
641 );
642 return Ok(crate::idempotency::persistence_failed_response());
643 }
644 }
645
646 if let Some(response) =
647 lapsed_lock_response(&settings, &key, &owner, lock_taken_at).await
648 {
649 return Ok(response);
650 }
651
652 let response = inner.call(req).await?;
653 Ok(cache_consumed_token_response(response, &settings, &key, &owner).await)
654 })
655 }
656}
657
658async fn lapsed_lock_response(
662 settings: &SubmitTokenSettings,
663 key: &str,
664 owner: &str,
665 taken_at: std::time::Instant,
666) -> Option<Response<Body>> {
667 use crate::idempotency::LockForHandler;
668 match crate::idempotency::keep_lock_for_handler(
669 settings.store.as_ref(),
670 key,
671 owner,
672 settings.in_flight_ttl,
673 taken_at,
674 )
675 .await
676 {
677 Ok(LockForHandler::Held | LockForHandler::Renewed) => None,
678 Ok(LockForHandler::Lost) => {
679 tracing::warn!("Submit-token lock lapsed before the handler ran; returning 409");
680 release_lock(&settings.store, key, owner).await;
681 Some(in_flight_conflict_response())
682 }
683 Err(error) => {
684 tracing::warn!(error = %error, "Submit-token lock renewal failed; failing closed");
685 Some(in_flight_conflict_response())
686 }
687 }
688}
689
690async fn cache_consumed_token_response(
695 response: Response<Body>,
696 settings: &SubmitTokenSettings,
697 key: &str,
698 owner: &str,
699) -> Response<Body> {
700 let (parts, body) = response.into_parts();
701 match collect_body(body, MAX_CACHEABLE_RESPONSE_BODY).await {
702 CollectedBody::Full(bytes) => {
703 let status = parts.status.as_u16();
704 if (200..400).contains(&status) {
707 let record = IdempotencyRecord {
708 status,
709 headers: replay_headers(&parts.headers),
710 body: bytes.to_vec(),
711 metadata: Vec::new(),
712 };
713 if let Err(error) = settings
722 .store
723 .set(key, owner, record, Vec::new(), settings.ttl)
724 .await
725 .and_then(crate::idempotency::stored)
726 {
727 tracing::error!(
728 error = %error,
729 "Submit-token persistence failed after handler success; failing closed"
730 );
731 return crate::idempotency::persistence_failed_response();
732 }
733 }
734 release_lock(&settings.store, key, owner).await;
735 Response::from_parts(parts, Body::from(bytes))
736 }
737 CollectedBody::Oversized { body, .. } => {
738 release_lock(&settings.store, key, owner).await;
742 Response::from_parts(parts, body)
743 }
744 CollectedBody::Errored(error) => {
745 let status = parts.status.as_u16();
746 tracing::error!(
762 error = %error,
763 "Submit-token response buffering failed on a read error; failing closed"
764 );
765 if !(200..400).contains(&status) {
766 release_lock(&settings.store, key, owner).await;
767 }
768 response_read_error_response()
769 }
770 }
771}
772
773#[cfg(test)]
774mod tests {
775 use super::*;
776 use crate::idempotency::{
777 IdempotencyEntry, IdempotencyFuture, IdempotencyStoreError, MemoryIdempotencyStore,
778 };
779 use axum::Router;
780 use axum::routing::{get, post};
781 use std::sync::atomic::{AtomicUsize, Ordering};
782 use tower::ServiceExt;
783
784 fn default_config() -> SubmitTokenConfig {
785 SubmitTokenConfig {
786 enabled: true,
787 ..Default::default()
788 }
789 }
790
791 fn layer_with_store(store: Arc<dyn IdempotencyStore>) -> SubmitTokenLayer {
792 SubmitTokenLayer::new(store, &default_config())
793 }
794
795 fn urlencoded_post(token: &str) -> Request<Body> {
796 Request::builder()
797 .method("POST")
798 .uri("/submit")
799 .header("Content-Type", "application/x-www-form-urlencoded")
800 .body(Body::from(format!("_submit_token={token}&title=hello")))
801 .unwrap()
802 }
803
804 fn multipart_post(content_type: &str, boundary: &str, token: &str) -> Request<Body> {
809 let body = format!(
810 "--{boundary}\r\n\
811 Content-Disposition: form-data; name=\"_submit_token\"\r\n\
812 \r\n\
813 {token}\r\n\
814 --{boundary}--\r\n"
815 );
816 Request::builder()
817 .method("POST")
818 .uri("/submit")
819 .header("Content-Type", content_type)
820 .body(Body::from(body))
821 .unwrap()
822 }
823
824 #[test]
825 fn submit_token_extractor_exposes_value() {
826 let token = SubmitToken::new("abc-123".to_owned());
827 assert_eq!(token.token(), "abc-123");
828 assert_eq!(token.to_string(), "abc-123");
829 }
830
831 #[tokio::test]
832 async fn mints_token_available_to_extractor() {
833 async fn handler(submit_token: SubmitToken) -> String {
834 submit_token.token().to_owned()
835 }
836 let store: Arc<dyn IdempotencyStore> =
837 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
838 let app = Router::new()
839 .route("/", get(handler))
840 .layer(layer_with_store(store));
841
842 let response = app
843 .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
844 .await
845 .unwrap();
846 assert_eq!(response.status(), StatusCode::OK);
847 let body = axum::body::to_bytes(response.into_body(), usize::MAX)
848 .await
849 .unwrap();
850 let token = String::from_utf8(body.to_vec()).unwrap();
851 assert!(
852 Uuid::parse_str(&token).is_ok(),
853 "minted token should be a uuid: {token}"
854 );
855 }
856
857 #[tokio::test]
858 async fn first_use_runs_handler_replay_short_circuits() {
859 let store: Arc<dyn IdempotencyStore> =
860 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
861 let count = Arc::new(AtomicUsize::new(0));
862 let count_inner = count.clone();
863 let app = Router::new()
864 .route(
865 "/submit",
866 post(move || {
867 let count = count_inner.clone();
868 async move {
869 count.fetch_add(1, Ordering::SeqCst);
870 "created"
871 }
872 }),
873 )
874 .layer(layer_with_store(store));
875
876 let token = "tok-replay";
877 let first = app.clone().oneshot(urlencoded_post(token)).await.unwrap();
879 assert_eq!(first.status(), StatusCode::OK);
880 assert!(first.headers().get(SUBMIT_TOKEN_REPLAYED).is_none());
881 let first_body = axum::body::to_bytes(first.into_body(), usize::MAX)
882 .await
883 .unwrap();
884 assert_eq!(&first_body[..], b"created");
885
886 let second = app.clone().oneshot(urlencoded_post(token)).await.unwrap();
888 assert_eq!(second.status(), StatusCode::OK);
889 assert_eq!(
890 second
891 .headers()
892 .get(SUBMIT_TOKEN_REPLAYED)
893 .map(|v| v.to_str().unwrap()),
894 Some("true")
895 );
896 let second_body = axum::body::to_bytes(second.into_body(), usize::MAX)
897 .await
898 .unwrap();
899 assert_eq!(&second_body[..], b"created");
900
901 assert_eq!(
902 count.load(Ordering::SeqCst),
903 1,
904 "handler must run exactly once"
905 );
906 }
907
908 #[tokio::test]
909 async fn lowercase_multipart_consumes_token_and_replay_short_circuits() {
910 let store: Arc<dyn IdempotencyStore> =
914 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
915 let count = Arc::new(AtomicUsize::new(0));
916 let count_inner = count.clone();
917 let app = Router::new()
918 .route(
919 "/submit",
920 post(move || {
921 let count = count_inner.clone();
922 async move {
923 count.fetch_add(1, Ordering::SeqCst);
924 "created"
925 }
926 }),
927 )
928 .layer(layer_with_store(store));
929
930 let token = "tok-mp-lower";
931 let ct = "multipart/form-data; boundary=simpleboundary123";
932 let boundary = "simpleboundary123";
933
934 let first = app
935 .clone()
936 .oneshot(multipart_post(ct, boundary, token))
937 .await
938 .unwrap();
939 assert_eq!(first.status(), StatusCode::OK);
940 assert!(first.headers().get(SUBMIT_TOKEN_REPLAYED).is_none());
941
942 let second = app
943 .clone()
944 .oneshot(multipart_post(ct, boundary, token))
945 .await
946 .unwrap();
947 assert_eq!(second.status(), StatusCode::OK);
948 assert_eq!(
949 second
950 .headers()
951 .get(SUBMIT_TOKEN_REPLAYED)
952 .map(|v| v.to_str().unwrap()),
953 Some("true")
954 );
955
956 assert_eq!(
957 count.load(Ordering::SeqCst),
958 1,
959 "handler must run exactly once"
960 );
961 }
962
963 #[tokio::test]
964 async fn mixed_case_multipart_consumes_token_and_replay_short_circuits() {
965 let store: Arc<dyn IdempotencyStore> =
971 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
972 let count = Arc::new(AtomicUsize::new(0));
973 let count_inner = count.clone();
974 let app = Router::new()
975 .route(
976 "/submit",
977 post(move || {
978 let count = count_inner.clone();
979 async move {
980 count.fetch_add(1, Ordering::SeqCst);
981 "created"
982 }
983 }),
984 )
985 .layer(layer_with_store(store));
986
987 let token = "tok-mp-mixed";
988 let ct = "Multipart/Form-Data; Boundary=BoUnDaRy-XyZ-123";
989 let boundary = "BoUnDaRy-XyZ-123";
990
991 let first = app
992 .clone()
993 .oneshot(multipart_post(ct, boundary, token))
994 .await
995 .unwrap();
996 assert_eq!(first.status(), StatusCode::OK);
997 assert!(first.headers().get(SUBMIT_TOKEN_REPLAYED).is_none());
998
999 let second = app
1000 .clone()
1001 .oneshot(multipart_post(ct, boundary, token))
1002 .await
1003 .unwrap();
1004 assert_eq!(second.status(), StatusCode::OK);
1005 assert_eq!(
1006 second
1007 .headers()
1008 .get(SUBMIT_TOKEN_REPLAYED)
1009 .map(|v| v.to_str().unwrap()),
1010 Some("true")
1011 );
1012
1013 assert_eq!(
1014 count.load(Ordering::SeqCst),
1015 1,
1016 "handler must run exactly once"
1017 );
1018 }
1019
1020 #[tokio::test]
1021 async fn quoted_semicolon_boundary_consumes_token_and_replay_short_circuits() {
1022 let store: Arc<dyn IdempotencyStore> =
1032 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
1033 let count = Arc::new(AtomicUsize::new(0));
1034 let count_inner = count.clone();
1035 let app = Router::new()
1036 .route(
1037 "/submit",
1038 post(move || {
1039 let count = count_inner.clone();
1040 async move {
1041 count.fetch_add(1, Ordering::SeqCst);
1042 "created"
1043 }
1044 }),
1045 )
1046 .layer(layer_with_store(store));
1047
1048 let token = "tok-mp-quoted-semi";
1049 let ct = "multipart/form-data; boundary=\"x;y\"";
1050 let boundary = "x;y";
1051
1052 let first = app
1053 .clone()
1054 .oneshot(multipart_post(ct, boundary, token))
1055 .await
1056 .unwrap();
1057 assert_eq!(first.status(), StatusCode::OK);
1058 assert!(first.headers().get(SUBMIT_TOKEN_REPLAYED).is_none());
1059
1060 let second = app
1061 .clone()
1062 .oneshot(multipart_post(ct, boundary, token))
1063 .await
1064 .unwrap();
1065 assert_eq!(second.status(), StatusCode::OK);
1066 assert_eq!(
1067 second
1068 .headers()
1069 .get(SUBMIT_TOKEN_REPLAYED)
1070 .map(|v| v.to_str().unwrap()),
1071 Some("true")
1072 );
1073
1074 assert_eq!(
1075 count.load(Ordering::SeqCst),
1076 1,
1077 "handler must run exactly once"
1078 );
1079 }
1080
1081 #[tokio::test]
1082 async fn uppercase_urlencoded_consumes_token_and_replay_short_circuits() {
1083 let store: Arc<dyn IdempotencyStore> =
1086 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
1087 let count = Arc::new(AtomicUsize::new(0));
1088 let count_inner = count.clone();
1089 let app = Router::new()
1090 .route(
1091 "/submit",
1092 post(move || {
1093 let count = count_inner.clone();
1094 async move {
1095 count.fetch_add(1, Ordering::SeqCst);
1096 "created"
1097 }
1098 }),
1099 )
1100 .layer(layer_with_store(store));
1101
1102 let token = "tok-ue-upper";
1103 let make_req = || {
1104 Request::builder()
1105 .method("POST")
1106 .uri("/submit")
1107 .header("Content-Type", "APPLICATION/X-WWW-FORM-URLENCODED")
1108 .body(Body::from(format!("_submit_token={token}&title=hello")))
1109 .unwrap()
1110 };
1111
1112 let first = app.clone().oneshot(make_req()).await.unwrap();
1113 assert_eq!(first.status(), StatusCode::OK);
1114 assert!(first.headers().get(SUBMIT_TOKEN_REPLAYED).is_none());
1115
1116 let second = app.clone().oneshot(make_req()).await.unwrap();
1117 assert_eq!(second.status(), StatusCode::OK);
1118 assert_eq!(
1119 second
1120 .headers()
1121 .get(SUBMIT_TOKEN_REPLAYED)
1122 .map(|v| v.to_str().unwrap()),
1123 Some("true")
1124 );
1125
1126 assert_eq!(
1127 count.load(Ordering::SeqCst),
1128 1,
1129 "handler must run exactly once"
1130 );
1131 }
1132
1133 #[tokio::test]
1134 async fn missing_token_passes_through() {
1135 let store: Arc<dyn IdempotencyStore> =
1136 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
1137 let count = Arc::new(AtomicUsize::new(0));
1138 let count_inner = count.clone();
1139 let app = Router::new()
1140 .route(
1141 "/submit",
1142 post(move || {
1143 let count = count_inner.clone();
1144 async move {
1145 count.fetch_add(1, Ordering::SeqCst);
1146 "ok"
1147 }
1148 }),
1149 )
1150 .layer(layer_with_store(store));
1151
1152 for _ in 0..2 {
1154 let req = Request::builder()
1155 .method("POST")
1156 .uri("/submit")
1157 .header("Content-Type", "application/x-www-form-urlencoded")
1158 .body(Body::from("title=hello"))
1159 .unwrap();
1160 let resp = app.clone().oneshot(req).await.unwrap();
1161 assert_eq!(resp.status(), StatusCode::OK);
1162 }
1163 assert_eq!(count.load(Ordering::SeqCst), 2);
1164 }
1165
1166 #[tokio::test]
1167 async fn exempt_path_is_exact_or_subtree_only() {
1168 let store: Arc<dyn IdempotencyStore> =
1173 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
1174 let count = Arc::new(AtomicUsize::new(0));
1175 let count_inner = count.clone();
1176 let layer =
1177 SubmitTokenLayer::new(store, &default_config()).with_exempt_path("/webhooks/stripe");
1178 let app = Router::new()
1179 .route(
1180 "/webhooks/stripe",
1181 post({
1182 let count = count_inner.clone();
1183 move || {
1184 let count = count.clone();
1185 async move {
1186 count.fetch_add(1, Ordering::SeqCst);
1187 "ok"
1188 }
1189 }
1190 }),
1191 )
1192 .route(
1193 "/webhooks/stripe/events",
1194 post({
1195 let count = count_inner.clone();
1196 move || {
1197 let count = count.clone();
1198 async move {
1199 count.fetch_add(1, Ordering::SeqCst);
1200 "ok"
1201 }
1202 }
1203 }),
1204 )
1205 .route(
1206 "/webhooks/stripe-admin",
1207 post(move || {
1208 let count = count_inner.clone();
1209 async move {
1210 count.fetch_add(1, Ordering::SeqCst);
1211 "ok"
1212 }
1213 }),
1214 )
1215 .layer(layer);
1216
1217 let submit = |uri: &'static str, token: &'static str| {
1218 let app = app.clone();
1219 async move {
1220 app.oneshot(
1221 Request::builder()
1222 .method("POST")
1223 .uri(uri)
1224 .header("Content-Type", "application/x-www-form-urlencoded")
1225 .body(Body::from(format!("_submit_token={token}")))
1226 .unwrap(),
1227 )
1228 .await
1229 .unwrap()
1230 }
1231 };
1232
1233 count.store(0, Ordering::SeqCst);
1236 submit("/webhooks/stripe", "tok-exact").await;
1237 submit("/webhooks/stripe", "tok-exact").await;
1238 assert_eq!(count.load(Ordering::SeqCst), 2);
1239
1240 count.store(0, Ordering::SeqCst);
1242 submit("/webhooks/stripe/events", "tok-subtree").await;
1243 submit("/webhooks/stripe/events", "tok-subtree").await;
1244 assert_eq!(count.load(Ordering::SeqCst), 2);
1245
1246 count.store(0, Ordering::SeqCst);
1250 submit("/webhooks/stripe-admin", "tok-adjacent").await;
1251 submit("/webhooks/stripe-admin", "tok-adjacent").await;
1252 assert_eq!(count.load(Ordering::SeqCst), 1);
1253 }
1254
1255 #[tokio::test]
1259 async fn dot_segments_cannot_reach_an_exempt_prefix() {
1260 let store: Arc<dyn IdempotencyStore> =
1261 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
1262 let count = Arc::new(AtomicUsize::new(0));
1263 let count_inner = count.clone();
1264 let layer = SubmitTokenLayer::new(store, &default_config()).with_exempt_path("/api/");
1265 let app = Router::new()
1266 .fallback(move || {
1267 let count = count_inner.clone();
1268 async move {
1269 count.fetch_add(1, Ordering::SeqCst);
1270 "ok"
1271 }
1272 })
1273 .layer(layer);
1274
1275 let submit = |uri: &'static str, token: &'static str| {
1276 let app = app.clone();
1277 async move {
1278 app.oneshot(
1279 Request::builder()
1280 .method("POST")
1281 .uri(uri)
1282 .header("Content-Type", "application/x-www-form-urlencoded")
1283 .body(Body::from(format!("_submit_token={token}")))
1284 .unwrap(),
1285 )
1286 .await
1287 .unwrap()
1288 }
1289 };
1290
1291 submit("/api/items", "tok-exempt").await;
1293 submit("/api/items", "tok-exempt").await;
1294 assert_eq!(count.load(Ordering::SeqCst), 2);
1295
1296 count.store(0, Ordering::SeqCst);
1299 submit("/api/%2e%2e/submit", "tok-dot").await;
1300 submit("/api/%2e%2e/submit", "tok-dot").await;
1301 assert_eq!(count.load(Ordering::SeqCst), 1);
1302 }
1303
1304 #[tokio::test]
1305 async fn expired_token_re_runs_after_ttl() {
1306 let store = MemoryIdempotencyStore::new(Duration::from_millis(10));
1309 let key = storage_key("ttl-token");
1310 store
1311 .set(
1312 &key,
1313 "",
1314 IdempotencyRecord {
1315 status: 200,
1316 headers: Vec::new(),
1317 body: b"first".to_vec(),
1318 metadata: Vec::new(),
1319 },
1320 Vec::new(),
1321 Duration::from_millis(10),
1322 )
1323 .await
1324 .unwrap();
1325 assert!(store.get(&key).await.unwrap().is_some());
1326 std::thread::sleep(Duration::from_millis(30));
1327 assert!(
1328 store.get(&key).await.unwrap().is_none(),
1329 "record must expire after its TTL"
1330 );
1331 }
1332
1333 #[tokio::test]
1334 async fn distinct_from_csrf_replayed_submit_token_short_circuits() {
1335 let store: Arc<dyn IdempotencyStore> =
1339 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
1340 let count = Arc::new(AtomicUsize::new(0));
1341 let count_inner = count.clone();
1342 let app = Router::new()
1343 .route(
1344 "/submit",
1345 post(move || {
1346 let count = count_inner.clone();
1347 async move {
1348 count.fetch_add(1, Ordering::SeqCst);
1349 "created"
1350 }
1351 }),
1352 )
1353 .layer(layer_with_store(store));
1354
1355 let token = "tok-csrf-distinct";
1356 let build = || {
1358 Request::builder()
1359 .method("POST")
1360 .uri("/submit")
1361 .header("Content-Type", "application/x-www-form-urlencoded")
1362 .body(Body::from(format!(
1363 "_csrf=valid-csrf&_submit_token={token}"
1364 )))
1365 .unwrap()
1366 };
1367 let first = app.clone().oneshot(build()).await.unwrap();
1368 assert_eq!(first.status(), StatusCode::OK);
1369 let second = app.clone().oneshot(build()).await.unwrap();
1372 assert_eq!(
1373 second
1374 .headers()
1375 .get(SUBMIT_TOKEN_REPLAYED)
1376 .map(|v| v.to_str().unwrap()),
1377 Some("true")
1378 );
1379 assert_eq!(count.load(Ordering::SeqCst), 1);
1380 }
1381
1382 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1386 async fn ten_concurrent_posts_persist_exactly_one_row() {
1387 let store: Arc<dyn IdempotencyStore> =
1388 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
1389 let rows = Arc::new(AtomicUsize::new(0));
1391 let rows_inner = rows.clone();
1392 let app = Router::new()
1393 .route(
1394 "/posts",
1395 post(move || {
1396 let rows = rows_inner.clone();
1397 async move {
1398 tokio::time::sleep(Duration::from_millis(20)).await;
1400 rows.fetch_add(1, Ordering::SeqCst);
1401 (StatusCode::SEE_OTHER, [("location", "/posts")], "")
1402 }
1403 }),
1404 )
1405 .layer(layer_with_store(store));
1406
1407 let token = "concurrent-token";
1408 let mut handles = Vec::new();
1409 for _ in 0..10 {
1410 let app = app.clone();
1411 handles.push(tokio::spawn(async move {
1412 let req = Request::builder()
1413 .method("POST")
1414 .uri("/posts")
1415 .header("Content-Type", "application/x-www-form-urlencoded")
1416 .body(Body::from(format!("_submit_token={token}&title=hello")))
1417 .unwrap();
1418 app.oneshot(req).await.unwrap().status()
1419 }));
1420 }
1421
1422 let mut succeeded = 0;
1423 let mut conflicts = 0;
1424 for h in handles {
1425 let status = h.await.unwrap();
1426 if status == StatusCode::SEE_OTHER {
1427 succeeded += 1;
1428 } else if status == StatusCode::CONFLICT {
1429 conflicts += 1;
1430 } else {
1431 panic!("unexpected status: {status}");
1432 }
1433 }
1434
1435 assert_eq!(
1436 rows.load(Ordering::SeqCst),
1437 1,
1438 "exactly one row must be persisted from 10 concurrent identical POSTs"
1439 );
1440 assert_eq!(
1441 succeeded + conflicts,
1442 10,
1443 "every request must either return the first response or be rejected in-flight"
1444 );
1445 assert!(
1446 succeeded >= 1,
1447 "at least the first submission must succeed and be replayable"
1448 );
1449 }
1450
1451 #[tokio::test]
1457 async fn over_limit_first_chunk_detects_leading_token_and_preserves_body() {
1458 let store: Arc<dyn IdempotencyStore> =
1459 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
1460 let count = Arc::new(AtomicUsize::new(0));
1461 let count_inner = count.clone();
1462 let seen_len = Arc::new(AtomicUsize::new(0));
1463 let seen_len_inner = seen_len.clone();
1464 let app = Router::new()
1467 .route(
1468 "/submit",
1469 post(move |body: Bytes| {
1470 let count = count_inner.clone();
1471 let seen_len = seen_len_inner.clone();
1472 async move {
1473 count.fetch_add(1, Ordering::SeqCst);
1474 seen_len.store(body.len(), Ordering::SeqCst);
1475 "created"
1476 }
1477 }),
1478 )
1479 .layer(layer_with_store(store).with_max_scan_bytes(64));
1481
1482 let token = "over-limit-token";
1483 let filler = "x".repeat(4096);
1486 let full_body = format!("_submit_token={token}&title={filler}");
1487 let full_len = full_body.len();
1488 assert!(full_len > 64, "body must exceed the scan cap for this test");
1489 let make = || {
1490 Request::builder()
1491 .method("POST")
1492 .uri("/submit")
1493 .header("Content-Type", "application/x-www-form-urlencoded")
1494 .body(Body::from(full_body.clone()))
1495 .unwrap()
1496 };
1497
1498 let first = app.clone().oneshot(make()).await.unwrap();
1501 assert_eq!(first.status(), StatusCode::OK);
1502 assert!(first.headers().get(SUBMIT_TOKEN_REPLAYED).is_none());
1503 assert_eq!(
1504 seen_len.load(Ordering::SeqCst),
1505 full_len,
1506 "handler must receive the complete body, not just the scanned prefix"
1507 );
1508
1509 let second = app.clone().oneshot(make()).await.unwrap();
1512 assert_eq!(
1513 second
1514 .headers()
1515 .get(SUBMIT_TOKEN_REPLAYED)
1516 .map(|v| v.to_str().unwrap()),
1517 Some("true"),
1518 "an over-limit first chunk with a leading token must still be guarded"
1519 );
1520 assert_eq!(
1521 count.load(Ordering::SeqCst),
1522 1,
1523 "handler must run exactly once"
1524 );
1525 }
1526
1527 struct FailingSetStore {
1531 inner: MemoryIdempotencyStore,
1532 }
1533
1534 impl FailingSetStore {
1535 fn new() -> Self {
1536 Self {
1537 inner: MemoryIdempotencyStore::new(Duration::from_secs(600)),
1538 }
1539 }
1540 }
1541
1542 impl IdempotencyStore for FailingSetStore {
1543 fn get<'a>(&'a self, key: &'a str) -> IdempotencyFuture<'a, Option<IdempotencyEntry>> {
1544 self.inner.get(key)
1545 }
1546
1547 fn set<'a>(
1548 &'a self,
1549 _key: &'a str,
1550 _owner: &'a str,
1551 _record: IdempotencyRecord,
1552 _body_hash: Vec<u8>,
1553 _ttl: Duration,
1554 ) -> IdempotencyFuture<'a, bool> {
1555 Box::pin(async {
1556 Err(IdempotencyStoreError::backend(
1557 "simulated consumed-token persistence failure",
1558 ))
1559 })
1560 }
1561
1562 fn try_lock<'a>(
1563 &'a self,
1564 key: &'a str,
1565 owner: &'a str,
1566 lock_ttl: Duration,
1567 ) -> IdempotencyFuture<'a, bool> {
1568 self.inner.try_lock(key, owner, lock_ttl)
1569 }
1570
1571 fn unlock<'a>(&'a self, key: &'a str, owner: &'a str) -> IdempotencyFuture<'a, ()> {
1572 self.inner.unlock(key, owner)
1573 }
1574 }
1575
1576 #[tokio::test]
1583 async fn persistence_failure_fails_closed_and_holds_lock() {
1584 let store: Arc<dyn IdempotencyStore> = Arc::new(FailingSetStore::new());
1585 let count = Arc::new(AtomicUsize::new(0));
1586 let count_inner = count.clone();
1587 let app = Router::new()
1588 .route(
1589 "/submit",
1590 post(move || {
1591 let count = count_inner.clone();
1592 async move {
1593 count.fetch_add(1, Ordering::SeqCst);
1594 "created"
1595 }
1596 }),
1597 )
1598 .layer(layer_with_store(store));
1599
1600 let token = "tok-persist-fail";
1601
1602 let first = app.clone().oneshot(urlencoded_post(token)).await.unwrap();
1605 assert_eq!(
1606 first.status(),
1607 StatusCode::SERVICE_UNAVAILABLE,
1608 "a persistence failure after handler success must fail closed, not return the success"
1609 );
1610
1611 let second = app.clone().oneshot(urlencoded_post(token)).await.unwrap();
1615 assert_eq!(
1616 second.status(),
1617 StatusCode::CONFLICT,
1618 "a retry after a persistence failure must be rejected in-flight, not re-run"
1619 );
1620
1621 assert_eq!(
1622 count.load(Ordering::SeqCst),
1623 1,
1624 "the handler must run at most once even when persistence fails"
1625 );
1626 }
1627
1628 struct FencedSetStore {
1631 inner: MemoryIdempotencyStore,
1632 }
1633
1634 impl IdempotencyStore for FencedSetStore {
1635 fn get<'a>(&'a self, key: &'a str) -> IdempotencyFuture<'a, Option<IdempotencyEntry>> {
1636 self.inner.get(key)
1637 }
1638
1639 fn set<'a>(
1640 &'a self,
1641 _key: &'a str,
1642 _owner: &'a str,
1643 _record: IdempotencyRecord,
1644 _body_hash: Vec<u8>,
1645 _ttl: Duration,
1646 ) -> IdempotencyFuture<'a, bool> {
1647 Box::pin(std::future::ready(Ok(false)))
1648 }
1649
1650 fn try_lock<'a>(
1651 &'a self,
1652 key: &'a str,
1653 owner: &'a str,
1654 lock_ttl: Duration,
1655 ) -> IdempotencyFuture<'a, bool> {
1656 self.inner.try_lock(key, owner, lock_ttl)
1657 }
1658
1659 fn unlock<'a>(&'a self, key: &'a str, owner: &'a str) -> IdempotencyFuture<'a, ()> {
1660 self.inner.unlock(key, owner)
1661 }
1662 }
1663
1664 #[tokio::test]
1667 async fn fenced_out_write_fails_closed_and_holds_lock() {
1668 let store: Arc<dyn IdempotencyStore> = Arc::new(FencedSetStore {
1669 inner: MemoryIdempotencyStore::new(Duration::from_secs(600)),
1670 });
1671 let count = Arc::new(AtomicUsize::new(0));
1672 let count_inner = count.clone();
1673 let app = Router::new()
1674 .route(
1675 "/submit",
1676 post(move || {
1677 let count = count_inner.clone();
1678 async move {
1679 count.fetch_add(1, Ordering::SeqCst);
1680 "created"
1681 }
1682 }),
1683 )
1684 .layer(layer_with_store(store));
1685
1686 let first = app
1687 .clone()
1688 .oneshot(urlencoded_post("tok-fenced"))
1689 .await
1690 .unwrap();
1691 assert_eq!(first.status(), StatusCode::SERVICE_UNAVAILABLE);
1692 let second = app
1693 .clone()
1694 .oneshot(urlencoded_post("tok-fenced"))
1695 .await
1696 .unwrap();
1697 assert_eq!(second.status(), StatusCode::CONFLICT);
1698 assert_eq!(count.load(Ordering::SeqCst), 1);
1699 }
1700
1701 #[tokio::test]
1710 async fn response_stream_error_fails_closed_and_holds_lock() {
1711 let store = Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
1712 let store_dyn: Arc<dyn IdempotencyStore> = store.clone();
1713 let count = Arc::new(AtomicUsize::new(0));
1714 let count_inner = count.clone();
1715 let app = Router::new()
1716 .route(
1717 "/submit",
1718 post(move || {
1719 let count = count_inner.clone();
1720 async move {
1721 count.fetch_add(1, Ordering::SeqCst);
1726 let chunks: Vec<Result<Bytes, std::io::Error>> = vec![
1727 Ok(Bytes::from("created")),
1728 Err(std::io::Error::other("simulated response read failure")),
1729 ];
1730 Response::new(Body::from_stream(futures::stream::iter(chunks)))
1731 }
1732 }),
1733 )
1734 .layer(layer_with_store(store_dyn));
1735
1736 let token = "tok-resp-stream-fail";
1737
1738 let first = app.clone().oneshot(urlencoded_post(token)).await.unwrap();
1742 assert_eq!(
1743 first.status(),
1744 StatusCode::INTERNAL_SERVER_ERROR,
1745 "a response-stream error after handler commit must fail closed, not return a truncated body"
1746 );
1747
1748 let key = storage_key(token);
1750 assert!(
1751 store.get(&key).await.unwrap().is_none(),
1752 "a response-stream error must not persist a consumed-token record"
1753 );
1754
1755 let second = app.clone().oneshot(urlencoded_post(token)).await.unwrap();
1759 assert_eq!(
1760 second.status(),
1761 StatusCode::CONFLICT,
1762 "a retry after a response-stream error must be rejected in-flight, not re-run"
1763 );
1764
1765 assert_eq!(
1766 count.load(Ordering::SeqCst),
1767 1,
1768 "the handler must run at most once even when the response stream errors"
1769 );
1770 }
1771
1772 #[tokio::test]
1780 async fn response_stream_error_on_non_success_releases_lock() {
1781 let store = Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
1782 let store_dyn: Arc<dyn IdempotencyStore> = store.clone();
1783 let count = Arc::new(AtomicUsize::new(0));
1784 let count_inner = count.clone();
1785 let app = Router::new()
1786 .route(
1787 "/submit",
1788 post(move || {
1789 let count = count_inner.clone();
1790 async move {
1791 count.fetch_add(1, Ordering::SeqCst);
1797 let chunks: Vec<Result<Bytes, std::io::Error>> = vec![
1798 Ok(Bytes::from("invalid")),
1799 Err(std::io::Error::other("simulated response read failure")),
1800 ];
1801 Response::builder()
1802 .status(StatusCode::UNPROCESSABLE_ENTITY)
1803 .body(Body::from_stream(futures::stream::iter(chunks)))
1804 .unwrap()
1805 }
1806 }),
1807 )
1808 .layer(layer_with_store(store_dyn));
1809
1810 let token = "tok-resp-stream-fail-422";
1811
1812 let first = app.clone().oneshot(urlencoded_post(token)).await.unwrap();
1816 assert_eq!(
1817 first.status(),
1818 StatusCode::INTERNAL_SERVER_ERROR,
1819 "a response-stream error must fail closed with 500 rather than a truncated body"
1820 );
1821
1822 let key = storage_key(token);
1825 assert!(
1826 store.get(&key).await.unwrap().is_none(),
1827 "a non-success response-stream error must not persist a consumed-token record"
1828 );
1829
1830 let second = app.clone().oneshot(urlencoded_post(token)).await.unwrap();
1833 assert_ne!(
1834 second.status(),
1835 StatusCode::CONFLICT,
1836 "a retry after a non-success response-stream error must be retryable, not a 409 in-flight conflict"
1837 );
1838
1839 assert_eq!(
1840 count.load(Ordering::SeqCst),
1841 2,
1842 "the handler must re-run on retry once the non-success lock is released"
1843 );
1844 }
1845
1846 struct FailingGetStore {
1851 inner: MemoryIdempotencyStore,
1852 fail_reads: std::sync::atomic::AtomicBool,
1853 }
1854
1855 impl FailingGetStore {
1856 fn new() -> Self {
1857 Self {
1858 inner: MemoryIdempotencyStore::new(Duration::from_secs(600)),
1859 fail_reads: std::sync::atomic::AtomicBool::new(false),
1860 }
1861 }
1862 }
1863
1864 impl IdempotencyStore for FailingGetStore {
1865 fn get<'a>(&'a self, key: &'a str) -> IdempotencyFuture<'a, Option<IdempotencyEntry>> {
1866 if self.fail_reads.load(Ordering::SeqCst) {
1867 Box::pin(async {
1868 Err(IdempotencyStoreError::backend(
1869 "simulated consumed-token lookup failure",
1870 ))
1871 })
1872 } else {
1873 self.inner.get(key)
1874 }
1875 }
1876
1877 fn set<'a>(
1878 &'a self,
1879 key: &'a str,
1880 owner: &'a str,
1881 record: IdempotencyRecord,
1882 body_hash: Vec<u8>,
1883 ttl: Duration,
1884 ) -> IdempotencyFuture<'a, bool> {
1885 self.inner.set(key, owner, record, body_hash, ttl)
1886 }
1887
1888 fn try_lock<'a>(
1889 &'a self,
1890 key: &'a str,
1891 owner: &'a str,
1892 lock_ttl: Duration,
1893 ) -> IdempotencyFuture<'a, bool> {
1894 self.inner.try_lock(key, owner, lock_ttl)
1895 }
1896
1897 fn unlock<'a>(&'a self, key: &'a str, owner: &'a str) -> IdempotencyFuture<'a, ()> {
1898 self.inner.unlock(key, owner)
1899 }
1900 }
1901
1902 #[tokio::test]
1908 async fn read_failure_fails_closed_and_does_not_rerun() {
1909 let store = Arc::new(FailingGetStore::new());
1910 let store_dyn: Arc<dyn IdempotencyStore> = store.clone();
1911 let count = Arc::new(AtomicUsize::new(0));
1912 let count_inner = count.clone();
1913 let app = Router::new()
1914 .route(
1915 "/submit",
1916 post(move || {
1917 let count = count_inner.clone();
1918 async move {
1919 count.fetch_add(1, Ordering::SeqCst);
1920 "created"
1921 }
1922 }),
1923 )
1924 .layer(layer_with_store(store_dyn));
1925
1926 let token = "tok-read-fail";
1927
1928 let first = app.clone().oneshot(urlencoded_post(token)).await.unwrap();
1930 assert_eq!(first.status(), StatusCode::OK);
1931 assert_eq!(count.load(Ordering::SeqCst), 1);
1932
1933 store.fail_reads.store(true, Ordering::SeqCst);
1936 let second = app.clone().oneshot(urlencoded_post(token)).await.unwrap();
1937 assert_eq!(
1938 second.status(),
1939 StatusCode::SERVICE_UNAVAILABLE,
1940 "a consumed-token lookup failure must fail closed, not re-run the handler"
1941 );
1942 assert_eq!(
1943 count.load(Ordering::SeqCst),
1944 1,
1945 "the handler must not re-run when the consumed-token lookup fails"
1946 );
1947 }
1948
1949 struct SlowLookupNoRenewStore {
1952 inner: MemoryIdempotencyStore,
1953 lookups: AtomicUsize,
1954 }
1955
1956 impl IdempotencyStore for SlowLookupNoRenewStore {
1957 fn get<'a>(&'a self, key: &'a str) -> IdempotencyFuture<'a, Option<IdempotencyEntry>> {
1958 let lookup = self.lookups.fetch_add(1, Ordering::SeqCst);
1959 Box::pin(async move {
1960 if lookup == 1 {
1961 tokio::time::sleep(Duration::from_millis(150)).await;
1962 }
1963 self.inner.get(key).await
1964 })
1965 }
1966
1967 fn set<'a>(
1968 &'a self,
1969 key: &'a str,
1970 owner: &'a str,
1971 record: IdempotencyRecord,
1972 body_hash: Vec<u8>,
1973 ttl: Duration,
1974 ) -> IdempotencyFuture<'a, bool> {
1975 self.inner.set(key, owner, record, body_hash, ttl)
1976 }
1977
1978 fn try_lock<'a>(
1979 &'a self,
1980 key: &'a str,
1981 owner: &'a str,
1982 lock_ttl: Duration,
1983 ) -> IdempotencyFuture<'a, bool> {
1984 self.inner.try_lock(key, owner, lock_ttl)
1985 }
1986
1987 fn unlock<'a>(&'a self, key: &'a str, owner: &'a str) -> IdempotencyFuture<'a, ()> {
1988 self.inner.unlock(key, owner)
1989 }
1990 }
1991
1992 #[tokio::test]
1996 async fn slow_post_lock_lookup_without_renewal_does_not_run_the_handler() {
1997 let store: Arc<dyn IdempotencyStore> = Arc::new(SlowLookupNoRenewStore {
1998 inner: MemoryIdempotencyStore::new(Duration::from_secs(600)),
1999 lookups: AtomicUsize::new(0),
2000 });
2001 let count = Arc::new(AtomicUsize::new(0));
2002 let count_inner = count.clone();
2003 let config = SubmitTokenConfig {
2004 in_flight_ttl_secs: 1,
2005 ..default_config()
2006 };
2007 let app = Router::new()
2008 .route(
2009 "/submit",
2010 post(move || {
2011 let count = count_inner.clone();
2012 async move {
2013 count.fetch_add(1, Ordering::SeqCst);
2014 "created"
2015 }
2016 }),
2017 )
2018 .layer(SubmitTokenLayer::new(store, &config));
2019
2020 let response = app.oneshot(urlencoded_post("tok-slow")).await.unwrap();
2021 assert_eq!(response.status(), StatusCode::CONFLICT);
2022 assert_eq!(count.load(Ordering::SeqCst), 0, "the handler must not run");
2023 }
2024
2025 #[tokio::test]
2031 async fn body_read_error_rejects_request_and_does_not_reach_handler() {
2032 let store: Arc<dyn IdempotencyStore> =
2033 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
2034 let count = Arc::new(AtomicUsize::new(0));
2035 let count_inner = count.clone();
2036 let app = Router::new()
2037 .route(
2038 "/submit",
2039 post(move |_body: Bytes| {
2040 let count = count_inner.clone();
2041 async move {
2042 count.fetch_add(1, Ordering::SeqCst);
2043 "created"
2044 }
2045 }),
2046 )
2047 .layer(layer_with_store(store));
2048
2049 let chunks: Vec<Result<Bytes, std::io::Error>> = vec![
2052 Ok(Bytes::from("_submit_token=tok-read-err&ti")),
2053 Err(std::io::Error::other("simulated body read failure")),
2054 ];
2055 let body = Body::from_stream(futures::stream::iter(chunks));
2056 let req = Request::builder()
2057 .method("POST")
2058 .uri("/submit")
2059 .header("Content-Type", "application/x-www-form-urlencoded")
2060 .body(body)
2061 .unwrap();
2062
2063 let resp = app.oneshot(req).await.unwrap();
2064 assert_eq!(
2065 resp.status(),
2066 StatusCode::BAD_REQUEST,
2067 "a mid-read body-stream error must reject the request, not forward a truncated form"
2068 );
2069 assert_eq!(
2070 count.load(Ordering::SeqCst),
2071 0,
2072 "the handler must never see a truncated body when the stream errors mid-read"
2073 );
2074 }
2075
2076 struct TtlRecordingStore {
2081 inner: MemoryIdempotencyStore,
2082 lock_ttl: std::sync::Mutex<Option<Duration>>,
2083 set_ttl: std::sync::Mutex<Option<Duration>>,
2084 }
2085
2086 impl TtlRecordingStore {
2087 fn new() -> Self {
2088 Self {
2089 inner: MemoryIdempotencyStore::new(Duration::from_secs(600)),
2090 lock_ttl: std::sync::Mutex::new(None),
2091 set_ttl: std::sync::Mutex::new(None),
2092 }
2093 }
2094 }
2095
2096 impl IdempotencyStore for TtlRecordingStore {
2097 fn get<'a>(&'a self, key: &'a str) -> IdempotencyFuture<'a, Option<IdempotencyEntry>> {
2098 self.inner.get(key)
2099 }
2100
2101 fn set<'a>(
2102 &'a self,
2103 key: &'a str,
2104 owner: &'a str,
2105 record: IdempotencyRecord,
2106 body_hash: Vec<u8>,
2107 ttl: Duration,
2108 ) -> IdempotencyFuture<'a, bool> {
2109 *self.set_ttl.lock().unwrap() = Some(ttl);
2110 self.inner.set(key, owner, record, body_hash, ttl)
2111 }
2112
2113 fn try_lock<'a>(
2114 &'a self,
2115 key: &'a str,
2116 owner: &'a str,
2117 lock_ttl: Duration,
2118 ) -> IdempotencyFuture<'a, bool> {
2119 *self.lock_ttl.lock().unwrap() = Some(lock_ttl);
2120 self.inner.try_lock(key, owner, lock_ttl)
2121 }
2122
2123 fn unlock<'a>(&'a self, key: &'a str, owner: &'a str) -> IdempotencyFuture<'a, ()> {
2124 self.inner.unlock(key, owner)
2125 }
2126 }
2127
2128 #[tokio::test]
2136 async fn in_flight_lock_ttl_is_decoupled_from_replay_ttl() {
2137 let store = Arc::new(TtlRecordingStore::new());
2138 let store_dyn: Arc<dyn IdempotencyStore> = store.clone();
2139 let config = SubmitTokenConfig {
2141 enabled: true,
2142 ttl_secs: 1,
2143 in_flight_ttl_secs: 86_400,
2144 ..Default::default()
2145 };
2146 let app = Router::new()
2147 .route("/submit", post(|| async { "created" }))
2148 .layer(SubmitTokenLayer::new(store_dyn, &config));
2149
2150 let resp = app.oneshot(urlencoded_post("tok-decoupled")).await.unwrap();
2151 assert_eq!(resp.status(), StatusCode::OK);
2152
2153 let lock_ttl = store
2156 .lock_ttl
2157 .lock()
2158 .unwrap()
2159 .expect("the guard must acquire an in-flight lock");
2160 assert_eq!(
2161 lock_ttl,
2162 Duration::from_secs(86_400),
2163 "the in-flight lock TTL must come from in_flight_ttl_secs, not ttl_secs"
2164 );
2165 let set_ttl = store
2167 .set_ttl
2168 .lock()
2169 .unwrap()
2170 .expect("the guard must record the consumed token");
2171 assert_eq!(
2172 set_ttl,
2173 Duration::from_secs(1),
2174 "the consumed-record replay TTL must still come from ttl_secs"
2175 );
2176 }
2177
2178 #[tokio::test]
2182 async fn in_flight_token_excludes_retry_with_small_replay_ttl() {
2183 let store = Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
2184 let store_dyn: Arc<dyn IdempotencyStore> = store.clone();
2185 let config = SubmitTokenConfig {
2186 enabled: true,
2187 ttl_secs: 1,
2188 in_flight_ttl_secs: 86_400,
2189 ..Default::default()
2190 };
2191 let count = Arc::new(AtomicUsize::new(0));
2192 let count_inner = count.clone();
2193 let app = Router::new()
2194 .route(
2195 "/submit",
2196 post(move || {
2197 let count = count_inner.clone();
2198 async move {
2199 count.fetch_add(1, Ordering::SeqCst);
2200 "created"
2201 }
2202 }),
2203 )
2204 .layer(SubmitTokenLayer::new(store_dyn, &config));
2205
2206 let token = "tok-inflight";
2210 let key = storage_key(token);
2211 assert!(
2212 store
2213 .try_lock(&key, "first-attempt", Duration::from_secs(86_400))
2214 .await
2215 .unwrap()
2216 );
2217
2218 let resp = app.oneshot(urlencoded_post(token)).await.unwrap();
2221 assert_eq!(
2222 resp.status(),
2223 StatusCode::CONFLICT,
2224 "a retry against a still-in-flight token must be excluded, not re-run"
2225 );
2226 assert_eq!(
2227 count.load(Ordering::SeqCst),
2228 0,
2229 "the handler must not run for a retry held out by the in-flight lock"
2230 );
2231 }
2232
2233 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
2238 async fn late_release_keeps_a_newer_attempts_lock() {
2239 let store: Arc<dyn IdempotencyStore> =
2240 Arc::new(MemoryIdempotencyStore::new(Duration::from_secs(600)));
2241 let config = SubmitTokenConfig {
2242 enabled: true,
2243 in_flight_ttl_secs: 1,
2244 ..Default::default()
2245 };
2246 let count = Arc::new(AtomicUsize::new(0));
2247 let count_inner = count.clone();
2248 let app = Router::new()
2249 .route(
2250 "/submit",
2251 post(move || {
2252 let count = count_inner.clone();
2253 async move {
2254 let run = count.fetch_add(1, Ordering::SeqCst);
2255 match run {
2256 0 => {
2258 tokio::time::sleep(Duration::from_millis(1_500)).await;
2259 (StatusCode::INTERNAL_SERVER_ERROR, "failed")
2260 }
2261 1 => {
2263 tokio::time::sleep(Duration::from_secs(2)).await;
2264 (StatusCode::CREATED, "created")
2265 }
2266 _ => (StatusCode::CREATED, "created"),
2267 }
2268 }
2269 }),
2270 )
2271 .layer(SubmitTokenLayer::new(store, &config));
2272
2273 let token = "tok-late-release";
2274 let a = tokio::spawn(app.clone().oneshot(urlencoded_post(token)));
2275 tokio::time::sleep(Duration::from_millis(1_200)).await;
2276 let b = tokio::spawn(app.clone().oneshot(urlencoded_post(token)));
2277 let a = a.await.unwrap().unwrap();
2279 assert_eq!(a.status(), StatusCode::INTERNAL_SERVER_ERROR);
2280 tokio::time::sleep(Duration::from_millis(200)).await;
2281
2282 let c = app.clone().oneshot(urlencoded_post(token)).await.unwrap();
2283 assert_eq!(
2284 c.status(),
2285 StatusCode::CONFLICT,
2286 "B still holds the token; A's release must not free B's lock"
2287 );
2288 assert_eq!(count.load(Ordering::SeqCst), 2, "C did not run the handler");
2289 let b = b.await.unwrap().unwrap();
2290 assert_eq!(b.status(), StatusCode::CREATED);
2291 }
2292
2293 struct LateRecordStore {
2298 inner: MemoryIdempotencyStore,
2299 gets: AtomicUsize,
2300 }
2301
2302 impl IdempotencyStore for LateRecordStore {
2303 fn get<'a>(&'a self, key: &'a str) -> IdempotencyFuture<'a, Option<IdempotencyEntry>> {
2304 if self.gets.fetch_add(1, Ordering::SeqCst) == 0 {
2305 return Box::pin(async { Ok(None) });
2306 }
2307 self.inner.get(key)
2308 }
2309
2310 fn set<'a>(
2311 &'a self,
2312 key: &'a str,
2313 owner: &'a str,
2314 record: IdempotencyRecord,
2315 body_hash: Vec<u8>,
2316 ttl: Duration,
2317 ) -> IdempotencyFuture<'a, bool> {
2318 self.inner.set(key, owner, record, body_hash, ttl)
2319 }
2320
2321 fn try_lock<'a>(
2322 &'a self,
2323 _key: &'a str,
2324 _owner: &'a str,
2325 _lock_ttl: Duration,
2326 ) -> IdempotencyFuture<'a, bool> {
2327 Box::pin(async { Ok(false) })
2328 }
2329
2330 fn unlock<'a>(&'a self, key: &'a str, owner: &'a str) -> IdempotencyFuture<'a, ()> {
2331 self.inner.unlock(key, owner)
2332 }
2333 }
2334
2335 #[tokio::test]
2336 async fn lock_refused_by_a_late_record_replays_it() {
2337 let token = "tok-late-record";
2338 let inner = MemoryIdempotencyStore::new(Duration::from_secs(600));
2339 let record = IdempotencyRecord {
2340 status: 201,
2341 headers: Vec::new(),
2342 body: b"first".to_vec(),
2343 metadata: Vec::new(),
2344 };
2345 inner
2346 .set(
2347 &storage_key(token),
2348 "racer",
2349 record,
2350 Vec::new(),
2351 Duration::from_secs(600),
2352 )
2353 .await
2354 .unwrap();
2355 let store: Arc<dyn IdempotencyStore> = Arc::new(LateRecordStore {
2356 inner,
2357 gets: AtomicUsize::new(0),
2358 });
2359 let count = Arc::new(AtomicUsize::new(0));
2360 let count_inner = count.clone();
2361 let app = Router::new()
2362 .route(
2363 "/submit",
2364 post(move || {
2365 let count = count_inner.clone();
2366 async move {
2367 count.fetch_add(1, Ordering::SeqCst);
2368 "created"
2369 }
2370 }),
2371 )
2372 .layer(layer_with_store(store));
2373
2374 let response = app.oneshot(urlencoded_post(token)).await.unwrap();
2375 assert_eq!(
2376 response.status(),
2377 StatusCode::CREATED,
2378 "the racing request's response is replayed, not a 409"
2379 );
2380 assert_eq!(count.load(Ordering::SeqCst), 0, "the handler did not run");
2381 }
2382
2383 #[test]
2388 fn storage_key_is_pinned_hex() {
2389 assert_eq!(
2390 storage_key("token-abc"),
2391 "submit:2a8a128a33f0f4cec8eb049c91fc3242d32de4ca2384fa3e7a3770b7ab952d1b"
2392 );
2393 }
2394}