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

pingap_core/
http_header.rs
