1use std::io::Read;
7
8pub const DECODE_MAX: usize = 64 * 1024;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum Sniffed {
14 Png { width: u32, height: u32 },
15 Jpeg { width: u32, height: u32 },
16 Gif { width: u32, height: u32 },
17 Pdf,
18 Gzip,
19 Zstd,
20 Zip,
21 Parquet,
22 Arrow,
23 Utf8,
24}
25
26impl Sniffed {
27 pub fn label(self) -> String {
29 match self {
30 Sniffed::Png { width, height } => format!("PNG {width}x{height}"),
31 Sniffed::Jpeg { width, height } if width > 0 => format!("JPEG {width}x{height}"),
32 Sniffed::Jpeg { .. } => "JPEG".to_string(),
33 Sniffed::Gif { width, height } => format!("GIF {width}x{height}"),
34 Sniffed::Pdf => "PDF".to_string(),
35 Sniffed::Gzip => "gzip".to_string(),
36 Sniffed::Zstd => "zstd".to_string(),
37 Sniffed::Zip => "zip".to_string(),
38 Sniffed::Parquet => "Parquet".to_string(),
39 Sniffed::Arrow => "Arrow".to_string(),
40 Sniffed::Utf8 => "UTF-8 text".to_string(),
41 }
42 }
43
44 pub fn extension(self) -> &'static str {
46 match self {
47 Sniffed::Png { .. } => "png",
48 Sniffed::Jpeg { .. } => "jpg",
49 Sniffed::Gif { .. } => "gif",
50 Sniffed::Pdf => "pdf",
51 Sniffed::Gzip => "gz",
52 Sniffed::Zstd => "zst",
53 Sniffed::Zip => "zip",
54 Sniffed::Parquet => "parquet",
55 Sniffed::Arrow => "arrow",
56 Sniffed::Utf8 => "txt",
57 }
58 }
59
60 pub fn is_document(self) -> bool {
62 matches!(
63 self,
64 Sniffed::Png { .. } | Sniffed::Jpeg { .. } | Sniffed::Gif { .. } | Sniffed::Pdf
65 )
66 }
67}
68
69fn be32(b: &[u8], at: usize) -> Option<u32> {
70 Some(u32::from_be_bytes(b.get(at..at + 4)?.try_into().ok()?))
71}
72
73fn le16(b: &[u8], at: usize) -> Option<u32> {
74 Some(u16::from_le_bytes(b.get(at..at + 2)?.try_into().ok()?) as u32)
75}
76
77fn be16(b: &[u8], at: usize) -> Option<u32> {
78 Some(u16::from_be_bytes(b.get(at..at + 2)?.try_into().ok()?) as u32)
79}
80
81fn jpeg_size(b: &[u8]) -> Option<(u32, u32)> {
83 let mut i = 2;
84 let end = b.len().min(64 * 1024);
85 while i + 9 < end {
86 if b[i] != 0xff {
87 i += 1;
88 continue;
89 }
90 let marker = b[i + 1];
91 let len = be16(b, i + 2)? as usize;
92 if matches!(marker, 0xc0..=0xcf) && !matches!(marker, 0xc4 | 0xc8 | 0xcc) {
93 return Some((be16(b, i + 7)?, be16(b, i + 5)?));
94 }
95 i += 2 + len;
96 }
97 None
98}
99
100pub fn sniff(b: &[u8]) -> Option<Sniffed> {
103 if b.starts_with(b"\x89PNG\r\n\x1a\n") {
104 return Some(Sniffed::Png {
105 width: be32(b, 16).unwrap_or(0),
106 height: be32(b, 20).unwrap_or(0),
107 });
108 }
109 if b.starts_with(&[0xff, 0xd8, 0xff]) {
110 let (width, height) = jpeg_size(b).unwrap_or((0, 0));
111 return Some(Sniffed::Jpeg { width, height });
112 }
113 if b.starts_with(b"GIF87a") || b.starts_with(b"GIF89a") {
114 return Some(Sniffed::Gif {
115 width: le16(b, 6).unwrap_or(0),
116 height: le16(b, 8).unwrap_or(0),
117 });
118 }
119 if b.starts_with(b"%PDF-") {
120 return Some(Sniffed::Pdf);
121 }
122 if b.starts_with(&[0x1f, 0x8b]) {
123 return Some(Sniffed::Gzip);
124 }
125 if b.starts_with(&[0x28, 0xb5, 0x2f, 0xfd]) {
126 return Some(Sniffed::Zstd);
127 }
128 if b.starts_with(b"PK\x03\x04") || b.starts_with(b"PK\x05\x06") {
129 return Some(Sniffed::Zip);
130 }
131 if b.len() >= 8 && b.starts_with(b"PAR1") && b.ends_with(b"PAR1") {
132 return Some(Sniffed::Parquet);
133 }
134 if b.starts_with(b"ARROW1") {
135 return Some(Sniffed::Arrow);
136 }
137 if !b.is_empty() && !b.contains(&0) && std::str::from_utf8(b).is_ok() {
138 return Some(Sniffed::Utf8);
139 }
140 None
141}
142
143#[derive(Debug, Clone, PartialEq, Eq)]
146pub struct Decoded {
147 pub text: String,
148 pub from: &'static str,
149 pub cut: bool,
150}
151
152fn text_of(mut out: Vec<u8>, cut: bool) -> Option<String> {
154 match String::from_utf8(out.clone()) {
155 Ok(text) if !text.contains('\0') => Some(text),
156 Ok(_) => None,
157 Err(e) if cut && e.utf8_error().error_len().is_none() => {
159 out.truncate(e.utf8_error().valid_up_to());
160 String::from_utf8(out).ok().filter(|t| !t.contains('\0'))
161 }
162 Err(_) => None,
163 }
164}
165
166pub fn decode_text(bytes: &[u8], kind: Option<Sniffed>) -> Option<Decoded> {
169 let (reader, from): (Box<dyn Read + '_>, &'static str) = match kind? {
170 Sniffed::Utf8 => {
171 let cut = bytes.len() > DECODE_MAX;
172 let head = &bytes[..bytes.len().min(DECODE_MAX)];
173 return text_of(head.to_vec(), cut).map(|text| Decoded {
174 text,
175 from: "UTF-8",
176 cut,
177 });
178 }
179 Sniffed::Gzip => (Box::new(flate2::read::GzDecoder::new(bytes)), "gzip"),
180 Sniffed::Zstd => (Box::new(zstd::Decoder::new(bytes).ok()?), "zstd"),
181 _ => return None,
182 };
183 let mut out = Vec::new();
184 reader
185 .take(DECODE_MAX as u64 + 1)
186 .read_to_end(&mut out)
187 .ok()?;
188 let cut = out.len() > DECODE_MAX;
189 out.truncate(DECODE_MAX);
190 text_of(out, cut).map(|text| Decoded { text, from, cut })
191}
192
193#[cfg(test)]
194mod tests {
195 use super::*;
196 use std::io::Write;
197
198 fn png(width: u32, height: u32) -> Vec<u8> {
199 let mut b = b"\x89PNG\r\n\x1a\n\0\0\0\x0dIHDR".to_vec();
200 b.extend(width.to_be_bytes());
201 b.extend(height.to_be_bytes());
202 b.extend([8, 2, 0, 0, 0]);
203 b
204 }
205
206 fn gzip(text: &str) -> Vec<u8> {
207 let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
208 e.write_all(text.as_bytes()).unwrap();
209 e.finish().unwrap()
210 }
211
212 #[test]
214 fn sniffs_the_common_kinds_from_their_magic() {
215 let jpeg = [
216 0xff, 0xd8, 0xff, 0xe0, 0x00, 0x04, 0x00, 0x00, 0xff, 0xc0, 0x00, 0x11, 0x08, 0x00,
217 0x20, 0x00, 0x40, 0x03, 0, 0, 0, 0, 0, 0,
218 ];
219 let zstd_bytes = zstd::encode_all(&b"hello zstd"[..], 3).unwrap();
220 let cases: Vec<(Vec<u8>, Option<&str>)> = vec![
221 (png(4, 4), Some("PNG 4x4")),
222 (jpeg.to_vec(), Some("JPEG 64x32")),
223 (b"GIF89a\x10\x00\x08\x00".to_vec(), Some("GIF 16x8")),
224 (b"%PDF-1.7\n".to_vec(), Some("PDF")),
225 (gzip("hello"), Some("gzip")),
226 (zstd_bytes, Some("zstd")),
227 (b"PK\x03\x04rest".to_vec(), Some("zip")),
228 (b"PAR1xxxxPAR1".to_vec(), Some("Parquet")),
229 (b"ARROW1\0\0".to_vec(), Some("Arrow")),
230 ("Grüße, world".as_bytes().to_vec(), Some("UTF-8 text")),
231 (vec![0xe9, 0x91, 0x23, 0xbe, 0x00], None),
232 (Vec::new(), None),
233 (b"a\0b".to_vec(), None),
234 ];
235 for (bytes, want) in cases {
236 assert_eq!(
237 sniff(&bytes).map(Sniffed::label).as_deref(),
238 want,
239 "{bytes:02x?}"
240 );
241 }
242 }
243
244 #[test]
245 fn decodes_utf8_and_compressed_text_within_the_cap() {
246 let d = decode_text("ünï\nline".as_bytes(), Some(Sniffed::Utf8)).unwrap();
247 assert_eq!(
248 (d.text.as_str(), d.from, d.cut),
249 ("ünï\nline", "UTF-8", false)
250 );
251 let d = decode_text(&gzip("hello gzip"), Some(Sniffed::Gzip)).unwrap();
252 assert_eq!(d.text, "hello gzip");
253 let z = zstd::encode_all(&b"hello zstd"[..], 3).unwrap();
254 assert_eq!(
255 decode_text(&z, Some(Sniffed::Zstd)).unwrap().text,
256 "hello zstd"
257 );
258 let big = "a".repeat(4 << 20);
260 let d = decode_text(&gzip(&big), Some(Sniffed::Gzip)).unwrap();
261 assert!(d.cut);
262 assert_eq!(d.text.len(), DECODE_MAX);
263 let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
265 e.write_all(&[0u8, 1, 2, 0xff]).unwrap();
266 assert!(decode_text(&e.finish().unwrap(), Some(Sniffed::Gzip)).is_none());
267 assert!(decode_text(&png(1, 1), sniff(&png(1, 1))).is_none());
268 }
269}