From f74da51a52b1f2290999df446d5316f8725f0286 Mon Sep 17 00:00:00 2001 From: zuisong Date: Sat, 13 Apr 2024 14:25:54 +0800 Subject: [PATCH] decode responses in zstd format --- Cargo.lock | 42 ++++++++++++++++++++ Cargo.toml | 1 + src/decoder.rs | 16 +++++++- src/main.rs | 2 +- tests/cli.rs | 47 +++++++++++++++++------ tests/fixtures/responses/README.md | 3 ++ tests/fixtures/responses/hello_world.zst | Bin 0 -> 25 bytes 7 files changed, 98 insertions(+), 13 deletions(-) create mode 100644 tests/fixtures/responses/hello_world.zst diff --git a/Cargo.lock b/Cargo.lock index 40b3ef60..21ff2156 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -222,6 +222,10 @@ name = "cc" version = "1.0.92" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2678b2e3449475e95b0aa6f9b506a28e61b3dc8996592b983695e8ebb58a8b41" +dependencies = [ + "jobserver", + "libc", +] [[package]] name = "cfg-if" @@ -883,6 +887,15 @@ version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" +[[package]] +name = "jobserver" +version = "0.1.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "685a7d121ee3f65ae4fddd72b25a04bb36b6af81bc0828f7d5434c0fe60fa3a2" +dependencies = [ + "libc", +] + [[package]] name = "js-sys" version = "0.3.69" @@ -2429,6 +2442,7 @@ dependencies = [ "tokio", "unicode-width", "url", + "zstd", ] [[package]] @@ -2445,3 +2459,31 @@ name = "zeroize" version = "1.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "525b4ec142c6b68a2d10f01f7bbf6755599ca3f81ea53b8431b7dd348f5fdb2d" + +[[package]] +name = "zstd" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d789b1514203a1120ad2429eae43a7bd32b90976a7bb8a05f7ec02fa88cc23a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1cd99b45c6bc03a018c8b8a86025678c87e55526064e38f9df301989dce7ec0a" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.10+zstd.1.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c253a4914af5bafc8fa8c86ee400827e83cf6ec01195ec1f1ed8441bf00d65aa" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/Cargo.toml b/Cargo.toml index 81623bc7..43114caa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -46,6 +46,7 @@ termcolor = "1.1.2" time = "0.3.16" unicode-width = "0.1.9" url = "2.2.2" +zstd = { version = "0.13.1", default-features = false } [dependencies.reqwest] version = "0.12.3" diff --git a/src/decoder.rs b/src/decoder.rs index 7d7541f9..e5ebe6b0 100644 --- a/src/decoder.rs +++ b/src/decoder.rs @@ -1,15 +1,17 @@ -use std::io::{self, Read}; +use std::io::{self, BufReader, Read}; use std::str::FromStr; use brotli::Decompressor as BrotliDecoder; use flate2::read::{GzDecoder, ZlibDecoder}; use reqwest::header::{HeaderMap, CONTENT_ENCODING, CONTENT_LENGTH, TRANSFER_ENCODING}; +use zstd::Decoder as ZstdDecoder; #[derive(Debug)] pub enum CompressionType { Gzip, Deflate, Brotli, + Zstd, } impl FromStr for CompressionType { @@ -19,6 +21,7 @@ impl FromStr for CompressionType { "gzip" => Ok(CompressionType::Gzip), "deflate" => Ok(CompressionType::Deflate), "br" => Ok(CompressionType::Brotli), + "zstd" => Ok(CompressionType::Zstd), _ => Err(anyhow::anyhow!("unknown compression type")), } } @@ -88,6 +91,7 @@ enum Decoder { Gzip(GzDecoder>), Deflate(ZlibDecoder>), Brotli(BrotliDecoder>), + Zstd(ZstdDecoder<'static, BufReader>>), } impl Read for Decoder { @@ -121,6 +125,15 @@ impl Read for Decoder { format!("error decoding brotli response body: {}", e), )), }, + Decoder::Zstd(decoder) => match decoder.read(buf) { + Ok(n) => Ok(n), + Err(e) if decoder.get_ref().get_ref().has_errored => Err(e), + Err(_) if !decoder.get_ref().get_ref().has_read_data => Ok(0), + Err(e) => Err(io::Error::new( + e.kind(), + format!("error decoding zstd response body: {}", e), + )), + }, } } } @@ -135,6 +148,7 @@ pub fn decompress( Some(CompressionType::Deflate) => Decoder::Deflate(ZlibDecoder::new(reader)), Some(CompressionType::Brotli) => Decoder::Brotli(BrotliDecoder::new(reader, 4096)), None => Decoder::PlainText(reader), + Some(CompressionType::Zstd) => Decoder::Zstd(ZstdDecoder::new(reader).unwrap()), } } diff --git a/src/main.rs b/src/main.rs index 01f5601a..3a74e52b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -365,7 +365,7 @@ fn run(args: Cli) -> Result { .request(method, url.clone()) .header( ACCEPT_ENCODING, - HeaderValue::from_static("gzip, deflate, br"), + HeaderValue::from_static("gzip, deflate, br, zstd"), ) .header(USER_AGENT, get_user_agent()); diff --git a/tests/cli.rs b/tests/cli.rs index 9e774f9f..e3b258b1 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -418,7 +418,7 @@ fn verbose() { .stdout(indoc! {r#" POST / HTTP/1.1 Accept: application/json, */*;q=0.5 - Accept-Encoding: gzip, deflate, br + Accept-Encoding: gzip, deflate, br, zstd Connection: keep-alive Content-Length: 9 Content-Type: application/json @@ -940,7 +940,7 @@ fn digest_auth_with_redirection() { .stdout(indoc! {r#" GET /login_page HTTP/1.1 Accept: */* - Accept-Encoding: gzip, deflate, br + Accept-Encoding: gzip, deflate, br, zstd Connection: keep-alive Host: http.mock User-Agent: xh/0.0.0 (test mode) @@ -954,7 +954,7 @@ fn digest_auth_with_redirection() { GET /login_page HTTP/1.1 Accept: */* - Accept-Encoding: gzip, deflate, br + Accept-Encoding: gzip, deflate, br, zstd Authorization: Digest username="ahmed", realm="me@xh.com", nonce="e5051361f053723a807674177fc7022f", uri="/login_page", qop=auth, nc=00000001, cnonce="f2/wE4q74E6zIJEtWaHKaf5wv/H5QzzpXusqGemxURZJ", response="894fd5ee1dcc702df7e4a6abed37fd56", opaque="9dcf562038f1ec1c8d02f218ef0e7a4b", algorithm=MD5 Connection: keep-alive Host: http.mock @@ -969,7 +969,7 @@ fn digest_auth_with_redirection() { GET /admin_page HTTP/1.1 Accept: */* - Accept-Encoding: gzip, deflate, br + Accept-Encoding: gzip, deflate, br, zstd Connection: keep-alive Host: http.mock User-Agent: xh/0.0.0 (test mode) @@ -2015,7 +2015,7 @@ fn can_unset_default_headers() { .stdout(indoc! {r#" GET / HTTP/1.1 Accept: */* - Accept-Encoding: gzip, deflate, br + Accept-Encoding: gzip, deflate, br, zstd Connection: keep-alive Host: http.mock @@ -2030,7 +2030,7 @@ fn can_unset_headers() { .stdout(indoc! {r#" GET / HTTP/1.1 Accept: */* - Accept-Encoding: gzip, deflate, br + Accept-Encoding: gzip, deflate, br, zstd Connection: keep-alive Hello: world Host: http.mock @@ -2047,7 +2047,7 @@ fn can_set_unset_header() { .stdout(indoc! {r#" GET / HTTP/1.1 Accept: */* - Accept-Encoding: gzip, deflate, br + Accept-Encoding: gzip, deflate, br, zstd Connection: keep-alive Hello: world Host: http.mock @@ -2780,7 +2780,7 @@ fn print_intermediate_requests_and_responses() { .stdout(indoc! {r#" GET /first_page HTTP/1.1 Accept: */* - Accept-Encoding: gzip, deflate, br + Accept-Encoding: gzip, deflate, br, zstd Connection: keep-alive Host: http.mock User-Agent: xh/0.0.0 (test mode) @@ -2794,7 +2794,7 @@ fn print_intermediate_requests_and_responses() { GET /second_page HTTP/1.1 Accept: */* - Accept-Encoding: gzip, deflate, br + Accept-Encoding: gzip, deflate, br, zstd Connection: keep-alive Host: http.mock User-Agent: xh/0.0.0 (test mode) @@ -2835,7 +2835,7 @@ fn history_print() { .stdout(indoc! {r#" GET /first_page HTTP/1.1 Accept: */* - Accept-Encoding: gzip, deflate, br + Accept-Encoding: gzip, deflate, br, zstd Connection: keep-alive Host: http.mock User-Agent: xh/0.0.0 (test mode) @@ -2847,7 +2847,7 @@ fn history_print() { GET /second_page HTTP/1.1 Accept: */* - Accept-Encoding: gzip, deflate, br + Accept-Encoding: gzip, deflate, br, zstd Connection: keep-alive Host: http.mock User-Agent: xh/0.0.0 (test mode) @@ -3378,6 +3378,31 @@ fn brotli() { "#}); } +#[test] +fn zstd() { + let server = server::http(|_req| async move { + let compressed_bytes = fs::read("./tests/fixtures/responses/hello_world.zst").unwrap(); + hyper::Response::builder() + .header("date", "N/A") + .header("content-encoding", "zstd") + .body(compressed_bytes.into()) + .unwrap() + }); + + get_command() + .arg(server.base_url()) + .assert() + .stdout(indoc! {r#" + HTTP/1.1 200 OK + Content-Encoding: zstd + Content-Length: 25 + Date: N/A + + Hello world + + "#}); +} + #[test] fn empty_response_with_content_encoding() { let server = server::http(|_req| async move { diff --git a/tests/fixtures/responses/README.md b/tests/fixtures/responses/README.md index dadf983a..9bd8669c 100644 --- a/tests/fixtures/responses/README.md +++ b/tests/fixtures/responses/README.md @@ -9,4 +9,7 @@ $ pigz -z hello_world # hello_world.zz $ echo "Hello world" > hello_world $ brotli hello_world # hello_world.br + +$ echo "Hello world" > hello_world +$ zstd hello_world # hello_world.zst ``` diff --git a/tests/fixtures/responses/hello_world.zst b/tests/fixtures/responses/hello_world.zst new file mode 100644 index 0000000000000000000000000000000000000000..a42154ef9c510b4ff47e5c4bcbaa87d50837c960 GIT binary patch literal 25 gcmdPcs{dDoCy{}{BQ+-{U!gp|C?|!hXMx2m0Bs)#-v9sr literal 0 HcmV?d00001