use std::collections::HashMap;
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::time::UNIX_EPOCH;
use serde::Serialize;
use crate::error::{Error, Result};
pub const HEAD_BYTES: usize = 64;
fn at(b: &[u8], start: usize, lit: &[u8]) -> bool {
b.len() >= start + lit.len() && &b[start..start + lit.len()] == lit
}
pub fn sniff(b: &[u8]) -> &'static str {
let byte = |i: usize| b.get(i).copied();
if b.len() >= 3 && b[0] == 0xff && b[1] == 0xd8 && b[2] == 0xff {
return "jpeg";
}
if b.len() >= 4 && b[0] == 0x89 && at(b, 1, b"PNG") {
return "png";
}
if at(b, 0, b"GIF8") {
return "gif";
}
if at(b, 0, b"RIFF") {
if at(b, 8, b"WEBP") {
return "webp";
}
if at(b, 8, b"WAVE") {
return "wav";
}
if at(b, 8, b"AVI ") {
return "avi";
}
}
if at(b, 0, b"BM") {
return "bmp";
}
if at(b, 0, &[0x49, 0x49, 0x2a, 0]) || at(b, 0, &[0x4d, 0x4d, 0, 0x2a]) {
return "tiff";
}
if at(b, 4, b"ftyp") {
let brand = &b[b.len().min(8)..b.len().min(12)];
const HEIC: [&[u8]; 8] = [b"heic", b"heix", b"hevc", b"hevx", b"mif1", b"msf1", b"heim", b"heis"];
if HEIC.contains(&brand) {
return "heic";
}
if brand == b"avif" || brand == b"avis" {
return "avif";
}
if brand == b"qt " {
return "mov";
}
if brand.starts_with(b"3g") {
return "3gp";
}
if brand == b"M4A " || brand == b"M4B " {
return "m4a";
}
return "mp4";
}
if at(b, 0, &[0x1a, 0x45, 0xdf, 0xa3]) {
let head = &b[..b.len().min(64)];
return if head.windows(4).any(|w| w == b"webm") { "webm" } else { "mkv" };
}
if at(b, 0, b"OggS") {
return "ogg";
}
if at(b, 0, b"#!AMR") {
return "amr";
}
if at(b, 0, b"ID3") {
return "mp3";
}
if byte(0) == Some(0xff) {
if let Some(second) = byte(1) {
if second & 0xf6 == 0xf0 {
return "aac"; }
if second & 0xe0 == 0xe0 {
return "mp3";
}
}
}
"unknown"
}
pub fn content_type(format: &str) -> Option<&'static str> {
Some(match format {
"jpeg" => "image/jpeg",
"png" => "image/png",
"gif" => "image/gif",
"webp" => "image/webp",
"mp4" => "video/mp4",
"mov" => "video/quicktime",
"3gp" => "video/3gpp",
"webm" => "video/webm",
"m4a" => "audio/mp4",
"aac" => "audio/aac",
"mp3" => "audio/mpeg",
"amr" => "audio/amr",
"ogg" => "audio/ogg",
"wav" => "audio/wav",
_ => return None,
})
}
fn other_content_type(format: &str) -> Option<&'static str> {
Some(match format {
"bmp" => "image/bmp",
"tiff" => "image/tiff",
"heic" => "image/heic",
"avif" => "image/avif",
"mkv" => "video/x-matroska",
"avi" => "video/x-msvideo",
_ => return None,
})
}
pub fn sniffed_content_type(format: &str) -> &'static str {
content_type(format).or_else(|| other_content_type(format)).unwrap_or("application/octet-stream")
}
fn content_type_by_ext(name: &str) -> Option<&'static str> {
let ext = name.rsplit_once('.').map(|(_, e)| e.to_ascii_lowercase())?;
Some(match ext.as_str() {
"txt" | "log" | "md" => "text/plain",
"csv" => "text/csv",
"json" => "application/json",
"xml" => "application/xml",
"pdf" => "application/pdf",
"zip" => "application/zip",
"gz" => "application/gzip",
"7z" => "application/x-7z-compressed",
"rar" => "application/vnd.rar",
"tar" => "application/x-tar",
"doc" => "application/msword",
"docx" => "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
"xls" => "application/vnd.ms-excel",
"xlsx" => "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
"ppt" => "application/vnd.ms-powerpoint",
"pptx" => "application/vnd.openxmlformats-officedocument.presentationml.presentation",
"svg" => "image/svg+xml",
"html" | "htm" => "text/html",
_ => return None,
})
}
#[derive(Debug, Clone, Serialize)]
pub struct FileInfo {
pub path: String,
pub name: String,
pub size: u64,
pub mtime: u64,
pub format: String,
pub content_type: String,
}
pub fn mtime_ms(meta: &fs::Metadata) -> u64 {
meta.modified().ok().and_then(|t| t.duration_since(UNIX_EPOCH).ok()).map(|d| d.as_millis() as u64).unwrap_or(0)
}
pub fn read_head(file: &mut File, n: usize) -> std::io::Result<Vec<u8>> {
let mut head = Vec::with_capacity(n);
Read::by_ref(file).take(n as u64).read_to_end(&mut head)?;
Ok(head)
}
pub fn stat(path: &str, real: &Path) -> Result<FileInfo> {
let meta = fs::metadata(real).map_err(|e| Error::from_io_read(&e))?;
if !meta.is_file() {
return Err(Error::validation("file_unreadable", "不是普通文件"));
}
let mut file = File::open(real).map_err(|e| Error::from_io_read(&e))?;
let head = read_head(&mut file, HEAD_BYTES).map_err(|e| Error::from_io_read(&e))?;
let format = sniff(&head);
let name = Path::new(path).file_name().map(|n| n.to_string_lossy().into_owned()).unwrap_or_default();
let content_type = content_type(format)
.or_else(|| other_content_type(format))
.or_else(|| content_type_by_ext(&name))
.unwrap_or("application/octet-stream");
Ok(FileInfo { path: path.to_string(), name, size: meta.len(), mtime: mtime_ms(&meta), format: format.to_string(), content_type: content_type.to_string() })
}
pub fn read_range(real: &Path, offset: u64, length: u64) -> Result<Vec<u8>> {
let mut file = File::open(real).map_err(|e| Error::from_io_read(&e))?;
let size = file.metadata().map_err(|e| Error::from_io_read(&e))?.len();
if offset >= size || length == 0 {
return Ok(Vec::new());
}
let length = length.min(size - offset);
file.seek(SeekFrom::Start(offset)).map_err(|e| Error::from_io_read(&e))?;
let mut buf = Vec::with_capacity(length as usize);
file.take(length).read_to_end(&mut buf).map_err(|e| Error::from_io_read(&e))?;
Ok(buf)
}
struct Pending {
label: String,
instance: u64,
path: PathBuf,
size: u64,
}
#[derive(Default)]
pub struct BlobWrites {
pending: Mutex<HashMap<String, Pending>>,
}
#[derive(Debug, Serialize)]
pub struct WriteResult {
pub handle: String,
pub path: String,
pub size: u64,
}
impl BlobWrites {
#[allow(clippy::too_many_arguments)]
pub fn write(&self, temp: &Path, label: &str, instance: u64, handle: &str, name: &str, last: bool, bytes: &[u8]) -> Result<WriteResult> {
let (handle, path, size) = if handle.is_empty() {
fs::create_dir_all(temp).map_err(|e| Error::from_io(&e))?;
let handle = random_hex();
let file_name = super::transfer::sanitize_file_name(name, "blob");
let path = temp.join(format!("{handle}-{file_name}"));
let mut file = OpenOptions::new().write(true).create_new(true).open(&path).map_err(|e| Error::from_io(&e))?;
if let Err(err) = file.write_all(bytes) {
let _ = fs::remove_file(&path);
return Err(Error::from_io(&err));
}
(handle, path, bytes.len() as u64)
} else {
let (path, size) = {
let pending = self.pending.lock().unwrap();
let entry = pending.get(handle).filter(|p| p.label == label).ok_or_else(|| Error::validation("invalid_handle", "写入句柄不存在"))?;
(entry.path.clone(), entry.size)
};
let mut file = OpenOptions::new().append(true).open(&path).map_err(|e| Error::from_io(&e))?;
if let Err(err) = file.write_all(bytes) {
self.abandon(handle);
return Err(Error::from_io(&err));
}
(handle.to_string(), path, size + bytes.len() as u64)
};
let mut pending = self.pending.lock().unwrap();
if last {
pending.remove(&handle);
} else {
pending.insert(handle.clone(), Pending { label: label.to_string(), instance, path: path.clone(), size });
}
Ok(WriteResult { handle, path: path.to_string_lossy().into_owned(), size })
}
fn abandon(&self, handle: &str) {
if let Some(p) = self.pending.lock().unwrap().remove(handle) {
let _ = fs::remove_file(&p.path);
}
}
pub fn release_instance(&self, label: &str, instance: u64) {
let gone: Vec<String> = self.pending.lock().unwrap().iter().filter(|(_, p)| p.label == label && p.instance == instance).map(|(h, _)| h.clone()).collect();
for handle in gone {
self.abandon(&handle);
}
}
}
pub fn random_hex() -> String {
hex::encode(rand::random::<[u8; 8]>())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn sniff_matches_web() {
assert_eq!(sniff(&[0xff, 0xd8, 0xff, 0xe0]), "jpeg");
assert_eq!(sniff(b"\x89PNG\r\n\x1a\n"), "png");
assert_eq!(sniff(b"GIF89a"), "gif");
assert_eq!(sniff(b"RIFF\0\0\0\0WEBPVP8 "), "webp");
assert_eq!(sniff(b"RIFF\0\0\0\0WAVEfmt "), "wav");
assert_eq!(sniff(b"RIFF\0\0\0\0AVI LIST"), "avi");
assert_eq!(sniff(b"BM\0\0"), "bmp");
assert_eq!(sniff(&[0x49, 0x49, 0x2a, 0]), "tiff");
assert_eq!(sniff(b"\0\0\0\x18ftypheic"), "heic");
assert_eq!(sniff(b"\0\0\0\x18ftypmif1"), "heic");
assert_eq!(sniff(b"\0\0\0\x18ftypavif"), "avif");
assert_eq!(sniff(b"\0\0\0\x14ftypqt "), "mov");
assert_eq!(sniff(b"\0\0\0\x14ftyp3gp4"), "3gp");
assert_eq!(sniff(b"\0\0\0\x20ftypM4A "), "m4a");
assert_eq!(sniff(b"\0\0\0\x20ftypisom"), "mp4");
assert_eq!(sniff(b"\0\0\0\x08ftyp"), "mp4");
let mut webm = vec![0x1a, 0x45, 0xdf, 0xa3, 0x9f, 0x42, 0x82, 0x84];
webm.extend_from_slice(b"webm");
assert_eq!(sniff(&webm), "webm");
assert_eq!(sniff(&[0x1a, 0x45, 0xdf, 0xa3, 0, 0]), "mkv");
assert_eq!(sniff(b"OggS\0"), "ogg");
assert_eq!(sniff(b"#!AMR\n"), "amr");
assert_eq!(sniff(b"ID3\x04"), "mp3");
assert_eq!(sniff(&[0xff, 0xf1, 0x50]), "aac");
assert_eq!(sniff(&[0xff, 0xfb, 0x90]), "mp3");
assert_eq!(sniff(&[0xff]), "unknown");
assert_eq!(sniff(b""), "unknown");
assert_eq!(sniff(b"%PDF-1.7"), "unknown");
}
#[test]
fn stat_and_read() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("photo.bin");
fs::write(&path, [0xff, 0xd8, 0xff, 0xe0, 1, 2, 3, 4, 5, 6]).unwrap();
let info = stat(path.to_str().unwrap(), &path).unwrap();
assert_eq!((info.name.as_str(), info.size, info.format.as_str(), info.content_type.as_str()), ("photo.bin", 10, "jpeg", "image/jpeg"));
assert!(info.mtime > 0);
let pdf = dir.path().join("a.pdf");
fs::write(&pdf, b"%PDF-1.7").unwrap();
let info = stat(pdf.to_str().unwrap(), &pdf).unwrap();
assert_eq!((info.format.as_str(), info.content_type.as_str()), ("unknown", "application/pdf"));
assert_eq!(read_range(&path, 2, 3).unwrap(), vec![0xff, 0xe0, 1]);
assert_eq!(read_range(&path, 8, 100).unwrap(), vec![5, 6]);
assert!(read_range(&path, 100, 1).unwrap().is_empty());
assert_eq!(stat("x", dir.path()).unwrap_err().reason.as_deref(), Some("file_unreadable"));
assert_eq!(sniffed_content_type("unknown"), "application/octet-stream");
}
#[test]
fn blob_chunks() {
let dir = tempfile::tempdir().unwrap();
let writes = BlobWrites::default();
let first = writes.write(dir.path(), "main", 1, "", "截图 1.png", false, b"abc").unwrap();
assert_eq!(first.size, 3);
assert!(first.path.ends_with("-截图 1.png"));
let second = writes.write(dir.path(), "main", 1, &first.handle, "", false, b"de").unwrap();
assert_eq!(second.size, 5);
assert!(writes.write(dir.path(), "other", 2, &first.handle, "", true, b"x").is_err());
let last = writes.write(dir.path(), "main", 1, &first.handle, "", true, b"f").unwrap();
assert_eq!(last.size, 6);
assert_eq!(fs::read(&last.path).unwrap(), b"abcdef");
assert!(writes.write(dir.path(), "main", 1, &first.handle, "", true, b"g").is_err());
let empty = writes.write(dir.path(), "main", 1, "", "", true, b"").unwrap();
assert_eq!(empty.size, 0);
assert!(empty.path.ends_with("-blob"));
let open = writes.write(dir.path(), "main", 1, "", "a", false, b"1").unwrap();
writes.release_instance("main", 1);
assert!(!Path::new(&open.path).exists());
}
}