pingap_core/
http_header.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
15// Import necessary modules and types from supervisors and external crates.
16use super::{Ctx, get_hostname};
17use ahash::AHashSet;
18use bytes::BytesMut;
19use http::header;
20use http::{HeaderName, HeaderValue};
21use ipnet::IpNet;
22use pingora::http::RequestHeader;
23use pingora::proxy::Session;
24use snafu::{ResultExt, Snafu};
25use std::borrow::Cow;
26use std::fmt::Write;
27use std::net::IpAddr;
28use std::str::FromStr;
29use std::sync::atomic::{AtomicBool, Ordering};
30use std::sync::{LazyLock, RwLock};
31
32// Define string constants for commonly used HTTP header names.
33const HTTP_HEADER_X_FORWARDED_FOR: &str = "x-forwarded-for";
34const HTTP_HEADER_X_REAL_IP: &str = "x-real-ip";
35
36// Define byte slice constants for special variable tags used in header value processing.
37// These are matched against the raw bytes of a header value.
38pub const HOST_NAME_TAG: &[u8] = b"$hostname";
39const HOST_TAG: &[u8] = b"$host";
40const SCHEME_TAG: &[u8] = b"$scheme";
41const REMOTE_ADDR_TAG: &[u8] = b"$remote_addr";
42const REMOTE_PORT_TAG: &[u8] = b"$remote_port";
43const SERVER_ADDR_TAG: &[u8] = b"$server_addr";
44const SERVER_PORT_TAG: &[u8] = b"$server_port";
45const PROXY_ADD_FORWARDED_TAG: &[u8] = b"$proxy_add_x_forwarded_for";
46const UPSTREAM_ADDR_TAG: &[u8] = b"$upstream_addr";
47
48// Define static HeaderValues for HTTP and HTTPS schemes to avoid re-creation.
49static SCHEME_HTTPS: HeaderValue = HeaderValue::from_static("https");
50static SCHEME_HTTP: HeaderValue = HeaderValue::from_static("http");
51
52/// Defines the custom error types for this module using the snafu crate.
53#[derive(Debug, Snafu)]
54pub enum Error {
55    /// Error for when a string cannot be parsed into a valid HeaderValue.
56    #[snafu(display("invalid header value: {value} - {source}"))]
57    InvalidHeaderValue {
58        value: String,
59        source: header::InvalidHeaderValue,
60    },
61    /// Error for when a string cannot be parsed into a valid HeaderName.
62    #[snafu(display("invalid header name: {value} - {source}"))]
63    InvalidHeaderName {
64        value: String,
65        source: header::InvalidHeaderName,
66    },
67}
68/// A convenient type alias for `Result` with the module's `Error` type.
69type Result<T, E = Error> = std::result::Result<T, E>;
70
71/// A type alias for a tuple representing an HTTP header.
72pub type HttpHeader = (HeaderName, HeaderValue);
73
74/// Gets the request host by checking the URI first, then falling back to the "Host" header.
75///
76/// This function follows the common practice of prioritizing the host from the absolute URI
77/// (e.g., in `GET http://example.com/path HTTP/1.1`) over the `Host` header field.
78pub fn get_host(header: &RequestHeader) -> Option<&str> {
79    // First, try to get the host directly from the parsed URI.
80    // http2 will always have a host in the uri
81    if let Some(host) = header.uri.host() {
82        return Some(host);
83    }
84    // If not in the URI, fall back to the "Host" header.
85    header
86        .headers
87        .get(http::header::HOST)
88        // Convert the header value to a string slice.
89        .and_then(|value| value.to_str().ok())
90        // The host header can include a port (e.g., "example.com:8080"), so we split and take the first part.
91        .and_then(|host| host.split(':').next())
92}
93
94/// Converts a single string in "name: value" format into an `HttpHeader` tuple.
95///
96/// This is a utility function for parsing header configurations. It trims whitespace
97/// from both the name and the value.
98pub fn convert_header(value: &str) -> Result<Option<HttpHeader>> {
99    // `split_once` is an efficient way to split the string into two parts at the first colon.
100    value
101        .split_once(':')
102        // If a colon exists, map the key and value parts.
103        .map(|(k, v)| {
104            // Parse the trimmed key into a HeaderName, wrapping errors.
105            let name = HeaderName::from_str(k.trim())
106                .context(InvalidHeaderNameSnafu { value: k })?;
107            // Parse the trimmed value into a HeaderValue, wrapping errors.
108            let value = HeaderValue::from_str(v.trim())
109                .context(InvalidHeaderValueSnafu { value: v })?;
110            // If both parsing steps succeed, return the header tuple.
111            Ok(Some((name, value)))
112        })
113        // If `split_once` returns None (no colon), default to `Ok(None)`.
114        .unwrap_or(Ok(None))
115}
116
117/// Converts a slice of strings into a `Vec` of `HttpHeader`s.
118///
119/// This function iterates over a list of header strings and uses `convert_header`
120/// on each, collecting the valid results into a vector.
121pub fn convert_headers(header_values: &[String]) -> Result<Vec<HttpHeader>> {
122    header_values
123        .iter()
124        // `filter_map` is used to iterate, convert, and filter out `None` results elegantly.
125        // `transpose` flips `Option<Result<T>>` to `Result<Option<T>>`, which is what `filter_map` expects.
126        .filter_map(|item| convert_header(item).transpose())
127        // `collect` gathers the `Result<HttpHeader>` items. If any item is an `Err`, `collect` will return that `Err`.
128        .collect()
129}
130
131// Define common, pre-built HTTP headers as static constants for reuse and performance.
132pub static HTTP_HEADER_NO_STORE: HttpHeader = (
133    header::CACHE_CONTROL,
134    HeaderValue::from_static("private, no-store"),
135);
136pub static HTTP_HEADER_NO_CACHE: HttpHeader = (
137    header::CACHE_CONTROL,
138    HeaderValue::from_static("private, no-cache"),
139);
140pub static HTTP_HEADER_CONTENT_JSON: HttpHeader = (
141    header::CONTENT_TYPE,
142    HeaderValue::from_static("application/json; charset=utf-8"),
143);
144pub static HTTP_HEADER_CONTENT_HTML: HttpHeader = (
145    header::CONTENT_TYPE,
146    HeaderValue::from_static("text/html; charset=utf-8"),
147);
148pub static HTTP_HEADER_CONTENT_TEXT: HttpHeader = (
149    header::CONTENT_TYPE,
150    HeaderValue::from_static("text/plain; charset=utf-8"),
151);
152pub static HTTP_HEADER_TRANSFER_CHUNKED: HttpHeader = (
153    header::TRANSFER_ENCODING,
154    HeaderValue::from_static("chunked"),
155);
156pub static HTTP_HEADER_NAME_X_REQUEST_ID: HeaderName =
157    HeaderName::from_static("x-request-id");
158
159/// Processes a `HeaderValue` that may contain a special dynamic variable (e.g., `$host`).
160/// It replaces the variable with its corresponding runtime value.
161#[inline]
162pub fn convert_header_value(
163    value: &HeaderValue,
164    session: &Session,
165    ctx: &Ctx,
166) -> Option<HeaderValue> {
167    // Work with the raw byte representation of the header value for efficient matching.
168    let buf = value.as_bytes();
169
170    // Perform a quick check for the special variable prefix ('$' or ':') to exit early
171    // for normal header values, which is the most common case.
172    if buf.is_empty() || !(buf[0] == b'$' || buf[0] == b':') {
173        return None;
174    }
175
176    // A helper closure to reduce boilerplate when converting a string slice to a HeaderValue.
177    let to_header_value = |s: &str| HeaderValue::from_str(s).ok();
178
179    // Match the entire byte slice against the predefined variable tags.
180    match buf {
181        HOST_TAG => get_host(session.req_header()).and_then(to_header_value),
182        SCHEME_TAG => Some(if ctx.conn.tls_version.is_some() {
183            SCHEME_HTTPS.clone()
184        } else {
185            SCHEME_HTTP.clone()
186        }),
187        HOST_NAME_TAG => to_header_value(get_hostname()),
188        REMOTE_ADDR_TAG => {
189            ctx.conn.remote_addr.as_deref().and_then(to_header_value)
190        },
191        REMOTE_PORT_TAG => ctx.conn.remote_port.and_then(|p| {
192            // Use `itoa` to format the integer directly into a valid header value
193            // without creating an intermediate `String`.
194            HeaderValue::from_str(itoa::Buffer::new().format(p)).ok()
195        }),
196        SERVER_ADDR_TAG => {
197            ctx.conn.server_addr.as_deref().and_then(to_header_value)
198        },
199        SERVER_PORT_TAG => ctx.conn.server_port.and_then(|p| {
200            HeaderValue::from_str(itoa::Buffer::new().format(p)).ok()
201        }),
202        UPSTREAM_ADDR_TAG => {
203            if !ctx.upstream.address.is_empty() {
204                to_header_value(&ctx.upstream.address)
205            } else {
206                None
207            }
208        },
209        PROXY_ADD_FORWARDED_TAG => {
210            ctx.conn.remote_addr.as_deref().and_then(|remote_addr| {
211                // Build the new `x-forwarded-for` value efficiently using `BytesMut` to avoid `format!`.
212                let existing = session.get_header(HTTP_HEADER_X_FORWARDED_FOR);
213                let capacity =
214                    existing.map(|v| v.as_bytes().len() + 2).unwrap_or(0)
215                        + remote_addr.len();
216                let mut value_buf = BytesMut::with_capacity(capacity);
217                if let Some(existing) = existing {
218                    value_buf.extend_from_slice(existing.as_bytes());
219                    value_buf.extend_from_slice(b", ");
220                }
221                value_buf.extend_from_slice(remote_addr.as_bytes());
222                HeaderValue::from_bytes(&value_buf).ok()
223            })
224        },
225        // If no predefined tag matches, it might be a different type of variable (e.g., `$http_...`).
226        _ => handle_special_headers(buf, session, ctx),
227    }
228}
229
230/// A helper function to handle more complex or less common special header variables.
231/// This function is called as a fallback from `convert_header_value`.
232#[inline]
233fn handle_special_headers(
234    buf: &[u8],
235    session: &Session,
236    ctx: &Ctx,
237) -> Option<HeaderValue> {
238    // Handle variables that reference other request headers, like `$http_user_agent`.
239    if buf.starts_with(b"$http_") {
240        // Attempt to parse the header name from the slice after the prefix.
241        let key = std::str::from_utf8(&buf[6..]).ok()?;
242        // Get the corresponding header from the request and clone its value.
243        return session.get_header(key).cloned();
244    }
245    // Handle variables that reference environment variables, like `$PATH`.
246    if buf.starts_with(b"$") {
247        let var_name = std::str::from_utf8(&buf[1..]).ok()?;
248        // Look up the environment variable and convert its value to a HeaderValue.
249        return std::env::var(var_name)
250            .ok()
251            .and_then(|v| HeaderValue::from_str(&v).ok());
252    }
253    // Handle variables that reference fields in the `Ctx` struct, like `:connection_id`.
254    if buf.starts_with(b":") {
255        let key = std::str::from_utf8(&buf[1..]).ok()?;
256        // Use `append_log_value` to get the string representation of the context field.
257        let mut value = BytesMut::with_capacity(20);
258        ctx.append_log_value(&mut value, key);
259        if !value.is_empty() {
260            // Convert the resulting bytes to a HeaderValue.
261            return HeaderValue::from_bytes(&value).ok();
262        }
263    }
264    // If no pattern matches, return None.
265    None
266}
267
268/// Gets the remote address (IP and port) from the session.
269pub fn get_remote_addr(session: &Session) -> Option<(String, u16)> {
270    session
271        .client_addr()
272        // Ensure the address is an IP address (v4 or v6).
273        .and_then(|addr| addr.as_inet())
274        // Map it to a tuple of (String, u16).
275        .map(|addr| (addr.ip().to_string(), addr.port()))
276}
277
278/// Parsed trusted downstream proxy addresses: individual IPs plus CIDR
279/// networks. Kept small and local since it is only used by `get_client_ip`.
280struct TrustedProxies {
281    nets: Vec<IpNet>,
282    ips: AHashSet<IpAddr>,
283}
284
285impl TrustedProxies {
286    /// Parses a list of IPs / CIDR ranges. Invalid entries are ignored (a
287    /// mistyped proxy simply fails closed: its forwarded headers are dropped).
288    fn parse(values: &[String]) -> Self {
289        let mut nets = Vec::new();
290        let mut ips = AHashSet::new();
291        for item in values {
292            if let Ok(net) = IpNet::from_str(item) {
293                nets.push(net);
294            } else if let Ok(ip) = IpAddr::from_str(item) {
295                ips.insert(ip);
296            }
297        }
298        Self { nets, ips }
299    }
300
301    /// Returns true if `peer` (an IP string) is one of the trusted proxies.
302    fn contains(&self, peer: &str) -> bool {
303        let Ok(addr) = peer.parse::<IpAddr>() else {
304            return false;
305        };
306        self.ips.contains(&addr)
307            || self.nets.iter().any(|net| net.contains(&addr))
308    }
309}
310
311// Trusted downstream proxies. When configured, the forwarded client-IP headers
312// (`X-Forwarded-For` / `X-Real-IP`) are only honoured for connections whose
313// direct TCP peer is one of these addresses; a client connecting directly must
314// not be able to spoof its IP for IP-based access control, rate limiting, etc.
315// A cheap atomic flag keeps the common "not configured" path lock-free.
316static TRUSTED_PROXIES_ENABLED: AtomicBool = AtomicBool::new(false);
317static TRUSTED_PROXIES: LazyLock<RwLock<Option<TrustedProxies>>> =
318    LazyLock::new(|| RwLock::new(None));
319
320/// Sets the trusted downstream proxy addresses (individual IPs or CIDR ranges).
321///
322/// When a non-empty list is configured, forwarded headers are only trusted for
323/// connections coming directly from one of these addresses. Passing `None` or
324/// an empty list restores the default behaviour of trusting forwarded headers
325/// unconditionally (backwards compatible). Safe to call repeatedly on reload.
326pub fn set_trusted_proxies(proxies: &Option<Vec<String>>) {
327    let parsed = match proxies {
328        Some(list) if !list.is_empty() => Some(TrustedProxies::parse(list)),
329        _ => None,
330    };
331    if let Ok(mut guard) = TRUSTED_PROXIES.write() {
332        TRUSTED_PROXIES_ENABLED.store(parsed.is_some(), Ordering::Relaxed);
333        *guard = parsed;
334    }
335}
336
337/// Returns true if the direct peer address is a configured trusted proxy.
338fn is_trusted_proxy(peer: &str) -> bool {
339    TRUSTED_PROXIES
340        .read()
341        .ok()
342        .and_then(|guard| guard.as_ref().map(|tp| tp.contains(peer)))
343        .unwrap_or(false)
344}
345
346/// Ensures `ctx.conn.client_ip` is populated and returns a borrowed reference.
347///
348/// Prefer this on the request path over calling [`get_client_ip`] repeatedly —
349/// the first call allocates once and subsequent callers reuse the cached value.
350#[inline]
351pub fn ensure_client_ip<'a>(session: &Session, ctx: &'a mut Ctx) -> &'a str {
352    if ctx.conn.client_ip.is_none() {
353        ctx.conn.client_ip = Some(get_client_ip(session));
354    }
355    // Just inserted or already present — never None after the block above.
356    ctx.conn.client_ip.as_deref().unwrap_or_default()
357}
358
359/// Gets the client's IP address.
360///
361/// When trusted proxies are configured, `X-Forwarded-For` / `X-Real-IP` are
362/// only honoured if the direct TCP peer is a trusted proxy; otherwise the
363/// peer's own address is returned. When no trusted proxies are configured the
364/// lookup order is:
365/// 1. `X-Forwarded-For` (taking the first IP in the list)
366/// 2. `X-Real-IP`
367/// 3. The remote address of the direct TCP connection
368pub fn get_client_ip(session: &Session) -> String {
369    // When trusted proxies are configured, a direct (untrusted) peer's
370    // forwarded headers must be ignored to prevent client-IP spoofing.
371    if TRUSTED_PROXIES_ENABLED.load(Ordering::Relaxed) {
372        let peer = get_remote_addr(session).map(|(addr, _)| addr);
373        let trusted = peer.as_deref().map(is_trusted_proxy).unwrap_or(false);
374        if !trusted {
375            return peer.unwrap_or_default();
376        }
377    }
378    // 1. Check `X-Forwarded-For`.
379    if let Some(value) = session.get_header(HTTP_HEADER_X_FORWARDED_FOR) {
380        // Efficiently take the first IP without creating an intermediate Vec.
381        if let Ok(s) = value.to_str()
382            && let Some(ip) = s.split(',').next()
383        {
384            let trimmed_ip = ip.trim();
385            if !trimmed_ip.is_empty() {
386                return trimmed_ip.to_string();
387            }
388        }
389    }
390    // 2. Check `X-Real-IP`.
391    if let Some(value) = session.get_header(HTTP_HEADER_X_REAL_IP) {
392        return value.to_str().unwrap_or_default().to_string();
393    }
394    // 3. Fall back to the direct connection's remote address.
395    if let Some((addr, _)) = get_remote_addr(session) {
396        return addr;
397    }
398    // If all checks fail, return an empty string.
399    "".to_string()
400}
401
402/// A convenient helper to get a header value as a `&str` from a `RequestHeader`.
403pub fn get_req_header_value<'a>(
404    req_header: &'a RequestHeader,
405    key: &str,
406) -> Option<&'a str> {
407    // Get the header by its key.
408    if let Some(value) = req_header.headers.get(key) {
409        // Try to convert it to a string slice. Fails if the value is not valid UTF-8.
410        if let Ok(value) = value.to_str() {
411            return Some(value);
412        }
413    }
414    None
415}
416
417/// Parses the "Cookie" header to find the value of a specific cookie.
418pub fn get_cookie_value<'a>(
419    req_header: &'a RequestHeader,
420    cookie_name: &str,
421) -> Option<&'a str> {
422    // First, get the entire "Cookie" header string. The '?' operator will short-circuit if it's not present.
423    get_req_header_value(req_header, "cookie")?
424        // Split the string into individual cookies.
425        .split(';')
426        // `find_map` is an efficient way to find the first cookie that matches our criteria.
427        .find_map(|item| {
428            // This chained logic attempts to quickly find a match.
429            // It's more complex to handle cases like "key=value" vs "key=" correctly.
430            item.trim()
431                .strip_prefix(cookie_name)?
432                .strip_prefix('=')
433                .or_else(|| {
434                    // Fallback logic to ensure the cookie name is an exact match.
435                    let (k, v) = item.split_once('=')?;
436                    if k.trim() == cookie_name {
437                        Some(v.trim())
438                    } else {
439                        None
440                    }
441                })
442        })
443}
444
445/// Gets the value of a specific query parameter from the request URI.
446pub fn get_query_value<'a>(
447    req_header: &'a RequestHeader,
448    name: &str,
449) -> Option<&'a str> {
450    // Get the query string from the URI, exiting if it doesn't exist.
451    req_header
452        .uri
453        .query()?
454        // Split the query string into key-value pairs.
455        .split('&')
456        // `find_map` efficiently searches for the first pair where the key matches.
457        .find_map(|item| {
458            // Split the pair into key and value.
459            let (k, v) = item.split_once('=')?;
460            // If the key matches, return the value.
461            if k == name { Some(v) } else { None }
462        })
463}
464
465/// Removes a specific query parameter from the request header's URI.
466///
467/// This function modifies the `req_header` in place.
468pub fn remove_query_from_header(
469    req_header: &mut RequestHeader,
470    name: &str,
471) -> Result<(), http::uri::InvalidUri> {
472    // If there is no query string, there is nothing to do.
473    let Some(query_str) = req_header.uri.query() else {
474        return Ok(());
475    };
476
477    // Pre-allocate a String with enough capacity to hold the new query string,
478    // which is a performance optimization to avoid reallocations.
479    let mut new_query = String::with_capacity(query_str.len());
480
481    // Iterate over each key-value pair in the original query string.
482    for item in query_str.split('&') {
483        // Get the key part of the pair.
484        let key = item.split('=').next().unwrap_or(item);
485
486        // If the key is not the one we want to remove, keep the item.
487        if key != name {
488            // If the new query string is not empty, add a separator first.
489            if !new_query.is_empty() {
490                new_query.push('&');
491            }
492            // Append the original "key=value" slice, which is allocation-free.
493            new_query.push_str(item);
494        }
495    }
496
497    // Reconstruct the URI from its path and the new query string.
498    let path = req_header.uri.path();
499    // Use `Cow` (Clone-on-Write) to avoid allocating a new String for the path if the query is empty.
500    let new_uri_str = if new_query.is_empty() {
501        // If the new query is empty, the new URI is just the path. Borrow it.
502        Cow::Borrowed(path)
503    } else {
504        // If the new query is not empty, build a new String. Own it.
505        let mut s = String::with_capacity(path.len() + 1 + new_query.len());
506        // `write!` is an efficient way to format into an existing String buffer.
507        let _ = write!(&mut s, "{path}?{new_query}");
508        Cow::Owned(s)
509    };
510
511    // Parse the newly constructed string into a `http::Uri`.
512    let new_uri = http::Uri::from_str(&new_uri_str)?;
513    // Update the request header with the new URI.
514    req_header.set_uri(new_uri);
515
516    Ok(())
517}
518
519#[cfg(test)]
520mod tests {
521    use super::*;
522    use crate::{ConnectionInfo, UpstreamInfo};
523    use pretty_assertions::assert_eq;
524    use tokio_test::io::Builder;
525
526    #[test]
527    fn test_convert_headers() {
528        let headers = convert_headers(&[
529            "Content-Type: application/octet-stream".to_string(),
530            "X-Server: $hostname".to_string(),
531            "X-User: $USER".to_string(),
532        ])
533        .unwrap();
534        assert_eq!(3, headers.len());
535        assert_eq!("content-type", headers[0].0.to_string());
536        assert_eq!("application/octet-stream", headers[0].1.to_str().unwrap());
537        assert_eq!("x-server", headers[1].0.to_string());
538        assert_eq!(false, headers[1].1.to_str().unwrap().is_empty());
539        assert_eq!("x-user", headers[2].0.to_string());
540        assert_eq!(false, headers[2].1.to_str().unwrap().is_empty());
541    }
542
543    #[test]
544    fn test_static_value() {
545        assert_eq!(
546            "cache-control: private, no-store",
547            format!(
548                "{}: {}",
549                HTTP_HEADER_NO_STORE.0.to_string(),
550                HTTP_HEADER_NO_STORE.1.to_str().unwrap_or_default()
551            )
552        );
553
554        assert_eq!(
555            "cache-control: private, no-cache",
556            format!(
557                "{}: {}",
558                HTTP_HEADER_NO_CACHE.0.to_string(),
559                HTTP_HEADER_NO_CACHE.1.to_str().unwrap_or_default()
560            )
561        );
562
563        assert_eq!(
564            "content-type: application/json; charset=utf-8",
565            format!(
566                "{}: {}",
567                HTTP_HEADER_CONTENT_JSON.0.to_string(),
568                HTTP_HEADER_CONTENT_JSON.1.to_str().unwrap_or_default()
569            )
570        );
571
572        assert_eq!(
573            "content-type: text/html; charset=utf-8",
574            format!(
575                "{}: {}",
576                HTTP_HEADER_CONTENT_HTML.0.to_string(),
577                HTTP_HEADER_CONTENT_HTML.1.to_str().unwrap_or_default()
578            )
579        );
580
581        assert_eq!(
582            "transfer-encoding: chunked",
583            format!(
584                "{}: {}",
585                HTTP_HEADER_TRANSFER_CHUNKED.0.to_string(),
586                HTTP_HEADER_TRANSFER_CHUNKED.1.to_str().unwrap_or_default()
587            )
588        );
589
590        assert_eq!("x-request-id", HTTP_HEADER_NAME_X_REQUEST_ID.to_string());
591
592        assert_eq!(
593            "content-type: text/plain; charset=utf-8",
594            format!(
595                "{}: {}",
596                HTTP_HEADER_CONTENT_TEXT.0.to_string(),
597                HTTP_HEADER_CONTENT_TEXT.1.to_str().unwrap_or_default()
598            )
599        );
600    }
601
602    #[tokio::test]
603    async fn test_convert_header_value() {
604        let headers = ["Host: pingap.io"].join("\r\n");
605        let input_header =
606            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
607        let mock_io = Builder::new().read(input_header.as_bytes()).build();
608        let mut session = Session::new_h1(Box::new(mock_io));
609        session.read_request().await.unwrap();
610        let default_state = Ctx {
611            upstream: UpstreamInfo {
612                address: "10.1.1.3:4123".to_string(),
613                ..Default::default()
614            },
615            conn: ConnectionInfo {
616                id: 102,
617                remote_addr: Some("10.1.1.1".to_string()),
618                remote_port: Some(6000),
619                server_addr: Some("10.1.1.2".to_string()),
620                server_port: Some(6001),
621                tls_version: Some("tls1.3".into()),
622                ..Default::default()
623            },
624            ..Default::default()
625        };
626
627        let value = convert_header_value(
628            &HeaderValue::from_str("$host").unwrap(),
629            &session,
630            &Ctx {
631                ..Default::default()
632            },
633        );
634        assert_eq!(true, value.is_some());
635        assert_eq!("pingap.io", value.unwrap().to_str().unwrap());
636
637        let value = convert_header_value(
638            &HeaderValue::from_str("$scheme").unwrap(),
639            &session,
640            &Ctx {
641                ..Default::default()
642            },
643        );
644        assert_eq!(true, value.is_some());
645        assert_eq!("http", value.unwrap().to_str().unwrap());
646        let value = convert_header_value(
647            &HeaderValue::from_str("$scheme").unwrap(),
648            &session,
649            &default_state,
650        );
651        assert_eq!(true, value.is_some());
652        assert_eq!("https", value.unwrap().to_str().unwrap());
653
654        let value = convert_header_value(
655            &HeaderValue::from_str("$remote_addr").unwrap(),
656            &session,
657            &default_state,
658        );
659        assert_eq!(true, value.is_some());
660        assert_eq!("10.1.1.1", value.unwrap().to_str().unwrap());
661
662        let value = convert_header_value(
663            &HeaderValue::from_str("$remote_port").unwrap(),
664            &session,
665            &default_state,
666        );
667        assert_eq!(true, value.is_some());
668        assert_eq!("6000", value.unwrap().to_str().unwrap());
669
670        let value = convert_header_value(
671            &HeaderValue::from_str("$server_addr").unwrap(),
672            &session,
673            &default_state,
674        );
675        assert_eq!(true, value.is_some());
676        assert_eq!("10.1.1.2", value.unwrap().to_str().unwrap());
677
678        let value = convert_header_value(
679            &HeaderValue::from_str("$server_port").unwrap(),
680            &session,
681            &default_state,
682        );
683        assert_eq!(true, value.is_some());
684        assert_eq!("6001", value.unwrap().to_str().unwrap());
685
686        let value = convert_header_value(
687            &HeaderValue::from_str("$upstream_addr").unwrap(),
688            &session,
689            &default_state,
690        );
691        assert_eq!(true, value.is_some());
692        assert_eq!("10.1.1.3:4123", value.unwrap().to_str().unwrap());
693
694        let value = convert_header_value(
695            &HeaderValue::from_str(":connection_id").unwrap(),
696            &session,
697            &default_state,
698        );
699        assert_eq!(true, value.is_some());
700        assert_eq!("102", value.unwrap().to_str().unwrap());
701
702        let headers = ["X-Forwarded-For: 1.1.1.1, 2.2.2.2"].join("\r\n");
703        let input_header =
704            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
705        let mock_io = Builder::new().read(input_header.as_bytes()).build();
706        let mut session = Session::new_h1(Box::new(mock_io));
707        session.read_request().await.unwrap();
708        let value = convert_header_value(
709            &HeaderValue::from_str("$proxy_add_x_forwarded_for").unwrap(),
710            &session,
711            &Ctx {
712                conn: ConnectionInfo {
713                    remote_addr: Some("10.1.1.1".to_string()),
714                    ..Default::default()
715                },
716                ..Default::default()
717            },
718        );
719        assert_eq!(true, value.is_some());
720        assert_eq!(
721            "1.1.1.1, 2.2.2.2, 10.1.1.1",
722            value.unwrap().to_str().unwrap()
723        );
724
725        let headers = [""].join("\r\n");
726        let input_header =
727            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
728        let mock_io = Builder::new().read(input_header.as_bytes()).build();
729        let mut session = Session::new_h1(Box::new(mock_io));
730        session.read_request().await.unwrap();
731        let value = convert_header_value(
732            &HeaderValue::from_str("$proxy_add_x_forwarded_for").unwrap(),
733            &session,
734            &Ctx {
735                conn: ConnectionInfo {
736                    remote_addr: Some("10.1.1.1".to_string()),
737                    ..Default::default()
738                },
739                ..Default::default()
740            },
741        );
742        assert_eq!(true, value.is_some());
743        assert_eq!("10.1.1.1", value.unwrap().to_str().unwrap());
744
745        let headers = [""].join("\r\n");
746        let input_header =
747            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
748        let mock_io = Builder::new().read(input_header.as_bytes()).build();
749        let mut session = Session::new_h1(Box::new(mock_io));
750        session.read_request().await.unwrap();
751        let value = convert_header_value(
752            &HeaderValue::from_str("$upstream_addr").unwrap(),
753            &session,
754            &Ctx {
755                upstream: UpstreamInfo {
756                    address: "10.1.1.1:8001".to_string(),
757                    ..Default::default()
758                },
759                ..Default::default()
760            },
761        );
762        assert_eq!(true, value.is_some());
763        assert_eq!("10.1.1.1:8001", value.unwrap().to_str().unwrap());
764
765        let headers = ["Origin: https://github.com"].join("\r\n");
766        let input_header =
767            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
768        let mock_io = Builder::new().read(input_header.as_bytes()).build();
769        let mut session = Session::new_h1(Box::new(mock_io));
770        session.read_request().await.unwrap();
771        let value = convert_header_value(
772            &HeaderValue::from_str("$http_origin").unwrap(),
773            &session,
774            &Ctx::default(),
775        );
776        assert_eq!(true, value.is_some());
777        assert_eq!("https://github.com", value.unwrap().to_str().unwrap());
778
779        let headers = ["Origin: https://github.com"].join("\r\n");
780        let input_header =
781            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
782        let mock_io = Builder::new().read(input_header.as_bytes()).build();
783        let mut session = Session::new_h1(Box::new(mock_io));
784        session.read_request().await.unwrap();
785        let value = convert_header_value(
786            &HeaderValue::from_str("$hostname").unwrap(),
787            &session,
788            &Ctx::default(),
789        );
790        assert_eq!(true, value.is_some());
791
792        let headers = ["Origin: https://github.com"].join("\r\n");
793        let input_header =
794            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
795        let mock_io = Builder::new().read(input_header.as_bytes()).build();
796        let mut session = Session::new_h1(Box::new(mock_io));
797        session.read_request().await.unwrap();
798        let value = convert_header_value(
799            &HeaderValue::from_str("$HOME").unwrap(),
800            &session,
801            &Ctx::default(),
802        );
803        assert_eq!(true, value.is_some());
804
805        let headers = ["Origin: https://github.com"].join("\r\n");
806        let input_header =
807            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
808        let mock_io = Builder::new().read(input_header.as_bytes()).build();
809        let mut session = Session::new_h1(Box::new(mock_io));
810        session.read_request().await.unwrap();
811        let value = convert_header_value(
812            &HeaderValue::from_str("UUID").unwrap(),
813            &session,
814            &Ctx::default(),
815        );
816        assert_eq!(false, value.is_some());
817    }
818
819    #[tokio::test]
820    async fn test_get_host() {
821        let headers = ["Host: pingap.io"].join("\r\n");
822        let input_header =
823            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
824        let mock_io = Builder::new().read(input_header.as_bytes()).build();
825        let mut session = Session::new_h1(Box::new(mock_io));
826        session.read_request().await.unwrap();
827        assert_eq!(get_host(session.req_header()), Some("pingap.io"));
828    }
829
830    #[test]
831    fn test_remove_query_from_header() {
832        let mut req =
833            RequestHeader::build("GET", b"/?apikey=123", None).unwrap();
834        remove_query_from_header(&mut req, "apikey").unwrap();
835        assert_eq!("/", req.uri.to_string());
836
837        let mut req =
838            RequestHeader::build("GET", b"/?apikey=123&name=pingap", None)
839                .unwrap();
840        remove_query_from_header(&mut req, "apikey").unwrap();
841        assert_eq!("/?name=pingap", req.uri.to_string());
842    }
843
844    #[tokio::test]
845    async fn test_get_client_ip() {
846        let headers = ["X-Forwarded-For:192.168.1.1"].join("\r\n");
847        let input_header =
848            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
849        let mock_io = Builder::new().read(input_header.as_bytes()).build();
850        let mut session = Session::new_h1(Box::new(mock_io));
851        session.read_request().await.unwrap();
852        assert_eq!(get_client_ip(&session), "192.168.1.1");
853
854        let headers = ["X-Real-Ip:192.168.1.2"].join("\r\n");
855        let input_header =
856            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
857        let mock_io = Builder::new().read(input_header.as_bytes()).build();
858        let mut session = Session::new_h1(Box::new(mock_io));
859        session.read_request().await.unwrap();
860        assert_eq!(get_client_ip(&session), "192.168.1.2");
861
862        // With trusted proxies configured, a forwarded header from an untrusted
863        // direct peer (the mock session has no trusted peer address) must be
864        // ignored instead of being taken at face value.
865        set_trusted_proxies(&Some(vec!["10.0.0.0/8".to_string()]));
866        let headers = ["X-Forwarded-For:192.168.1.1"].join("\r\n");
867        let input_header =
868            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
869        let mock_io = Builder::new().read(input_header.as_bytes()).build();
870        let mut session = Session::new_h1(Box::new(mock_io));
871        session.read_request().await.unwrap();
872        assert_ne!(get_client_ip(&session), "192.168.1.1");
873
874        // Restoring the default trusts forwarded headers again.
875        set_trusted_proxies(&None);
876        let headers = ["X-Forwarded-For:192.168.1.1"].join("\r\n");
877        let input_header =
878            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
879        let mock_io = Builder::new().read(input_header.as_bytes()).build();
880        let mut session = Session::new_h1(Box::new(mock_io));
881        session.read_request().await.unwrap();
882        assert_eq!(get_client_ip(&session), "192.168.1.1");
883    }
884
885    #[tokio::test]
886    async fn test_get_header_value() {
887        let headers = ["Host: pingap.io"].join("\r\n");
888        let input_header =
889            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
890        let mock_io = Builder::new().read(input_header.as_bytes()).build();
891        let mut session = Session::new_h1(Box::new(mock_io));
892        session.read_request().await.unwrap();
893        assert_eq!(
894            get_req_header_value(session.req_header(), "Host"),
895            Some("pingap.io")
896        );
897    }
898
899    #[tokio::test]
900    async fn test_get_cookie_value() {
901        let headers = ["Cookie: name=pingap"].join("\r\n");
902        let input_header =
903            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
904        let mock_io = Builder::new().read(input_header.as_bytes()).build();
905        let mut session = Session::new_h1(Box::new(mock_io));
906        session.read_request().await.unwrap();
907        assert_eq!(
908            get_cookie_value(session.req_header(), "name"),
909            Some("pingap")
910        );
911    }
912
913    #[tokio::test]
914    async fn test_get_query_value() {
915        let headers = ["X-Forwarded-For:192.168.1.1"].join("\r\n");
916        let input_header =
917            format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
918        let mock_io = Builder::new().read(input_header.as_bytes()).build();
919        let mut session = Session::new_h1(Box::new(mock_io));
920        session.read_request().await.unwrap();
921        assert_eq!(get_query_value(session.req_header(), "size"), Some("1"));
922    }
923
924    /// Tests `convert_header` with edge cases like empty strings,
925    /// strings without colons, and invalid header names/values.
926    #[test]
927    fn test_convert_header_edge_cases() {
928        // Empty string should result in Ok(None).
929        assert!(convert_header("").unwrap().is_none());
930        // String without a colon should result in Ok(None).
931        assert!(convert_header("no-colon").unwrap().is_none());
932        // Invalid header name should result in an error.
933        assert!(convert_header("Invalid Name: value").is_err());
934        // Invalid header value (with newline) should result in an error.
935        assert!(convert_header("Valid-Name: invalid\r\nvalue").is_err());
936    }
937
938    /// Tests `get_host` logic with different request formats.
939    #[test]
940    fn test_get_host_variants() {
941        // Case 1: Host is in the URI authority.
942        let uri_string = "http://user:pass@authority.com/path";
943
944        // 使用 .parse() 或 from_str 来创建 Uri
945        let uri = http::Uri::from_str(uri_string).unwrap();
946        let mut req_with_authority =
947            RequestHeader::build("GET", b"/path", None).unwrap();
948        req_with_authority.set_uri(uri);
949        assert_eq!(get_host(&req_with_authority), Some("authority.com"));
950
951        // Case 2: Host is in the "Host" header.
952        let mut req_with_host_header =
953            RequestHeader::build("GET", b"/path", None).unwrap();
954        req_with_host_header
955            .insert_header("Host", "header-host.com:8080")
956            .unwrap();
957        assert_eq!(get_host(&req_with_host_header), Some("header-host.com"));
958
959        // Case 3: No host information available.
960        let req_no_host = RequestHeader::build("GET", b"/path", None).unwrap();
961        assert_eq!(get_host(&req_no_host), None);
962    }
963
964    /// Tests `get_cookie_value` with multiple cookies and edge cases.
965    #[test]
966    fn test_get_cookie_value_advanced() {
967        let mut req = RequestHeader::build("GET", b"/", None).unwrap();
968        req.insert_header("Cookie", "id=123; session=abc; theme=dark")
969            .unwrap();
970
971        assert_eq!(get_cookie_value(&req, "session"), Some("abc"));
972        assert_eq!(get_cookie_value(&req, "id"), Some("123"));
973        assert_eq!(get_cookie_value(&req, "theme"), Some("dark"));
974        // Test for a non-existent cookie.
975        assert_eq!(get_cookie_value(&req, "lang"), None);
976        // Test for a cookie name that is a prefix of another.
977        assert_eq!(get_cookie_value(&req, "the"), None);
978    }
979
980    #[test]
981    fn test_remove_query_from_header_variants() {
982        // Case 1: Remove the only query param.
983        let mut req =
984            RequestHeader::build("GET", b"/path?key=val", None).unwrap();
985        remove_query_from_header(&mut req, "key").unwrap();
986        assert_eq!(req.uri.to_string(), "/path");
987
988        // Case 2: Remove the first of multiple params.
989        let mut req =
990            RequestHeader::build("GET", b"/path?key1=val1&key2=val2", None)
991                .unwrap();
992        remove_query_from_header(&mut req, "key1").unwrap();
993        assert_eq!(req.uri.to_string(), "/path?key2=val2");
994
995        // Case 3: Remove the last of multiple params.
996        let mut req =
997            RequestHeader::build("GET", b"/path?key1=val1&key2=val2", None)
998                .unwrap();
999        remove_query_from_header(&mut req, "key2").unwrap();
1000        assert_eq!(req.uri.to_string(), "/path?key1=val1");
1001
1002        // Case 4: Remove a middle param.
1003        let mut req =
1004            RequestHeader::build("GET", b"/path?key1=v1&key2=v2&key3=v3", None)
1005                .unwrap();
1006        remove_query_from_header(&mut req, "key2").unwrap();
1007        assert_eq!(req.uri.to_string(), "/path?key1=v1&key3=v3");
1008
1009        // Case 5: Param to remove is not present.
1010        let mut req =
1011            RequestHeader::build("GET", b"/path?key=val", None).unwrap();
1012        remove_query_from_header(&mut req, "nonexistent").unwrap();
1013        assert_eq!(req.uri.to_string(), "/path?key=val");
1014
1015        // Case 6: No query string to begin with.
1016        let mut req = RequestHeader::build("GET", b"/path", None).unwrap();
1017        remove_query_from_header(&mut req, "key").unwrap();
1018        assert_eq!(req.uri.to_string(), "/path");
1019    }
1020}
pingap_core/http_header.rs

pingap_core/
http_header.rs
