pingap_upstream/
hash_strategy.rs

1// Copyright 2024-2025 Tree xie.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use pingap_core::{
16    get_client_ip, get_cookie_value, get_query_value, get_req_header_value,
17};
18use pingora::proxy::Session;
19use std::borrow::Cow;
20
21/// A pre-parsed, efficient representation of the hash strategy.
22#[derive(PartialEq)]
23pub enum HashStrategy {
24    Url,
25    Ip,
26    Header(String),
27    Cookie(String),
28    Query(String),
29    Path, // Default
30}
31
32impl HashStrategy {
33    /// Gets the value to use for consistent hashing.
34    /// This is optimized to avoid allocations where possible.
35    pub fn get_value<'a>(
36        &self,
37        session: &'a Session,
38        client_ip: &'a Option<String>,
39    ) -> Cow<'a, str> {
40        match self {
41            HashStrategy::Url => {
42                Cow::Owned(session.req_header().uri.to_string())
43            },
44            HashStrategy::Ip => {
45                if let Some(ip) = client_ip {
46                    Cow::Borrowed(ip)
47                } else {
48                    Cow::Owned(get_client_ip(session))
49                }
50            },
51            HashStrategy::Header(key) => {
52                get_req_header_value(session.req_header(), key)
53                    .map(Cow::Borrowed)
54                    .unwrap_or(Cow::Borrowed(""))
55            },
56            HashStrategy::Cookie(key) => {
57                get_cookie_value(session.req_header(), key)
58                    .map(Cow::Borrowed)
59                    .unwrap_or(Cow::Borrowed(""))
60            },
61            HashStrategy::Query(key) => {
62                get_query_value(session.req_header(), key)
63                    .map(Cow::Borrowed)
64                    .unwrap_or(Cow::Borrowed(""))
65            },
66            HashStrategy::Path => {
67                Cow::Borrowed(session.req_header().uri.path())
68            },
69        }
70    }
71}
72
73impl From<(&str, &str)> for HashStrategy {
74    fn from(tuple: (&str, &str)) -> Self {
75        match tuple.0 {
76            "url" => HashStrategy::Url,
77            "ip" => HashStrategy::Ip,
78            "header" => HashStrategy::Header(tuple.1.to_string()),
79            "cookie" => HashStrategy::Cookie(tuple.1.to_string()),
80            "query" => HashStrategy::Query(tuple.1.to_string()),
81            _ => HashStrategy::Path,
82        }
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::HashStrategy;
89    use pingora::proxy::Session;
90    use pretty_assertions::assert_eq;
91    use tokio_test::io::Builder;
92
93    #[test]
94    fn test_new_hash_strategy() {
95        assert!(HashStrategy::Url == HashStrategy::from(("url", "")));
96        assert!(HashStrategy::Ip == HashStrategy::from(("ip", "")));
97        assert!(
98            HashStrategy::Header("User-Agent".to_string())
99                == HashStrategy::from(("header", "User-Agent"))
100        );
101        assert!(
102            HashStrategy::Cookie("deviceId".to_string())
103                == HashStrategy::from(("cookie", "deviceId"))
104        );
105        assert!(
106            HashStrategy::Query("id".to_string())
107                == HashStrategy::from(("query", "id"))
108        );
109        assert!(HashStrategy::Path == HashStrategy::from(("", "")));
110    }
111
112    #[tokio::test]
113    async fn test_get_hash_key_value() {
114        let headers = [
115            "Host: github.com",
116            "Referer: https://github.com/",
117            "User-Agent: pingap/0.1.1",
118            "Cookie: deviceId=abc",
119            "Accept: application/json",
120            "X-Forwarded-For: 1.1.1.1",
121        ]
122        .join("\r\n");
123        let input_header = format!(
124            "GET /vicanso/pingap?id=1234 HTTP/1.1\r\n{headers}\r\n\r\n"
125        );
126        let mock_io = Builder::new().read(input_header.as_bytes()).build();
127
128        let mut session = Session::new_h1(Box::new(mock_io));
129        session.read_request().await.unwrap();
130
131        assert_eq!(
132            "/vicanso/pingap?id=1234",
133            HashStrategy::Url.get_value(&session, &None)
134        );
135
136        assert_eq!("1.1.1.1", HashStrategy::Ip.get_value(&session, &None));
137        assert_eq!(
138            "2.2.2.2",
139            HashStrategy::Ip.get_value(&session, &Some("2.2.2.2".to_string()))
140        );
141
142        assert_eq!(
143            "pingap/0.1.1",
144            HashStrategy::Header("User-Agent".to_string())
145                .get_value(&session, &None)
146        );
147
148        assert_eq!(
149            "abc",
150            HashStrategy::Cookie("deviceId".to_string())
151                .get_value(&session, &None)
152        );
153        assert_eq!(
154            "1234",
155            HashStrategy::Query("id".to_string()).get_value(&session, &None)
156        );
157        assert_eq!(
158            "/vicanso/pingap",
159            HashStrategy::Path.get_value(&session, &None)
160        );
161    }
162}
pingap_upstream/hash_strategy.rs

pingap_upstream/
hash_strategy.rs
