From f6c68fa3c56a5af3ba996174d9432f80384f470e Mon Sep 17 00:00:00 2001
From: vicanso <tree.xie@outlook.com>
Date: Sat, 13 Apr 2024 09:03:29 +0800
Subject: [PATCH] fix: fix limit value parse

---
 TODO.md            |  3 +-
 src/proxy/limit.rs | 79 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/TODO.md b/TODO.md
index 459bccfb..5b142654 100644
--- a/TODO.md
+++ b/TODO.md
@@ -1,12 +1,13 @@
 # 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
 - [ ] tls cert auto update
+- [ ] support validate config before save(web)
+- [x] support alpn for location
 - [x] support add header for location
 - [x] support x-forwarded-for
 - [x] error page
diff --git a/src/proxy/limit.rs b/src/proxy/limit.rs
index 1ab8f959..41574764 100644
--- a/src/proxy/limit.rs
+++ b/src/proxy/limit.rs
@@ -31,7 +31,7 @@ pub enum Error {
 }
 type Result<T, E = Error> = std::result::Result<T, E>;
 
-#[derive(PartialEq)]
+#[derive(PartialEq, Debug)]
 pub enum LimitTag {
     Ip,
     RequestHeader,
@@ -57,7 +57,6 @@ impl Limiter {
                 message: key.to_string(),
             });
         }
-        let key = key.substring(1, key.len() - 1);
         let ch = key.substring(0, 1);
         let value = key.substring(1, key.len());
         let tag = match ch {
@@ -75,7 +74,7 @@ impl Limiter {
         })
     }
     /// Increment `key` by 1. If value gt max, an error will be return.
-    /// Otherwise returns a Guard.
+    /// Otherwise returns a Guard. It may set the client ip to context.
     pub fn incr(&self, session: &Session, ctx: &mut State) -> Result<()> {
         let key = match self.tag {
             LimitTag::Query => utils::get_query_value(session.req_header(), &self.value)
@@ -109,3 +108,77 @@ impl Limiter {
         Ok(())
     }
 }
+#[cfg(test)]
+mod tests {
+    use crate::state::State;
+
+    use super::{LimitTag, Limiter};
+    use pingora::proxy::Session;
+    use pretty_assertions::assert_eq;
+    use tokio_test::io::Builder;
+    async fn new_session() -> Session {
+        let headers = vec![
+            "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
+    }
+    #[tokio::test]
+    async fn test_new_cookie_limiter() {
+        let limiter = Limiter::new("~deviceId 10").unwrap();
+        assert_eq!(LimitTag::Cookie, limiter.tag);
+        let mut ctx = State {
+            ..Default::default()
+        };
+        let session = new_session().await;
+
+        limiter.incr(&session, &mut ctx).unwrap();
+        assert_eq!(true, ctx.guard.is_some());
+    }
+    #[tokio::test]
+    async fn test_new_req_header_limiter() {
+        let limiter = Limiter::new(">X-Uuid 10").unwrap();
+        assert_eq!(LimitTag::RequestHeader, limiter.tag);
+        let mut ctx = State {
+            ..Default::default()
+        };
+        let session = new_session().await;
+
+        limiter.incr(&session, &mut ctx).unwrap();
+        assert_eq!(true, ctx.guard.is_some());
+    }
+    #[tokio::test]
+    async fn test_new_query_limiter() {
+        let limiter = Limiter::new("?key 10").unwrap();
+        assert_eq!(LimitTag::Query, limiter.tag);
+        let mut ctx = State {
+            ..Default::default()
+        };
+        let session = new_session().await;
+
+        limiter.incr(&session, &mut ctx).unwrap();
+        assert_eq!(true, ctx.guard.is_some());
+    }
+    #[tokio::test]
+    async fn test_new_ip_limiter() {
+        let limiter = Limiter::new("ip 10").unwrap();
+        assert_eq!(LimitTag::Ip, limiter.tag);
+        let mut ctx = State {
+            ..Default::default()
+        };
+        let session = new_session().await;
+
+        limiter.incr(&session, &mut ctx).unwrap();
+        assert_eq!(true, ctx.guard.is_some());
+    }
+}