tourmaline/src/task/chrooting/mapping.rs

use std::path::{Path, PathBuf};

use sys_mount::{FilesystemType, Mount, MountBuilder, MountFlags, UnmountDrop, UnmountFlags};
use tokio::fs;

use crate::error::ChrootError;

pub enum Mapping {
    Device(PathBuf),
    Dir(PathBuf, DirOpts),
    Special(PathBuf, PathBuf, DirOpts),
    Link(PathBuf, PathBuf),
    Copy(PathBuf, PathBuf),
}

#[derive(Default)]
pub struct DirOpts {
    fs_type: Option<String>,
    mount_flags: Option<MountFlags>,
    required: bool,
    unmount_flags: Option<UnmountFlags>,
}

impl DirOpts {
    pub fn fs_type<S: Into<String>>(mut self, fs_type: S) -> Self {
        self.fs_type = Some(fs_type.into());

        self
    }

    pub fn flags(mut self, flags: MountFlags) -> Self {
        self.mount_flags = Some(flags);

        self
    }

    pub fn optional(mut self) -> Self {
        self.required = false;

        self
    }

    pub fn unmount_flags(mut self, flags: UnmountFlags) -> Self {
        self.unmount_flags = Some(flags);

        self
    }
}

pub enum MappingHandle {
    Mount(UnmountDrop<Mount>),
    Link(LinkDrop),
    None,
}

pub struct LinkDrop {
    path: PathBuf,
}

impl Drop for LinkDrop {
    fn drop(&mut self) {
        std::fs::remove_file(&self.path).unwrap();
    }
}

impl Mapping {
    fn dev<P: Into<PathBuf>>(path: P) -> Self {
        Self::Device(path.into())
    }

    fn dir<P: Into<PathBuf>>(path: P, opt: DirOpts) -> Self {
        Self::Dir(path.into(), opt)
    }

    fn special<P1: Into<PathBuf>, P2: Into<PathBuf>>(src: P1, dst: P2, opt: DirOpts) -> Self {
        Self::Special(src.into(), dst.into(), opt)
    }

    fn link<P1: Into<PathBuf>, P2: Into<PathBuf>>(src: P1, dst: P2) -> Self {
        Self::Link(src.into(), dst.into())
    }

    fn copy<P1: Into<PathBuf>, P2: Into<PathBuf>>(src: P1, dst: P2) -> Self {
        Self::Copy(src.into(), dst.into())
    }

    pub async fn create_mapping(&self, root_path: &Path) -> Result<MappingHandle, ChrootError> {
        match &self {
            Mapping::Device(d) => self.map_dev(d, &root_path.join(d)).await,
            Mapping::Dir(d, opt) => self.map_dir(d, &root_path.join(d), opt).await,
            Mapping::Special(src, dst, opt) => self.map_dir(src, &root_path.join(dst), opt).await,
            Mapping::Link(src, dst) => self.map_link(src, &root_path.join(dst)).await,
            Mapping::Copy(src, dst) => self.copy_file(src, &root_path.join(dst)).await,
        }
    }

    async fn map_dev(&self, src: &Path, dst: &Path) -> Result<MappingHandle, ChrootError> {
        if !dst.exists() {
            if let Some(parent) = dst.parent() {
                fs::create_dir_all(&parent)
                    .await
                    .map_err(|e| ChrootError::Mount(dst.to_owned(), e))?;
            }
            fs::write(&dst, "")
                .await
                .map_err(|e| ChrootError::Mount(dst.to_owned(), e))?;
        }
        let mount = MountBuilder::default()
            .flags(MountFlags::BIND)
            .mount_autodrop(src, dst, UnmountFlags::empty())
            .map_err(|e| ChrootError::Mount(dst.to_owned(), e))?;

        Ok(MappingHandle::Mount(mount))
    }

    async fn map_dir(
        &self,
        src: &Path,
        dst: &Path,
        opt: &DirOpts,
    ) -> Result<MappingHandle, ChrootError> {
        if !dst.exists() {
            fs::create_dir_all(&dst)
                .await
                .map_err(|e| ChrootError::Mount(dst.to_owned(), e))?;
        }
        let fs_type = FilesystemType::Manual(opt.fs_type.as_deref().unwrap_or(""));
        let flags = opt.mount_flags.unwrap_or_else(MountFlags::empty);

        let mount = MountBuilder::default()
            .fstype(fs_type)
            .flags(flags)
            .mount_autodrop(src, dst, opt.unmount_flags.unwrap_or(UnmountFlags::empty()))
            .map_err(|e| ChrootError::Mount(src.to_owned(), e))?;

        Ok(MappingHandle::Mount(mount))
    }

    async fn map_link(&self, src: &Path, dst: &Path) -> Result<MappingHandle, ChrootError> {
        if dst.exists() && dst.is_symlink() {
            fs::remove_file(dst).await.map_err(ChrootError::Unlink)?;
        }
        fs::symlink(src, dst)
            .await
            .map_err(|e| ChrootError::Link(src.to_owned(), e))?;
        Ok(MappingHandle::Link(LinkDrop {
            path: dst.to_owned(),
        }))
    }

    async fn copy_file(&self, src: &Path, dst: &Path) -> Result<MappingHandle, ChrootError> {
        if dst.exists() && dst.is_file() {
            fs::remove_file(dst)
                .await
                .map_err(|e| ChrootError::Copy(src.to_owned(), e))?;
        }

        if src.exists() {
            fs::copy(src, dst)
                .await
                .map_err(|e| ChrootError::Copy(src.to_owned(), e))?;
        }

        Ok(MappingHandle::None)
    }
}

pub fn default_mappings() -> Vec<Mapping> {
    static NOSUID: MountFlags = MountFlags::NOSUID;
    static NOEXEC: MountFlags = MountFlags::NOEXEC;
    static NODEV: MountFlags = MountFlags::NODEV;
    static STRICTATIME: MountFlags = MountFlags::STRICTATIME;

    vec![
        Mapping::dir(
            "proc",
            DirOpts::default()
                .fs_type("proc")
                .flags(NOSUID | NOEXEC | NODEV),
        ),
        Mapping::dir(
            "sys",
            DirOpts::default()
                .fs_type("sysfs")
                .flags(NOSUID | NOEXEC | NODEV)
                .unmount_flags(UnmountFlags::DETACH),
        ),
        Mapping::dir(
            "sys/firmware/efi/efivars",
            DirOpts::default()
                .fs_type("efivarfs")
                .flags(NOSUID | NOEXEC | NODEV)
                .optional(),
        ),
        Mapping::special(
            "udev",
            "dev",
            DirOpts::default().fs_type("devtmpfs").flags(NOSUID),
        ),
        Mapping::special(
            "devpts",
            "dev/pts",
            DirOpts::default().fs_type("devpts").flags(NOSUID | NOEXEC),
        ),
        Mapping::special(
            "shm",
            "dev/shm",
            DirOpts::default().fs_type("tmpfs").flags(NOSUID | NODEV),
        ),
        Mapping::dev("/dev/full"),
        Mapping::dev("/dev/null"),
        Mapping::dev("/dev/random"),
        Mapping::dev("/dev/tty"),
        Mapping::dev("/dev/urandom"),
        Mapping::dev("/dev/zero"),
        Mapping::link("/proc/self/fd", "dev/fd"),
        Mapping::link("/proc/self/fd/0", "dev/stdin"),
        Mapping::link("/proc/self/fd/1", "dev/stdout"),
        Mapping::link("/proc/self/fd/2", "dev/stderr"),
        Mapping::copy("/etc/resolv.conf", "etc/resolv.conf"),
        Mapping::dir(
            "/run",
            DirOpts::default().fs_type("tmpfs").flags(NOSUID | NODEV),
        ),
        Mapping::dir(
            "/tmp",
            DirOpts::default()
                .fs_type("tmpfs")
                .flags(STRICTATIME | NODEV | NOSUID),
        ),
    ]
}