From 496e51a92c436f7586a0d230708aa3375891a1d8 Mon Sep 17 00:00:00 2001
From: shenlan <manbuzhe2009@qq.com>
Date: Sun, 3 Aug 2025 20:21:37 +0800
Subject: [PATCH] Add Askai limiter Proxy-Wasm module

---
 deploy/nginx/askai-limiter.conf | 24 +++++++++++++++
 wasm/askai_limiter/Cargo.toml   | 10 ++++++
 wasm/askai_limiter/README.md    | 50 ++++++++++++++++++++++++++++++
 wasm/askai_limiter/src/lib.rs   | 54 +++++++++++++++++++++++++++++++++
 4 files changed, 138 insertions(+)
 create mode 100644 deploy/nginx/askai-limiter.conf
 create mode 100644 wasm/askai_limiter/Cargo.toml
 create mode 100644 wasm/askai_limiter/README.md
 create mode 100644 wasm/askai_limiter/src/lib.rs

diff --git a/deploy/nginx/askai-limiter.conf b/deploy/nginx/askai-limiter.conf
new file mode 100644
index 0000000..89596e2
--- /dev/null
+++ b/deploy/nginx/askai-limiter.conf
@@ -0,0 +1,24 @@
+load_module modules/ngx_http_wasm_module.so;
+
+http {
+    wasm {
+        module limiter /etc/nginx/wasm/askai_limiter.wasm;
+    }
+
+    server {
+        listen 443 ssl;
+        server_name cn-homepage.svc.plus;
+
+        ssl_certificate /etc/ssl/svc.plus.pem;
+        ssl_certificate_key /etc/ssl/svc.plus.rsa.key;
+
+        location /api/askai {
+            wasm_call limiter;
+            proxy_pass http://127.0.0.1:8080/api/askai;
+        }
+
+        location / {
+            try_files $uri $uri/ /index.html;
+        }
+    }
+}
diff --git a/wasm/askai_limiter/Cargo.toml b/wasm/askai_limiter/Cargo.toml
new file mode 100644
index 0000000..becc049
--- /dev/null
+++ b/wasm/askai_limiter/Cargo.toml
@@ -0,0 +1,10 @@
+[package]
+name = "askai_limiter"
+version = "0.1.0"
+edition = "2021"
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+proxy-wasm = "0.2"
diff --git a/wasm/askai_limiter/README.md b/wasm/askai_limiter/README.md
new file mode 100644
index 0000000..77ef37a
--- /dev/null
+++ b/wasm/askai_limiter/README.md
@@ -0,0 +1,50 @@
+# Askai Limiter Proxy-Wasm Module
+
+This module provides a simple API rate limiter for Nginx using the experimental
+`ngx_http_wasm_module` and the [proxy-wasm-rust-sdk](https://github.com/proxy-wasm/proxy-wasm-rust-sdk).
+It enforces a global daily limit of **200** requests per API endpoint.
+
+## Build
+
+```bash
+rustup target add wasm32-wasip1
+cargo build --release --target wasm32-wasip1
+```
+
+The compiled module will be located at
+`target/wasm32-wasip1/release/askai_limiter.wasm`.
+
+## Nginx Configuration
+
+Example snippet that loads the compiled module and applies it to the
+`/api/askai` route:
+
+```nginx
+load_module modules/ngx_http_wasm_module.so;
+
+http {
+    wasm {
+        module limiter /etc/nginx/wasm/askai_limiter.wasm;
+    }
+
+    server {
+        listen 443 ssl;
+        server_name cn-homepage.svc.plus;
+
+        ssl_certificate /etc/ssl/svc.plus.pem;
+        ssl_certificate_key /etc/ssl/svc.plus.rsa.key;
+
+        location /api/askai {
+            wasm_call limiter;
+            proxy_pass http://127.0.0.1:8080/api/askai;
+        }
+
+        location / {
+            try_files $uri $uri/ /index.html;
+        }
+    }
+}
+```
+
+Requests beyond the first 200 in a single day will return HTTP 429 with the
+body `{"error":"API daily limit reached"}`.
diff --git a/wasm/askai_limiter/src/lib.rs b/wasm/askai_limiter/src/lib.rs
new file mode 100644
index 0000000..134bc57
--- /dev/null
+++ b/wasm/askai_limiter/src/lib.rs
@@ -0,0 +1,54 @@
+use proxy_wasm::traits::*;
+use proxy_wasm::types::*;
+use std::time::{SystemTime, UNIX_EPOCH};
+
+struct AskaiLimiter;
+
+impl Context for AskaiLimiter {}
+
+impl HttpContext for AskaiLimiter {
+    fn on_http_request_headers(&mut self, _num_headers: usize, _end_of_stream: bool) -> Action {
+        // Use the day since UNIX epoch as the key
+        let today = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_secs()
+            / 86_400;
+        let key = format!("askai:{}", today);
+
+        // Read the current count from shared data
+        let (data, _cas) = self.get_shared_data(&key);
+        let mut count = data
+            .and_then(|d| String::from_utf8(d).ok())
+            .and_then(|s| s.parse::<u32>().ok())
+            .unwrap_or(0);
+
+        if count >= 200 {
+            self.send_http_response(
+                429,
+                vec![("Content-Type", "application/json")],
+                Some(b"{\"error\":\"API daily limit reached\"}"),
+            );
+            return Action::Pause;
+        }
+
+        // Increment and store the updated count
+        count += 1;
+        let _ = self.set_shared_data(&key, Some(count.to_string().as_bytes()), None);
+        Action::Continue
+    }
+}
+
+impl RootContext for AskaiLimiter {
+    fn on_configure(&mut self, _configuration_size: usize) -> bool {
+        true
+    }
+
+    fn create_http_context(&self, _context_id: u32) -> Option<Box<dyn HttpContext>> {
+        Some(Box::new(AskaiLimiter))
+    }
+}
+
+proxy_wasm::main! {{
+    proxy_wasm::set_http_context(|_, _| Box::new(AskaiLimiter));
+}}