shell_rs/

truncate.rs

// Copyright (c) 2021 Xu Shaohua <shaohua@biofan.org>. All rights reserved.
// Use of this source is governed by General Public License that can be found
// in the LICENSE file.

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

use crate::error::{Error, ErrorKind};

// TODO(Shaohua): Read block size info from system.
const BLK_SIZE: nc::blksize_t = 512;

#[derive(Debug, Clone)]
pub struct Options {
    /// Do not create any files.
    pub no_create: bool,

    /// Treat `size` as number of IO blocks instead of bytes.
    pub io_blocks: bool,

    /// Base size on another file.
    pub reference: Option<PathBuf>,

    /// Set or adjust the file size by `size` bytes.
    ///
    /// The `size` argument is an integer and optional unit (example: 10K is 10*1024).
    /// Units are K,M,G,T,P,E, (powers of  1024) or KB,MB,... (powers of 1000).
    /// Binary prefixes can be used, too: KiB=K, MiB=M, and so on.
    ///
    /// `size` may also be prefixed by one of the following modifying characters:
    ///   * '+' extend by
    ///   * '-' reduce by
    ///   * '<' at most
    ///   * '>' at least
    ///   * '/' round down to multiple of
    ///   * '%' round up to multiple of
    pub size: Option<String>,
}

impl Default for Options {
    fn default() -> Self {
        Self {
            no_create: false,
            io_blocks: false,
            reference: None,
            size: None,
        }
    }
}

/// Shrink or extend the size of a file to the specified size.
///
/// Shrink or extend the size of file to the specified size.
/// `file` argument that does not exist is created.
/// If `file` is larger than the specified size, the extra data is lost.
/// If `file` is shorter, it is extended and the sparse extended part (hole) reads as zero bytes.
/// Mandatory arguments to long options are mandatory for short options too.
pub fn truncate<P: AsRef<Path>>(file: P, options: &Options) -> Result<u64, Error> {
    let new_size: u64 = if let Some(size) = &options.size {
        parse_size(file.as_ref(), size, options.io_blocks)?
    } else if let Some(ref_file) = &options.reference {
        let fd = unsafe { nc::openat(nc::AT_FDCWD, ref_file, nc::O_RDONLY, 0)? };
        let mut statbuf = nc::stat_t::default();
        unsafe { nc::fstat(fd, &mut statbuf)? };
        unsafe { nc::close(fd)? };
        statbuf.st_size as u64
    } else {
        return Err(Error::new(
            ErrorKind::ParameterError,
            "Please specify either`size` or `reference`",
        ));
    };

    let file = file.as_ref();
    if let Err(err) = unsafe { nc::access(file, nc::R_OK | nc::W_OK) } {
        if options.no_create {
            return Err(err.into());
        }
    }

    let fd = unsafe { nc::openat(nc::AT_FDCWD, file, nc::O_CREAT | nc::O_WRONLY, 0o644)? };
    unsafe { nc::ftruncate(fd, new_size as nc::off_t)? };
    unsafe { nc::close(fd)? };
    Ok(new_size as u64)
}

fn parse_size<P: AsRef<Path>>(file: P, size: &str, io_blocks: bool) -> Result<u64, Error> {
    let pattern =
        regex::Regex::new(r"^(?P<prefix>[+\-<>/%]*)(?P<num>\d+)(?P<suffix>\w*)$").unwrap();
    let matches = pattern.captures(size).unwrap();
    let size_error = Error::from_string(
        ErrorKind::ParameterError,
        format!("Invalid size: {:?}", size),
    );

    let mut num: u64 = matches
        .name("num")
        .ok_or_else(|| size_error.clone())?
        .as_str()
        .parse()?;
    if let Some(suffix) = matches.name("suffix") {
        let suffix = suffix.as_str();
        if suffix.is_empty() {
            // Do nothing.
        } else if suffix == "K" {
            num *= 1024;
        } else if suffix == "KB" {
            num *= 1000;
        } else if suffix == "M" {
            num *= 1024 * 1024;
        } else if suffix == "MB" {
            num *= 1000 * 1000;
        } else if suffix == "G" {
            num *= 1024 * 1024 * 1024;
        } else if suffix == "GB" {
            num *= 1000 * 1000 * 1000;
        } else if suffix == "T" {
            num *= 1024 * 1024 * 1024 * 1024;
        } else if suffix == "TB" {
            num *= 1000 * 1000 * 1000 * 1000;
        } else if suffix == "P" {
            num *= 1024 * 1024 * 1024 * 1024 * 1024;
        } else if suffix == "PB" {
            num *= 1000 * 1000 * 1000 * 1000 * 1000;
        } else if suffix == "E" {
            num *= 1024 * 1024 * 1024 * 1024 * 1024 * 1024;
        } else if suffix == "EB" {
            num *= 1000 * 1000 * 1000 * 1000 * 1000 * 1000;
        } else {
            return Err(size_error);
        }
    }

    let num = if io_blocks {
        num * BLK_SIZE as u64
    } else {
        num
    };
    let mut new_size = num;
    if let Some(prefix) = matches.name("prefix") {
        let prefix = prefix.as_str();
        if !prefix.is_empty() {
            let fd = unsafe { nc::openat(nc::AT_FDCWD, file.as_ref(), nc::O_RDONLY, 0)? };
            let mut statbuf = nc::stat_t::default();
            unsafe { nc::fstat(fd, &mut statbuf)? };
            unsafe { nc::close(fd)? };
            let old_size = statbuf.st_size as u64;
            if prefix == "+" {
                // Extend
                new_size = old_size + num;
            } else if prefix == "-" {
                // Shrink
                new_size = old_size - num;
            } else if prefix == "<" {
                // At most
                new_size = u64::min(old_size, num);
            } else if prefix == ">" {
                // At least
                new_size = u64::max(old_size, num);
            } else if prefix == "/" {
                // Round down
                let rem = old_size.rem_euclid(num);
                println!("rem: {}", rem);
                new_size = old_size - rem;
            } else if prefix == "%" {
                // Round up
                let rem = old_size.rem_euclid(num);
                println!("rem: {}", rem);
                new_size = old_size - rem + num;
            } else {
                return Err(size_error);
            }
        }
    }

    Ok(new_size)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_size() {
        let file = "tests/Rust_Wikipedia.pdf";
        const OLD_SIZE: u64 = 455977;

        assert_eq!(parse_size(file, "1024", false), Ok(1024));
        assert_eq!(parse_size(file, "1K", false), Ok(1024));
        assert_eq!(parse_size(file, "1M", false), Ok(1024 * 1024));
        assert_eq!(parse_size(file, "1G", false), Ok(1024 * 1024 * 1024));
        assert_eq!(parse_size(file, "1T", false), Ok(1024 * 1024 * 1024 * 1024));
        assert_eq!(parse_size(file, "1KB", false), Ok(1000));
        assert_eq!(parse_size(file, "1MB", false), Ok(1000 * 1000));
        assert_eq!(parse_size(file, "1GB", false), Ok(1000 * 1000 * 1000));
        assert_eq!(
            parse_size(file, "1TB", false),
            Ok(1000 * 1000 * 1000 * 1000)
        );

        assert_eq!(parse_size(file, "+1024", false), Ok(OLD_SIZE + 1024));
        assert_eq!(parse_size(file, "-1024", false), Ok(OLD_SIZE - 1024));
        assert_eq!(parse_size(file, "<1024", false), Ok(1024));
        assert_eq!(parse_size(file, ">1024", false), Ok(OLD_SIZE));
        assert_eq!(parse_size(file, "/1024", false), Ok(455680));
        assert_eq!(parse_size(file, "%1024", false), Ok(456704));
    }

    #[test]
    fn test_truncate() {
        let file = "/tmp/truncate.shell-rs";
        assert!(truncate(file, &Options::default()).is_err());

        assert_eq!(
            truncate(
                file,
                &Options {
                    size: Some("1M".to_string()),
                    ..Options::default()
                },
            ),
            Ok(1024 * 1024)
        );

        assert_eq!(
            truncate(
                file,
                &Options {
                    size: Some("1K".to_string()),
                    ..Options::default()
                },
            ),
            Ok(1024)
        );

        let _ = nc::unlink(file);
    }
}