HMC/handmade-revolt-backend
6
0
Fork 0
forked from jmug/stoatchat
Code Pull Requests Actions Packages Releases 2 Activity

fix(delta): set ratelimit headers on 429 responses

Browse Source
This commit is contained in:
Bob Bobs
2024-02-09 16:26:06 -07:00
committed by Paul Makles
parent 226dbca6e0
commit ea00f0fec1
1 changed files with 42 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
crates/delta/src/util/ratelimiter.rs
View File

@@ -180,7 +180,7 @@ impl Ratelimiter {
pub fn from( pub fn from(
identifier: &str, identifier: &str,
(bucket, resource): (&str, Option<&str>), (bucket, resource): (&str, Option<&str>),
) -> Result<Ratelimiter, u128> { ) -> Result<Ratelimiter, Ratelimiter> {
let mut key = DefaultHasher::new(); let mut key = DefaultHasher::new();
key.write(identifier.as_bytes()); key.write(identifier.as_bytes());
key.write(bucket.as_bytes()); key.write(bucket.as_bytes());
@@ -194,27 +194,29 @@ impl Ratelimiter {
let mut entry = Entry::from(key); let mut entry = Entry::from(key);
let remaining = entry.get_remaining(limit); let remaining = entry.get_remaining(limit);
if remaining > 0 { let reset = entry.left_until_reset();
entry.deduct(); let mut ratelimiter = Ratelimiter {
key,
let reset = entry.left_until_reset(); limit,
entry.save(key); remaining,
reset,
Ok(Ratelimiter { };
key, if remaining == 0 {
limit, return Err(ratelimiter);
remaining: remaining - 1,
reset,
})
} else {
Err(entry.left_until_reset())
} }
entry.deduct();
entry.save(key);
ratelimiter.remaining -= 1;
ratelimiter.reset = entry.left_until_reset();
Ok(ratelimiter)
} }
} }
#[async_trait] #[async_trait]
impl<'r> FromRequest<'r> for Ratelimiter { impl<'r> FromRequest<'r> for Ratelimiter {
type Error = u128; type Error = Ratelimiter;
async fn from_request(request: &'r rocket::Request<'_>) -> Outcome<Self, Self::Error> { async fn from_request(request: &'r rocket::Request<'_>) -> Outcome<Self, Self::Error> {
let ratelimiter = request let ratelimiter = request
@@ -233,7 +235,7 @@ impl<'r> FromRequest<'r> for Ratelimiter {
match ratelimiter { match ratelimiter {
Ok(ratelimiter) => Outcome::Success(*ratelimiter), Ok(ratelimiter) => Outcome::Success(*ratelimiter),
Err(retry_after) => Outcome::Failure((Status::TooManyRequests, *retry_after)), Err(ratelimiter) => Outcome::Failure((Status::TooManyRequests, *ratelimiter)),
} }
} }
} }
@@ -275,23 +277,24 @@ impl Fairing for RatelimitFairing {
} }
async fn on_response<'r>(&self, request: &'r Request<'_>, response: &mut Response<'r>) { async fn on_response<'r>(&self, request: &'r Request<'_>, response: &mut Response<'r>) {
use rocket::outcome::Outcome; let guard = request.guard::<Ratelimiter>().await;
match request.guard::<Ratelimiter>().await { let (Outcome::Success(ratelimiter) | Outcome::Failure((_, ratelimiter))) = guard else {
Outcome::Success(ratelimiter) => { unreachable!()
let Ratelimiter { };
key, let Ratelimiter {
limit, key,
remaining, limit,
reset, remaining,
} = ratelimiter; reset,
} = ratelimiter;
response.set_raw_header("X-RateLimit-Limit", limit.to_string()); response.set_raw_header("X-RateLimit-Limit", limit.to_string());
response.set_raw_header("X-RateLimit-Bucket", key.to_string()); response.set_raw_header("X-RateLimit-Bucket", key.to_string());
response.set_raw_header("X-RateLimit-Remaining", remaining.to_string()); response.set_raw_header("X-RateLimit-Remaining", remaining.to_string());
response.set_raw_header("X-RateLimit-Reset-After", reset.to_string()); response.set_raw_header("X-RateLimit-Reset-After", reset.to_string());
}
Outcome::Failure(_) => response.set_status(Status::TooManyRequests), if guard.is_failure() {
Outcome::Forward(_) => unreachable!(), response.set_status(Status::TooManyRequests);
} }
} }
} }
@@ -308,11 +311,14 @@ impl<'r> FromRequest<'r> for RatelimitInformation {
type Error = u128; type Error = u128;
async fn from_request(request: &'r rocket::Request<'_>) -> Outcome<Self, Self::Error> { async fn from_request(request: &'r rocket::Request<'_>) -> Outcome<Self, Self::Error> {
Outcome::Success(match request.guard::<Ratelimiter>().await { let info = match request.guard::<Ratelimiter>().await {
Outcome::Success(ratelimiter) => RatelimitInformation::Success(ratelimiter), Outcome::Success(ratelimiter) => RatelimitInformation::Success(ratelimiter),
Outcome::Failure((_, retry_after)) => RatelimitInformation::Failure { retry_after }, Outcome::Failure((_, ratelimiter)) => RatelimitInformation::Failure {
retry_after: ratelimiter.reset,
},
_ => unreachable!(), _ => unreachable!(),
}) };
Outcome::Success(info)
} }
} }