pingap-plugin 0.15.0

// Copyright 2024-2025 Tree xie.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use super::{
    Error, get_hash_key, get_int_conf, get_int_conf_or_default,
    get_step_conf_in, get_str_conf,
};
use async_trait::async_trait;
use http::header::RETRY_AFTER;
use http::{HeaderValue, StatusCode};
use humantime::parse_duration;
use pingap_config::{PluginCategory, PluginConf};
use pingap_core::{
    Ctx, HttpHeader, HttpResponse, Inflight, Plugin, PluginStep, Rate,
    RequestPluginResult,
};
use pingap_core::{
    ensure_client_ip, get_cookie_value, get_query_value, get_req_header_value,
};
use pingora::proxy::Session;
use std::borrow::Cow;
use std::time::Duration;
use tracing::debug;

type Result<T, E = Error> = std::result::Result<T, E>;

// LimitTag determines what value will be used as the rate limiting key
#[derive(PartialEq, Eq, Debug)]
pub enum LimitTag {
    Ip,            // Use client IP (from X-Forwarded-For or direct connection)
    RequestHeader, // Use value from a specified HTTP request header
    Cookie,        // Use value from a specified cookie
    Query,         // Use value from a specified URL query parameter
}

// Limiter implements rate limiting and concurrent request limiting
// It can be configured via TOML with settings like:
// ```toml
// type = "rate"          # or "inflight"
// tag = "cookie"         # or "header", "query", "ip"
// key = "session_id"     # name of header/cookie/query param to use
// max = 100             # maximum requests allowed
// interval = "60s"      # time window for rate limiting
// ```
/// A rate limiter or concurrent request limiter that can be configured to limit based on
/// different request attributes (IP, headers, cookies, query params)
pub struct Limiter {
    /// Determines what value will be used as the rate limiting key (IP, header, cookie, or query param)
    tag: LimitTag,

    /// Maximum number of requests/connections allowed within the interval (for rate limiting)
    /// or at the same time (for inflight limiting)
    max: f64,

    /// The name of the header/cookie/query parameter to use as the limiting key
    /// Only used when tag is not LimitTag::Ip
    key: String,

    /// Tracks concurrent requests using atomic counters.
    /// When a request completes, the counter automatically decrements via RAII guard.
    /// Only used when configured as an inflight limiter (type = "inflight")
    inflight: Option<Inflight>,

    /// Tracks request counts over a sliding time window.
    /// Automatically expires old requests based on configured interval.
    /// Only used when configured as a rate limiter (type = "rate")
    rate: Option<Rate>,

    /// When to apply the limiting logic:
    /// - PluginStep::Request: During initial request processing
    /// - PluginStep::ProxyUpstream: Before forwarding to upstream server
    plugin_step: PluginStep,

    /// Unique identifier for this limiter instance, used to distinguish between
    /// different limiters in the same application
    hash_value: String,

    /// The weight of current slot
    weight: f64,

    /// `Retry-After` for a rate limiter's 429: the window length, the
    /// soonest the budget can have moved on.
    retry_after: Option<HttpHeader>,
}

/// Converts a plugin configuration into a Limiter instance
///
/// # Arguments
/// * `value` - Plugin configuration containing limiter settings
///
/// # Returns
/// * `Result<Self>` - New Limiter instance or error if configuration is invalid
///
/// # Configuration Options
/// * `type` - "rate" or "inflight"
/// * `tag` - "ip", "cookie", "header", or "query"
/// * `key` - Name of header/cookie/query parameter to use
/// * `max` - Maximum allowed requests/connections
/// * `interval` - Time window for rate limiting (e.g. "60s")
impl TryFrom<&PluginConf> for Limiter {
    type Error = Error;
    fn try_from(value: &PluginConf) -> Result<Self> {
        let hash_value = get_hash_key(value);
        let category = PluginCategory::Limit.to_string();
        let invalid = |message: String| Error::Invalid {
            category: category.clone(),
            message,
        };
        // Limiting only makes sense before the upstream is involved.
        let step = get_step_conf_in(
            value,
            &category,
            PluginStep::Request,
            &[PluginStep::Request, PluginStep::ProxyUpstream],
        )?;

        // Every setting here decides who gets limited, so a value that is
        // not one of the documented ones is an error rather than a silent
        // fallback: a misspelt `tag` used to limit by ip, a misspelt `type`
        // rate-limited, and a forgotten `max` rejected almost everything.
        let tag = match get_str_conf(value, "tag").as_str() {
            "" | "ip" => LimitTag::Ip,
            "cookie" => LimitTag::Cookie,
            "header" => LimitTag::RequestHeader,
            "query" => LimitTag::Query,
            other => {
                return Err(invalid(format!(
                    "Invalid tag({other}), expect ip, header, cookie or query"
                )));
            },
        };
        let key = get_str_conf(value, "key");
        if tag != LimitTag::Ip && key.is_empty() {
            return Err(invalid(
                "key is required for a header, cookie or query tag".to_string(),
            ));
        }
        let is_inflight = match get_str_conf(value, "type").as_str() {
            "" | "rate" => false,
            "inflight" => true,
            other => {
                return Err(invalid(format!(
                    "Invalid type({other}), expect rate or inflight"
                )));
            },
        };
        if !value.contains_key("max") {
            return Err(invalid("max is required".to_string()));
        }
        let max = get_int_conf(value, "max");
        if max < 0 {
            return Err(invalid("max must not be negative".to_string()));
        }

        // Parse time interval for rate limiting
        // Format examples: "10s", "1m", "2h"
        // Default: 10 seconds if not specified
        let interval = get_str_conf(value, "interval");
        let interval = if !interval.is_empty() {
            parse_duration(&interval).map_err(|e| invalid(e.to_string()))?
        } else {
            Duration::from_secs(10)
        };
        if interval.is_zero() {
            return Err(invalid(
                "interval must be greater than zero".to_string(),
            ));
        }

        // Create either inflight or rate limiter based on config
        let mut inflight = None;
        let mut rate = None;
        let mut retry_after = None;
        let mut max = max as f64;
        if is_inflight {
            // Inflight limiter uses atomic counters to track concurrent requests
            inflight = Some(Inflight::new());
        } else {
            // convert it to rps
            max /= interval.as_secs_f64().max(1.0);
            // Rate limiter uses time-bucketed counters
            rate = Some(Rate::new(interval));
            retry_after = Some((
                RETRY_AFTER,
                HeaderValue::from(interval.as_secs().max(1)),
            ));
        }

        let weight = get_int_conf_or_default(value, "weight", 50).clamp(0, 100)
            as f64
            / 100.0;

        Ok(Self {
            hash_value,
            tag,
            key,
            max,
            inflight,
            rate,
            plugin_step: step,
            weight,
            retry_after,
        })
    }
}

impl Limiter {
    /// Creates a new Limiter instance from plugin configuration
    ///
    /// # Arguments
    /// * `params` - Plugin configuration containing limiter settings like type, tag, key, max, etc.
    ///
    /// # Returns
    /// * `Result<Self>` - New Limiter instance or error if configuration is invalid
    ///
    /// # Example Configuration
    /// ```toml
    /// type = "rate"          # or "inflight"
    /// tag = "cookie"         # or "header", "query", "ip"
    /// key = "session_id"     # name of header/cookie/query param to use
    /// max = 100             # maximum requests allowed
    /// interval = "60s"      # time window for rate limiting
    /// ```
    pub fn new(params: &PluginConf) -> Result<Self> {
        debug!(params = params.to_string(), "new limit plugin");
        Self::try_from(params)
    }
    /// Increments and checks the limit counter for the current request
    ///
    /// # Arguments
    /// * `session` - The HTTP session containing request details
    /// * `ctx` - Mutable state context for storing request data
    ///
    /// # Returns
    /// * `Result<()>` - Ok if within limits, Error if limit exceeded
    ///
    /// # Effects
    /// * For rate limiting: Records request in time window
    /// * For inflight limiting: Increments counter and stores RAII guard in context
    /// * For IP-based limiting: Stores client IP in context
    pub fn incr(&self, session: &Session, ctx: &mut Ctx) -> Result<()> {
        // Extract the key value based on configured tag type.
        // Borrow where possible — Rate/Inflight only need `Hash`, not an owned String.
        let key: Cow<'_, str> = match self.tag {
            LimitTag::Query => Cow::Borrowed(
                get_query_value(session.req_header(), &self.key)
                    .unwrap_or_default(),
            ),
            LimitTag::RequestHeader => Cow::Borrowed(
                get_req_header_value(session.req_header(), &self.key)
                    .unwrap_or_default(),
            ),
            LimitTag::Cookie => Cow::Borrowed(
                get_cookie_value(session.req_header(), &self.key)
                    .unwrap_or_default(),
            ),
            _ => Cow::Borrowed(ensure_client_ip(session, ctx)),
        };

        // Skip limiting if no key found (e.g., missing header/cookie)
        if key.is_empty() {
            return Ok(());
        }

        // Track request based on limiter type.
        // Pass `&Cow` (Sized) rather than `&str` — pingora-limits requires `T: Hash + Sized`.
        let value = if let Some(rate) = &self.rate {
            // For rate limiting:
            rate.observe(&key, 1); // Record this request
            if self.weight > 0.0 {
                rate.rate_with(&key, |rate_info| {
                    let prev =
                        rate_info.prev_samples as f64 * (1. - self.weight);
                    let curr = rate_info.curr_samples as f64 * self.weight;
                    (prev + curr) / rate_info.interval.as_secs_f64()
                })
            } else {
                rate.rate(&key) // get the per second rate estimation of previous time window
            }
        } else if let Some(inflight) = &self.inflight {
            // For inflight limiting:
            // Increment counter
            // Store guard in context - when guard is dropped, counter auto-decrements
            let (guard, value) = inflight.incr(&key, 1);
            ctx.state.guard = Some(guard);
            value as f64
        } else {
            0.0
        };

        // Check if limit exceeded
        if value > self.max {
            return Err(Error::Exceed {
                category: PluginCategory::Limit.to_string(),
                max: self.max,
                value,
            });
        }
        Ok(())
    }
}

#[async_trait]
impl Plugin for Limiter {
    /// Returns unique identifier for this limiter instance
    #[inline]
    fn config_key(&self) -> Cow<'_, str> {
        Cow::Borrowed(&self.hash_value)
    }

    /// Handles incoming HTTP requests by applying configured limits
    ///
    /// # Arguments
    /// * `step` - Current plugin execution step
    /// * `session` - Mutable HTTP session
    /// * `ctx` - Mutable state context
    ///
    /// # Returns
    /// * `pingora::Result<Option<HttpResponse>>` - None to continue processing,
    ///   Some(response) with 429 status if limit exceeded
    ///
    /// # Effects
    /// * Increments and checks appropriate limit counter
    /// * Returns 429 Too Many Requests if limit exceeded
    #[inline]
    async fn handle_request(
        &self,
        step: PluginStep,
        session: &mut Session,
        ctx: &mut Ctx,
    ) -> pingora::Result<RequestPluginResult> {
        // Only run at configured plugin step
        if step != self.plugin_step {
            return Ok(RequestPluginResult::Skipped);
        }

        // Try to increment counter
        if let Err(e) = self.incr(session, ctx) {
            // If limit exceeded, return 429 Too Many Requests
            return Ok(RequestPluginResult::Respond(HttpResponse {
                status: StatusCode::TOO_MANY_REQUESTS,
                headers: self.retry_after.clone().map(|header| vec![header]),
                body: e.to_string().into(),
                ..Default::default()
            }));
        }

        // Continue normal request processing if within limits
        Ok(RequestPluginResult::Continue)
    }
}

register_plugin!("limit", Limiter);

#[cfg(test)]
mod tests {
    use super::*;
    use http::StatusCode;
    use pingap_config::PluginConf;
    use pingap_core::{Ctx, PluginStep};
    use pingora::proxy::Session;
    use pretty_assertions::assert_eq;
    use std::time::Duration;
    use tokio_test::io::Builder;

    async fn new_session() -> Session {
        let headers = [
            "Host: github.com",
            "Referer: https://github.com/",
            "User-Agent: pingap/0.1.1",
            "Cookie: deviceId=abc",
            "Accept: application/json",
            "X-Uuid: 138q71",
            "X-Forwarded-For: 1.1.1.1, 192.168.1.2",
        ]
        .join("\r\n");
        let input_header =
            format!("GET /vicanso/pingap?key=1 HTTP/1.1\r\n{headers}\r\n\r\n");
        let mock_io = Builder::new().read(input_header.as_bytes()).build();
        let mut session = Session::new_h1(Box::new(mock_io));
        session.read_request().await.unwrap();
        session
    }

    #[test]
    fn test_limit_params() {
        let params = Limiter::try_from(
            &toml::from_str::<PluginConf>(
                r###"
type = "inflight"
tag = "cookie"
key = "deviceId"
max = 10
"###,
            )
            .unwrap(),
        )
        .unwrap();
        assert_eq!("request", params.plugin_step.to_string());
        assert_eq!(true, params.inflight.is_some());
        assert_eq!(LimitTag::Cookie, params.tag);
        assert_eq!("deviceId", params.key);

        let result = Limiter::try_from(
            &toml::from_str::<PluginConf>(
                r###"
step = "response"
type = "inflight"
tag = "cookie"
key = "deviceId"
max = 10
"###,
            )
            .unwrap(),
        );
        assert_eq!(
            "Plugin limit invalid, message: Invalid step(response), expect one of: request, proxy_upstream",
            result.err().unwrap().to_string()
        );

        // Settings that decide who is limited are checked, not defaulted.
        for (conf, expect) in [
            (
                "tag = \"cookies\"\nkey = \"a\"\nmax = 1",
                "Invalid tag(cookies)",
            ),
            ("tag = \"header\"\nmax = 1", "key is required"),
            ("type = \"inflght\"\nmax = 1", "Invalid type(inflght)"),
            ("type = \"rate\"", "max is required"),
            ("max = -1", "max must not be negative"),
            ("max = 1\ninterval = \"0s\"", "interval must be greater"),
        ] {
            let err =
                Limiter::try_from(&toml::from_str::<PluginConf>(conf).unwrap())
                    .err()
                    .unwrap()
                    .to_string();
            assert_eq!(true, err.contains(expect), "{conf}: {err}");
        }

        // A rate limiter tells the client when to come back.
        let params = Limiter::try_from(
            &toml::from_str::<PluginConf>("max = 10\ninterval = \"30s\"")
                .unwrap(),
        )
        .unwrap();
        assert_eq!(
            Some("30"),
            params
                .retry_after
                .as_ref()
                .and_then(|(_, value)| value.to_str().ok())
        );
        let params = Limiter::try_from(
            &toml::from_str::<PluginConf>("type = \"inflight\"\nmax = 10")
                .unwrap(),
        )
        .unwrap();
        assert_eq!(true, params.retry_after.is_none());
    }

    #[tokio::test]
    async fn test_new_cookie_limiter() {
        let limiter = Limiter::new(
            &toml::from_str::<PluginConf>(
                r###"
type = "inflight"
tag = "cookie"
key = "deviceId"
max = 10
    "###,
            )
            .unwrap(),
        )
        .unwrap();

        assert_eq!(LimitTag::Cookie, limiter.tag);
        let mut ctx = Ctx {
            ..Default::default()
        };
        let session = new_session().await;

        limiter.incr(&session, &mut ctx).unwrap();
        assert_eq!(true, ctx.state.guard.is_some());
    }
    #[tokio::test]
    async fn test_new_req_header_limiter() {
        let limiter = Limiter::new(
            &toml::from_str::<PluginConf>(
                r###"
type = "inflight"
tag = "header"
key = "X-Uuid"
max = 10
    "###,
            )
            .unwrap(),
        )
        .unwrap();
        assert_eq!(LimitTag::RequestHeader, limiter.tag);
        let mut ctx = Ctx {
            ..Default::default()
        };
        let session = new_session().await;

        limiter.incr(&session, &mut ctx).unwrap();
        assert_eq!(true, ctx.state.guard.is_some());
    }
    #[tokio::test]
    async fn test_new_query_limiter() {
        let limiter = Limiter::new(
            &toml::from_str::<PluginConf>(
                r###"
type = "inflight"
tag = "query"
key = "key"
max = 10
    "###,
            )
            .unwrap(),
        )
        .unwrap();
        assert_eq!(LimitTag::Query, limiter.tag);
        let mut ctx = Ctx {
            ..Default::default()
        };
        let session = new_session().await;

        limiter.incr(&session, &mut ctx).unwrap();
        assert_eq!(true, ctx.state.guard.is_some());
    }
    #[tokio::test]
    async fn test_new_ip_limiter() {
        let limiter = Limiter::new(
            &toml::from_str::<PluginConf>(
                r###"
type = "inflight"
max = 10
"###,
            )
            .unwrap(),
        )
        .unwrap();
        assert_eq!(LimitTag::Ip, limiter.tag);
        let mut ctx = Ctx {
            ..Default::default()
        };
        let session = new_session().await;

        limiter.incr(&session, &mut ctx).unwrap();
        assert_eq!(true, ctx.state.guard.is_some());
    }
    #[tokio::test]
    async fn test_inflight_limit() {
        let limiter = Limiter::new(
            &toml::from_str::<PluginConf>(
                r###"
type = "inflight"
max = 0
"###,
            )
            .unwrap(),
        )
        .unwrap();

        let headers = ["X-Forwarded-For: 1.1.1.1"].join("\r\n");
        let input_header =
            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
        let mock_io = Builder::new().read(input_header.as_bytes()).build();
        let mut session = Session::new_h1(Box::new(mock_io));
        session.read_request().await.unwrap();
        let result = limiter
            .handle_request(
                PluginStep::Request,
                &mut session,
                &mut Ctx::default(),
            )
            .await
            .unwrap();

        let RequestPluginResult::Respond(resp) = result else {
            panic!("result is not Respond");
        };
        assert_eq!(StatusCode::TOO_MANY_REQUESTS, resp.status);

        let limiter = Limiter::new(
            &toml::from_str::<PluginConf>(
                r###"
type = "inflight"
max = 1
"###,
            )
            .unwrap(),
        )
        .unwrap();
        let result = limiter
            .handle_request(
                PluginStep::Request,
                &mut session,
                &mut Ctx::default(),
            )
            .await
            .unwrap();

        assert_eq!(true, result == RequestPluginResult::Continue);
    }

    #[tokio::test]
    async fn test_rate_limit() {
        let limiter = Limiter::new(
            &toml::from_str::<PluginConf>(
                r###"
type = "rate"
max = 1
interval = "1s"
"###,
            )
            .unwrap(),
        )
        .unwrap();

        let headers = ["X-Forwarded-For: 1.1.1.1"].join("\r\n");
        let input_header =
            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
        let mock_io = Builder::new().read(input_header.as_bytes()).build();
        let mut session = Session::new_h1(Box::new(mock_io));
        session.read_request().await.unwrap();
        let result = limiter
            .handle_request(
                PluginStep::Request,
                &mut session,
                &mut Ctx::default(),
            )
            .await
            .unwrap();

        assert_eq!(true, result == RequestPluginResult::Continue);

        let _ = limiter
            .handle_request(
                PluginStep::Request,
                &mut session,
                &mut Ctx::default(),
            )
            .await
            .unwrap();

        // wait for the next loop
        tokio::time::sleep(Duration::from_secs(1)).await;
        let result = limiter
            .handle_request(
                PluginStep::Request,
                &mut session,
                &mut Ctx::default(),
            )
            .await
            .unwrap();
        let RequestPluginResult::Respond(resp) = result else {
            panic!("result is not Respond");
        };
        assert_eq!(StatusCode::TOO_MANY_REQUESTS, resp.status);
        assert_eq!(
            Some("1"),
            resp.headers
                .as_ref()
                .and_then(|headers| headers.first())
                .and_then(|(_, value)| value.to_str().ok())
        );

        // wait for rate limiter to reset
        tokio::time::sleep(Duration::from_secs(1)).await;
        let result = limiter
            .handle_request(
                PluginStep::Request,
                &mut session,
                &mut Ctx::default(),
            )
            .await
            .unwrap();
        assert_eq!(true, result == RequestPluginResult::Continue);
    }
}