Skip to content

Commit

Permalink
Refactored Counter _out_ the public crate API
Browse files Browse the repository at this point in the history
  • Loading branch information
@alexsnaps
alexsnaps committed Jun 3, 2024
1 parent d040c0d commit eefd5d4
Show file tree
Hide file tree
Showing 3 changed files with 74 additions and 98 deletions.
75 changes: 17 additions & 58 deletions limitador-server/src/envoy_rls/server.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,12 @@ use opentelemetry::propagation::Extractor;
use std::collections::HashMap;
use std::sync::Arc;

use limitador::CheckResult;
use tonic::codegen::http::HeaderMap;
use tonic::{transport, transport::Server, Request, Response, Status};
use tracing::Span;
use tracing_opentelemetry::OpenTelemetrySpanExt;

use limitador::counter::Counter;

use crate::envoy_rls::server::envoy::config::core::v3::HeaderValue;
use crate::envoy_rls::server::envoy::service::ratelimit::v3::rate_limit_response::Code;
use crate::envoy_rls::server::envoy::service::ratelimit::v3::rate_limit_service_server::{
Expand All @@ -29,6 +28,21 @@ pub enum RateLimitHeaders {
DraftVersion03,
}

impl RateLimitHeaders {
pub fn headers(&self, response: &mut CheckResult) -> Vec<HeaderValue> {
let mut headers = match self {
RateLimitHeaders::None => Vec::default(),
RateLimitHeaders::DraftVersion03 => response
.response_header()
.into_iter()
.map(|(key, value)| HeaderValue { key, value })
.collect(),
};
headers.sort_by(|a, b| a.key.cmp(&b.key));
headers
}
}

pub struct MyRateLimiter {
limiter: Arc<Limiter>,
rate_limit_headers: RateLimitHeaders,
Expand Down Expand Up @@ -142,10 +156,7 @@ impl RateLimitService for MyRateLimiter {
overall_code: resp_code.into(),
statuses: vec![],
request_headers_to_add: vec![],
response_headers_to_add: to_response_header(
&self.rate_limit_headers,
&mut rate_limited_resp.counters,
),
response_headers_to_add: self.rate_limit_headers.headers(&mut rate_limited_resp),
raw_body: vec![],
dynamic_metadata: None,
quota: None,
Expand All @@ -155,58 +166,6 @@ impl RateLimitService for MyRateLimiter {
}
}

pub fn to_response_header(
rate_limit_headers: &RateLimitHeaders,
counters: &mut [Counter],
) -> Vec<HeaderValue> {
let mut headers = Vec::new();
match rate_limit_headers {
RateLimitHeaders::None => {}

// creates response headers per https://datatracker.ietf.org/doc/id/draft-polli-ratelimit-headers-03.html
RateLimitHeaders::DraftVersion03 => {
// sort by the limit remaining..
counters.sort_by(|a, b| {
let a_remaining = a.remaining().unwrap_or(a.max_value());
let b_remaining = b.remaining().unwrap_or(b.max_value());
a_remaining.cmp(&b_remaining)
});

let mut all_limits_text = String::with_capacity(20 * counters.len());
counters.iter_mut().for_each(|counter| {
all_limits_text.push_str(
format!(", {};w={}", counter.max_value(), counter.window().as_secs()).as_str(),
);
if let Some(name) = counter.limit().name() {
all_limits_text
.push_str(format!(";name=\"{}\"", name.replace('"', "'")).as_str());
}
});

if let Some(counter) = counters.first() {
headers.push(HeaderValue {
key: "X-RateLimit-Limit".to_string(),
value: format!("{}{}", counter.max_value(), all_limits_text),
});

let remaining = counter.remaining().unwrap_or(counter.max_value());
headers.push(HeaderValue {
key: "X-RateLimit-Remaining".to_string(),
value: format!("{}", remaining),
});

if let Some(duration) = counter.expires_in() {
headers.push(HeaderValue {
key: "X-RateLimit-Reset".to_string(),
value: format!("{}", duration.as_secs()),
});
}
}
}
};
headers
}

struct RateLimitRequestHeaders {
inner: HeaderMap,
}
Expand Down
54 changes: 14 additions & 40 deletions limitador-server/src/http_api/server.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use crate::prometheus_metrics::PrometheusMetrics;
use crate::Limiter;
use actix_web::{http::StatusCode, HttpResponse, HttpResponseBuilder, ResponseError};
use actix_web::{App, HttpServer};
use limitador::CheckResult;
use paperclip::actix::{
api_v2_errors,
api_v2_operation,
Expand Down Expand Up @@ -209,7 +210,7 @@ async fn check_and_report(
add_response_header(
&mut resp,
response_headers.as_str(),
&mut is_rate_limited.counters,
&mut is_rate_limited,
);
resp.json(())
}
Expand All @@ -224,7 +225,7 @@ async fn check_and_report(
add_response_header(
&mut resp,
response_headers.as_str(),
&mut is_rate_limited.counters,
&mut is_rate_limited,
);
resp.json(())
}
Expand All @@ -238,48 +239,21 @@ async fn check_and_report(
pub fn add_response_header(
resp: &mut HttpResponseBuilder,
rate_limit_headers: &str,
counters: &mut [limitador::counter::Counter],
result: &mut CheckResult,
) {
match rate_limit_headers {
if rate_limit_headers == "DraftVersion03" {
// creates response headers per https://datatracker.ietf.org/doc/id/draft-polli-ratelimit-headers-03.html
"DraftVersion03" => {
// sort by the limit remaining..
counters.sort_by(|a, b| {
let a_remaining = a.remaining().unwrap_or(a.max_value());
let b_remaining = b.remaining().unwrap_or(b.max_value());
a_remaining.cmp(&b_remaining)
});

let mut all_limits_text = String::with_capacity(20 * counters.len());
counters.iter_mut().for_each(|counter| {
all_limits_text.push_str(
format!(", {};w={}", counter.max_value(), counter.window().as_secs()).as_str(),
);
if let Some(name) = counter.limit().name() {
all_limits_text
.push_str(format!(";name=\"{}\"", name.replace('"', "'")).as_str());
}
});

if let Some(counter) = counters.first() {
resp.insert_header((
"X-RateLimit-Limit",
format!("{}{}", counter.max_value(), all_limits_text),
));

let remaining = counter.remaining().unwrap_or(counter.max_value());
resp.insert_header((
"X-RateLimit-Remaining".to_string(),
format!("{}", remaining),
));

if let Some(duration) = counter.expires_in() {
resp.insert_header(("X-RateLimit-Reset", format!("{}", duration.as_secs())));
}
let headers = result.response_header();
if let Some(limit) = headers.get("X-RateLimit-Limit") {
resp.insert_header(("X-RateLimit-Limit", limit.clone()));
}
if let Some(remaining) = headers.get("X-RateLimit-Remaining") {
resp.insert_header(("X-RateLimit-Remaining".to_string(), remaining.clone()));
if let Some(duration) = headers.get("X-RateLimit-Reset") {
resp.insert_header(("X-RateLimit-Reset", duration.clone()));
}
}
_default => {}
};
}
}

pub async fn run_http_server(
Expand Down
43 changes: 43 additions & 0 deletions limitador/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,49 @@ pub struct CheckResult {
pub limit_name: Option<String>,
}

impl CheckResult {
pub fn response_header(&mut self) -> HashMap<String, String> {
let mut headers = HashMap::new();
// sort by the limit remaining..
self.counters.sort_by(|a, b| {
let a_remaining = a.remaining().unwrap_or(a.max_value());
let b_remaining = b.remaining().unwrap_or(b.max_value());
a_remaining.cmp(&b_remaining)
});

let mut all_limits_text = String::with_capacity(20 * self.counters.len());
self.counters.iter_mut().for_each(|counter| {
all_limits_text.push_str(
format!(", {};w={}", counter.max_value(), counter.window().as_secs()).as_str(),
);
if let Some(name) = counter.limit().name() {
all_limits_text.push_str(format!(";name=\"{}\"", name.replace('"', "'")).as_str());
}
});

if let Some(counter) = self.counters.first() {
headers.insert(
"X-RateLimit-Limit".to_string(),
format!("{}{}", counter.max_value(), all_limits_text),
);

let remaining = counter.remaining().unwrap_or(counter.max_value());
headers.insert(
"X-RateLimit-Remaining".to_string(),
format!("{}", remaining),
);

if let Some(duration) = counter.expires_in() {
headers.insert(
"X-RateLimit-Reset".to_string(),
format!("{}", duration.as_secs()),
);
}
}
headers
}
}

impl From<CheckResult> for bool {
fn from(value: CheckResult) -> Self {
value.limited
Expand Down

0 comments on commit eefd5d4

Please sign in to comment.