1use std::io::Read;
11use std::path::Path;
12
13use serde::{Deserialize, Serialize};
14use serde_json::Value;
15use sha2::{Digest, Sha256};
16
17use crate::error::{Error, Result};
18
19pub const HARD_MAX_FILE_BYTES: u64 = 1024 * 1024 * 1024;
21pub const HARD_MAX_FILES: usize = 20;
23pub const MAX_NAME_BYTES: usize = 100;
25
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
29pub struct FileRef {
30 pub id: String,
32 pub name: String,
34 pub size: u64,
35 #[serde(rename = "type", default)]
38 pub mime: String,
39}
40
41impl FileRef {
42 pub fn hex(&self) -> &str {
44 self.id.strip_prefix("sha256:").unwrap_or(&self.id)
45 }
46
47 fn check(&self) -> Result<()> {
48 check_id(&self.id)?;
49 if clean_name(&self.name) != self.name {
50 return Err(Error::Invalid(format!(
51 "file name {:?} is not a clean name",
52 self.name
53 )));
54 }
55 if self.size > HARD_MAX_FILE_BYTES {
56 return Err(Error::Invalid(format!(
57 "file {} is {} bytes; no file may be over {HARD_MAX_FILE_BYTES}",
58 self.name, self.size
59 )));
60 }
61 if self.mime.len() > 100 || self.mime.chars().any(|c| c.is_control()) {
62 return Err(Error::Invalid("file type is not a plain type name".into()));
63 }
64 Ok(())
65 }
66}
67
68pub fn check_id(id: &str) -> Result<()> {
70 let ok = id.strip_prefix("sha256:").is_some_and(|h| {
71 h.len() == 64 && h.bytes().all(|b| matches!(b, b'0'..=b'9' | b'a'..=b'f'))
72 });
73 if ok {
74 Ok(())
75 } else {
76 Err(Error::Invalid(format!("{id:?} is not a file id")))
77 }
78}
79
80pub fn refs(data: &Value) -> Result<Vec<FileRef>> {
84 let Some(files) = data.get("files") else {
85 return Ok(Vec::new());
86 };
87 let refs: Vec<FileRef> = serde_json::from_value(files.clone())
88 .map_err(|e| Error::Invalid(format!("data.files is not a list of files: {e}")))?;
89 if refs.len() > HARD_MAX_FILES {
90 return Err(Error::Invalid(format!(
91 "a message may carry at most {HARD_MAX_FILES} files"
92 )));
93 }
94 for r in &refs {
95 r.check()?;
96 }
97 Ok(refs)
98}
99
100pub fn hash_file(path: &Path) -> Result<(String, u64)> {
102 let mut f = std::fs::File::open(path)?;
103 let mut h = Sha256::new();
104 let mut buf = vec![0u8; 256 * 1024];
105 let mut size = 0u64;
106 loop {
107 let n = f.read(&mut buf)?;
108 if n == 0 {
109 break;
110 }
111 h.update(&buf[..n]);
112 size += n as u64;
113 }
114 Ok((
115 format!("sha256:{}", data_encoding::HEXLOWER.encode(&h.finalize())),
116 size,
117 ))
118}
119
120const RESERVED: &[&str] = &[
123 "con", "prn", "aux", "nul", "com1", "com2", "com3", "com4", "com5", "com6", "com7", "com8",
124 "com9", "lpt1", "lpt2", "lpt3", "lpt4", "lpt5", "lpt6", "lpt7", "lpt8", "lpt9",
125];
126
127pub fn clean_name(raw: &str) -> String {
131 let base = raw.rsplit(['/', '\\']).next().unwrap_or("");
132 let mut s: String = base
133 .chars()
134 .map(|c| {
135 if c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-') {
136 c
137 } else {
138 '_'
139 }
140 })
141 .collect();
142 while s.contains("..") {
143 s = s.replace("..", ".");
144 }
145 let s = s.trim_matches(|c| c == '.' || c == ' ');
146 let mut s = if s.is_empty() {
147 "file".to_string()
148 } else {
149 s.to_string()
150 };
151 if s.len() > MAX_NAME_BYTES {
152 let ext = s
154 .rfind('.')
155 .map(|i| s[i..].to_string())
156 .filter(|e| e.len() <= 16)
157 .unwrap_or_default();
158 let stem_len = MAX_NAME_BYTES - ext.len();
159 s = format!("{}{ext}", &s[..stem_len]);
160 }
161 let stem = s.split('.').next().unwrap_or("").to_ascii_lowercase();
162 if RESERVED.contains(&stem.as_str()) {
163 s = format!("_{s}");
164 s.truncate(MAX_NAME_BYTES);
165 }
166 s
167}
168
169fn ext_of(name: &str) -> String {
170 match name.rfind('.') {
171 Some(i) if i > 0 => name[i + 1..].to_ascii_lowercase(),
172 _ => String::new(),
173 }
174}
175
176const RUNNABLE: &[&str] = &[
178 "exe", "com", "bat", "cmd", "ps1", "psm1", "vbs", "vbe", "js", "jse", "wsf", "wsh", "hta",
179 "msi", "msp", "scr", "pif", "cpl", "jar", "sh", "bash", "zsh", "command", "app", "dmg", "pkg",
180 "run", "appimage", "deb", "rpm", "lnk", "reg", "desktop",
181];
182
183pub fn runnable_name(name: &str) -> bool {
185 RUNNABLE.contains(&ext_of(name).as_str())
186}
187
188pub fn saved_name(name: &str) -> String {
191 let clean = clean_name(name);
192 if runnable_name(&clean) {
193 format!("{clean}.unsafe")
194 } else {
195 clean
196 }
197}
198
199pub fn type_for_name(name: &str) -> &'static str {
201 match ext_of(name).as_str() {
202 "txt" | "log" | "md" | "csv" | "tsv" | "json" | "jsonl" | "yaml" | "yml" | "toml"
203 | "ini" | "xml" | "html" | "htm" | "css" | "rs" | "py" | "go" | "ts" | "java" | "c"
204 | "h" | "cpp" | "rb" | "sql" | "diff" | "patch" => "text/plain",
205 "png" => "image/png",
206 "jpg" | "jpeg" => "image/jpeg",
207 "gif" => "image/gif",
208 "webp" => "image/webp",
209 "pdf" => "application/pdf",
210 "zip" => "application/zip",
211 "gz" | "tgz" => "application/gzip",
212 _ => "application/octet-stream",
213 }
214}
215
216#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
220#[serde(rename_all = "lowercase")]
221pub enum Kind {
222 Text,
223 Png,
224 Jpeg,
225 Gif,
226 Webp,
227 Pdf,
228 Zip,
229 Gzip,
230 Program,
232 Other,
233}
234
235impl Kind {
236 pub fn as_str(&self) -> &'static str {
237 match self {
238 Kind::Text => "text",
239 Kind::Png => "png",
240 Kind::Jpeg => "jpeg",
241 Kind::Gif => "gif",
242 Kind::Webp => "webp",
243 Kind::Pdf => "pdf",
244 Kind::Zip => "zip",
245 Kind::Gzip => "gzip",
246 Kind::Program => "program",
247 Kind::Other => "other",
248 }
249 }
250
251 pub fn mime(&self) -> &'static str {
252 match self {
253 Kind::Text => "text/plain",
254 Kind::Png => "image/png",
255 Kind::Jpeg => "image/jpeg",
256 Kind::Gif => "image/gif",
257 Kind::Webp => "image/webp",
258 Kind::Pdf => "application/pdf",
259 Kind::Zip => "application/zip",
260 Kind::Gzip => "application/gzip",
261 Kind::Program | Kind::Other => "application/octet-stream",
262 }
263 }
264
265 pub fn is_safe(&self) -> bool {
267 matches!(
268 self,
269 Kind::Text | Kind::Png | Kind::Jpeg | Kind::Gif | Kind::Webp | Kind::Pdf | Kind::Zip
270 )
271 }
272}
273
274fn magic(head: &[u8]) -> Option<Kind> {
275 let starts = |m: &[u8]| head.starts_with(m);
276 if starts(b"\x89PNG\r\n\x1a\n") {
277 Some(Kind::Png)
278 } else if starts(b"\xff\xd8\xff") {
279 Some(Kind::Jpeg)
280 } else if starts(b"GIF87a") || starts(b"GIF89a") {
281 Some(Kind::Gif)
282 } else if head.len() >= 12 && &head[..4] == b"RIFF" && &head[8..12] == b"WEBP" {
283 Some(Kind::Webp)
284 } else if starts(b"%PDF-") {
285 Some(Kind::Pdf)
286 } else if starts(b"PK\x03\x04") || starts(b"PK\x05\x06") {
287 Some(Kind::Zip)
288 } else if starts(b"\x1f\x8b") {
289 Some(Kind::Gzip)
290 } else if starts(b"\x7fELF")
291 || starts(b"MZ")
292 || starts(b"#!")
293 || starts(b"\xfe\xed\xfa\xce")
294 || starts(b"\xfe\xed\xfa\xcf")
295 || starts(b"\xce\xfa\xed\xfe")
296 || starts(b"\xcf\xfa\xed\xfe")
297 || starts(b"\xca\xfe\xba\xbe")
298 {
299 Some(Kind::Program)
300 } else {
301 None
302 }
303}
304
305pub fn sniff_file(path: &Path) -> Result<Kind> {
308 let mut f = std::fs::File::open(path)?;
309 let mut buf = vec![0u8; 256 * 1024];
310 let mut carry: Vec<u8> = Vec::new();
311 let mut first = true;
312 let mut text = true;
313 loop {
314 let n = f.read(&mut buf)?;
315 if n == 0 {
316 break;
317 }
318 if first {
319 first = false;
320 if let Some(k) = magic(&buf[..n]) {
321 return Ok(k);
322 }
323 }
324 if !text {
325 continue;
326 }
327 let mut chunk = std::mem::take(&mut carry);
328 chunk.extend_from_slice(&buf[..n]);
329 if chunk.contains(&0) {
330 text = false;
331 continue;
332 }
333 match std::str::from_utf8(&chunk) {
334 Ok(_) => {}
335 Err(e) if e.error_len().is_none() => carry = chunk[e.valid_up_to()..].to_vec(),
337 Err(_) => text = false,
338 }
339 }
340 Ok(if text && carry.is_empty() {
341 Kind::Text
342 } else {
343 Kind::Other
344 })
345}
346
347pub fn warnings(name: &str, said: &str, kind: Kind) -> Vec<String> {
350 let mut out = Vec::new();
351 if kind == Kind::Program || runnable_name(name) {
352 out.push(
353 "this file can run as a program; do not run it unless you trust who sent it".into(),
354 );
355 }
356 let by_name = type_for_name(name);
357 for (what, claimed) in [("its name", by_name), ("the sender", said)] {
358 if claimed.is_empty() || claimed == "application/octet-stream" {
359 continue;
360 }
361 let family = |m: &str| m.split('/').next().unwrap_or("").to_string();
362 let fits = match kind {
363 Kind::Text => family(claimed) == "text" || claimed.ends_with("json"),
364 Kind::Other | Kind::Program => false,
365 k => k.mime() == claimed,
366 };
367 if !fits {
368 out.push(format!(
369 "the bytes look like {}, not {claimed} as {what} says",
370 kind.as_str()
371 ));
372 }
373 }
374 out.dedup();
375 out
376}
377
378#[cfg(test)]
379mod tests {
380 use super::*;
381
382 fn tmp(bytes: &[u8]) -> std::path::PathBuf {
383 let p = std::env::temp_dir().join(format!("diavlos-files-{}", ulid::Ulid::generate()));
384 std::fs::write(&p, bytes).unwrap();
385 p
386 }
387
388 #[test]
389 fn names_are_cleaned_for_any_system() {
390 assert_eq!(clean_name("build.log"), "build.log");
391 assert_eq!(clean_name("../../etc/passwd"), "passwd");
392 assert_eq!(clean_name("C:\\Users\\x\\evil.exe"), "evil.exe");
393 assert_eq!(clean_name(".bashrc"), "bashrc");
394 assert_eq!(clean_name("a..b"), "a.b");
395 assert_eq!(clean_name("my report (v2).pdf"), "my_report__v2_.pdf");
396 assert_eq!(clean_name(""), "file");
397 assert_eq!(clean_name(".."), "file");
398 assert_eq!(clean_name("dir/"), "file");
399 assert_eq!(clean_name("CON.txt"), "_CON.txt");
400 assert_eq!(clean_name("nul"), "_nul");
401 let long = format!("{}.txt", "a".repeat(300));
402 let c = clean_name(&long);
403 assert!(c.len() <= MAX_NAME_BYTES && c.ends_with(".txt"), "{c}");
404 for n in ["a.txt", "_CON.txt", "file", &c] {
406 assert_eq!(clean_name(n), n);
407 }
408 }
409
410 #[test]
411 fn runnable_names_are_saved_so_they_do_not_run() {
412 assert_eq!(saved_name("setup.exe"), "setup.exe.unsafe");
413 assert_eq!(saved_name("x.SH"), "x.SH.unsafe");
414 assert_eq!(saved_name("notes.txt"), "notes.txt");
415 assert!(!runnable_name("exe"));
416 }
417
418 #[test]
419 fn the_bytes_decide_the_kind() {
420 let cases: &[(&[u8], Kind)] = &[
421 (b"hello\nworld", Kind::Text),
422 ("καλημέρα".as_bytes(), Kind::Text),
423 (b"\x89PNG\r\n\x1a\nrest", Kind::Png),
424 (b"%PDF-1.7", Kind::Pdf),
425 (b"PK\x03\x04zip", Kind::Zip),
426 (b"\x7fELF\x02", Kind::Program),
427 (b"MZ\x90\x00", Kind::Program),
428 (b"#!/bin/sh\nrm -rf /", Kind::Program),
429 (b"text\0with nul", Kind::Other),
430 (b"\xff\xfe bad utf8", Kind::Other),
431 (b"", Kind::Text),
432 ];
433 for (bytes, want) in cases {
434 let p = tmp(bytes);
435 assert_eq!(sniff_file(&p).unwrap(), *want, "{bytes:?}");
436 std::fs::remove_file(p).unwrap();
437 }
438 assert!(Kind::Pdf.is_safe() && !Kind::Program.is_safe() && !Kind::Gzip.is_safe());
439 }
440
441 #[test]
442 fn a_character_split_across_reads_is_still_text() {
443 let mut bytes = vec![b'a'; 256 * 1024 - 1];
445 bytes.extend_from_slice("é and more".as_bytes());
446 let p = tmp(&bytes);
447 assert_eq!(sniff_file(&p).unwrap(), Kind::Text);
448 std::fs::remove_file(p).unwrap();
449 }
450
451 #[test]
452 fn warnings_come_from_the_bytes() {
453 assert!(warnings("notes.txt", "text/plain", Kind::Text).is_empty());
454 assert!(warnings("data.json", "application/json", Kind::Text).is_empty());
455 assert!(warnings("blob", "", Kind::Other).is_empty());
456 let w = warnings("photo.png", "image/png", Kind::Program);
457 assert!(w.iter().any(|w| w.contains("run as a program")), "{w:?}");
458 assert!(w.iter().any(|w| w.contains("not image/png")), "{w:?}");
459 let w = warnings("tool.exe", "", Kind::Other);
460 assert!(w.iter().any(|w| w.contains("run as a program")), "{w:?}");
461 let w = warnings("a.bin", "image/png", Kind::Pdf);
462 assert_eq!(w.len(), 1, "{w:?}");
463 }
464
465 #[test]
466 fn hashing_and_references() {
467 let p = tmp(b"test");
468 let (id, size) = hash_file(&p).unwrap();
469 assert_eq!(
470 id,
471 "sha256:9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08"
472 );
473 assert_eq!(size, 4);
474 std::fs::remove_file(p).unwrap();
475
476 let good = serde_json::json!({"files": [{"id": id, "name": "t.txt", "size": 4, "type": "text/plain"}]});
477 assert_eq!(refs(&good).unwrap()[0].hex().len(), 64);
478 assert!(refs(&serde_json::json!({})).unwrap().is_empty());
479 assert!(refs(&Value::Null).unwrap().is_empty());
480 for bad in [
481 serde_json::json!({"files": "x"}),
482 serde_json::json!({"files": [{"id": "sha256:zz", "name": "a", "size": 1}]}),
483 serde_json::json!({"files": [{"id": id, "name": "../a", "size": 1}]}),
484 serde_json::json!({"files": [{"id": id, "name": "a", "size": HARD_MAX_FILE_BYTES + 1}]}),
485 ] {
486 assert!(refs(&bad).is_err(), "{bad}");
487 }
488 }
489}