Skip to content

Commit

Permalink
Implement Rate Limiter Binding (#603)
Browse files Browse the repository at this point in the history
* Implement Rate Limiter binding

* Make success field public
  • Loading branch information
kflansburg authored Sep 26, 2024
1 parent 566166f commit 9fc2fbd
Show file tree
Hide file tree
Showing 4 changed files with 82 additions and 0 deletions.
2 changes: 2 additions & 0 deletions worker-sys/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ mod incoming_request_cf_properties;
#[cfg(feature = "queue")]
mod queue;
mod r2;
mod rate_limit;
mod schedule;
mod socket;
mod tls_client_auth;
Expand All @@ -30,6 +31,7 @@ pub use incoming_request_cf_properties::*;
#[cfg(feature = "queue")]
pub use queue::*;
pub use r2::*;
pub use rate_limit::*;
pub use schedule::*;
pub use socket::*;
pub use tls_client_auth::*;
Expand Down
12 changes: 12 additions & 0 deletions worker-sys/src/types/rate_limit.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
use js_sys::Promise;
use wasm_bindgen::JsValue;

#[wasm_bindgen::prelude::wasm_bindgen]
extern "C" {
#[wasm_bindgen(extends=js_sys::Object)]
#[derive(Clone, PartialEq, Eq)]
pub type RateLimiter;

#[wasm_bindgen(method, catch)]
pub fn limit(this: &RateLimiter, arg: js_sys::Object) -> Result<Promise, JsValue>;
}
2 changes: 2 additions & 0 deletions worker/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,7 @@ pub use crate::hyperdrive::*;
#[cfg(feature = "queue")]
pub use crate::queue::*;
pub use crate::r2::*;
pub use crate::rate_limit::RateLimiter;
pub use crate::request::{FromRequest, Request};
pub use crate::request_init::*;
pub use crate::response::{EncodeBody, IntoResponse, Response, ResponseBody, ResponseBuilder};
Expand Down Expand Up @@ -218,6 +219,7 @@ mod hyperdrive;
#[cfg(feature = "queue")]
mod queue;
mod r2;
mod rate_limit;
mod request;
mod request_init;
mod response;
Expand Down
66 changes: 66 additions & 0 deletions worker/src/rate_limit.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
use crate::{send::SendFuture, EnvBinding, Result};
use serde::{Deserialize, Serialize};
use wasm_bindgen::{JsCast, JsValue};
use wasm_bindgen_futures::JsFuture;
use worker_sys::RateLimiter as RateLimiterSys;

pub struct RateLimiter(RateLimiterSys);

#[derive(Serialize, Deserialize)]
struct RateLimitOptions {
key: String,
}

#[derive(Serialize, Deserialize)]
pub struct RateLimitOutcome {
pub success: bool,
}

unsafe impl Send for RateLimiter {}
unsafe impl Sync for RateLimiter {}

impl EnvBinding for RateLimiter {
const TYPE_NAME: &'static str = "RateLimiter";
}
impl RateLimiter {
pub async fn limit(&self, key: String) -> Result<RateLimitOutcome> {
let arg = serde_wasm_bindgen::to_value(&RateLimitOptions { key })?;
let promise = self.0.limit(arg.into())?;
let fut = SendFuture::new(JsFuture::from(promise));
let result = fut.await?;
let outcome = serde_wasm_bindgen::from_value(result)?;
Ok(outcome)
}
}

impl JsCast for RateLimiter {
fn instanceof(val: &JsValue) -> bool {
val.is_instance_of::<RateLimiterSys>()
}

fn unchecked_from_js(val: JsValue) -> Self {
Self(val.into())
}

fn unchecked_from_js_ref(val: &JsValue) -> &Self {
unsafe { &*(val as *const JsValue as *const Self) }
}
}

impl From<RateLimiter> for JsValue {
fn from(limiter: RateLimiter) -> Self {
JsValue::from(limiter.0)
}
}

impl AsRef<JsValue> for RateLimiter {
fn as_ref(&self) -> &JsValue {
&self.0
}
}

impl From<RateLimiterSys> for RateLimiter {
fn from(inner: RateLimiterSys) -> Self {
Self(inner)
}
}

0 comments on commit 9fc2fbd

Please sign in to comment.