refactor: adjust limit handle function

TODO.md

@@ -1,22 +1,24 @@
 # TODO
 
+- [ ] support alpn for location
+- [ ] custom error for pingora error
+- [ ] authentication for admin page
+- [ ] support etcd or other storage for config
+- [ ] better error handler
+- [ ] log rotate
 - [x] support add header for location
 - [x] support x-forwarded-for
 - [x] error page
 - [x] http peer option
 - [x] access log
 - [x] support format for env logger(or tokio tracing)
-- [ ] better error handler
 - [x] config validate
 - [x] support add tls
-- [ ] log rotate
 - [x] stats of server
 - [x] start without config
 - [x] static serve for admin
 - [x] status:499 for client abort
 - [x] support get pingap start time
-- [ ] custom error for pingora error
-- [ ] authentication for admin page
 - [x] static file serve
 - [x] set priority for location
 - [x] mock response for upstream
@@ -25,12 +27,11 @@
 - [x] support set upstream_keepalive_pool_size
 - [x] graceful restart for admin web
 - [x] use stable pingora
-- [ ] support etcd or other storage for config
 - [x] support web hook for backend health check
 - [x] sentry uri config
 - [x] charset for static file
 - [x] web hook for wecom, dingtalk robot
 - [x] verify_cert option for http peer
 - [x] compression: zstd, br, gzip
 - [x] support set threads for each server
-- [ ] limit of request: ip or custom field
+- [x] location limit of request: ip or custom field

conf/pingap.toml

@@ -116,4 +116,4 @@ locations = ["lo"]
 stats_path = "/stats"
 
 # Admin path for admin server. Default `None`
-admin_path = "/admin"
+admin_path = "/pingap"

src/config/load.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::utils;
+use crate::{proxy::Limiter, utils};
 use base64::{engine::general_purpose::STANDARD, Engine};
 use glob::glob;
 use http::HeaderValue;
@@ -126,6 +126,7 @@ pub struct LocationConf {
     pub gzip_level: Option<u32>,
     pub br_level: Option<u32>,
     pub zstd_level: Option<u32>,
+    pub limit: Option<String>,
     pub remark: Option<String>,
 }
 
@@ -156,6 +157,11 @@ impl LocationConf {
                 message: format!("{} upstream is not found(location:{name})", self.upstream),
             });
         }
+        if let Some(limit) = &self.limit {
+            Limiter::new(limit).map_err(|err| Error::Invalid {
+                message: format!("{err}({limit})"),
+            })?;
+        }
         Ok(())
     }
 
@@ -197,7 +203,6 @@ pub struct ServerConf {
     pub stats_path: Option<String>,
     pub admin_path: Option<String>,
     pub threads: Option<usize>,
-    pub limit: Option<String>,
     pub remark: Option<String>,
 }
 

src/proxy/limit.rs

@@ -33,10 +33,10 @@ type Result<T, E = Error> = std::result::Result<T, E>;
 
 #[derive(PartialEq)]
 pub enum LimitTag {
-    Cookie,
+    Ip,
     RequestHeader,
+    Cookie,
     Query,
-    Ip,
 }
 
 pub struct Limiter {
@@ -78,17 +78,22 @@ impl Limiter {
     /// Otherwise returns a Guard.
     pub fn incr(&self, session: &Session, ctx: &mut State) -> Result<()> {
         let key = match self.tag {
-            LimitTag::Ip => {
-                let client_ip = utils::get_client_ip(session);
-                ctx.client_ip = Some(client_ip.clone());
-                client_ip
-            }
+            LimitTag::Query => utils::get_query_value(session.req_header(), &self.value)
+                .unwrap_or_default()
+                .to_string(),
             LimitTag::RequestHeader => {
                 utils::get_req_header_value(session.req_header(), &self.value)
                     .unwrap_or_default()
                     .to_string()
             }
-            _ => "".to_string(),
+            LimitTag::Cookie => utils::get_cookie_value(session.req_header(), &self.value)
+                .unwrap_or_default()
+                .to_string(),
+            _ => {
+                let client_ip = utils::get_client_ip(session);
+                ctx.client_ip = Some(client_ip.clone());
+                client_ip
+            }
         };
         if key.is_empty() {
             return Ok(());

src/proxy/location.rs

@@ -12,12 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use super::Upstream;
+use super::{Limiter, Upstream};
 use crate::config::LocationConf;
 use crate::http_extra::{convert_headers, HttpHeader};
+use crate::state::State;
 use http::header::HeaderValue;
+use log::error;
 use once_cell::sync::Lazy;
 use pingora::http::{RequestHeader, ResponseHeader};
+use pingora::proxy::Session;
 use regex::Regex;
 use snafu::{ResultExt, Snafu};
 use std::sync::Arc;
@@ -89,6 +92,7 @@ pub struct Location {
     gzip_level: u32,
     br_level: u32,
     zstd_level: u32,
+    limiter: Option<Limiter>,
     pub support_compression: bool,
     pub upstream: Arc<Upstream>,
 }
@@ -138,6 +142,17 @@ impl Location {
         let zstd_level = conf.zstd_level.unwrap_or_default();
         let support_compression = gzip_level + br_level + zstd_level > 0;
         let path = conf.path.clone().unwrap_or_default();
+        let limiter = if let Some(limit) = &conf.limit {
+            match Limiter::new(limit) {
+                Ok(l) => Some(l),
+                Err(e) => {
+                    error!("New limiter fail: {e}");
+                    None
+                }
+            }
+        } else {
+            None
+        };
         Ok(Location {
             // name: conf.name.clone(),
             path_selector: new_path_selector(&path)?,
@@ -151,6 +166,7 @@ impl Location {
             br_level,
             zstd_level,
             support_compression,
+            limiter,
         })
     }
     /// Returns `true` if the host and path match location.
@@ -233,6 +249,17 @@ impl Location {
         }
         None
     }
+    #[inline]
+    /// Validate the request count less than max limit, and set the guard to ctx.
+    pub fn validate_limit(&self, session: &Session, ctx: &mut State) -> pingora::Result<()> {
+        if let Some(limiter) = &self.limiter {
+            return limiter.incr(session, ctx).map_err(|_e| {
+                let e = pingora::Error::new(pingora::ErrorType::InternalError);
+                pingora::Error::because(pingora::ErrorType::HTTPStatus(429), "Limit eceed", e)
+            });
+        }
+        Ok(())
+    }
 }
 
 #[cfg(test)]

src/proxy/logger.rs

@@ -335,15 +335,10 @@ impl Parser {
                     buf.push_str(&format!("{:?}", Duration::from_millis(ms as u64)));
                 }
                 TagCategory::Cookie => {
-                    let cookie_name = tag.data.clone().unwrap_or_default();
-                    let cookie_value =
-                        utils::get_req_header_value(req_header, "Cookie").unwrap_or_default();
-                    for item in cookie_value.split(';') {
-                        if let Some((k, v)) = item.split_once('=') {
-                            if k == cookie_name {
-                                buf.push_str(v.trim());
-                            }
-                        }
+                    if let Some(value) =
+                        utils::get_cookie_value(req_header, &tag.data.clone().unwrap_or_default())
+                    {
+                        buf.push_str(value);
                     }
                 }
                 TagCategory::RequestHeader => {

src/proxy/server.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use super::logger::Parser;
-use super::{Limiter, Location, Upstream};
+use super::{Location, Upstream};
 use crate::config::{LocationConf, PingapConf, UpstreamConf};
 use crate::http_extra::{HttpResponse, HTTP_HEADER_CONTENT_JSON};
 use crate::serve::Serve;
@@ -77,7 +77,6 @@ pub struct ServerConf {
     pub tls_cert: Option<Vec<u8>>,
     pub tls_key: Option<Vec<u8>>,
     pub threads: Option<usize>,
-    pub limit: Option<String>,
     pub error_template: String,
 }
 
@@ -143,7 +142,6 @@ impl From<PingapConf> for Vec<ServerConf> {
                 upstreams: filter_upstreams,
                 locations: filter_locations,
                 threads: item.threads,
-                limit: item.limit,
                 error_template,
             });
         }
@@ -172,7 +170,6 @@ pub struct Server {
     threads: Option<usize>,
     tls_cert: Option<Vec<u8>>,
     tls_key: Option<Vec<u8>>,
-    limiter: Option<Limiter>,
 }
 
 #[derive(Serialize)]
@@ -223,17 +220,6 @@ impl Server {
         if let Some(access_log) = conf.access_log {
             p = Some(Parser::from(access_log.as_str()));
         }
-        let limiter = if let Some(limit) = &conf.limit {
-            match Limiter::new(limit) {
-                Ok(limit) => Some(limit),
-                Err(e) => {
-                    error!("Parse limit fail: {e}");
-                    None
-                }
-            }
-        } else {
-            None
-        };
 
         Ok(Server {
             admin: conf.admin,
@@ -248,7 +234,6 @@ impl Server {
             tls_key: conf.tls_key,
             tls_cert: conf.tls_cert,
             threads: conf.threads,
-            limiter,
         })
     }
     pub fn run(self, conf: &Arc<configuration::ServerConf>) -> Result<ServerServices> {
@@ -392,12 +377,6 @@ impl ProxyHttp for Server {
     where
         Self::CTX: Send + Sync,
     {
-        if let Some(limiter) = &self.limiter {
-            limiter.incr(session, ctx).map_err(|_e| {
-                let e = pingora::Error::new(pingora::ErrorType::InternalError);
-                pingora::Error::because(pingora::ErrorType::HTTPStatus(429), "Limit eceed", e)
-            })?;
-        }
         ctx.processing = self.processing.fetch_add(1, Ordering::Relaxed);
         self.accepted.fetch_add(1, Ordering::Relaxed);
         // session.cache.enable(storage, eviction, predictor, cache_lock)
@@ -440,6 +419,7 @@ impl ProxyHttp for Server {
             let result = mock.handle(session, ctx).await?;
             return Ok(result);
         }
+        lo.validate_limit(session, ctx)?;
         // TODO get response from cache
         // check location support cache
 

src/serve/admin.rs

@@ -21,6 +21,7 @@ use crate::state::State;
 use crate::state::{get_start_time, restart};
 use crate::utils::get_pkg_version;
 use async_trait::async_trait;
+use bytes::{BufMut, BytesMut};
 use bytesize::ByteSize;
 use http::{Method, StatusCode};
 use log::error;
@@ -125,7 +126,10 @@ impl AdminServe {
         category: &str,
         name: &str,
     ) -> pingora::Result<HttpResponse> {
-        let buf = session.read_request_body().await?.unwrap_or_default();
+        let mut buf = BytesMut::with_capacity(4096);
+        while let Some(value) = session.read_request_body().await? {
+            buf.put(value.as_ref());
+        }
         let key = name.to_string();
         let mut conf = self.load_config()?;
         match category {

src/utils/mod.rs

@@ -92,3 +92,29 @@ pub fn get_req_header_value<'a>(req_header: &'a RequestHeader, key: &str) -> Opt
     }
     None
 }
+
+pub fn get_cookie_value<'a>(req_header: &'a RequestHeader, cookie_name: &str) -> Option<&'a str> {
+    if let Some(cookie_value) = get_req_header_value(req_header, "Cookie") {
+        for item in cookie_value.split(';') {
+            if let Some((k, v)) = item.split_once('=') {
+                if k == cookie_name {
+                    return Some(v.trim());
+                }
+            }
+        }
+    }
+    None
+}
+
+pub fn get_query_value<'a>(req_header: &'a RequestHeader, name: &str) -> Option<&'a str> {
+    if let Some(query) = req_header.uri.query() {
+        for item in query.split('&') {
+            if let Some((k, v)) = item.split_once('=') {
+                if k == name {
+                    return Some(v.trim());
+                }
+            }
+        }
+    }
+    None
+}

web/src/pages/location-info.tsx

@@ -100,6 +100,13 @@ export default function LocationInfo() {
       span: 4,
       category: FormItemCategory.NUMBER,
     },
+    {
+      id: "limit",
+      label: "Limit",
+      defaultValue: location.limit,
+      span: 6,
+      category: FormItemCategory.TEXT,
+    },
     {
       id: "remark",
       label: "Remark",