use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use crate::entities::chat_file::{
ChatFile, FileOrigin, mime_for, same_name, sha256_hex, versioned,
};
use crate::shared::os_open;
const MAX_VERSIONS: u32 = 10_000;
#[derive(Debug, Clone, PartialEq)]
pub enum Stored {
New(ChatFile),
Unchanged(ChatFile),
Restored(ChatFile),
}
pub fn store(dir: &Path, listed: &[ChatFile], name: &str, content: &[u8]) -> io::Result<Stored> {
store_as(
dir,
listed,
name,
content,
FileOrigin::Sandbox,
Some(os_open::FROM_ELSEWHERE),
)
}
pub fn store_as(
dir: &Path,
listed: &[ChatFile],
name: &str,
content: &[u8],
origin: FileOrigin,
zone: Option<&[u8]>,
) -> io::Result<Stored> {
confined(dir, name)?;
std::fs::create_dir_all(dir)?;
let sha256 = sha256_hex(content);
let on_disk = names_in(dir);
for n in 1..=MAX_VERSIONS {
let candidate = versioned(name, n);
if let Some(existing) = listed.iter().find(|f| same_name(&f.name, &candidate)) {
if existing.sha256 == sha256 {
if on_disk.iter().any(|d| same_name(d, &candidate)) {
return Ok(Stored::Unchanged(existing.clone()));
}
match write_new(&dir.join(&candidate), content) {
Err(e) if e.kind() == io::ErrorKind::AlreadyExists => {
return Ok(Stored::Unchanged(existing.clone()));
}
other => other?,
}
mark(&dir.join(&candidate), zone);
return Ok(Stored::Restored(existing.clone()));
}
continue;
}
if on_disk.iter().any(|d| same_name(d, &candidate)) {
continue;
}
match write_new(&dir.join(&candidate), content) {
Ok(()) => {}
Err(e) if e.kind() == io::ErrorKind::AlreadyExists => continue,
Err(e) => return Err(e),
}
mark(&dir.join(&candidate), zone);
return Ok(Stored::New(ChatFile {
id: uuid::Uuid::new_v4(),
mime: mime_for(&candidate, content).to_string(),
name: candidate,
origin,
bytes: content.len() as u64,
sha256,
added_at: chrono::Utc::now(),
}));
}
Err(io::Error::other("no free version of the name"))
}
fn write_new(path: &Path, content: &[u8]) -> io::Result<()> {
let mut file = std::fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(path)?;
if let Err(e) = file.write_all(content).and_then(|()| file.sync_all()) {
drop(file);
let _ = std::fs::remove_file(path);
return Err(e);
}
Ok(())
}
pub fn mark(path: &Path, zone: Option<&[u8]>) {
let Some(zone) = zone else {
return;
};
if let Err(e) = os_open::set_zone(path, zone) {
tracing::warn!(
path = %path.display(),
error = %e,
"stored files: could not mark the file as come from elsewhere"
);
}
}
pub fn unlisted(dir: &Path, listed: &[ChatFile]) -> Vec<ChatFile> {
let Ok(entries) = std::fs::read_dir(dir) else {
return Vec::new();
};
let mut found: Vec<ChatFile> = Vec::new();
for entry in entries.filter_map(Result::ok) {
let Some(name) = entry.file_name().to_str().map(str::to_string) else {
continue;
};
if listed
.iter()
.chain(&found)
.any(|f| same_name(&f.name, &name))
{
continue;
}
if !std::fs::symlink_metadata(entry.path()).is_ok_and(|m| m.is_file()) {
continue;
}
let Ok((sha256, head, bytes)) = hash_file(&entry.path()) else {
continue;
};
found.push(ChatFile {
id: uuid::Uuid::new_v4(),
mime: mime_for(&name, &head).to_string(),
name,
origin: FileOrigin::Recovered,
bytes,
sha256,
added_at: chrono::Utc::now(),
});
}
found.sort_by(|a, b| a.name.cmp(&b.name));
found
}
pub fn remove(dir: &Path, name: &str) -> io::Result<()> {
match std::fs::remove_file(confined(dir, name)?) {
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
other => other,
}
}
pub fn exists(dir: &Path, name: &str) -> bool {
confined(dir, name).is_ok_and(|path| path.is_file())
}
pub fn confined(dir: &Path, name: &str) -> io::Result<PathBuf> {
let plain = !name.is_empty() && name != "." && name != ".." && !name.contains(['/', '\\', ':']);
if plain {
Ok(dir.join(name))
} else {
Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!("not a plain file name: {name:?}"),
))
}
}
fn names_in(dir: &Path) -> Vec<String> {
std::fs::read_dir(dir)
.map(|entries| {
entries
.filter_map(Result::ok)
.map(|e| e.file_name().to_string_lossy().into_owned())
.collect()
})
.unwrap_or_default()
}
fn hash_file(path: &Path) -> io::Result<(String, Vec<u8>, u64)> {
use sha2::{Digest, Sha256};
const HEAD: usize = 64;
let mut file = std::fs::File::open(path)?;
let mut hasher = Sha256::new();
let mut head = Vec::with_capacity(HEAD);
let mut buf = vec![0u8; 64 * 1024];
let mut total = 0u64;
loop {
let n = file.read(&mut buf)?;
if n == 0 {
break;
}
if head.len() < HEAD {
head.extend_from_slice(&buf[..n.min(HEAD - head.len())]);
}
hasher.update(&buf[..n]);
total += n as u64;
}
let hex = hasher
.finalize()
.iter()
.map(|b| format!("{b:02x}"))
.collect();
Ok((hex, head, total))
}
#[cfg(test)]
mod tests {
use super::*;
const PNG: &[u8] = b"\x89PNG\r\n\x1a\n-a-chart-";
fn listing(name: &str, content: &[u8]) -> ChatFile {
ChatFile::new(name, FileOrigin::Sandbox, content)
}
fn new_name(stored: Stored) -> String {
match stored {
Stored::New(file) => file.name,
Stored::Unchanged(file) => panic!("expected a new file, got unchanged {}", file.name),
Stored::Restored(file) => panic!("expected a new file, got restored {}", file.name),
}
}
#[test]
fn stores_under_the_name_and_lists_what_it_wrote() {
let dir = tempfile::tempdir().unwrap();
let Stored::New(file) = store(dir.path(), &[], "chart.png", PNG).unwrap() else {
panic!("expected a new file");
};
assert_eq!(file.name, "chart.png");
assert_eq!(file.mime, "image/png");
assert_eq!(file.bytes, PNG.len() as u64);
assert_eq!(file.sha256, sha256_hex(PNG));
assert_eq!(file.origin, FileOrigin::Sandbox);
assert_eq!(std::fs::read(dir.path().join("chart.png")).unwrap(), PNG);
}
#[test]
fn a_listed_name_with_other_bytes_gets_the_next_version_case_insensitively() {
let dir = tempfile::tempdir().unwrap();
let listed = [listing("Chart.PNG", b"older")];
std::fs::write(dir.path().join("Chart.PNG"), b"older").unwrap();
let name = new_name(store(dir.path(), &listed, "chart.png", PNG).unwrap());
assert_eq!(name, "chart (2).png");
assert_eq!(
std::fs::read(dir.path().join("Chart.PNG")).unwrap(),
b"older"
);
}
#[test]
fn the_same_bytes_under_a_listed_name_of_the_family_write_nothing() {
let dir = tempfile::tempdir().unwrap();
let listed = [
listing("chart.png", b"first"),
listing("chart (2).png", PNG),
];
std::fs::write(dir.path().join("chart.png"), b"first").unwrap();
std::fs::write(dir.path().join("chart (2).png"), PNG).unwrap();
let stored = store(dir.path(), &listed, "chart.png", PNG).unwrap();
assert_eq!(stored, Stored::Unchanged(listed[1].clone()));
assert_eq!(names_in(dir.path()).len(), 2, "nothing was written");
}
#[test]
fn the_same_bytes_put_a_listed_copy_back_when_it_has_gone_missing() {
let dir = tempfile::tempdir().unwrap();
let listed = [listing("chart.png", PNG)];
assert_eq!(names_in(dir.path()).len(), 0);
let stored = store(dir.path(), &listed, "chart.png", PNG).unwrap();
assert_eq!(
stored,
Stored::Restored(listed[0].clone()),
"the listing keeps its identity; only the bytes came back"
);
assert_eq!(
std::fs::read(dir.path().join("chart.png")).unwrap(),
PNG,
"the copy has to be readable again, which is the point"
);
let again = store(dir.path(), &listed, "chart.png", PNG).unwrap();
assert_eq!(again, Stored::Unchanged(listed[0].clone()));
}
#[cfg(windows)]
#[test]
fn what_a_call_wrote_is_marked_and_a_users_copy_carries_its_own_mark() {
use crate::shared::os_open::{FROM_ELSEWHERE, zone_of};
let dir = tempfile::tempdir().unwrap();
let name = new_name(store(dir.path(), &[], "sales.csv", b"=1+1\n").unwrap());
assert_eq!(
zone_of(&dir.path().join(&name)).as_deref(),
Some(FROM_ELSEWHERE),
"a new output"
);
let listed = [listing("chart.png", PNG)];
let restored = store(dir.path(), &listed, "chart.png", PNG).unwrap();
assert!(matches!(restored, Stored::Restored(_)), "{restored:?}");
assert_eq!(
zone_of(&dir.path().join("chart.png")).as_deref(),
Some(FROM_ELSEWHERE),
"an output written back"
);
let own = store_as(
dir.path(),
&[],
"book.xlsx",
b"PK",
FileOrigin::Attached,
None,
);
let own = new_name(own.unwrap());
assert_eq!(
zone_of(&dir.path().join(&own)),
None,
"the user's file came with no mark"
);
let zone: &[u8] = b"[ZoneTransfer]\r\nZoneId=3\r\nHostUrl=https://example.com/a.docx\r\n";
let got = store_as(
dir.path(),
&[],
"a.docx",
b"PK",
FileOrigin::Attached,
Some(zone),
);
let got = new_name(got.unwrap());
assert_eq!(
zone_of(&dir.path().join(&got)).as_deref(),
Some(zone),
"a download's copy keeps the download's mark"
);
}
#[test]
fn a_file_on_disk_is_never_overwritten() {
let dir = tempfile::tempdir().unwrap();
std::fs::write(dir.path().join("chart.png"), b"kept").unwrap();
let name = new_name(store(dir.path(), &[], "chart.png", PNG).unwrap());
assert_eq!(name, "chart (2).png");
assert_eq!(
std::fs::read(dir.path().join("chart.png")).unwrap(),
b"kept"
);
}
#[test]
fn a_name_that_is_not_one_component_is_refused_and_nothing_is_written() {
let root = tempfile::tempdir().unwrap();
let dir = root.path().join("chat");
for name in ["../escape.txt", r"..\escape.txt", "..", "", "C:escape.txt"] {
let err = store(&dir, &[], name, b"x").unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::InvalidInput, "{name:?}");
assert!(remove(&dir, name).is_err(), "{name:?}");
}
assert!(!root.path().join("escape.txt").exists());
assert!(!dir.exists(), "a refused store creates nothing");
}
#[test]
fn unlisted_files_are_found_with_their_hash_and_type_and_nothing_else_is() {
let dir = tempfile::tempdir().unwrap();
std::fs::write(dir.path().join("listed.csv"), b"a,b").unwrap();
std::fs::write(dir.path().join("orphan.png"), PNG).unwrap();
std::fs::create_dir(dir.path().join("a directory")).unwrap();
let listed = [listing("LISTED.csv", b"a,b")];
let found = unlisted(dir.path(), &listed);
assert_eq!(found.len(), 1);
assert_eq!(found[0].name, "orphan.png");
assert_eq!(found[0].origin, FileOrigin::Recovered);
assert_eq!(found[0].mime, "image/png");
assert_eq!(found[0].sha256, sha256_hex(PNG));
assert_eq!(found[0].bytes, PNG.len() as u64);
assert!(
dir.path().join("orphan.png").exists(),
"adopting deletes nothing"
);
}
#[test]
fn an_absent_folder_holds_nothing_unlisted() {
let dir = tempfile::tempdir().unwrap();
assert!(unlisted(&dir.path().join("none"), &[]).is_empty());
}
#[test]
fn remove_deletes_our_copy_and_a_missing_one_counts_as_deleted() {
let dir = tempfile::tempdir().unwrap();
std::fs::write(dir.path().join("chart.png"), PNG).unwrap();
assert!(exists(dir.path(), "chart.png"));
remove(dir.path(), "chart.png").unwrap();
assert!(!exists(dir.path(), "chart.png"));
remove(dir.path(), "chart.png").unwrap();
}
}