Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(delta): set ratelimit headers on 429 responses #302

Merged
merged 1 commit into from
Feb 10, 2024
Merged

fix(delta): set ratelimit headers on 429 responses #302

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 42 additions & 36 deletions crates/delta/src/util/ratelimiter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ impl Ratelimiter {
pub fn from(
identifier: &str,
(bucket, resource): (&str, Option<&str>),
) -> Result<Ratelimiter, u128> {
) -> Result<Ratelimiter, Ratelimiter> {
let mut key = DefaultHasher::new();
key.write(identifier.as_bytes());
key.write(bucket.as_bytes());
Expand All @@ -194,27 +194,29 @@ impl Ratelimiter {
let mut entry = Entry::from(key);

let remaining = entry.get_remaining(limit);
if remaining > 0 {
entry.deduct();
let reset = entry.left_until_reset();
let mut ratelimiter = Ratelimiter {
key,
limit,
remaining,
reset,
};
if remaining == 0 {
return Err(ratelimiter);
}

let reset = entry.left_until_reset();
entry.save(key);
entry.deduct();
entry.save(key);
ratelimiter.remaining -= 1;
ratelimiter.reset = entry.left_until_reset();

Ok(Ratelimiter {
key,
limit,
remaining: remaining - 1,
reset,
})
} else {
Err(entry.left_until_reset())
}
Ok(ratelimiter)
}
}

#[async_trait]
impl<'r> FromRequest<'r> for Ratelimiter {
type Error = u128;
type Error = Ratelimiter;

async fn from_request(request: &'r rocket::Request<'_>) -> Outcome<Self, Self::Error> {
let ratelimiter = request
Expand All @@ -233,7 +235,7 @@ impl<'r> FromRequest<'r> for Ratelimiter {

match ratelimiter {
Ok(ratelimiter) => Outcome::Success(*ratelimiter),
Err(retry_after) => Outcome::Failure((Status::TooManyRequests, *retry_after)),
Err(ratelimiter) => Outcome::Failure((Status::TooManyRequests, *ratelimiter)),
}
}
}
Expand Down Expand Up @@ -275,23 +277,24 @@ impl Fairing for RatelimitFairing {
}

async fn on_response<'r>(&self, request: &'r Request<'_>, response: &mut Response<'r>) {
use rocket::outcome::Outcome;
match request.guard::<Ratelimiter>().await {
Outcome::Success(ratelimiter) => {
let Ratelimiter {
key,
limit,
remaining,
reset,
} = ratelimiter;

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-Remaining", remaining.to_string());
response.set_raw_header("X-RateLimit-Reset-After", reset.to_string());
}
Outcome::Failure(_) => response.set_status(Status::TooManyRequests),
Outcome::Forward(_) => unreachable!(),
let guard = request.guard::<Ratelimiter>().await;
let (Outcome::Success(ratelimiter) | Outcome::Failure((_, ratelimiter))) = guard else {
unreachable!()
};
let Ratelimiter {
key,
limit,
remaining,
reset,
} = ratelimiter;

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-Remaining", remaining.to_string());
response.set_raw_header("X-RateLimit-Reset-After", reset.to_string());

if guard.is_failure() {
response.set_status(Status::TooManyRequests);
}
}
}
Expand All @@ -308,11 +311,14 @@ impl<'r> FromRequest<'r> for RatelimitInformation {
type Error = u128;

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::Failure((_, retry_after)) => RatelimitInformation::Failure { retry_after },
Outcome::Failure((_, ratelimiter)) => RatelimitInformation::Failure {
retry_after: ratelimiter.reset,
},
_ => unreachable!(),
})
};
Outcome::Success(info)
}
}

Expand Down
Loading