From 8a4cb649e34378f442a6c69b82c57366f01a4718 Mon Sep 17 00:00:00 2001 From: shenlan Date: Mon, 4 Aug 2025 00:00:49 +0800 Subject: [PATCH] feat: add flexible askai rate limiter --- .github/workflows/build-askai-limiter.yml | 46 +++++++++++++ Makefile | 84 +++++++++++++---------- wasm/askai_limiter/README.md | 27 +++++++- wasm/askai_limiter/src/lib.rs | 79 +++++++++++++++++---- 4 files changed, 182 insertions(+), 54 deletions(-) create mode 100644 .github/workflows/build-askai-limiter.yml diff --git a/.github/workflows/build-askai-limiter.yml b/.github/workflows/build-askai-limiter.yml new file mode 100644 index 0000000..95368d3 --- /dev/null +++ b/.github/workflows/build-askai-limiter.yml @@ -0,0 +1,46 @@ +name: Build and Release Askai Limiter + +on: + workflow_dispatch: + push: + paths: + - 'wasm/askai_limiter/**' + - '.github/workflows/build-askai-limiter.yml' + - 'Makefile' + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions-rs/toolchain@v1 + with: + toolchain: stable + target: wasm32-wasip1 + profile: minimal + override: true + - name: Build Wasm Module + run: make wasm-askai-limiter + - name: Upload artifact + uses: actions/upload-artifact@v4 + with: + name: askai_limiter.wasm + path: build/askai_limiter.wasm + + release: + runs-on: ubuntu-latest + needs: build + steps: + - uses: actions/checkout@v4 + - uses: actions/download-artifact@v4 + with: + name: askai_limiter.wasm + path: release + - name: Publish GitHub Release + uses: softprops/action-gh-release@v1 + with: + tag_name: askai-limiter-${{ github.run_number }} + name: Askai Limiter ${{ github.run_number }} + files: release/askai_limiter.wasm + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/Makefile b/Makefile index 81174ae..20d0eb0 100644 --- a/Makefile +++ b/Makefile @@ -2,72 +2,82 @@ GO=go APP_NAME=xcontrol +WASM_TARGET=wasm32-wasip1 +WASM_MODULE=askai_limiter + UNAME_S := $(shell uname -s) UNAME_M := $(shell uname -m) .PHONY: all build agent vet test clean release \ - macos-x64 macos-arm64 windows-x64 linux-x64 linux-arm64 macos linux + macos-x64 macos-arm64 windows-x64 linux-x64 linux-arm64 macos linux \ + wasm-askai-limiter all: vet test build build: -$(GO) build -o bin/$(APP_NAME) ./cmd/api + $(GO) build -o bin/$(APP_NAME) ./cmd/api agent: -$(GO) build -o bin/$(APP_NAME)-agent ./cmd/agent + $(GO) build -o bin/$(APP_NAME)-agent ./cmd/agent vet: -$(GO) vet ./... + $(GO) vet ./... test: -$(GO) test ./... + $(GO) test ./... clean: -rm -rf bin build + rm -rf bin build wasm/$(WASM_MODULE)/target + +wasm-askai-limiter: + rustup target add $(WASM_TARGET) + cargo build --release --target $(WASM_TARGET) --manifest-path wasm/$(WASM_MODULE)/Cargo.toml + mkdir -p build + cp wasm/$(WASM_MODULE)/target/$(WASM_TARGET)/release/$(WASM_MODULE).wasm build/$(WASM_MODULE).wasm macos-x64: -@if [ "$(UNAME_S)" = "Darwin" ] && [ "$(UNAME_M)" = "x86_64" ]; then \ -GOOS=darwin GOARCH=amd64 $(GO) build -o build/macos-x64/$(APP_NAME) ./cmd/api; \ -else \ -echo "macos-x64 build requires macOS x86_64"; \ -fi + @if [ "$(UNAME_S)" = "Darwin" ] && [ "$(UNAME_M)" = "x86_64" ]; then \ + GOOS=darwin GOARCH=amd64 $(GO) build -o build/macos-x64/$(APP_NAME) ./cmd/api; \ + else \ + echo "macos-x64 build requires macOS x86_64"; \ + fi macos-arm64: -@if [ "$(UNAME_S)" = "Darwin" ] && [ "$(UNAME_M)" = "arm64" ]; then \ -GOOS=darwin GOARCH=arm64 $(GO) build -o build/macos-arm64/$(APP_NAME) ./cmd/api; \ -else \ -echo "macos-arm64 build requires macOS arm64"; \ -fi + @if [ "$(UNAME_S)" = "Darwin" ] && [ "$(UNAME_M)" = "arm64" ]; then \ + GOOS=darwin GOARCH=arm64 $(GO) build -o build/macos-arm64/$(APP_NAME) ./cmd/api; \ + else \ + echo "macos-arm64 build requires macOS arm64"; \ + fi windows-x64: -GOOS=windows GOARCH=amd64 $(GO) build -o build/windows-x64/$(APP_NAME).exe ./cmd/api + GOOS=windows GOARCH=amd64 $(GO) build -o build/windows-x64/$(APP_NAME).exe ./cmd/api linux-x64: -GOOS=linux GOARCH=amd64 $(GO) build -o build/linux-x64/$(APP_NAME) ./cmd/api + GOOS=linux GOARCH=amd64 $(GO) build -o build/linux-x64/$(APP_NAME) ./cmd/api linux-arm64: -GOOS=linux GOARCH=arm64 $(GO) build -o build/linux-arm64/$(APP_NAME) ./cmd/api + GOOS=linux GOARCH=arm64 $(GO) build -o build/linux-arm64/$(APP_NAME) ./cmd/api -release: macos-x64 macos-arm64 windows-x64 linux-x64 linux-arm64 +release: macos-x64 macos-arm64 windows-x64 linux-x64 linux-arm64 wasm-askai-limiter macos: -@if [ "$(UNAME_S)" = "Darwin" ]; then \ -if [ "$(UNAME_M)" = "arm64" ]; then \ -GOOS=darwin GOARCH=arm64 $(GO) build -o build/macos-arm64/$(APP_NAME) ./cmd/api; \ -else \ -GOOS=darwin GOARCH=amd64 $(GO) build -o build/macos-x64/$(APP_NAME) ./cmd/api; \ -fi; \ -else \ -echo "macos build requires macOS host"; \ -fi + @if [ "$(UNAME_S)" = "Darwin" ]; then \ + if [ "$(UNAME_M)" = "arm64" ]; then \ + GOOS=darwin GOARCH=arm64 $(GO) build -o build/macos-arm64/$(APP_NAME) ./cmd/api; \ + else \ + GOOS=darwin GOARCH=amd64 $(GO) build -o build/macos-x64/$(APP_NAME) ./cmd/api; \ + fi; \ + else \ + echo "macos build requires macOS host"; \ + fi linux: -@if [ "$(UNAME_S)" = "Linux" ]; then \ -if [ "$(UNAME_M)" = "aarch64" ] || [ "$(UNAME_M)" = "arm64" ]; then \ -GOOS=linux GOARCH=arm64 $(GO) build -o build/linux-arm64/$(APP_NAME) ./cmd/api; \ -else \ -GOOS=linux GOARCH=amd64 $(GO) build -o build/linux-x64/$(APP_NAME) ./cmd/api; \ -fi; \ -else \ -echo "linux build requires Linux host"; \ + @if [ "$(UNAME_S)" = "Linux" ]; then \ + if [ "$(UNAME_M)" = "aarch64" ] || [ "$(UNAME_M)" = "arm64" ]; then \ + GOOS=linux GOARCH=arm64 $(GO) build -o build/linux-arm64/$(APP_NAME) ./cmd/api; \ + else \ + GOOS=linux GOARCH=amd64 $(GO) build -o build/linux-x64/$(APP_NAME) ./cmd/api; \ + fi; \ + else \ + echo "linux build requires Linux host"; \ fi diff --git a/wasm/askai_limiter/README.md b/wasm/askai_limiter/README.md index 77ef37a..c5c4103 100644 --- a/wasm/askai_limiter/README.md +++ b/wasm/askai_limiter/README.md @@ -2,7 +2,28 @@ 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. +It supports two counting modes: + +- **Rolling** – requests are counted within a sliding time window. The default + configuration limits to **200** requests every **24** hours. +- **Unlimited** – a fixed quota that never resets until the limit is exhausted. + +### Configuration + +The module accepts a simple comma‑separated configuration string when loaded: + +``` +limit=200,window=86400 +``` + +- `limit` – maximum allowed requests. +- `window` – (optional) duration of the rolling window in seconds. Omit this + field to enable the unlimited mode. + +Examples: + +- `limit=1000` – unlimited counting with a quota of 1000 total requests. +- `limit=1400,window=604800` – 1400 requests allowed every 7 days (168 hours). ## Build @@ -46,5 +67,5 @@ http { } ``` -Requests beyond the first 200 in a single day will return HTTP 429 with the -body `{"error":"API daily limit reached"}`. +Requests exceeding the configured quota will return HTTP 429 with the body +`{"error":"API limit reached"}`. diff --git a/wasm/askai_limiter/src/lib.rs b/wasm/askai_limiter/src/lib.rs index 134bc57..913cf66 100644 --- a/wasm/askai_limiter/src/lib.rs +++ b/wasm/askai_limiter/src/lib.rs @@ -2,53 +2,104 @@ use proxy_wasm::traits::*; use proxy_wasm::types::*; use std::time::{SystemTime, UNIX_EPOCH}; -struct AskaiLimiter; +#[derive(Clone, Copy)] +struct Config { + limit: u32, + window: Option, // seconds; None means unlimited +} + +impl Default for Config { + fn default() -> Self { + Self { + limit: 200, + window: Some(86_400), + } + } +} + +struct AskaiLimiter { + config: Config, +} + +struct AskaiLimiterRoot { + config: Config, +} 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() + let now = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() - .as_secs() - / 86_400; - let key = format!("askai:{}", today); + .as_secs(); + let key = match self.config.window { + Some(window) => format!("askai:{}", now / window), + None => "askai:total".to_string(), + }; - // 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::().ok()) .unwrap_or(0); - if count >= 200 { + if count >= self.config.limit { self.send_http_response( 429, vec![("Content-Type", "application/json")], - Some(b"{\"error\":\"API daily limit reached\"}"), + Some(b"{\"error\":\"API 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 { +impl Context for AskaiLimiterRoot {} + +impl RootContext for AskaiLimiterRoot { + fn on_configure(&mut self, _size: usize) -> bool { + if let Some(config_bytes) = self.get_plugin_configuration() { + if let Ok(text) = String::from_utf8(config_bytes) { + self.config = parse_config(&text, self.config); + } + } true } fn create_http_context(&self, _context_id: u32) -> Option> { - Some(Box::new(AskaiLimiter)) + Some(Box::new(AskaiLimiter { config: self.config })) } } +fn parse_config(text: &str, mut cfg: Config) -> Config { + for part in text.split(',') { + let mut kv = part.splitn(2, '='); + let key = kv.next().unwrap_or("").trim(); + let val = kv.next().unwrap_or("").trim(); + match key { + "limit" => { + if let Ok(v) = val.parse::() { + cfg.limit = v; + } + } + "window" | "period" => { + if let Ok(v) = val.parse::() { + cfg.window = Some(v); + } + } + _ => {} + } + } + cfg +} + proxy_wasm::main! {{ - proxy_wasm::set_http_context(|_, _| Box::new(AskaiLimiter)); + proxy_wasm::set_root_context(|_| Box::new(AskaiLimiterRoot { + config: Config::default(), + })); }}