use anyhow::{Context, Result, bail, ensure};
use std::{
fs,
io::{Read, Write},
path::{Path, PathBuf},
};
use tempfile::NamedTempFile;
pub const MAX_IMAGE_BYTES: usize = 20 * 1024 * 1024;
#[derive(Clone)]
pub struct ImageReference {
pub id: i64,
pub name: String,
pub media_type: String,
}
impl ImageReference {
pub fn markdown(&self, task_id: i64) -> Result<String> {
ensure!(task_id > 0 && self.id > 0, "Invalid image identity");
let alt = self
.name
.chars()
.map(|ch| match ch {
'\\' | '[' | ']' | '*' | '_' | '`' | '&' | '~' => format!("\\{ch}"),
'\r' | '\n' | '\t' => " ".to_owned(),
_ => ch.to_string(),
})
.collect::<String>();
Ok(format!(
"",
self.id,
image_extension(&self.media_type)?
))
}
}
fn image_extension(media_type: &str) -> Result<&'static str> {
Ok(match media_type {
"image/png" => "png",
"image/jpeg" => "jpg",
"image/gif" => "gif",
"image/webp" => "webp",
_ => bail!("Unsupported stored image media type: {media_type}"),
})
}
pub struct ImageStore {
root: PathBuf,
}
pub struct PendingFiles {
created: Vec<PathBuf>,
active: bool,
}
impl PendingFiles {
pub fn new() -> Self {
Self {
created: Vec::new(),
active: true,
}
}
pub fn keep(&mut self) {
self.active = false;
}
pub fn discard(&mut self) -> Result<()> {
self.active = false;
let mut failure = None;
for path in self.created.iter().rev() {
if let Err(error) = fs::remove_file(path) {
if error.kind() != std::io::ErrorKind::NotFound && failure.is_none() {
failure = Some((path.clone(), error));
}
}
}
self.created.clear();
match failure {
Some((path, error)) => {
Err(error).with_context(|| format!("Cannot remove {}", path.display()))
}
None => Ok(()),
}
}
}
impl Drop for PendingFiles {
fn drop(&mut self) {
if self.active {
let _ = self.discard();
}
}
}
impl ImageStore {
pub fn new(root: PathBuf) -> Self {
Self { root }
}
pub fn path(&self, task_id: i64, image_id: i64, media_type: &str) -> Result<PathBuf> {
ensure!(task_id > 0 && image_id > 0, "Invalid image identity");
let ext = image_extension(media_type)?;
Ok(self
.root
.join(task_id.to_string())
.join(format!("{image_id}.{ext}")))
}
pub fn write(
&self,
pending: &mut PendingFiles,
task_id: i64,
image_id: i64,
media_type: &str,
data: &[u8],
) -> Result<()> {
let path = self.path(task_id, image_id, media_type)?;
let directory = path.parent().context("Missing image directory")?;
let parent = self.root.parent().context("Missing image storage parent")?;
verify_directory(parent)?;
create_real_directory(&self.root)?;
create_real_directory(directory)?;
sync_directory(parent)?;
sync_directory(&self.root)?;
let mut temporary = NamedTempFile::new_in(directory)?;
temporary.write_all(data)?;
temporary.as_file().sync_all()?;
match temporary.persist_noclobber(&path) {
Ok(_) => {
pending.created.push(path.clone());
sync_directory(directory)?;
Ok(())
}
Err(error) if error.error.kind() == std::io::ErrorKind::AlreadyExists => {
let existing = self
.read(task_id, image_id, media_type, data.len() as i64)
.context("Stored image path conflicts with existing bytes")?;
ensure!(
existing == data,
"Stored image path conflicts with existing bytes"
);
fs::File::open(&path)?.sync_all()?;
sync_directory(directory)?;
Ok(())
}
Err(error) => {
Err(error.error).with_context(|| format!("Cannot store {}", path.display()))
}
}
}
pub fn read(
&self,
task_id: i64,
image_id: i64,
media_type: &str,
expected_bytes: i64,
) -> Result<Vec<u8>> {
let path = self.path(task_id, image_id, media_type)?;
verify_directory(self.root.parent().context("Missing image storage parent")?)?;
verify_directory(&self.root)?;
verify_directory(path.parent().context("Missing image directory")?)?;
let metadata = fs::symlink_metadata(&path)
.with_context(|| format!("Cannot read stored image {}", path.display()))?;
ensure!(
metadata.file_type().is_file(),
"Stored image is not a regular file"
);
let data = fs::read(&path)
.with_context(|| format!("Cannot read stored image {}", path.display()))?;
ensure!(
expected_bytes >= 0 && data.len() as i64 == expected_bytes,
"Stored image byte count differs from database"
);
Ok(data)
}
}
fn verify_directory(path: &Path) -> Result<()> {
let metadata = fs::symlink_metadata(path)
.with_context(|| format!("Cannot inspect image directory {}", path.display()))?;
ensure!(
!metadata.file_type().is_symlink(),
"Stored image directory is a symlink: {}",
path.display()
);
ensure!(
metadata.is_dir(),
"Stored image directory is not a directory"
);
Ok(())
}
fn create_real_directory(path: &Path) -> Result<()> {
match fs::create_dir(path) {
Ok(()) => (),
Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => (),
Err(error) => {
return Err(error).with_context(|| format!("Cannot create {}", path.display()));
}
}
verify_directory(path)
}
#[cfg(unix)]
fn sync_directory(path: &Path) -> Result<()> {
fs::File::open(path)?.sync_all()?;
Ok(())
}
#[cfg(not(unix))]
fn sync_directory(_path: &Path) -> Result<()> {
Ok(())
}
#[derive(Clone, Debug)]
pub struct ImageInput {
pub name: String,
pub data: Vec<u8>,
}
impl ImageInput {
pub fn read(path: &Path) -> Result<Self> {
let metadata =
std::fs::metadata(path).with_context(|| format!("Cannot read {}", path.display()))?;
ensure!(metadata.is_file(), "Image must be a regular file");
let file =
std::fs::File::open(path).with_context(|| format!("Cannot read {}", path.display()))?;
ensure!(file.metadata()?.is_file(), "Image must be a regular file");
let mut data = Vec::new();
file.take(MAX_IMAGE_BYTES as u64 + 1)
.read_to_end(&mut data)?;
let input = Self {
name: path
.file_name()
.context("Missing image filename")?
.to_string_lossy()
.into_owned(),
data,
};
input.media_type()?;
Ok(input)
}
pub fn media_type(&self) -> Result<&'static str> {
ensure!(
self.data.len() <= MAX_IMAGE_BYTES,
"Image exceeds 20 MiB limit"
);
sniff_media_type(&self.data)
}
}
pub fn sniff_media_type(data: &[u8]) -> Result<&'static str> {
if data.starts_with(b"\x89PNG\r\n\x1a\n") {
Ok("image/png")
} else if data.starts_with(b"\xff\xd8\xff") {
Ok("image/jpeg")
} else if data.starts_with(b"GIF87a") || data.starts_with(b"GIF89a") {
Ok("image/gif")
} else if data.starts_with(b"RIFF") && data.get(8..12) == Some(b"WEBP") {
Ok("image/webp")
} else {
bail!("Unsupported image signature; expected PNG, JPEG, GIF or WebP")
}
}