sausage/vfs/src/memory.rs

//! Shared in-memory directories and identity-preserving open files.
use crate::*;
use std::{collections::BTreeMap, io::ErrorKind as E, sync::Mutex};

#[derive(Clone)]
enum Node {
    Dir(Arc<Mutex<Data>>),
    File(Arc<Mutex<Data>>),
}
#[derive(Default)]
struct Data {
    bytes: Vec<u8>,
    accessed: Option<SystemTime>,
    modified: Option<SystemTime>,
}
/// A shareable memory filesystem. Clones refer to the same tree.
#[derive(Clone)]
pub struct Memory {
    nodes: Arc<Mutex<BTreeMap<String, Node>>>,
}
impl Default for Memory {
    fn default() -> Self {
        Self {
            nodes: Arc::new(Mutex::new(BTreeMap::from([(
                "/".into(),
                Node::Dir(Arc::default()),
            )]))),
        }
    }
}
fn parent(path: &str) -> &str {
    path.rsplit_once('/')
        .map(|(p, _)| if p.is_empty() { "/" } else { p })
        .unwrap_or("/")
}
fn validate(nodes: &BTreeMap<String, Node>, path: &str) -> Result<()> {
    let mut prefix = String::new();
    let parts: Vec<_> = path
        .trim_end_matches('/')
        .split('/')
        .filter(|p| !p.is_empty())
        .collect();
    for part in parts.iter().take(parts.len().saturating_sub(1)) {
        prefix.push('/');
        prefix.push_str(part);
        directory(nodes, &prefix)?;
    }
    if path.ends_with('/') && path != "/" {
        directory(nodes, path.trim_end_matches('/'))?;
    }
    Ok(())
}
fn clean(path: &str) -> &str {
    if path == "/" {
        path
    } else {
        path.trim_end_matches('/')
    }
}
fn directory(nodes: &BTreeMap<String, Node>, path: &str) -> Result<()> {
    match nodes.get(path) {
        Some(Node::Dir(_)) => Ok(()),
        Some(_) => Err(error(E::NotADirectory)),
        None => Err(error(E::NotFound)),
    }
}
fn metadata(node: &Node) -> Stat {
    match node {
        Node::Dir(data) => {
            let data = data.lock().unwrap();
            Stat {
                directory: true,
                symlink: false,
                size: 0,
                accessed: data.accessed,
                modified: data.modified,
            }
        }
        Node::File(data) => {
            let d = data.lock().unwrap();
            Stat {
                directory: false,
                symlink: false,
                size: d.bytes.len() as u64,
                accessed: d.accessed,
                modified: d.modified,
            }
        }
    }
}
struct Handle {
    node: Node,
    write: bool,
}
#[async_trait]
impl Layer for Memory {
    async fn stat(&self, path: &str, follow: bool) -> Result<Stat> {
        let _ = follow;
        let nodes = self.nodes.lock().unwrap();
        validate(&nodes, path)?;
        nodes
            .get(clean(path))
            .map(metadata)
            .ok_or_else(|| error(E::NotFound))
    }
    async fn list(&self, path: &str) -> Result<Vec<Entry>> {
        let nodes = self.nodes.lock().unwrap();
        validate(&nodes, path)?;
        let path = clean(path);
        directory(&nodes, path)?;
        Ok(nodes
            .iter()
            .filter(|(p, _)| p.as_str() != path && parent(p) == path)
            .map(|(p, n)| Entry {
                name: p.rsplit('/').next().unwrap().into(),
                directory: matches!(n, Node::Dir(_)),
                symlink: false,
            })
            .collect())
    }
    async fn open(&self, path: &str, o: Open) -> Result<Arc<dyn File>> {
        let mut nodes = self.nodes.lock().unwrap();
        validate(&nodes, path)?;
        let path = clean(path);
        if o.truncate && !o.write {
            return Err(error(E::InvalidInput));
        }
        if nodes.contains_key(path) && o.create && o.exclusive {
            return Err(error(E::AlreadyExists));
        }
        if !nodes.contains_key(path) {
            if !o.create || o.directory {
                return Err(error(E::NotFound));
            }
            directory(&nodes, parent(path))?;
            nodes.insert(
                path.into(),
                Node::File(Arc::new(Mutex::new(Data {
                    bytes: vec![],
                    accessed: None,
                    modified: None,
                }))),
            );
        }
        let node = nodes[path].clone();
        match &node {
            Node::Dir(_) if o.write || o.truncate => return Err(error(E::IsADirectory)),
            Node::File(_) if o.directory => return Err(error(E::NotADirectory)),
            Node::File(d) if o.truncate => d.lock().unwrap().bytes.clear(),
            _ => {}
        }
        Ok(Arc::new(Handle {
            node,
            write: o.write,
        }))
    }
    async fn mkdir(&self, path: &str) -> Result<()> {
        let mut nodes = self.nodes.lock().unwrap();
        validate(&nodes, path)?;
        let path = clean(path);
        if nodes.contains_key(path) {
            return Err(error(E::AlreadyExists));
        }
        directory(&nodes, parent(path))?;
        nodes.insert(path.into(), Node::Dir(Arc::default()));
        Ok(())
    }
    async fn remove_file(&self, path: &str) -> Result<()> {
        let mut nodes = self.nodes.lock().unwrap();
        validate(&nodes, path)?;
        let path = clean(path);
        match nodes.get(path) {
            None => return Err(error(E::NotFound)),
            Some(Node::Dir(_)) => return Err(error(E::IsADirectory)),
            _ => {}
        }
        nodes.remove(path);
        Ok(())
    }
    async fn remove_dir(&self, path: &str) -> Result<()> {
        let mut nodes = self.nodes.lock().unwrap();
        validate(&nodes, path)?;
        let path = clean(path);
        directory(&nodes, path)?;
        if path == "/" {
            return Err(error(E::ResourceBusy));
        }
        if nodes.keys().any(|p| p != path && parent(p) == path) {
            return Err(error(E::DirectoryNotEmpty));
        }
        nodes.remove(path);
        Ok(())
    }
    async fn rename(&self, from: &str, to: &str) -> Result<()> {
        let mut nodes = self.nodes.lock().unwrap();
        validate(&nodes, from)?;
        validate(&nodes, to)?;
        let from = clean(from);
        let to = clean(to);
        let node = nodes.get(from).cloned().ok_or_else(|| error(E::NotFound))?;
        directory(&nodes, parent(to))?;
        if from == to {
            return Ok(());
        }
        if from == "/" || to == "/" {
            return Err(error(E::ResourceBusy));
        }
        if to.starts_with(&format!("{from}/")) {
            return Err(error(E::InvalidInput));
        }
        if let Some(target) = nodes.get(to) {
            match (&node, target) {
                (Node::Dir(_), Node::File(_)) => return Err(error(E::NotADirectory)),
                (Node::File(_), Node::Dir(_)) if from.starts_with(&format!("{to}/")) => {
                    return Err(error(E::DirectoryNotEmpty));
                }
                (Node::File(_), Node::Dir(_)) => return Err(error(E::IsADirectory)),
                (Node::Dir(_), Node::Dir(_))
                    if nodes.keys().any(|p| p != to && parent(p) == to) =>
                {
                    return Err(error(E::DirectoryNotEmpty));
                }
                _ => {}
            }
        }
        let prefix = format!("{from}/");
        let moved: Vec<_> = nodes
            .keys()
            .filter(|p| p.as_str() == from || p.starts_with(&prefix))
            .cloned()
            .collect();
        for p in moved {
            let n = nodes.remove(&p).unwrap();
            nodes.insert(format!("{to}{}", &p[from.len()..]), n);
        }
        Ok(())
    }
}
impl Handle {
    fn data(&self) -> Result<&Mutex<Data>> {
        match &self.node {
            Node::File(d) => Ok(d),
            _ => Err(error(E::IsADirectory)),
        }
    }
    fn writable(&self) -> Result<&Mutex<Data>> {
        if !self.write {
            return Err(error(E::PermissionDenied));
        }
        self.data()
    }
}
#[async_trait]
impl File for Handle {
    async fn stat(&self) -> Result<Stat> {
        Ok(metadata(&self.node))
    }
    async fn read(&self, offset: u64, length: usize) -> Result<Vec<u8>> {
        let d = self.data()?.lock().unwrap();
        let start = usize::try_from(offset)
            .unwrap_or(usize::MAX)
            .min(d.bytes.len());
        Ok(d.bytes[start..start.saturating_add(length).min(d.bytes.len())].to_vec())
    }
    async fn write(&self, offset: u64, bytes: Vec<u8>) -> Result<usize> {
        let mut d = self.writable()?.lock().unwrap();
        let start = usize::try_from(offset).map_err(|_| error(E::FileTooLarge))?;
        let end = start
            .checked_add(bytes.len())
            .ok_or_else(|| error(E::FileTooLarge))?;
        if !bytes.is_empty() {
            if end > d.bytes.len() {
                let additional = end - d.bytes.len();
                d.bytes
                    .try_reserve(additional)
                    .map_err(|_| error(E::FileTooLarge))?;
                d.bytes.resize(end, 0);
            }
            d.bytes[start..end].copy_from_slice(&bytes);
            d.modified = Some(SystemTime::now());
        }
        Ok(bytes.len())
    }
    async fn append(&self, bytes: Vec<u8>) -> Result<usize> {
        let mut d = self.writable()?.lock().unwrap();
        d.bytes
            .try_reserve(bytes.len())
            .map_err(|_| error(E::FileTooLarge))?;
        d.bytes.extend_from_slice(&bytes);
        d.modified = Some(SystemTime::now());
        Ok(bytes.len())
    }
    async fn set_size(&self, size: u64) -> Result<()> {
        let mut d = self.writable()?.lock().unwrap();
        let size = usize::try_from(size).map_err(|_| error(E::FileTooLarge))?;
        let additional = size.saturating_sub(d.bytes.len());
        d.bytes
            .try_reserve(additional)
            .map_err(|_| error(E::FileTooLarge))?;
        d.bytes.resize(size, 0);
        d.modified = Some(SystemTime::now());
        Ok(())
    }
    async fn sync(&self) -> Result<()> {
        Ok(())
    }
    async fn set_times(
        &self,
        accessed: Option<SystemTime>,
        modified: Option<SystemTime>,
    ) -> Result<()> {
        let data = match &self.node {
            Node::Dir(d) | Node::File(d) => d,
        };
        let mut d = data.lock().unwrap();
        if accessed.is_some() {
            d.accessed = accessed;
        }
        if modified.is_some() {
            d.modified = modified;
        }
        Ok(())
    }
}