MusicFS/src/origins/mod.rs

pub mod attrs;
pub mod local;
pub mod network;

use std::{
    collections::BTreeMap,
    io,
    path::Path,
    sync::{Arc, Mutex},
    time::Duration,
};

use fuser::{Errno, FileAttr, Filesystem, Generation, INodeNo, Request};
use sea_orm::{ActiveModelTrait, DatabaseConnection};
use tracing::{debug, error, trace};

use crate::item::{FileType, Item};
use crate::music::db::save_music_metadata;
use crate::origins::local::file_io;
use crate::virtual_dirs::parent_inode_from_path;

/// Transport-agnostic byte-range reader. `locator` is whatever string the
/// origin treats as a file key: a filesystem path for `LocalOrigin`, a remote
/// key for the future network origin. Both produce bytes at `(offset, len)`.
///
/// `inode` is the FUSE-level inode; NetworkOrigin uses it as the server-side
/// file id (and as the Postgres cache key). LocalOrigin ignores it.
pub trait ByteSource: Send + Sync {
    fn read_at(
        &self,
        inode: INodeNo,
        locator: &Path,
        offset: u64,
        len: usize,
    ) -> io::Result<Vec<u8>>;
}

pub trait FileWatcher: Send + Sync {
    fn watch(&self, files: Arc<Mutex<BTreeMap<INodeNo, Item>>>) -> WatcherHandle;
}

/// Returned by [`FileWatcher::watch`] so the caller can stop a background
/// watcher before the process (and its tokio runtime) shuts down. Without
/// this, an async watcher blocked on a timer/stream panics when the runtime is
/// torn down out from under it.
pub struct WatcherHandle {
    stop: Option<Box<dyn FnOnce() + Send>>,
}

impl WatcherHandle {
    pub fn new(stop: impl FnOnce() + Send + 'static) -> Self {
        return WatcherHandle {
            stop: Some(Box::new(stop)),
        };
    }

    /// A watcher with no shutdown work — its thread exits on its own.
    pub fn detached() -> Self {
        return WatcherHandle { stop: None };
    }

    /// Signal the watcher to stop and wait for it to finish.
    pub fn stop(mut self) {
        if let Some(stop) = self.stop.take() {
            stop();
        }
    }
}

pub trait Origin: Send + Sync {
    fn snapshot(&self) -> io::Result<BTreeMap<INodeNo, Item>>;
    fn byte_source(&self) -> Arc<dyn ByteSource>;
    fn watcher(&self) -> Box<dyn FileWatcher>;
}

pub struct FuseFs {
    pub files: Arc<Mutex<BTreeMap<INodeNo, Item>>>,
    pub bytes: Arc<dyn ByteSource>,
    pub client: DatabaseConnection,
    pub runtime_handle: tokio::runtime::Handle,
}

pub(crate) fn file_to_attr(item: &Item) -> FileAttr {
    let attrs = &item.attrs;
    let kind = match item.file_type {
        FileType::Directory => fuser::FileType::Directory,
        FileType::File => fuser::FileType::RegularFile,
    };

    let size = match &item.music_metadata {
        Some(mm) if !mm.header.is_empty() => mm.virtual_size(attrs.size),
        _ => attrs.size,
    };

    return FileAttr {
        ino: item.inode,
        size,
        blocks: attrs.blocks,
        atime: attrs.atime,
        mtime: attrs.mtime,
        ctime: attrs.ctime,
        crtime: attrs.crtime,
        kind,
        perm: attrs.perm,
        nlink: attrs.nlink,
        uid: attrs.uid,
        gid: attrs.gid,
        rdev: attrs.rdev,
        blksize: attrs.blksize,
        flags: 0,
    };
}

impl Filesystem for FuseFs {
    fn open(
        &self,
        _req: &Request,
        ino: INodeNo,
        _flags: fuser::OpenFlags,
        reply: fuser::ReplyOpen,
    ) {
        trace!(%ino, "open");
        if self.files.lock().unwrap().contains_key(&ino) {
            reply.opened(fuser::FileHandle(ino.0), fuser::FopenFlags::empty());
        } else {
            debug!(%ino, "open: not found");
            reply.error(Errno::ENOENT);
        }
    }

    fn setattr(
        &self,
        _req: &Request,
        ino: INodeNo,
        _mode: Option<u32>,
        _uid: Option<u32>,
        _gid: Option<u32>,
        _size: Option<u64>,
        _atime: Option<fuser::TimeOrNow>,
        _mtime: Option<fuser::TimeOrNow>,
        _ctime: Option<std::time::SystemTime>,
        _fh: Option<fuser::FileHandle>,
        _crtime: Option<std::time::SystemTime>,
        _chgtime: Option<std::time::SystemTime>,
        _bkuptime: Option<std::time::SystemTime>,
        _flags: Option<fuser::BsdFileFlags>,
        reply: fuser::ReplyAttr,
    ) {
        trace!(%ino, "setattr");
        match self.files.lock().unwrap().get(&ino) {
            Some(file) => reply.attr(&Duration::new(1, 0), &file_to_attr(file)),
            None => {
                debug!(%ino, "setattr: not found");
                reply.error(Errno::ENOENT);
            }
        }
    }

    fn write(
        &self,
        _req: &Request,
        ino: INodeNo,
        _fh: fuser::FileHandle,
        offset: u64,
        data: &[u8],
        _write_flags: fuser::WriteFlags,
        _flags: fuser::OpenFlags,
        _lock_owner: Option<fuser::LockOwner>,
        reply: fuser::ReplyWrite,
    ) {
        trace!(%ino, offset, len = data.len(), "write");
        let written = data.len() as u32;
        let write_start = offset;
        let write_end = write_start + data.len() as u64;

        let updated_music_metadata = {
            let mut files = self.files.lock().unwrap();
            let item = match files.get_mut(&ino) {
                Some(item) => item,
                None => {
                    debug!(%ino, "write: not found");
                    reply.written(written);
                    return;
                }
            };

            match &mut item.music_metadata {
                Some(mm) if mm.vorbis_comment_length > 0 => {
                    let vc_data_offset = mm.vorbis_comment_offset;
                    let vc_hdr_offset = vc_data_offset - 4;
                    if write_start <= vc_hdr_offset && write_end >= vc_data_offset {
                        let hdr_from = (vc_hdr_offset - write_start) as usize;
                        let new_length = u32::from_be_bytes([
                            0,
                            data[hdr_from + 1],
                            data[hdr_from + 2],
                            data[hdr_from + 3],
                        ]) as u64;
                        let vc_data_end = vc_data_offset + new_length;
                        if write_end >= vc_data_end {
                            let from = (vc_data_offset - write_start) as usize;
                            let to = (vc_data_end - write_start) as usize;
                            mm.update_from_vorbis_comment_data(&data[from..to]);
                            debug!(%ino, "write: vorbis comment tag update detected");
                            Some(mm.clone())
                        } else {
                            None
                        }
                    } else {
                        None
                    }
                }
                Some(mm)
                    if !mm.header.is_empty()
                        && write_start == 0
                        && data.len() >= 3
                        && &data[0..3] == b"ID3" =>
                {
                    mm.update_from_id3_data(data);
                    debug!(%ino, "write: ID3v2 tag update detected");
                    Some(mm.clone())
                }
                Some(mm) if mm.header.is_empty() && data.len() == 128 && &data[0..3] == b"TAG" => {
                    mm.update_from_id3v1_data(data);
                    debug!(%ino, "write: ID3v1 tag update detected");
                    Some(mm.clone())
                }
                _ => None,
            }
        };

        if let Some(music_metadata) = updated_music_metadata {
            let client = self.client.clone();
            let ino_i64 = ino.0 as i64;
            self.runtime_handle.block_on(async move {
                if let Err(e) = save_music_metadata(ino_i64, &music_metadata, &client).await {
                    error!(ino = ino_i64, error = %e, "write: save_music_metadata failed");
                } else {
                    debug!(ino = ino_i64, "write: persisted updated music metadata");
                }
            });
        }

        reply.written(written);
    }

    fn getattr(
        &self,
        _req: &fuser::Request,
        ino: INodeNo,
        _fh: Option<fuser::FileHandle>,
        reply: fuser::ReplyAttr,
    ) {
        trace!(%ino, "getattr");
        match self.files.lock().unwrap().get(&ino) {
            Some(file) => {
                let ttl = Duration::new(1, 0);
                let attr = file_to_attr(file);
                reply.attr(&ttl, &attr);
            }
            None => {
                debug!(%ino, "getattr: not found");
                reply.error(Errno::ENOENT);
            }
        }
    }

    fn readdir(
        &self,
        _req: &fuser::Request,
        ino: INodeNo,
        fh: fuser::FileHandle,
        offset: u64,
        mut reply: fuser::ReplyDirectory,
    ) {
        trace!(%ino, %fh, offset, "readdir");

        let files = self.files.lock().unwrap();

        let parent_inode = match files.get(&ino) {
            Some(dir) => dir.parent_inode,
            None => {
                debug!(%ino, "readdir: not found");
                reply.error(Errno::ENOENT);
                return;
            }
        };

        if offset < 1 {
            if reply.add(ino, 1, fuser::FileType::Directory, ".") {
                reply.ok();
                return;
            }
        }

        if offset < 2 {
            if reply.add(parent_inode, 2, fuser::FileType::Directory, "..") {
                reply.ok();
                return;
            }
        }

        let skip_count = if offset <= 2 {
            0
        } else {
            (offset - 2) as usize
        };

        for (i, (key, value)) in files
            .iter()
            .filter(|(_, v)| v.parent_inode == ino && v.inode != ino)
            .skip(skip_count)
            .enumerate()
        {
            let entry_offset = (skip_count + i + 3) as u64;
            let file_type = match value.file_type {
                FileType::Directory => fuser::FileType::Directory,
                FileType::File => fuser::FileType::RegularFile,
            };
            if reply.add(*key, entry_offset, file_type, &value.name) {
                break;
            }
        }
        reply.ok();
    }

    fn rename(
        &self,
        _req: &Request,
        parent: INodeNo,
        name: &std::ffi::OsStr,
        newparent: INodeNo,
        newname: &std::ffi::OsStr,
        _flags: fuser::RenameFlags,
        reply: fuser::ReplyEmpty,
    ) {
        trace!(
            %parent,
            name = %name.display(),
            %newparent,
            newname = %newname.display(),
            "rename"
        );
        let name_str = match name.to_str() {
            Some(s) => s,
            None => {
                error!(%parent, "rename: source name is not valid UTF-8");
                reply.error(Errno::EINVAL);
                return;
            }
        };
        let newname_str = match newname.to_str() {
            Some(s) => s,
            None => {
                error!(%newparent, "rename: target name is not valid UTF-8");
                reply.error(Errno::EINVAL);
                return;
            }
        };

        let mut files = self.files.lock().unwrap();

        let item_ino = match files
            .iter()
            .find(|(_, v)| v.parent_inode == parent && v.name == name_str)
        {
            Some((ino, _)) => *ino,
            None => {
                debug!(%parent, name = %name_str, "rename: source not found");
                reply.error(Errno::ENOENT);
                return;
            }
        };

        let new_local_path = if newparent == INodeNo::ROOT {
            std::path::PathBuf::from(newname_str)
        } else {
            match files.get(&newparent) {
                Some(dir) => dir.local_path.join(newname_str),
                None => {
                    debug!(%newparent, "rename: target parent not found");
                    reply.error(Errno::ENOENT);
                    return;
                }
            }
        };

        debug!(ino = %item_ino, from = %name_str, to = %newname_str, "rename");

        let new_parent_inode = parent_inode_from_path(&new_local_path);

        let item = files.get_mut(&item_ino).unwrap();
        item.name = newname_str.to_string();
        item.local_path = new_local_path.clone();
        item.parent_inode = new_parent_inode;

        let inode_i64 = item_ino.0 as i64;
        let new_name_owned = newname_str.to_string();
        let new_local_path_str = new_local_path.to_string_lossy().into_owned();
        let client = self.client.clone();

        drop(files);

        self.runtime_handle.block_on(async move {
            use sea_orm::ActiveValue::Set;
            if let Err(e) = (crate::db::entities::ActiveModel {
                inode: Set(inode_i64),
                name: Set(new_name_owned),
                local_path: Set(new_local_path_str),
                ..Default::default()
            }
            .update(&client)
            .await)
            {
                error!(ino = inode_i64, error = %e, "rename: db update failed");
            }
        });

        reply.ok();
    }

    fn read(
        &self,
        _req: &Request,
        ino: INodeNo,
        _fh: fuser::FileHandle,
        offset: u64,
        size: u32,
        _flags: fuser::OpenFlags,
        _lock_owner: Option<fuser::LockOwner>,
        reply: fuser::ReplyData,
    ) {
        trace!(%ino, offset, size, "read");
        let (inode, locator, virtual_layout) = {
            let files = self.files.lock().unwrap();
            let item = match files.get(&ino) {
                Some(f) => f,
                None => {
                    debug!(%ino, "read: not found");
                    reply.error(Errno::ENOENT);
                    return;
                }
            };
            let inode = item.inode;
            let locator = item.original_path.clone();
            let virtual_layout = item
                .music_metadata
                .as_ref()
                .filter(|mm| !mm.header.is_empty())
                .map(|mm| {
                    (
                        mm.header.starts_with(b"ID3"),
                        mm.header.clone(),
                        mm.picture_block_headers.clone(),
                        mm.picture_data_ranges.clone(),
                        mm.real_audio_start,
                    )
                });
            (inode, locator, virtual_layout)
        };

        let Some((is_mp3, header, pic_hdrs, pic_ranges, real_audio_start)) = virtual_layout else {
            match self.bytes.read_at(inode, &locator, offset, size as usize) {
                Ok(bytes) => reply.data(&bytes),
                Err(e) => {
                    error!(%inode, offset, size, error = %e, "read: read_at failed; returning EIO");
                    reply.error(Errno::EIO);
                }
            }
            return;
        };

        let bytes = &self.bytes;
        let reader = |off: u64, len: usize| bytes.read_at(inode, &locator, off, len);
        let result = if is_mp3 {
            file_io::assemble_mp3_read(
                &reader,
                &header,
                &pic_hdrs,
                &pic_ranges,
                real_audio_start,
                offset,
                size,
            )
        } else {
            file_io::assemble_flac_read(
                &reader,
                &header,
                &pic_hdrs,
                &pic_ranges,
                real_audio_start,
                offset,
                size,
            )
        };
        match result {
            Ok(bytes) => reply.data(&bytes),
            Err(e) => {
                error!(%inode, offset, size, error = %e, "read: assembly failed; returning EIO");
                reply.error(Errno::EIO);
            }
        }
    }

    fn lookup(
        &self,
        _req: &Request,
        parent: INodeNo,
        name: &std::ffi::OsStr,
        reply: fuser::ReplyEntry,
    ) {
        trace!(%parent, name = %name.display(), "lookup");

        match self
            .files
            .lock()
            .unwrap()
            .iter()
            .find(|item| item.1.parent_inode == parent && item.1.name == name.to_str().unwrap())
        {
            Some(item) => {
                let ttl = Duration::new(1, 0);
                let attr = file_to_attr(item.1);

                reply.entry(&ttl, &attr, Generation(0));
            }
            None => {
                debug!(%parent, name = %name.display(), "lookup: not found");
                reply.error(Errno::ENOENT);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::item::FileType;
    use crate::music::metadata::MusicMetadata;
    use crate::origins::attrs::FileAttrs;
    use fuser::INodeNo;
    use std::os::unix::fs::MetadataExt;
    use std::path::PathBuf;

    #[test]
    fn file_to_attr_uses_real_size_for_non_flac() {
        let tmp = tempfile::tempdir().unwrap();
        let metadata = std::fs::metadata(tmp.path()).unwrap();
        let attrs = FileAttrs::from(&metadata);
        let item = Item::new(
            INodeNo(42),
            INodeNo::ROOT,
            "test".to_string(),
            tmp.path().to_path_buf(),
            PathBuf::from("test"),
            FileType::File,
            attrs.clone(),
            None,
        );

        let attr = file_to_attr(&item);
        assert_eq!(attr.size, metadata.size());
    }

    #[test]
    fn file_to_attr_uses_virtual_size_for_flac() {
        let tmp = tempfile::tempdir().unwrap();
        let test_file = tmp.path().join("test.flac");
        std::fs::write(&test_file, vec![0u8; 1000]).unwrap();
        let metadata = std::fs::metadata(&test_file).unwrap();
        let attrs = FileAttrs::from(&metadata);
        let mm = MusicMetadata {
            real_audio_start: 500,
            header: vec![0u8; 100],
            picture_data_ranges: vec![(0, 50)],
            ..MusicMetadata::default()
        };
        let item = Item::new(
            INodeNo(43),
            INodeNo::ROOT,
            "test_flac".to_string(),
            test_file.clone(),
            PathBuf::from("test_flac"),
            FileType::File,
            attrs.clone(),
            Some(mm.clone()),
        );

        let attr = file_to_attr(&item);
        let expected_virtual_size = mm.virtual_size(attrs.size);
        assert_eq!(attr.size, expected_virtual_size);
    }

    #[test]
    fn file_to_attr_kind_directory() {
        let tmp = tempfile::tempdir().unwrap();
        let metadata = std::fs::metadata(tmp.path()).unwrap();
        let item = Item::new(
            INodeNo(44),
            INodeNo::ROOT,
            "test_dir".to_string(),
            tmp.path().to_path_buf(),
            PathBuf::from("test_dir"),
            FileType::Directory,
            FileAttrs::from(&metadata),
            None,
        );

        let attr = file_to_attr(&item);
        assert_eq!(attr.kind, fuser::FileType::Directory);
    }

    #[test]
    fn file_to_attr_kind_file() {
        let tmp = tempfile::tempdir().unwrap();
        let metadata = std::fs::metadata(tmp.path()).unwrap();
        let item = Item::new(
            INodeNo(45),
            INodeNo::ROOT,
            "test_file".to_string(),
            tmp.path().to_path_buf(),
            PathBuf::from("test_file"),
            FileType::File,
            FileAttrs::from(&metadata),
            None,
        );

        let attr = file_to_attr(&item);
        assert_eq!(attr.kind, fuser::FileType::RegularFile);
    }
}