use std::fs::{self, OpenOptions};
use std::io::{ErrorKind, Write};
use std::path::{Component, Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};
use anyhow::{Context, Result, bail};
use vtcode_commons::{canonicalize, validate_no_escaping_symlink_ancestors};
pub(crate) fn canonical_storage_dir(dir: &Path) -> Result<PathBuf> {
fs::create_dir_all(dir).with_context(|| format!("could not create storage directory {}", dir.display()))?;
resolve_canonical_storage_dir(dir)
}
pub(crate) fn canonical_existing_storage_dir(dir: &Path) -> Result<PathBuf> {
resolve_canonical_storage_dir(dir)
}
fn resolve_canonical_storage_dir(dir: &Path) -> Result<PathBuf> {
validate_no_escaping_symlink_ancestors(dir, false)
.with_context(|| format!("could not validate storage directory {}", dir.display()))?;
let metadata =
fs::symlink_metadata(dir).with_context(|| format!("could not inspect storage directory {}", dir.display()))?;
if metadata.file_type().is_symlink() {
bail!("refusing to use symlinked directory {}", dir.display());
}
if !metadata.is_dir() {
bail!("storage path {} is not a directory", dir.display());
}
canonicalize(dir).with_context(|| format!("could not resolve storage directory {}", dir.display()))
}
pub(crate) fn validate_record_key(key: &str) -> Result<()> {
if key.contains('\0') {
bail!("invalid record key: NUL byte");
}
let mut components = Path::new(key).components();
match (components.next(), components.next()) {
(Some(Component::Normal(_)), None) => Ok(()),
_ => bail!("invalid record key {key:?}: expected a single file name"),
}
}
pub(crate) fn contained_record_path(dir: &Path, key: &str, extension: &str) -> Result<PathBuf> {
validate_record_key(key)?;
Ok(canonical_storage_dir(dir)?.join(format!("{key}.{extension}")))
}
pub(crate) fn contained_existing_record_path(dir: &Path, key: &str, extension: &str) -> Result<PathBuf> {
validate_record_key(key)?;
Ok(canonical_existing_storage_dir(dir)?.join(format!("{key}.{extension}")))
}
pub(crate) fn reject_symlink_destination(path: &Path) -> Result<()> {
match fs::symlink_metadata(path) {
Ok(metadata) if metadata.file_type().is_symlink() => {
bail!("refusing to replace symlinked file {}", path.display())
}
Ok(metadata) if !metadata.is_file() => bail!("{} is not a regular file", path.display()),
Ok(_) => Ok(()),
Err(error) if error.kind() == ErrorKind::NotFound => Ok(()),
Err(error) => Err(error).with_context(|| format!("could not inspect {}", path.display())),
}
}
pub(crate) fn write_atomic(dir: &Path, file_name: &str, contents: &[u8]) -> Result<()> {
let destination = dir.join(file_name);
reject_symlink_destination(&destination)?;
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_nanos())
.unwrap_or_default();
for attempt in 0..32u32 {
let temporary = dir.join(format!(".{file_name}.{}-{timestamp}.{attempt}.tmp", std::process::id()));
let mut file = match OpenOptions::new().write(true).create_new(true).open(&temporary) {
Ok(file) => file,
Err(error) if error.kind() == ErrorKind::AlreadyExists => continue,
Err(error) => return Err(error).with_context(|| format!("could not create {}", temporary.display())),
};
if let Err(error) = file.write_all(contents).and_then(|()| file.sync_all()) {
drop(file);
drop(fs::remove_file(&temporary));
return Err(error).with_context(|| format!("could not stage {}", destination.display()));
}
drop(file);
if let Err(error) = fs::rename(&temporary, &destination) {
drop(fs::remove_file(&temporary));
return Err(error).with_context(|| format!("could not publish {}", destination.display()));
}
return Ok(());
}
bail!("could not allocate a unique temporary file in {}", dir.display())
}