use std::convert::Infallible;
use std::task::{Context, Poll};
use std::time::Duration;
use axum::body::Body;
use axum::http::{HeaderName, HeaderValue, Method, Request, Response, StatusCode};
use base64::Engine;
use futures::future::BoxFuture;
use http_body_util::BodyExt;
#[cfg(feature = "multi-tenant")]
use hwhkit_core::TenantId;
use hwhkit_integration_redis::RedisHandle;
use redis_client::AsyncCommands;
use serde::{Deserialize, Serialize};
use serde_json::json;
use sha2::{Digest, Sha256};
use tower::{Layer, Service};
const HEADER_IDEMPOTENCY_KEY: &str = "idempotency-key";
const NAMESPACE: &str = "hwhkit:idem";
#[derive(Clone)]
pub struct IdempotencyLayer {
handle: RedisHandle,
ttl: Duration,
namespace: String,
}
impl IdempotencyLayer {
pub fn new(handle: RedisHandle) -> Self {
Self {
handle,
ttl: Duration::from_secs(24 * 60 * 60),
namespace: NAMESPACE.to_string(),
}
}
#[must_use]
pub fn with_ttl(mut self, ttl: Duration) -> Self {
self.ttl = ttl;
self
}
#[must_use]
pub fn with_namespace(mut self, ns: impl Into<String>) -> Self {
self.namespace = ns.into();
self
}
}
impl<S> Layer<S> for IdempotencyLayer {
type Service = Idempotency<S>;
fn layer(&self, inner: S) -> Self::Service {
Idempotency {
inner,
handle: self.handle.clone(),
ttl: self.ttl,
namespace: self.namespace.clone(),
}
}
}
#[derive(Clone)]
pub struct Idempotency<S> {
inner: S,
handle: RedisHandle,
ttl: Duration,
namespace: String,
}
#[derive(Serialize, Deserialize)]
struct CachedResponse {
status: u16,
headers: Vec<(String, String)>,
body_b64: String,
fingerprint: String,
}
impl<S> Service<Request<Body>> for Idempotency<S>
where
S: Service<Request<Body>, Response = Response<Body>, Error = Infallible>
+ Clone
+ Send
+ 'static,
S::Future: Send + 'static,
{
type Response = Response<Body>;
type Error = Infallible;
type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, req: Request<Body>) -> Self::Future {
let mutating = matches!(
req.method(),
&Method::POST | &Method::PUT | &Method::PATCH | &Method::DELETE
);
let key = req
.headers()
.get(HEADER_IDEMPOTENCY_KEY)
.and_then(|v| v.to_str().ok())
.map(|s| s.to_string());
let clone = self.inner.clone();
let mut inner = std::mem::replace(&mut self.inner, clone);
if !mutating || key.is_none() {
return Box::pin(async move { inner.call(req).await });
}
let key = key.expect("verified Some above");
#[cfg(feature = "multi-tenant")]
let tenant_prefix: String = req
.extensions()
.get::<TenantId>()
.map(|t| format!("t:{}:", t.as_str()))
.unwrap_or_default();
#[cfg(not(feature = "multi-tenant"))]
let tenant_prefix: String = String::new();
let mut conn = self.handle.manager();
let ttl_secs = self.ttl.as_secs().max(1);
let redis_key = format!("{}:{}{}", self.namespace, tenant_prefix, key);
let method = req.method().clone();
let path = req.uri().path().to_string();
let (parts, body) = req.into_parts();
Box::pin(async move {
let collected = match body.collect().await {
Ok(c) => c.to_bytes(),
Err(err) => {
tracing::warn!(error = %err, "idempotency: failed to buffer request body");
return Ok(bad_gateway(
"failed to read request body for idempotency processing",
));
}
};
let fingerprint = fingerprint_request(&method, &path, &collected);
let cached: Option<String> = AsyncCommands::get(&mut conn, &redis_key)
.await
.unwrap_or(None);
if let Some(payload) = cached {
if let Ok(parsed) = serde_json::from_str::<CachedResponse>(&payload) {
if parsed.fingerprint == fingerprint {
return Ok(replay(parsed));
} else {
return Ok(idempotency_conflict());
}
}
}
let req = Request::from_parts(parts, Body::from(collected));
let response = inner.call(req).await?;
let (parts, body) = response.into_parts();
let collected = match body.collect().await {
Ok(c) => c.to_bytes(),
Err(err) => {
tracing::warn!(
error = %err,
idempotency_key = %key,
"failed to buffer response body for cache; serving empty body"
);
let resp = Response::from_parts(parts, Body::empty());
return Ok(resp);
}
};
let mut headers_vec = Vec::new();
for (k, v) in parts.headers.iter() {
if let Ok(val) = v.to_str() {
headers_vec.push((k.as_str().to_string(), val.to_string()));
}
}
let cached = CachedResponse {
status: parts.status.as_u16(),
headers: headers_vec,
body_b64: encode_b64(&collected),
fingerprint,
};
if let Ok(serialized) = serde_json::to_string(&cached) {
let _: redis_client::RedisResult<()> =
AsyncCommands::set_ex(&mut conn, &redis_key, serialized, ttl_secs).await;
}
Ok(Response::from_parts(parts, Body::from(collected)))
})
}
}
fn replay(cached: CachedResponse) -> Response<Body> {
let mut builder =
Response::builder().status(StatusCode::from_u16(cached.status).unwrap_or(StatusCode::OK));
for (k, v) in &cached.headers {
if let (Ok(name), Ok(value)) = (HeaderName::try_from(k.as_str()), HeaderValue::from_str(v))
{
builder = builder.header(name, value);
}
}
builder = builder.header("x-idempotent-replay", "1");
let bytes = decode_b64(&cached.body_b64).unwrap_or_default();
builder
.body(Body::from(bytes))
.unwrap_or_else(|_| Response::new(Body::empty()))
}
fn idempotency_conflict() -> Response<Body> {
let body = json!({
"type": "about:blank",
"title": "Idempotency-Key Conflict",
"status": 409,
"detail": "Idempotency-Key reused across requests with different fingerprints",
});
let bytes = serde_json::to_vec(&body).unwrap_or_default();
Response::builder()
.status(StatusCode::CONFLICT)
.header("content-type", "application/problem+json")
.body(Body::from(bytes))
.unwrap_or_else(|_| Response::new(Body::empty()))
}
fn bad_gateway(detail: &str) -> Response<Body> {
let body = json!({
"type": "about:blank",
"title": "Bad Gateway",
"status": 502,
"detail": detail,
});
let bytes = serde_json::to_vec(&body).unwrap_or_default();
Response::builder()
.status(StatusCode::BAD_GATEWAY)
.header("content-type", "application/problem+json")
.body(Body::from(bytes))
.unwrap_or_else(|_| Response::new(Body::empty()))
}
pub fn fingerprint_request(method: &Method, path: &str, body: &[u8]) -> String {
let mut h = Sha256::new();
h.update(method.as_str().as_bytes());
h.update(b"\n");
h.update(path.as_bytes());
h.update(b"\n");
h.update(body);
let digest = h.finalize();
let mut hex = String::with_capacity(digest.len() * 2);
for b in digest {
use std::fmt::Write;
let _ = write!(hex, "{b:02x}");
}
hex
}
fn encode_b64(bytes: &[u8]) -> String {
base64::engine::general_purpose::STANDARD.encode(bytes)
}
fn decode_b64(s: &str) -> Option<Vec<u8>> {
base64::engine::general_purpose::STANDARD.decode(s).ok()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn b64_roundtrip() {
let data = b"hello, world";
let enc = encode_b64(data);
let dec = decode_b64(&enc).unwrap();
assert_eq!(dec, data);
}
#[test]
fn fingerprint_changes_when_body_changes() {
let a = fingerprint_request(&Method::POST, "/x", b"a");
let b = fingerprint_request(&Method::POST, "/x", b"b");
assert_ne!(a, b);
}
#[test]
fn fingerprint_changes_when_path_changes() {
let a = fingerprint_request(&Method::POST, "/a", b"x");
let b = fingerprint_request(&Method::POST, "/b", b"x");
assert_ne!(a, b);
}
#[test]
fn fingerprint_changes_when_method_changes() {
let a = fingerprint_request(&Method::POST, "/x", b"x");
let b = fingerprint_request(&Method::PUT, "/x", b"x");
assert_ne!(a, b);
}
#[test]
fn fingerprint_stable_for_same_inputs() {
let a = fingerprint_request(&Method::POST, "/x", b"y");
let b = fingerprint_request(&Method::POST, "/x", b"y");
assert_eq!(a, b);
}
}