Skip to main content

aria2_protocol/http/
encoding.rs

1use flate2::read::{DeflateDecoder, GzDecoder};
2use std::io::Read;
3use tracing::debug;
4
5pub struct HttpEncoding;
6
7impl HttpEncoding {
8    pub fn decode(body: &[u8], content_encoding: Option<&str>) -> Result<Vec<u8>, String> {
9        let encoding = content_encoding.unwrap_or("").to_lowercase();
10        match encoding.as_str() {
11            "" | "identity" => Ok(body.to_vec()),
12            "gzip" | "x-gzip" => Self::decode_gzip(body),
13            "deflate" => Self::decode_deflate(body),
14            "br" => Err("Brotli compression not supported yet".to_string()),
15            "compress" => Err("compress compression format is deprecated".to_string()),
16            other => Err(format!("Unsupported Content-Encoding: {}", other)),
17        }
18    }
19
20    fn decode_gzip(data: &[u8]) -> Result<Vec<u8>, String> {
21        if data.is_empty() {
22            return Ok(Vec::new());
23        }
24        let mut decoder = GzDecoder::new(data);
25        let mut decompressed = Vec::with_capacity(data.len() * 4);
26        decoder
27            .read_to_end(&mut decompressed)
28            .map_err(|e| format!("gzip decompression failed: {}", e))?;
29        debug!(
30            "gzip decompression complete: {} -> {} bytes",
31            data.len(),
32            decompressed.len()
33        );
34        Ok(decompressed)
35    }
36
37    fn decode_deflate(data: &[u8]) -> Result<Vec<u8>, String> {
38        if data.is_empty() {
39            return Ok(Vec::new());
40        }
41        let mut decoder = DeflateDecoder::new(data);
42        let mut decompressed = Vec::with_capacity(data.len() * 4);
43        decoder
44            .read_to_end(&mut decompressed)
45            .map_err(|e| format!("deflate decompression failed: {}", e))?;
46        debug!(
47            "deflate decompression complete: {} -> {} bytes",
48            data.len(),
49            decompressed.len()
50        );
51        Ok(decompressed)
52    }
53
54    pub fn detect_best_accept_encoding() -> &'static str {
55        "gzip, deflate"
56    }
57
58    pub fn is_compressed(content_encoding: Option<&str>) -> bool {
59        matches!(
60            content_encoding.map(|e| e.to_lowercase()).as_deref(),
61            Some("gzip" | "deflate" | "br" | "compress")
62        )
63    }
64}
65
66pub struct ChunkedDecoder;
67
68impl ChunkedDecoder {
69    pub fn decode(data: &[u8]) -> Result<Vec<u8>, String> {
70        let mut result = Vec::new();
71        let mut pos = 0;
72
73        while pos < data.len() {
74            let line_end = data[pos..]
75                .iter()
76                .position(|&b| b == b'\r' || b == b'\n')
77                .ok_or("Chunked encoding format error: cannot find chunk size line")?;
78
79            let size_str = unsafe { std::str::from_utf8_unchecked(&data[pos..pos + line_end]) };
80            let size_str = size_str.trim();
81            let chunk_size: usize = usize::from_str_radix(size_str, 16)
82                .map_err(|e| format!("Chunk size parsing failed: {}", e))?;
83
84            pos += line_end;
85            if pos < data.len() && data[pos] == b'\r' {
86                pos += 1;
87            }
88            if pos < data.len() && data[pos] == b'\n' {
89                pos += 1;
90            }
91
92            if chunk_size == 0 {
93                break;
94            }
95
96            if pos + chunk_size > data.len() {
97                return Err("Chunked data truncated".to_string());
98            }
99
100            result.extend_from_slice(&data[pos..pos + chunk_size]);
101            pos += chunk_size;
102
103            if pos + 1 < data.len() && data[pos] == b'\r' && data[pos + 1] == b'\n' {
104                pos += 2;
105            } else if pos < data.len() && data[pos] == b'\n' {
106                pos += 1;
107            }
108        }
109
110        Ok(result)
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn test_decode_identity() {
120        let data = b"hello world";
121        let result = HttpEncoding::decode(data, None).unwrap();
122        assert_eq!(result, b"hello world");
123
124        let result = HttpEncoding::decode(data, Some("identity")).unwrap();
125        assert_eq!(result, b"hello world");
126    }
127
128    #[test]
129    fn test_gzip_roundtrip() {
130        use flate2::Compression;
131        use flate2::write::GzEncoder;
132        use std::io::Write;
133
134        let original = b"Hello, this is a test string that should compress well!";
135        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
136        encoder.write_all(original).unwrap();
137        let compressed = encoder.finish().unwrap();
138
139        let decompressed = HttpEncoding::decode(&compressed, Some("gzip")).unwrap();
140        assert_eq!(decompressed, original);
141    }
142
143    #[test]
144    fn test_detect_compression() {
145        assert!(HttpEncoding::is_compressed(Some("gzip")));
146        assert!(HttpEncoding::is_compressed(Some("deflate")));
147        assert!(HttpEncoding::is_compressed(Some("br")));
148        assert!(!HttpEncoding::is_compressed(None));
149        assert!(!HttpEncoding::is_compressed(Some("identity")));
150    }
151
152    #[test]
153    fn test_chunked_decoder() {
154        let raw_data = "5\r\nhello\r\n6\r\n world\r\n0\r\n\r\n";
155        let decoded = ChunkedDecoder::decode(raw_data.as_bytes()).unwrap();
156        assert_eq!(decoded, b"hello world");
157    }
158}