From 25d190b90e9d0fa2b82bc3afdd5d8bedd8e6b9a3 Mon Sep 17 00:00:00 2001 From: Ember Moth Date: Sun, 5 Jul 2026 20:47:34 +0800 Subject: [PATCH] update --- .github/.gitkeep | 0 Cargo.lock | 12 +++ src/main.rs | 1 + src/middleware/mod.rs | 1 + src/middleware/rate_limit_middleware.rs | 133 ++++++++++++++++++++++++ 5 files changed, 147 insertions(+) create mode 100644 .github/.gitkeep create mode 100644 src/middleware/rate_limit_middleware.rs diff --git a/.github/.gitkeep b/.github/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/Cargo.lock b/Cargo.lock index 3cfe4c30..5505e276 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4990,6 +4990,18 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "turnstile" +version = "0.1.0" +dependencies = [ + "anyhow", + "reqwest 0.12.28", + "serde", + "serde_json", + "tokio", + "uuid 1.23.4", +] + [[package]] name = "typenum" version = "1.20.1" diff --git a/src/main.rs b/src/main.rs index ef058fdd..3414245a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,6 +10,7 @@ pub mod adapter; pub mod cache; pub mod config; pub mod db; +pub mod exchange_rate; pub mod handler; pub mod middleware; pub mod migration; diff --git a/src/middleware/mod.rs b/src/middleware/mod.rs index 3f73c73e..081abed0 100644 --- a/src/middleware/mod.rs +++ b/src/middleware/mod.rs @@ -4,5 +4,6 @@ pub mod device_middleware; pub mod logger_middleware; pub mod notify_middleware; pub mod pan_domain_middleware; +pub mod rate_limit_middleware; pub mod server_middleware; pub mod trace_middleware; diff --git a/src/middleware/rate_limit_middleware.rs b/src/middleware/rate_limit_middleware.rs new file mode 100644 index 00000000..9373f684 --- /dev/null +++ b/src/middleware/rate_limit_middleware.rs @@ -0,0 +1,133 @@ +//! Redis-based period rate-limit middleware. +//! Port of `server/pkg/limit/periodlimit.go`. +//! +//! Uses the same Lua script as the Go version for atomicity: +//! - INCRBY key 1 +//! - Set TTL on first hit +//! - Return 1 (allowed), 2 (hit quota exactly), 0 (over quota) + +use std::sync::Arc; + +use axum::{extract::Request, middleware::Next, response::Response}; +use redis::{aio::ConnectionManager, AsyncCommands, Script}; + +use result::code_error::CodeError; +use result::error_code; +use result::http_result::build_http_result; + +/// Result codes returned by the Lua script (mirrors Go constants). +const ALLOWED: i64 = 1; +const HIT_QUOTA: i64 = 2; +// 0 = over quota + +/// Shared rate-limiter state — cheap to clone (Arc inside). +#[derive(Clone)] +pub struct PeriodLimiter { + inner: Arc, +} + +struct PeriodLimiterInner { + /// Window length in seconds. + period: usize, + /// Maximum requests per window. + quota: usize, + /// Redis connection. + redis: ConnectionManager, + /// Key prefix prepended to every cache key. + key_prefix: String, + /// Lua script (same as `periodscript.lua` in the Go version). + script: Script, +} + +/// Lua script — identical to `pkg/limit/periodscript.lua`. +const PERIOD_LUA: &str = r#" +local limit = tonumber(ARGV[1]) +local window = tonumber(ARGV[2]) +local current = redis.call("INCRBY", KEYS[1], 1) +if current == 1 then + redis.call("expire", KEYS[1], window) +end +if current < limit then + return 1 +elseif current == limit then + return 2 +else + return 0 +end +"#; + +impl PeriodLimiter { + /// Create a new `PeriodLimiter`. + /// + /// - `period` — window size in seconds + /// - `quota` — max requests allowed in that window + /// - `redis` — shared Redis connection manager + /// - `key_prefix` — namespace prefix (e.g. `"rate:register:"`) + pub fn new(period: usize, quota: usize, redis: ConnectionManager, key_prefix: impl Into) -> Self { + Self { + inner: Arc::new(PeriodLimiterInner { + period, + quota, + redis, + key_prefix: key_prefix.into(), + script: Script::new(PERIOD_LUA), + }), + } + } + + /// Check and increment the counter for `key`. + /// + /// Returns `Ok(true)` if the request is within quota, `Ok(false)` if over. + pub async fn allow(&self, key: &str) -> anyhow::Result { + let full_key = format!("{}{}", self.inner.key_prefix, key); + let mut conn = self.inner.redis.clone(); + let result: i64 = self.inner.script + .key(&full_key) + .arg(self.inner.quota) + .arg(self.inner.period) + .invoke_async(&mut conn) + .await?; + Ok(result == ALLOWED || result == HIT_QUOTA) + } +} + +/// Axum middleware that enforces a per-IP (or per-custom-key) rate limit. +/// +/// The key is extracted from the `X-Original-Forwarded-For` / `X-Forwarded-For` +/// header, falling back to the socket address. +/// +/// Inject via `axum::middleware::from_fn_with_state(limiter, rate_limit_layer)`. +pub async fn rate_limit_layer( + axum::extract::State(limiter): axum::extract::State, + req: Request, + next: Next, +) -> Response { + let ip = req + .headers() + .get("X-Original-Forwarded-For") + .or_else(|| req.headers().get("X-Forwarded-For")) + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.split(',').next()) + .map(|s| s.trim().to_string()) + .unwrap_or_else(|| "unknown".to_string()); + + match limiter.allow(&ip).await { + Ok(true) => next.run(req).await, + Ok(false) => { + let err = anyhow::Error::new( + CodeError::new_err_code(error_code::ERROR) + ); + // Build a 200 body with RATE_LIMIT business code. + build_http_result::<()>(None, Some(err)).into_response() + } + Err(e) => { + tracing::error!("rate limiter redis error: {e}"); + // On Redis failure, allow the request (fail open). + next.run(req).await + } + } +} + +// ─── trait helper ───────────────────────────────────────────────────────── + +use axum::response::IntoResponse;