use std::net::IpAddr;
use axum::extract::ConnectInfo;
use axum::http::request::Parts;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use super::class::IdempotencyClass;
use crate::error::AppError;
use crate::store::KeyspaceHandle;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Principal {
AuthToken(Vec<u8>),
Ip(IpAddr),
Anonymous,
}
impl Principal {
pub fn hash(&self) -> [u8; 32] {
let mut hasher = Sha256::new();
match self {
Principal::AuthToken(bytes) => {
hasher.update(b"auth-token:");
hasher.update(bytes);
}
Principal::Ip(ip) => {
hasher.update(b"ip:");
hasher.update(ip.to_string().as_bytes());
}
Principal::Anonymous => {
hasher.update(b"anonymous");
}
}
hasher.finalize().into()
}
}
pub fn principal_from_request(parts: &Parts) -> Principal {
if let Some(auth) = parts.headers.get(axum::http::header::AUTHORIZATION) {
return Principal::AuthToken(auth.as_bytes().to_vec());
}
if let Some(ConnectInfo(addr)) = parts.extensions.get::<ConnectInfo<std::net::SocketAddr>>() {
return Principal::Ip(addr.ip());
}
Principal::Anonymous
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct CacheEntry {
pub idempotency_key: String,
pub request_hash: [u8; 32],
pub response_status: u16,
pub response_headers: Vec<(String, String)>,
pub response_body: Vec<u8>,
pub class: IdempotencyClass,
pub created_at: DateTime<Utc>,
pub expires_at: DateTime<Utc>,
}
impl CacheEntry {
pub fn is_expired(&self, now: DateTime<Utc>) -> bool {
self.expires_at <= now
}
}
#[derive(Clone)]
pub struct IdempotencyStore {
ks: KeyspaceHandle,
}
impl IdempotencyStore {
pub fn new(ks: KeyspaceHandle) -> Self {
Self { ks }
}
pub async fn get(
&self,
principal_hash: &[u8; 32],
key: &str,
) -> Result<Option<CacheEntry>, AppError> {
let storage_key = storage_key(principal_hash, key);
let entry: Option<CacheEntry> = self.ks.get(storage_key).await?;
let now = Utc::now();
Ok(entry.filter(|e| !e.is_expired(now)))
}
pub async fn put(&self, principal_hash: &[u8; 32], entry: &CacheEntry) -> Result<(), AppError> {
let storage_key = storage_key(principal_hash, &entry.idempotency_key);
self.ks.insert(storage_key, entry).await
}
}
fn storage_key(principal_hash: &[u8; 32], key: &str) -> Vec<u8> {
let mut out = Vec::with_capacity(64 + key.len() + 5);
out.extend_from_slice(b"idem:");
out.extend_from_slice(hex::encode(principal_hash).as_bytes());
out.push(b':');
out.extend_from_slice(key.as_bytes());
out
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::StoreConfig;
use crate::store::Store;
use chrono::Duration;
fn temp_store() -> (IdempotencyStore, tempfile::TempDir) {
let dir = tempfile::tempdir().expect("tempdir");
let cfg = StoreConfig {
data_dir: dir.path().to_path_buf(),
};
let store = Store::open(&cfg).expect("store");
let ks = store.keyspace("idempotency-test").expect("ks");
(IdempotencyStore::new(ks), dir)
}
fn sample_entry() -> CacheEntry {
let now = Utc::now();
CacheEntry {
idempotency_key: "key-1".into(),
request_hash: [0xAB; 32],
response_status: 201,
response_headers: vec![("content-type".into(), "application/json".into())],
response_body: br#"{"ok":true}"#.to_vec(),
class: IdempotencyClass::NonDestructive,
created_at: now,
expires_at: now
+ Duration::seconds(IdempotencyClass::NonDestructive.ttl_seconds() as i64),
}
}
#[test]
fn principal_hash_is_stable_and_distinct_across_kinds() {
let a = Principal::AuthToken(b"Bearer abc".to_vec());
let a_again = Principal::AuthToken(b"Bearer abc".to_vec());
let b = Principal::AuthToken(b"Bearer xyz".to_vec());
let ip = Principal::Ip(IpAddr::V4("127.0.0.1".parse().unwrap()));
let anon = Principal::Anonymous;
assert_eq!(a.hash(), a_again.hash());
assert_ne!(a.hash(), b.hash());
assert_ne!(a.hash(), ip.hash());
assert_ne!(a.hash(), anon.hash());
assert_ne!(ip.hash(), anon.hash());
}
#[tokio::test]
async fn put_then_get_returns_entry() {
let (store, _dir) = temp_store();
let principal = Principal::AuthToken(b"Bearer t".to_vec()).hash();
let entry = sample_entry();
store.put(&principal, &entry).await.unwrap();
let got = store.get(&principal, &entry.idempotency_key).await.unwrap();
assert_eq!(got.as_ref(), Some(&entry));
}
#[tokio::test]
async fn entries_are_scoped_by_principal() {
let (store, _dir) = temp_store();
let a = Principal::AuthToken(b"alice".to_vec()).hash();
let b = Principal::AuthToken(b"bob".to_vec()).hash();
let entry = sample_entry();
store.put(&a, &entry).await.unwrap();
let got_a = store.get(&a, &entry.idempotency_key).await.unwrap();
let got_b = store.get(&b, &entry.idempotency_key).await.unwrap();
assert!(got_a.is_some());
assert!(got_b.is_none(), "principal scoping leaked");
}
#[tokio::test]
async fn expired_entries_are_filtered_at_read_time() {
let (store, _dir) = temp_store();
let principal = Principal::AuthToken(b"Bearer t".to_vec()).hash();
let mut entry = sample_entry();
entry.expires_at = Utc::now() - Duration::seconds(1);
store.put(&principal, &entry).await.unwrap();
let got = store.get(&principal, &entry.idempotency_key).await.unwrap();
assert!(got.is_none(), "stale entry served");
}
#[tokio::test]
async fn put_overwrites_existing_entry_under_same_key() {
let (store, _dir) = temp_store();
let principal = Principal::AuthToken(b"Bearer t".to_vec()).hash();
let first = sample_entry();
store.put(&principal, &first).await.unwrap();
let mut second = first.clone();
second.response_status = 204;
second.response_body = b"updated".to_vec();
store.put(&principal, &second).await.unwrap();
let got = store.get(&principal, &first.idempotency_key).await.unwrap();
assert_eq!(got.unwrap().response_status, 204);
}
}