libbtrfs 0.0.20

use super::{
    btrfs_ioctl, btrfs_ioctl_get_subvol_info_args, btrfs_root_ref, io, SubvolInfo, TreeSearch,
};
use crate::{
    bindings::{
        btrfs_ioctl_get_subvol_rootref_args, BTRFS_IOC_GET_SUBVOL_ROOTREF, BTRFS_ROOT_ITEM_KEY,
        BTRFS_ROOT_REF_KEY, BTRFS_ROOT_TREE_OBJECTID, BTRFS_VOL_NAME_MAX,
    },
    lookup, subvol, tree_search,
    util::{root_item_to_subvol_info_args, OptionFd},
    Opt,
};
use std::{
    collections::VecDeque,
    fs::File,
    os::unix::{
        fs::OpenOptionsExt,
        io::{AsRawFd, RawFd},
    },
    path::{Path, MAIN_SEPARATOR_STR as SEP},
};

fn subvol_info_name_from_bytes(bytes: &[u8]) -> [libc::c_char; BTRFS_VOL_NAME_MAX + 1] {
    let mut uninit = std::mem::MaybeUninit::<[libc::c_char; BTRFS_VOL_NAME_MAX + 1]>::uninit();
    let p = uninit.as_mut_ptr().cast::<libc::c_char>();
    unsafe {
        bytes.as_ptr().copy_to_nonoverlapping(p.cast(), bytes.len());
        p.add(bytes.len()).write(0);
        uninit.assume_init()
    }
}

/// Entries returned by [`Iter`]
pub struct SubvolEntry {
    treeid: u64,
    name: String,
    parent_id: u64,
    dirid: u64,
    path: Option<String>,
    info: Option<Box<SubvolInfo>>,
}

impl SubvolEntry {
    /// Id of this subvolume
    pub fn treeid(&self) -> u64 {
        self.treeid
    }

    /// Name for this subvolume
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Id of the subvolume which contains this subvolume
    pub fn parent_id(&self) -> u64 {
        self.parent_id
    }

    /// Inode of the directory containing this subvolume
    pub fn dirid(&self) -> u64 {
        self.dirid
    }

    /// Path for this subvoume relative to toplevel subvolume given to the iterator. This is `top`
    /// for privileged or the root subvolume referenced by `pathname` for unprivileged
    ///
    /// # Panics
    ///
    /// This function panics if [`Opt::GET_PATH`] flag was no provided
    ///
    pub fn path(&self) -> &str {
        self.path.as_ref().expect("GET_PATH Flag not provided")
    }

    /// [`SubvolInfo`] for this subvolume
    ///
    /// # Panics
    ///
    /// This function panics if [`Opt::GET_INFO`] flag was not provided
    ///
    pub fn info(&self) -> &SubvolInfo {
        self.info.as_ref().expect("GET_INFO Flag not provided")
    }
}

/// Iterator over subvolumes in a btrfs filesystem
///
/// Returned by the [`walk`] function and yields instances of <code>[io::Result]<[SubvolEntry]></code>
///
/// # Notes
///
/// **Requires CAP_SYS_ADIMN capabilities**
///
/// For an iterator that can be called by unprivileged processes see [`IterUser`]
///
/// # Panics
///
/// The iterator will panic if it encounters invalid UTF-8
///
pub struct Iter {
    fd: OptionFd,
    stack: Vec<SubvolEntry>,
    ts: tree_search::TreeSearch,
    flags: Opt,
    insert_ref: fn(&mut VecDeque<SubvolEntry>, SubvolEntry),
}

impl Iter {
    const fn opt_ordering(flags: &Opt) -> fn(&mut VecDeque<SubvolEntry>, SubvolEntry) {
        if flags.contains(Opt::DESCENDING) {
            VecDeque::push_back
        } else {
            VecDeque::push_front
        }
    }

    fn new_internal(top: u64, fd: OptionFd, flags: Opt) -> io::Result<Self> {
        let insert_ref = Self::opt_ordering(&flags);
        let parent_id = top;
        let mut ts = TreeSearch::new(fd.as_raw_fd(), |key| {
            key.tree_id = BTRFS_ROOT_TREE_OBJECTID;
            key.min_objectid = top;
            key.max_objectid = top;
            key.min_type = BTRFS_ROOT_REF_KEY;
            key.max_type = BTRFS_ROOT_REF_KEY;
        });

        let stack = ts.search()?.map_or_else(Vec::new, |items| {
            let mut ref_buf = VecDeque::with_capacity(items.nr_items as usize);

            for item in items {
                let rref = item.get::<&btrfs_root_ref>();
                let name = item.name_as_str(rref).expect("Invalid UTF-8").to_string();
                let dirid = u64::from_le(rref.dirid);
                let treeid = item.offset();

                let path = flags.contains(Opt::GET_PATH).then(|| name.clone());

                let info = flags.contains(Opt::GET_INFO).then(||
                    Box::new(SubvolInfo(btrfs_ioctl_get_subvol_info_args {
                        treeid, parent_id, dirid, name: subvol_info_name_from_bytes(name.as_bytes()),
                        ..Default::default()
                    }))
                );
                insert_ref(&mut ref_buf, SubvolEntry { treeid, name, parent_id, dirid, path, info })
            }
            ref_buf.into()
        });

        Ok(Self { fd, ts, stack, flags, insert_ref })
    }

    fn fill_root_item_info(&mut self, info: &mut SubvolInfo) -> io::Result<()> {
        self.ts
            .min_type(BTRFS_ROOT_ITEM_KEY)
            .max_type(BTRFS_ROOT_ITEM_KEY)
            .nr_items(1)
            .search()?
            .map_or(Err(io::ErrorKind::NotFound.into()), |mut item| {
                root_item_to_subvol_info_args(&mut info.0, item.next().unwrap().get());

                Ok(())
            })
    }
}

impl Iterator for Iter {
    type Item = io::Result<SubvolEntry>;

    fn next(&mut self) -> Option<Self::Item> {
        let mut top = self.stack.pop()?;
        let parent_id = top.treeid;

        match self
            .ts
            .min_type(BTRFS_ROOT_REF_KEY)
            .max_type(BTRFS_ROOT_REF_KEY)
            .min_objectid(parent_id)
            .max_objectid(parent_id)
            .nr_items(u32::MAX)
            .search()
        {
            Err(e) => return Some(Err(e)),
            Ok(items) => if let Some(items) = items {
                let mut ref_buf = VecDeque::with_capacity(items.nr_items as usize);

                for item in items {
                    let rref = item.get::<&btrfs_root_ref>();
                    let name = item.name_as_str(rref).expect("Invalid UTF-8").to_string();
                    let dirid = u64::from_le(rref.dirid);
                    let treeid = item.offset();

                    let path = if let Some(ref top_path) = top.path {
                        Some(top_path.clone()
                            + SEP
                            + match lookup::fd::path_as_str(self.fd.as_raw_fd(), dirid, parent_id) {
                                Err(e) => return Some(Err(e)),
                                Ok(lookup_path) => lookup_path
                            }
                            + &name)
                    } else {
                        None
                    };
                    let info = self.flags.contains(Opt::GET_INFO).then(||
                        Box::new(SubvolInfo(btrfs_ioctl_get_subvol_info_args {
                            treeid, parent_id, dirid, name: subvol_info_name_from_bytes(name.as_bytes()),
                            ..Default::default()
                        })
                    ));
                    (self.insert_ref)(&mut ref_buf, SubvolEntry {
                        treeid, name, parent_id, dirid, path, info
                    })
                }
                self.stack.extend(ref_buf)
            }
        }

        if let Some(ref mut info) = top.info {
            if let Err(e) = self.fill_root_item_info(info) {
                return Some(Err(e))
            }
        }

        Some(Ok(top))
    }
}

/// Iterator over subvolumes in a btrfs filesystem
///
/// Returned by the [`walk_user`] function and yields instances of
/// <code>[io::Result]<[SubvolEntry]></code>. Unlike [`Iter`] can be called by unprivileged
/// processes
///
/// # Panics
///
/// The iterator will panic if it encounters invalid UTF-8
///
pub struct IterUser {
    fd: OptionFd,
    args: btrfs_ioctl_get_subvol_rootref_args,
    stack: Vec<SubvolEntry>,
    flags: Opt,
    insert_ref: fn(&mut VecDeque<SubvolEntry>, SubvolEntry),
}

impl IterUser {
    fn new_internal(fd: OptionFd, flags: Opt) -> io::Result<IterUser> {
        let insert_ref = Iter::opt_ordering(&flags);
        let parent_id = lookup::fd::treeid(fd.as_raw_fd())?;
        let mut args = btrfs_ioctl_get_subvol_rootref_args::default();

        btrfs_ioctl(fd.as_raw_fd(), BTRFS_IOC_GET_SUBVOL_ROOTREF, &mut args)?;
        let num_items = args.num_items as usize;
        let mut ref_buf = VecDeque::with_capacity(num_items);

        for rref in args.rootref.iter().take(num_items) {
            let (lookup, name) = {
                match lookup::fd::user_path_as_str(fd.as_raw_fd(), rref.dirid, rref.treeid) {
                    // The IOC_LOOKUP_USER ioctl will fail with EACCES if `treeid` and `dirid` are
                    // not contained within the directory that `fd` refers to.
                    //
                    // This only happens if `fd` does not refer to a subvolume root. Ignoring this
                    // error allow findings subvolumes below any directory not just subvolumes
                    //
                    Err(e) if e.raw_os_error() == Some(libc::EACCES) => continue,
                    ret => ret?,
                }
            };
            // Add a null byte to the end of each path so it can be used for openat.
            // Remove it after the call.
            let path = Some(lookup.to_string() + name + "\0");
            let name = name.to_string();

            insert_ref(&mut ref_buf, SubvolEntry {
                name, parent_id, path, treeid: rref.treeid, dirid: rref.dirid, info: None,
            })
        }

        Ok(Self { stack: ref_buf.into(), fd, args, flags, insert_ref })
    }
}

impl Iterator for IterUser {
    type Item = io::Result<SubvolEntry>;

    fn next(&mut self) -> Option<Self::Item> {
        let mut top = self.stack.pop()?;
        let parent_id = top.treeid;
        let parent_path = top.path.as_mut().unwrap();

        let _close_res = match syscall!(unsafe {
            openat(self.fd.as_raw_fd(), parent_path.as_ptr().cast(), libc::O_RDONLY | libc::O_DIRECTORY)
        }) {
            Err(e) => return Some(Err(e)),
            Ok(fd) => {
                // remove null byte regardless of path_opt so path of refs are correct
                parent_path.truncate(parent_path.len() - 1);

                top.info =
                    self.flags
                        .contains(Opt::GET_INFO)
                        .then_some(match subvol::fd::get_info(fd) {
                            Err(e) => return Some(Err(e)),
                            Ok(info) => Box::new(info),
                        });

                self.args.min_treeid = 0;
                if let Err(e) = btrfs_ioctl(fd, BTRFS_IOC_GET_SUBVOL_ROOTREF, &mut self.args) {
                    return Some(Err(e))
                }
                let num_items = self.args.num_items as usize;
                let mut ref_buf = VecDeque::with_capacity(num_items);

                for rref in self.args.rootref.iter().take(num_items) {
                    let (lookup, name) = match lookup::fd::user_path_as_str(fd, rref.dirid, rref.treeid) {
                        Err(e) => return Some(Err(e)),
                        Ok(tup) => tup,
                    };
                    let path = Some(parent_path.clone() + SEP + lookup + name + "\0");
                    let name = name.to_string();

                    (self.insert_ref)(&mut ref_buf, SubvolEntry {
                        name, parent_id, path, dirid: rref.dirid, treeid: rref.treeid, info: None,
                    })
                }
                self.stack.extend(ref_buf);

                syscall!(unsafe { close(fd) })
            },
        };
        debug_assert!(_close_res.is_ok(), "close failed");

        if !self.flags.contains(Opt::GET_PATH) {
            top.path = None
        }

        Some(Ok(top))
    }
}

/// Returns an iterator that will walk subvolumes in a btrfs filesystem tree
///
/// The iterator returns all subvolumes below the subvolume referenced by `top`, wich must be a
/// subvolume in a btrfs filesystem referenced by `fs`. The iterator can be customized with the
/// options provided to `flags`. Calls to next will yeild instances of
/// <code>[io::Result]<[SubvolEntry]></code>
///
/// # Flags
///
/// Full list of available flags:
///
/// * [`Opt::ASCENDING `]
///
/// Subvolumes are returned in ascending order by treeid for each subvolume referencing a given root.
/// This is the default ordering and exists for documentation purposes. Passing this flag has no
/// effect
///
/// * [`Opt::DESCENDING `]
///
/// Subvolumes are returned in descending order by treeid for each subvolume referencing a given root
///
/// * [`Opt::GET_PATH`]
///
/// Get the subvolume path for each subvolume. [`SubvolEntry::path`] will panic if this flag is not
/// provided
///
/// * [`Opt::GET_INFO `]
///
/// Get [`SubvolInfo`] for each subvolume. [`SubvolEntry::info`] will panic if this flag is not
/// provided
///
/// Note that this option will likley change in the future because the [`SubvolInfo`] structure stores
/// the name field as an array not a dynamically allocated object which is not memory efficent
/// for members of a collection. Additionally some fields are redundant.
///
/// # Notes
///
/// **Requires CAP_SYS_ADIMN capabilities**
///
/// # Panics
///
/// The iterator will panic if it encounters invalid UTF-8
///
pub fn walk<P: AsRef<Path>>(top: u64, fs: P, flags: Opt) -> io::Result<Iter> {
    let fd = File::open(fs)?;

    Iter::new_internal(top, OptionFd::File(fd), flags)
}

/// Returns an iterator that will walk subvolumes in a btrfs filesystem tree
///
/// The iterator returns all subvolumes below the directory referenced by `pathname`. The
/// iterator can be customized with the options provided to `flags`. Calls to next will yeild
/// instances of <code>[io::Result]<[SubvolEntry]></code>
///
/// # Flags
///
/// For a full list of available flags see: [`walk`]
///
/// # Errors
///
/// * [`io::ErrorKind::NotADirectory`]
///
/// `pathname` does not refer to a directory
///
/// Note that in the raw file descriptor version of this function (see [`fd::walk_user`]) it is not
/// an error if the raw file descriptor does not refer to a directory, however no subvolumes will
/// be returned
///
/// # Panics
///
/// The iterator will panic if it encounters invalid UTF-8
///
pub fn walk_user<P: AsRef<Path>>(pathname: P, flags: Opt) -> io::Result<IterUser> {
    let fd = File::options()
        .read(true)
        .custom_flags(libc::O_DIRECTORY)
        .open(pathname)?;

    IterUser::new_internal(OptionFd::File(fd), flags)
}

pub mod fd {
    use super::*;

    /// See [super::walk()]
    pub fn walk(top: u64, fd: RawFd, flags: Opt) -> io::Result<Iter> {
        Iter::new_internal(top, OptionFd::Raw(fd), flags)
    }

    /// See [super::walk_user()]
    pub fn walk_user(fd: RawFd, flags: Opt) -> io::Result<IterUser> {
        IterUser::new_internal(OptionFd::Raw(fd), flags)
    }
}