revolt_database/util/
idempotency.rs

1use std::num::NonZeroUsize;
2
3use revolt_result::{create_error, Result};
4
5#[cfg(feature = "rocket-impl")]
6use revolt_result::Error;
7
8use async_std::sync::Mutex;
9use once_cell::sync::Lazy;
10use serde::{Deserialize, Serialize};
11
12#[derive(Serialize, Deserialize)]
13pub struct IdempotencyKey {
14    key: String,
15}
16
17static TOKEN_CACHE: Lazy<Mutex<lru::LruCache<String, ()>>> =
18    Lazy::new(|| Mutex::new(lru::LruCache::new(NonZeroUsize::new(1000).unwrap())));
19
20impl IdempotencyKey {
21    pub fn unchecked_from_string(key: String) -> Self {
22        Self { key }
23    }
24
25    // Backwards compatibility.
26    // Issue #109
27    pub async fn consume_nonce(&mut self, v: Option<String>) -> Result<()> {
28        if let Some(v) = v {
29            let mut cache = TOKEN_CACHE.lock().await;
30            if cache.get(&v).is_some() {
31                return Err(create_error!(DuplicateNonce));
32            }
33
34            cache.put(v.clone(), ());
35            self.key = v;
36        }
37
38        Ok(())
39    }
40
41    pub fn into_key(self) -> String {
42        self.key
43    }
44}
45
46#[cfg(feature = "rocket-impl")]
47use revolt_rocket_okapi::{
48    gen::OpenApiGenerator,
49    request::{OpenApiFromRequest, RequestHeaderInput},
50    revolt_okapi::openapi3::{Parameter, ParameterValue},
51};
52
53#[cfg(feature = "rocket-impl")]
54use schemars::schema::{InstanceType, SchemaObject, SingleOrVec};
55
56#[cfg(feature = "rocket-impl")]
57impl OpenApiFromRequest<'_> for IdempotencyKey {
58    fn from_request_input(
59        _gen: &mut OpenApiGenerator,
60        _name: String,
61        _required: bool,
62    ) -> revolt_rocket_okapi::Result<RequestHeaderInput> {
63        Ok(RequestHeaderInput::Parameter(Parameter {
64            name: "Idempotency-Key".to_string(),
65            description: Some("Unique key to prevent duplicate requests".to_string()),
66            allow_empty_value: false,
67            required: false,
68            deprecated: false,
69            extensions: schemars::Map::new(),
70            location: "header".to_string(),
71            value: ParameterValue::Schema {
72                allow_reserved: false,
73                example: None,
74                examples: None,
75                explode: None,
76                style: None,
77                schema: SchemaObject {
78                    instance_type: Some(SingleOrVec::Single(Box::new(InstanceType::String))),
79                    ..Default::default()
80                },
81            },
82        }))
83    }
84}
85
86#[cfg(feature = "rocket-impl")]
87use rocket::{
88    http::Status,
89    request::{FromRequest, Outcome},
90};
91
92#[cfg(feature = "rocket-impl")]
93#[async_trait]
94impl<'r> FromRequest<'r> for IdempotencyKey {
95    type Error = Error;
96
97    async fn from_request(request: &'r rocket::Request<'_>) -> Outcome<Self, Self::Error> {
98        if let Some(key) = request
99            .headers()
100            .get("Idempotency-Key")
101            .next()
102            .map(|k| k.to_string())
103        {
104            if key.len() > 64 {
105                return Outcome::Error((
106                    Status::BadRequest,
107                    create_error!(FailedValidation {
108                        error: "idempotency key too long".to_string(),
109                    }),
110                ));
111            }
112
113            let idempotency = IdempotencyKey { key };
114            let mut cache = TOKEN_CACHE.lock().await;
115            if cache.get(&idempotency.key).is_some() {
116                return Outcome::Error((Status::Conflict, create_error!(DuplicateNonce)));
117            }
118
119            cache.put(idempotency.key.clone(), ());
120            return Outcome::Success(idempotency);
121        }
122
123        Outcome::Success(IdempotencyKey {
124            key: ulid::Ulid::new().to_string(),
125        })
126    }
127}
revolt_database/util/idempotency.rs

revolt_database/util/
idempotency.rs
