RustProxy/src/compression/tests.rs

mod tests {
    use super::*;
    use http_body_util::BodyExt;
    fn response() -> Response {
        let mut response =
            Response::new(Body::from("RustProxy compressible payload. ".repeat(500)));
        response
            .headers_mut()
            .insert("content-type", "text/plain".parse().unwrap());
        response
            .headers_mut()
            .insert("etag", "\"origin\"".parse().unwrap());
        response
    }
    #[tokio::test]
    async fn large_thresholds_do_not_truncate_or_panic() {
        let mut h = HeaderMap::new();
        h.insert("accept-encoding", "gzip".parse().unwrap());
        for threshold in [65_535, 65_536, 1_048_576] {
            let settings = Settings {
                gzip: true,
                min_bytes: threshold,
                ..Default::default()
            };
            settings.validate().unwrap();
            for (size, expected) in [(threshold - 1, false), (threshold, true)] {
                let mut r = Response::new(Body::from("a".repeat(size as usize)));
                r.headers_mut()
                    .insert("content-type", "text/plain".parse().unwrap());
                assert_eq!(
                    encode(r, &h, &Method::GET, &settings)
                        .await
                        .headers()
                        .contains_key("content-encoding"),
                    expected,
                    "threshold={threshold}, size={size}"
                );
            }
        }
    }
    #[tokio::test]
    async fn negotiation_and_body_streaming() {
        let s = Settings {
            gzip: true,
            br: true,
            zstd: true,
            deflate: true,
            ..Default::default()
        };
        for encoding in ["gzip", "br", "zstd", "deflate"] {
            let mut h = HeaderMap::new();
            h.insert("accept-encoding", encoding.parse().unwrap());
            let r = encode(response(), &h, &Method::GET, &s).await;
            assert_eq!(r.headers()["content-encoding"], encoding);
            let expected = match encoding {
                "gzip" => AppliedEncoding::Gzip,
                "br" => AppliedEncoding::Brotli,
                "zstd" => AppliedEncoding::Zstd,
                "deflate" => AppliedEncoding::Deflate,
                _ => unreachable!(),
            };
            assert_eq!(
                r.extensions().get::<AppliedEncoding>().copied(),
                Some(expected)
            );
            assert_eq!(r.headers()["etag"], "W/\"origin\"");
            assert!(r
                .headers()
                .get_all("vary")
                .iter()
                .any(|v| v.to_str().unwrap().contains("accept-encoding")));
            assert!(r.into_body().collect().await.unwrap().to_bytes().len() < 1000);
        }
        let mut h = HeaderMap::new();
        h.insert("accept-encoding", "gzip;q=0, identity;q=1".parse().unwrap());
        assert!(!encode(response(), &h, &Method::GET, &s)
            .await
            .headers()
            .contains_key("content-encoding"));
    }
    #[tokio::test]
    async fn quality_weights_can_select_each_enabled_codec() {
        let settings = Settings {
            gzip: true,
            br: true,
            zstd: true,
            deflate: true,
            ..Default::default()
        };
        let cases = [
            (
                "gzip;q=1, br;q=0.2, zstd;q=0.2, deflate;q=0.2",
                "gzip",
                AppliedEncoding::Gzip,
            ),
            (
                "gzip;q=0.2, br;q=1, zstd;q=0.2, deflate;q=0.2",
                "br",
                AppliedEncoding::Brotli,
            ),
            (
                "gzip;q=0.2, br;q=0.2, zstd;q=1, deflate;q=0.2",
                "zstd",
                AppliedEncoding::Zstd,
            ),
            (
                "gzip;q=0.2, br;q=0.2, zstd;q=0.2, deflate;q=1",
                "deflate",
                AppliedEncoding::Deflate,
            ),
        ];
        for (accept, expected_header, expected_marker) in cases {
            let mut headers = HeaderMap::new();
            headers.insert("accept-encoding", accept.parse().unwrap());
            let response = encode(response(), &headers, &Method::GET, &settings).await;
            assert_eq!(response.headers()["content-encoding"], expected_header);
            assert_eq!(
                response.extensions().get::<AppliedEncoding>().copied(),
                Some(expected_marker),
                "accept-encoding={accept}"
            );
        }
    }

    #[tokio::test]
    async fn no_transform_personal_and_sse_are_not_compressed() {
        let s = Settings {
            gzip: true,
            ..Default::default()
        };
        let mut h = HeaderMap::new();
        h.insert("accept-encoding", "gzip".parse().unwrap());
        for name in ["authorization", "cookie", "range"] {
            let mut private = h.clone();
            private.insert(
                http::HeaderName::from_bytes(name.as_bytes()).unwrap(),
                "value".parse().unwrap(),
            );
            assert!(!encode(response(), &private, &Method::GET, &s)
                .await
                .headers()
                .contains_key("content-encoding"));
        }
        let mut r = response();
        r.headers_mut()
            .insert("cache-control", "no-transform".parse().unwrap());
        assert!(!encode(r, &h, &Method::GET, &s)
            .await
            .headers()
            .contains_key("content-encoding"));
        let mut r = response();
        r.headers_mut()
            .insert("content-type", "text/event-stream".parse().unwrap());
        assert!(!encode(r, &h, &Method::GET, &s)
            .await
            .headers()
            .contains_key("content-encoding"));
    }
}