From 9fc2fbd0f3dd282e9ac7ac5ed203e635d6f98a1d Mon Sep 17 00:00:00 2001
From: Kevin Flansburg <kflansburg@cloudflare.com>
Date: Thu, 26 Sep 2024 08:10:55 -0400
Subject: [PATCH] Implement Rate Limiter Binding (#603)

* Implement Rate Limiter binding

* Make success field public
---
 worker-sys/src/types.rs            |  2 +
 worker-sys/src/types/rate_limit.rs | 12 ++++++
 worker/src/lib.rs                  |  2 +
 worker/src/rate_limit.rs           | 66 ++++++++++++++++++++++++++++++
 4 files changed, 82 insertions(+)
 create mode 100644 worker-sys/src/types/rate_limit.rs
 create mode 100644 worker/src/rate_limit.rs

diff --git a/worker-sys/src/types.rs b/worker-sys/src/types.rs
index 76235254..e808e984 100644
--- a/worker-sys/src/types.rs
+++ b/worker-sys/src/types.rs
@@ -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;
@@ -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::*;
diff --git a/worker-sys/src/types/rate_limit.rs b/worker-sys/src/types/rate_limit.rs
new file mode 100644
index 00000000..0cfda3c2
--- /dev/null
+++ b/worker-sys/src/types/rate_limit.rs
@@ -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>;
+}
diff --git a/worker/src/lib.rs b/worker/src/lib.rs
index bbaac220..ddeccf28 100644
--- a/worker/src/lib.rs
+++ b/worker/src/lib.rs
@@ -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};
@@ -218,6 +219,7 @@ mod hyperdrive;
 #[cfg(feature = "queue")]
 mod queue;
 mod r2;
+mod rate_limit;
 mod request;
 mod request_init;
 mod response;
diff --git a/worker/src/rate_limit.rs b/worker/src/rate_limit.rs
new file mode 100644
index 00000000..64dd8513
--- /dev/null
+++ b/worker/src/rate_limit.rs
@@ -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)
+    }
+}