ureq-proto 0.7.1

use http::HeaderName;
use http::request::Builder;

/// Append a header while preserving its name's casing when writing the request.
///
/// This is an escape hatch for servers that compare header names case-sensitively.
/// Names and values retain the same validation as [`Builder::header`]. Values may
/// contain non-UTF-8 bytes, but cannot contain line breaks.
///
/// Distinct names follow the request header map's arbitrary iteration order.
/// Values sharing a name retain their insertion order.
///
/// Casing is associated with the appended value's position among values with
/// the same case-insensitive name. Subsequent header mutation can change which
/// value occupies that position. Removed headers are never restored, and headers
/// generated by the protocol use their normal casing.
pub fn header_raw(mut builder: Builder, name: &str, value: &[u8]) -> Builder {
    let index = builder
        .headers_ref()
        .map(|h| h.get_all(name).iter().count());
    builder = builder.header(name, value);
    let (Some(index), Some(extensions)) = (index, builder.extensions_mut()) else {
        return builder;
    };
    if extensions.get::<HeaderCasing>().is_none() {
        extensions.insert(HeaderCasing::default());
    }
    let header_casing = extensions.get_mut::<HeaderCasing>().unwrap();
    // The successful builder.header call above validated the name.
    let key = HeaderName::from_bytes(name.as_bytes()).unwrap();
    if let Some((_, _, casing)) = header_casing
        .0
        .iter_mut()
        .find(|(n, i, _)| *n == key && *i == index)
    {
        *casing = name.into();
    } else {
        header_casing.0.push((key, index, name.into()));
    }
    builder
}

#[derive(Clone, Default)]
pub(super) struct HeaderCasing(Vec<(HeaderName, usize, String)>);

impl HeaderCasing {
    pub(super) fn get<'a>(&'a self, name: &'a HeaderName, index: usize) -> &'a str {
        self.0
            .iter()
            .find(|(n, i, _)| n == name && *i == index)
            .map(|(_, _, casing)| casing.as_str())
            .unwrap_or_else(|| name.as_str())
    }
}

#[cfg(test)]
mod tests {
    use super::{HeaderCasing, header_raw};
    use crate::Error;
    use crate::client::state::Redirect;
    use crate::client::{Call, RedirectAuthHeaders};
    use http::{HeaderValue, Request, StatusCode, Version, header};

    fn output(request: Request<()>, size: usize) -> Vec<u8> {
        let mut call = Call::new(request).unwrap().proceed();
        let mut result = Vec::new();
        while !call.can_proceed() {
            let mut buffer = vec![0; size];
            let n = call.write(&mut buffer).unwrap();
            result.extend_from_slice(&buffer[..n]);
        }
        result
    }

    #[test]
    fn typed_and_raw_headers_keep_map_order_incrementally() {
        for raw in [false, true] {
            let mut builder = Request::get("http://q.test");
            for (name, value) in [
                ("x-first", "one"),
                ("x-remove", "gone"),
                ("x-second", "middle"),
                ("x-first", "two"),
                ("x-third", "last"),
                ("x-second", "end"),
            ] {
                builder = if raw {
                    header_raw(builder, name, value.as_bytes())
                } else {
                    builder.header(name, value)
                };
            }
            builder.headers_mut().unwrap().remove("x-remove");
            let request = builder.body(()).unwrap();
            let mut expected = b"GET / HTTP/1.1\r\nhost: q.test\r\n".to_vec();
            for (name, value) in request.headers() {
                expected.extend_from_slice(name.as_str().as_bytes());
                expected.extend_from_slice(b": ");
                expected.extend_from_slice(value.as_bytes());
                expected.extend_from_slice(b"\r\n");
            }
            expected.extend_from_slice(b"\r\n");
            assert_eq!(output(request, 24), expected);
        }
    }

    #[test]
    fn preserves_mixed_duplicates_and_typed_headers_incrementally() {
        let builder = Request::get("http://q.test").header("X-SECRET", "typed");
        let builder = header_raw(builder, "X-SECRET", b"one");
        let builder = header_raw(builder, "x-SeCrEt", &[0xff]);
        let request = builder.header("X-SECRET", "last").body(()).unwrap();
        assert_eq!(output(request, 24), b"GET / HTTP/1.1\r\nhost: q.test\r\nx-secret: typed\r\nX-SECRET: one\r\nx-SeCrEt: \xff\r\nx-secret: last\r\n\r\n");
    }

    #[test]
    fn remove_and_reappend_uses_latest_explicit_casing() {
        let mut builder = header_raw(Request::get("http://q.test"), "X-SECRET", b"old");
        for name in ["x-SeCrEt", "X-secret", "X-SeCrEt"] {
            builder.headers_mut().unwrap().remove("x-secret");
            builder = header_raw(builder, name, b"new");
            let casing = builder
                .extensions_ref()
                .unwrap()
                .get::<HeaderCasing>()
                .unwrap();
            assert_eq!(casing.0.len(), 1);
            assert_eq!(casing.0[0].1, 0);
            assert_eq!(casing.0[0].2, name);
        }
        assert_eq!(
            output(builder.body(()).unwrap(), 24),
            b"GET / HTTP/1.1\r\nhost: q.test\r\nX-SeCrEt: new\r\n\r\n"
        );
    }

    #[test]
    fn validates_names_and_values() {
        for name in ["", "X Bad", "X:Bad", "X\r\nInjected", "é"] {
            assert!(
                header_raw(Request::get("http://q.test"), name, b"ok")
                    .body(())
                    .is_err()
            );
        }
        for value in [b"bad\r\nInjected: yes".as_slice(), b"bad\n", b"bad\0"] {
            assert!(
                header_raw(Request::get("http://q.test"), "X-Good", value)
                    .body(())
                    .is_err()
            );
        }
        let builder = Request::get("bad uri");
        assert!(header_raw(builder, "X-Good", b"ok").body(()).is_err());
    }

    #[test]
    fn amendments_do_not_consume_original_ordinals() {
        let request = header_raw(Request::get("http://q.test"), "X-SECRET", b"original")
            .body(())
            .unwrap();
        let mut call = Call::new(request).unwrap();
        call.header("X-SECRET", "amended").unwrap();
        let mut call = call.proceed();
        let mut buffer = [0; 256];
        let n = call.write(&mut buffer).unwrap();
        assert_eq!(
            &buffer[..n],
            b"GET / HTTP/1.1\r\nx-secret: amended\r\nhost: q.test\r\nX-SECRET: original\r\n\r\n"
        );
    }

    #[test]
    fn raw_headers_still_participate_in_protocol_validation() {
        let builder = header_raw(Request::get("http://q.test"), "HOST", b"q.test");
        let request = builder.header("host", "other.test").body(()).unwrap();
        let mut call = Call::new(request).unwrap().proceed();
        assert_eq!(
            call.write(&mut [0; 256]).unwrap_err(),
            Error::TooManyHostHeaders
        );
        let builder = header_raw(Request::post("http://q.test"), "Content-Length", b"1");
        let request = builder.header("content-length", "2").body(()).unwrap();
        let mut call = Call::new(request).unwrap().proceed();
        assert_eq!(
            call.write(&mut [0; 256]).unwrap_err(),
            Error::TooManyContentLengthHeaders
        );
        let request = header_raw(Request::post("http://q.test"), "CONTENT-LENGTH", b"1, 2")
            .body(())
            .unwrap();
        let mut call = Call::new(request).unwrap().proceed();
        assert_eq!(
            call.write(&mut [0; 256]).unwrap_err(),
            Error::BadContentLengthHeader
        );
        let request = header_raw(Request::get("http://q.test"), "X-SECRET", b"ok")
            .version(Version::HTTP_2)
            .body(())
            .unwrap();
        let mut call = Call::new(request).unwrap().proceed();
        assert_eq!(
            call.write(&mut [0; 256]).unwrap_err(),
            Error::UnsupportedVersion
        );
    }

    #[test]
    fn redirect_removes_sensitive_values_and_keeps_remaining_casing() {
        let mut builder = Request::get("http://q.test");
        for (name, value) in [
            ("AUTHORIZATION", b"secret".as_slice()),
            ("COOKIE", b"secret"),
            ("CONTENT-LENGTH", b"0"),
            ("X-SECRET", b"kept"),
        ] {
            builder = header_raw(builder, name, value);
        }
        let call = Call::new(builder.body(()).unwrap()).unwrap();
        let mut redirect: Call<Redirect> = Call::wrap(call.inner);
        redirect.inner.status = Some(StatusCode::FOUND);
        redirect.inner.location = Some(HeaderValue::from_static("http://other.test/"));
        let mut next = redirect
            .as_new_call(RedirectAuthHeaders::Never)
            .unwrap()
            .unwrap()
            .proceed();
        let mut buffer = [0; 256];
        let n = next.write(&mut buffer).unwrap();
        assert_eq!(
            &buffer[..n],
            b"GET / HTTP/1.1\r\nhost: other.test\r\nX-SECRET: kept\r\n\r\n"
        );
    }

    #[test]
    fn removed_metadata_does_not_affect_regenerated_headers() {
        let mut request = header_raw(Request::get("http://q.test"), "HOST", b"old.test")
            .body(())
            .unwrap();
        request.headers_mut().remove(header::HOST);
        assert_eq!(
            output(request, 24),
            b"GET / HTTP/1.1\r\nhost: q.test\r\n\r\n"
        );
    }
}