<commit_before>// This file is part of the uutils coreutils package.
//
// (c) Jian Zeng <anonymousknight96@gmail.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

// spell-checker:ignore (ToDO) COMFOLLOW Chowner Passwd RFILE RFILE's derefer dgid duid

#[macro_use]
extern crate uucore;
pub use uucore::entries::{self, Group, Locate, Passwd};
use uucore::fs::resolve_relative_path;
use uucore::libc::{gid_t, uid_t};
use uucore::perms::{wrap_chown, Verbosity};

use clap::{App, Arg};

use walkdir::WalkDir;

use std::fs::{self, Metadata};
use std::os::unix::fs::MetadataExt;

use std::convert::AsRef;
use std::path::Path;

static ABOUT: &str = "change file owner and group";
static VERSION: &str = env!("CARGO_PKG_VERSION");

pub mod options {
    pub mod verbosity {
        pub static CHANGES: &str = "changes";
        pub static QUIET: &str = "quiet";
        pub static SILENT: &str = "silent";
        pub static VERBOSE: &str = "verbose";
    }
    pub mod preserve_root {
        pub static PRESERVE: &str = "preserve-root";
        pub static NO_PRESERVE: &str = "no-preserve-root";
    }
    pub mod dereference {
        pub static DEREFERENCE: &str = "dereference";
        pub static NO_DEREFERENCE: &str = "no-dereference"; // not actually read?
    }
    pub static FROM: &str = "from";
    pub static RECURSIVE: &str = "recursive"; // TODO(felipel) recursive, reference, traverse related?
    pub mod traverse {
        pub static TRAVERSE: &str = "H";
        pub static NO_TRAVERSE: &str = "P";
        pub static EVERY: &str = "L";
    }
    pub static REFERENCE: &str = "reference";
}

static ARG_OWNER: &str = "owner";
static ARG_FILES: &str = "files";

const FTS_COMFOLLOW: u8 = 1;
const FTS_PHYSICAL: u8 = 1 << 1;
const FTS_LOGICAL: u8 = 1 << 2;

fn get_usage() -> String {
    format!(
        "{0} [OPTION]... [OWNER][:[GROUP]] FILE...\n{0} [OPTION]... --reference=RFILE FILE...",
        executable!()
    )
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let args = args.collect_str();

    let usage = get_usage();

    let matches = App::new(executable!())
        .version(VERSION)
        .about(ABOUT)
        .usage(&usage[..])
        .arg(
            Arg::with_name(options::verbosity::CHANGES)
                .short("c")
                .long(options::verbosity::CHANGES)
                .help("like verbose but report only when a change is made"),
        )
        .arg(Arg::with_name(options::dereference::DEREFERENCE).long(options::dereference::DEREFERENCE).help(
            "affect the referent of each symbolic link (this is the default), rather than the symbolic link itself",
        ))
        .arg(
            Arg::with_name(options::dereference::NO_DEREFERENCE)
                .short("h")
                .long(options::dereference::NO_DEREFERENCE)
                .help(
                    "affect symbolic links instead of any referenced file (useful only on systems that can change the ownership of a symlink)",
                ),
        )
        .arg(
            Arg::with_name(options::FROM)
                .long(options::FROM)
                .help(
                    "change the owner and/or group of each file only if its current owner and/or group match those specified here. Either may be omitted, in which case a match is not required for the omitted attribute",
                )
                .value_name("CURRENT_OWNER:CURRENT_GROUP"),
        )
        .arg(
            Arg::with_name(options::preserve_root::PRESERVE)
                .long(options::preserve_root::PRESERVE)
                .help("fail to operate recursively on '/'"),
        )
        .arg(
            Arg::with_name(options::preserve_root::NO_PRESERVE)
                .long(options::preserve_root::NO_PRESERVE)
                .help("do not treat '/' specially (the default)"),
        )
        .arg(
            Arg::with_name(options::verbosity::QUIET)
                .long(options::verbosity::QUIET)
                .help("suppress most error messages"),
        )
        .arg(
            Arg::with_name(options::RECURSIVE)
                .short("R")
                .long(options::RECURSIVE)
                .help("operate on files and directories recursively"),
        )
        .arg(
            Arg::with_name(options::REFERENCE)
                .long(options::REFERENCE)
                .help("use RFILE's owner and group rather than specifying OWNER:GROUP values")
                .value_name("RFILE")
                .min_values(1),
        )
        .arg(Arg::with_name(options::verbosity::SILENT).short("f").long(options::verbosity::SILENT))
        .arg(
            Arg::with_name(options::traverse::TRAVERSE)
                .short(options::traverse::TRAVERSE)
                .help("if a command line argument is a symbolic link to a directory, traverse it")
                .overrides_with_all(&[options::traverse::EVERY, options::traverse::NO_TRAVERSE]),
        )
        .arg(
            Arg::with_name(options::traverse::EVERY)
                .short(options::traverse::EVERY)
                .help("traverse every symbolic link to a directory encountered")
                .overrides_with_all(&[options::traverse::TRAVERSE, options::traverse::NO_TRAVERSE]),
        )
        .arg(
            Arg::with_name(options::traverse::NO_TRAVERSE)
                .short(options::traverse::NO_TRAVERSE)
                .help("do not traverse any symbolic links (default)")
                .overrides_with_all(&[options::traverse::TRAVERSE, options::traverse::EVERY]),
        )
        .arg(
            Arg::with_name(options::verbosity::VERBOSE)
                .long(options::verbosity::VERBOSE)
                .help("output a diagnostic for every file processed"),
        )
        .arg(
            Arg::with_name(ARG_OWNER)
                .multiple(false)
                .takes_value(true)
                .required(true),
        )
        .arg(
            Arg::with_name(ARG_FILES)
                .multiple(true)
                .takes_value(true)
                .required(true)
                .min_values(1),
        )
        .get_matches_from(args);

    /* First arg is the owner/group */
    let owner = matches.value_of(ARG_OWNER).unwrap();

    /* Then the list of files */
    let files: Vec<String> = matches
        .values_of(ARG_FILES)
        .map(|v| v.map(ToString::to_string).collect())
        .unwrap_or_default();

    let preserve_root = matches.is_present(options::preserve_root::PRESERVE);

    let mut derefer = if matches.is_present(options::dereference::NO_DEREFERENCE) {
        1
    } else {
        0
    };

    let mut bit_flag = if matches.is_present(options::traverse::TRAVERSE) {
        FTS_COMFOLLOW | FTS_PHYSICAL
    } else if matches.is_present(options::traverse::EVERY) {
        FTS_LOGICAL
    } else {
        FTS_PHYSICAL
    };

    let recursive = matches.is_present(options::RECURSIVE);
    if recursive {
        if bit_flag == FTS_PHYSICAL {
            if derefer == 1 {
                show_info!("-R --dereference requires -H or -L");
                return 1;
            }
            derefer = 0;
        }
    } else {
        bit_flag = FTS_PHYSICAL;
    }

    let verbosity = if matches.is_present(options::verbosity::CHANGES) {
        Verbosity::Changes
    } else if matches.is_present(options::verbosity::SILENT)
        || matches.is_present(options::verbosity::QUIET)
    {
        Verbosity::Silent
    } else if matches.is_present(options::verbosity::VERBOSE) {
        Verbosity::Verbose
    } else {
        Verbosity::Normal
    };

    let filter = if let Some(spec) = matches.value_of(options::FROM) {
        match parse_spec(&spec) {
            Ok((Some(uid), None)) => IfFrom::User(uid),
            Ok((None, Some(gid))) => IfFrom::Group(gid),
            Ok((Some(uid), Some(gid))) => IfFrom::UserGroup(uid, gid),
            Ok((None, None)) => IfFrom::All,
            Err(e) => {
                show_info!("{}", e);
                return 1;
            }
        }
    } else {
        IfFrom::All
    };

    let dest_uid: Option<u32>;
    let dest_gid: Option<u32>;
    if let Some(file) = matches.value_of(options::REFERENCE) {
        match fs::metadata(&file) {
            Ok(meta) => {
                dest_gid = Some(meta.gid());
                dest_uid = Some(meta.uid());
            }
            Err(e) => {
                show_info!("failed to get attributes of '{}': {}", file, e);
                return 1;
            }
        }
    } else {
        match parse_spec(&owner) {
            Ok((u, g)) => {
                dest_uid = u;
                dest_gid = g;
            }
            Err(e) => {
                show_info!("{}", e);
                return 1;
            }
        }
    }
    let executor = Chowner {
        bit_flag,
        dest_uid,
        dest_gid,
        verbosity,
        recursive,
        dereference: derefer != 0,
        filter,
        preserve_root,
        files,
    };
    executor.exec()
}

fn parse_spec(spec: &str) -> Result<(Option<u32>, Option<u32>), String> {
    let args = spec.split(':').collect::<Vec<_>>();
    let usr_only = args.len() == 1;
    let grp_only = args.len() == 2 && args[0].is_empty() && !args[1].is_empty();
    let usr_grp = args.len() == 2 && !args[0].is_empty() && !args[1].is_empty();

    if usr_only {
        Ok((
            Some(match Passwd::locate(args[0]) {
                Ok(v) => v.uid(),
                _ => return Err(format!("invalid user: '{}'", spec)),
            }),
            None,
        ))
    } else if grp_only {
        Ok((
            None,
            Some(match Group::locate(args[1]) {
                Ok(v) => v.gid(),
                _ => return Err(format!("invalid group: '{}'", spec)),
            }),
        ))
    } else if usr_grp {
        Ok((
            Some(match Passwd::locate(args[0]) {
                Ok(v) => v.uid(),
                _ => return Err(format!("invalid user: '{}'", spec)),
            }),
            Some(match Group::locate(args[1]) {
                Ok(v) => v.gid(),
                _ => return Err(format!("invalid group: '{}'", spec)),
            }),
        ))
    } else {
        Ok((None, None))
    }
}

enum IfFrom {
    All,
    User(u32),
    Group(u32),
    UserGroup(u32, u32),
}

struct Chowner {
    dest_uid: Option<u32>,
    dest_gid: Option<u32>,
    bit_flag: u8,
    verbosity: Verbosity,
    filter: IfFrom,
    files: Vec<String>,
    recursive: bool,
    preserve_root: bool,
    dereference: bool,
}

macro_rules! unwrap {
    ($m:expr, $e:ident, $err:block) => {
        match $m {
            Ok(meta) => meta,
            Err($e) => $err,
        }
    };
}

impl Chowner {
    fn exec(&self) -> i32 {
        let mut ret = 0;
        for f in &self.files {
            ret |= self.traverse(f);
        }
        ret
    }

    fn traverse<P: AsRef<Path>>(&self, root: P) -> i32 {
        let follow_arg = self.dereference || self.bit_flag != FTS_PHYSICAL;
        let path = root.as_ref();
        let meta = match self.obtain_meta(path, follow_arg) {
            Some(m) => m,
            _ => return 1,
        };

        // Prohibit only if:
        // (--preserve-root and -R present) &&
        // (
        //     (argument is not symlink && resolved to be '/') ||
        //     (argument is symlink && should follow argument && resolved to be '/')
        // )
        if self.recursive && self.preserve_root {
            let may_exist = if follow_arg {
                path.canonicalize().ok()
            } else {
                let real = resolve_relative_path(path);
                if real.is_dir() {
                    Some(real.canonicalize().expect("failed to get real path"))
                } else {
                    Some(real.into_owned())
                }
            };

            if let Some(p) = may_exist {
                if p.parent().is_none() {
                    show_info!("it is dangerous to operate recursively on '/'");
                    show_info!("use --no-preserve-root to override this failsafe");
                    return 1;
                }
            }
        }

        let ret = if self.matched(meta.uid(), meta.gid()) {
            match wrap_chown(
                path,
                &meta,
                self.dest_uid,
                self.dest_gid,
                follow_arg,
                self.verbosity.clone(),
            ) {
                Ok(n) => {
                    if n != "" {
                        show_info!("{}", n);
                    }
                    0
                }
                Err(e) => {
                    if self.verbosity != Verbosity::Silent {
                        show_info!("{}", e);
                    }
                    1
                }
            }
        } else {
            0
        };

        if !self.recursive {
            ret
        } else {
            ret | self.dive_into(&root)
        }
    }

    fn dive_into<P: AsRef<Path>>(&self, root: P) -> i32 {
        let mut ret = 0;
        let root = root.as_ref();
        let follow = self.dereference || self.bit_flag & FTS_LOGICAL != 0;
        for entry in WalkDir::new(root).follow_links(follow).min_depth(1) {
            let entry = unwrap!(entry, e, {
                ret = 1;
                show_info!("{}", e);
                continue;
            });
            let path = entry.path();
            let meta = match self.obtain_meta(path, follow) {
                Some(m) => m,
                _ => {
                    ret = 1;
                    continue;
                }
            };

            if !self.matched(meta.uid(), meta.gid()) {
                continue;
            }

            ret = match wrap_chown(
                path,
                &meta,
                self.dest_uid,
                self.dest_gid,
                follow,
                self.verbosity.clone(),
            ) {
                Ok(n) => {
                    if n != "" {
                        show_info!("{}", n);
                    }
                    0
                }
                Err(e) => {
                    if self.verbosity != Verbosity::Silent {
                        show_info!("{}", e);
                    }
                    1
                }
            }
        }
        ret
    }

    fn obtain_meta<P: AsRef<Path>>(&self, path: P, follow: bool) -> Option<Metadata> {
        use self::Verbosity::*;
        let path = path.as_ref();
        let meta = if follow {
            unwrap!(path.metadata(), e, {
                match self.verbosity {
                    Silent => (),
                    _ => show_info!("cannot access '{}': {}", path.display(), e),
                }
                return None;
            })
        } else {
            unwrap!(path.symlink_metadata(), e, {
                match self.verbosity {
                    Silent => (),
                    _ => show_info!("cannot dereference '{}': {}", path.display(), e),
                }
                return None;
            })
        };
        Some(meta)
    }

    #[inline]
    fn matched(&self, uid: uid_t, gid: gid_t) -> bool {
        match self.filter {
            IfFrom::All => true,
            IfFrom::User(u) => u == uid,
            IfFrom::Group(g) => g == gid,
            IfFrom::UserGroup(u, g) => u == uid && g == gid,
        }
    }
}
<commit_msg>remove markers<commit_after>// This file is part of the uutils coreutils package.
//
// (c) Jian Zeng <anonymousknight96@gmail.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

// spell-checker:ignore (ToDO) COMFOLLOW Chowner Passwd RFILE RFILE's derefer dgid duid

#[macro_use]
extern crate uucore;
pub use uucore::entries::{self, Group, Locate, Passwd};
use uucore::fs::resolve_relative_path;
use uucore::libc::{gid_t, uid_t};
use uucore::perms::{wrap_chown, Verbosity};

use clap::{App, Arg};

use walkdir::WalkDir;

use std::fs::{self, Metadata};
use std::os::unix::fs::MetadataExt;

use std::convert::AsRef;
use std::path::Path;

static ABOUT: &str = "change file owner and group";
static VERSION: &str = env!("CARGO_PKG_VERSION");

pub mod options {
    pub mod verbosity {
        pub static CHANGES: &str = "changes";
        pub static QUIET: &str = "quiet";
        pub static SILENT: &str = "silent";
        pub static VERBOSE: &str = "verbose";
    }
    pub mod preserve_root {
        pub static PRESERVE: &str = "preserve-root";
        pub static NO_PRESERVE: &str = "no-preserve-root";
    }
    pub mod dereference {
        pub static DEREFERENCE: &str = "dereference";
        pub static NO_DEREFERENCE: &str = "no-dereference";
    }
    pub static FROM: &str = "from";
    pub static RECURSIVE: &str = "recursive";
    pub mod traverse {
        pub static TRAVERSE: &str = "H";
        pub static NO_TRAVERSE: &str = "P";
        pub static EVERY: &str = "L";
    }
    pub static REFERENCE: &str = "reference";
}

static ARG_OWNER: &str = "owner";
static ARG_FILES: &str = "files";

const FTS_COMFOLLOW: u8 = 1;
const FTS_PHYSICAL: u8 = 1 << 1;
const FTS_LOGICAL: u8 = 1 << 2;

fn get_usage() -> String {
    format!(
        "{0} [OPTION]... [OWNER][:[GROUP]] FILE...\n{0} [OPTION]... --reference=RFILE FILE...",
        executable!()
    )
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let args = args.collect_str();

    let usage = get_usage();

    let matches = App::new(executable!())
        .version(VERSION)
        .about(ABOUT)
        .usage(&usage[..])
        .arg(
            Arg::with_name(options::verbosity::CHANGES)
                .short("c")
                .long(options::verbosity::CHANGES)
                .help("like verbose but report only when a change is made"),
        )
        .arg(Arg::with_name(options::dereference::DEREFERENCE).long(options::dereference::DEREFERENCE).help(
            "affect the referent of each symbolic link (this is the default), rather than the symbolic link itself",
        ))
        .arg(
            Arg::with_name(options::dereference::NO_DEREFERENCE)
                .short("h")
                .long(options::dereference::NO_DEREFERENCE)
                .help(
                    "affect symbolic links instead of any referenced file (useful only on systems that can change the ownership of a symlink)",
                ),
        )
        .arg(
            Arg::with_name(options::FROM)
                .long(options::FROM)
                .help(
                    "change the owner and/or group of each file only if its current owner and/or group match those specified here. Either may be omitted, in which case a match is not required for the omitted attribute",
                )
                .value_name("CURRENT_OWNER:CURRENT_GROUP"),
        )
        .arg(
            Arg::with_name(options::preserve_root::PRESERVE)
                .long(options::preserve_root::PRESERVE)
                .help("fail to operate recursively on '/'"),
        )
        .arg(
            Arg::with_name(options::preserve_root::NO_PRESERVE)
                .long(options::preserve_root::NO_PRESERVE)
                .help("do not treat '/' specially (the default)"),
        )
        .arg(
            Arg::with_name(options::verbosity::QUIET)
                .long(options::verbosity::QUIET)
                .help("suppress most error messages"),
        )
        .arg(
            Arg::with_name(options::RECURSIVE)
                .short("R")
                .long(options::RECURSIVE)
                .help("operate on files and directories recursively"),
        )
        .arg(
            Arg::with_name(options::REFERENCE)
                .long(options::REFERENCE)
                .help("use RFILE's owner and group rather than specifying OWNER:GROUP values")
                .value_name("RFILE")
                .min_values(1),
        )
        .arg(Arg::with_name(options::verbosity::SILENT).short("f").long(options::verbosity::SILENT))
        .arg(
            Arg::with_name(options::traverse::TRAVERSE)
                .short(options::traverse::TRAVERSE)
                .help("if a command line argument is a symbolic link to a directory, traverse it")
                .overrides_with_all(&[options::traverse::EVERY, options::traverse::NO_TRAVERSE]),
        )
        .arg(
            Arg::with_name(options::traverse::EVERY)
                .short(options::traverse::EVERY)
                .help("traverse every symbolic link to a directory encountered")
                .overrides_with_all(&[options::traverse::TRAVERSE, options::traverse::NO_TRAVERSE]),
        )
        .arg(
            Arg::with_name(options::traverse::NO_TRAVERSE)
                .short(options::traverse::NO_TRAVERSE)
                .help("do not traverse any symbolic links (default)")
                .overrides_with_all(&[options::traverse::TRAVERSE, options::traverse::EVERY]),
        )
        .arg(
            Arg::with_name(options::verbosity::VERBOSE)
                .long(options::verbosity::VERBOSE)
                .help("output a diagnostic for every file processed"),
        )
        .arg(
            Arg::with_name(ARG_OWNER)
                .multiple(false)
                .takes_value(true)
                .required(true),
        )
        .arg(
            Arg::with_name(ARG_FILES)
                .multiple(true)
                .takes_value(true)
                .required(true)
                .min_values(1),
        )
        .get_matches_from(args);

    /* First arg is the owner/group */
    let owner = matches.value_of(ARG_OWNER).unwrap();

    /* Then the list of files */
    let files: Vec<String> = matches
        .values_of(ARG_FILES)
        .map(|v| v.map(ToString::to_string).collect())
        .unwrap_or_default();

    let preserve_root = matches.is_present(options::preserve_root::PRESERVE);

    let mut derefer = if matches.is_present(options::dereference::NO_DEREFERENCE) {
        1
    } else {
        0
    };

    let mut bit_flag = if matches.is_present(options::traverse::TRAVERSE) {
        FTS_COMFOLLOW | FTS_PHYSICAL
    } else if matches.is_present(options::traverse::EVERY) {
        FTS_LOGICAL
    } else {
        FTS_PHYSICAL
    };

    let recursive = matches.is_present(options::RECURSIVE);
    if recursive {
        if bit_flag == FTS_PHYSICAL {
            if derefer == 1 {
                show_info!("-R --dereference requires -H or -L");
                return 1;
            }
            derefer = 0;
        }
    } else {
        bit_flag = FTS_PHYSICAL;
    }

    let verbosity = if matches.is_present(options::verbosity::CHANGES) {
        Verbosity::Changes
    } else if matches.is_present(options::verbosity::SILENT)
        || matches.is_present(options::verbosity::QUIET)
    {
        Verbosity::Silent
    } else if matches.is_present(options::verbosity::VERBOSE) {
        Verbosity::Verbose
    } else {
        Verbosity::Normal
    };

    let filter = if let Some(spec) = matches.value_of(options::FROM) {
        match parse_spec(&spec) {
            Ok((Some(uid), None)) => IfFrom::User(uid),
            Ok((None, Some(gid))) => IfFrom::Group(gid),
            Ok((Some(uid), Some(gid))) => IfFrom::UserGroup(uid, gid),
            Ok((None, None)) => IfFrom::All,
            Err(e) => {
                show_info!("{}", e);
                return 1;
            }
        }
    } else {
        IfFrom::All
    };

    let dest_uid: Option<u32>;
    let dest_gid: Option<u32>;
    if let Some(file) = matches.value_of(options::REFERENCE) {
        match fs::metadata(&file) {
            Ok(meta) => {
                dest_gid = Some(meta.gid());
                dest_uid = Some(meta.uid());
            }
            Err(e) => {
                show_info!("failed to get attributes of '{}': {}", file, e);
                return 1;
            }
        }
    } else {
        match parse_spec(&owner) {
            Ok((u, g)) => {
                dest_uid = u;
                dest_gid = g;
            }
            Err(e) => {
                show_info!("{}", e);
                return 1;
            }
        }
    }
    let executor = Chowner {
        bit_flag,
        dest_uid,
        dest_gid,
        verbosity,
        recursive,
        dereference: derefer != 0,
        filter,
        preserve_root,
        files,
    };
    executor.exec()
}

fn parse_spec(spec: &str) -> Result<(Option<u32>, Option<u32>), String> {
    let args = spec.split(':').collect::<Vec<_>>();
    let usr_only = args.len() == 1;
    let grp_only = args.len() == 2 && args[0].is_empty() && !args[1].is_empty();
    let usr_grp = args.len() == 2 && !args[0].is_empty() && !args[1].is_empty();

    if usr_only {
        Ok((
            Some(match Passwd::locate(args[0]) {
                Ok(v) => v.uid(),
                _ => return Err(format!("invalid user: '{}'", spec)),
            }),
            None,
        ))
    } else if grp_only {
        Ok((
            None,
            Some(match Group::locate(args[1]) {
                Ok(v) => v.gid(),
                _ => return Err(format!("invalid group: '{}'", spec)),
            }),
        ))
    } else if usr_grp {
        Ok((
            Some(match Passwd::locate(args[0]) {
                Ok(v) => v.uid(),
                _ => return Err(format!("invalid user: '{}'", spec)),
            }),
            Some(match Group::locate(args[1]) {
                Ok(v) => v.gid(),
                _ => return Err(format!("invalid group: '{}'", spec)),
            }),
        ))
    } else {
        Ok((None, None))
    }
}

enum IfFrom {
    All,
    User(u32),
    Group(u32),
    UserGroup(u32, u32),
}

struct Chowner {
    dest_uid: Option<u32>,
    dest_gid: Option<u32>,
    bit_flag: u8,
    verbosity: Verbosity,
    filter: IfFrom,
    files: Vec<String>,
    recursive: bool,
    preserve_root: bool,
    dereference: bool,
}

macro_rules! unwrap {
    ($m:expr, $e:ident, $err:block) => {
        match $m {
            Ok(meta) => meta,
            Err($e) => $err,
        }
    };
}

impl Chowner {
    fn exec(&self) -> i32 {
        let mut ret = 0;
        for f in &self.files {
            ret |= self.traverse(f);
        }
        ret
    }

    fn traverse<P: AsRef<Path>>(&self, root: P) -> i32 {
        let follow_arg = self.dereference || self.bit_flag != FTS_PHYSICAL;
        let path = root.as_ref();
        let meta = match self.obtain_meta(path, follow_arg) {
            Some(m) => m,
            _ => return 1,
        };

        // Prohibit only if:
        // (--preserve-root and -R present) &&
        // (
        //     (argument is not symlink && resolved to be '/') ||
        //     (argument is symlink && should follow argument && resolved to be '/')
        // )
        if self.recursive && self.preserve_root {
            let may_exist = if follow_arg {
                path.canonicalize().ok()
            } else {
                let real = resolve_relative_path(path);
                if real.is_dir() {
                    Some(real.canonicalize().expect("failed to get real path"))
                } else {
                    Some(real.into_owned())
                }
            };

            if let Some(p) = may_exist {
                if p.parent().is_none() {
                    show_info!("it is dangerous to operate recursively on '/'");
                    show_info!("use --no-preserve-root to override this failsafe");
                    return 1;
                }
            }
        }

        let ret = if self.matched(meta.uid(), meta.gid()) {
            match wrap_chown(
                path,
                &meta,
                self.dest_uid,
                self.dest_gid,
                follow_arg,
                self.verbosity.clone(),
            ) {
                Ok(n) => {
                    if n != "" {
                        show_info!("{}", n);
                    }
                    0
                }
                Err(e) => {
                    if self.verbosity != Verbosity::Silent {
                        show_info!("{}", e);
                    }
                    1
                }
            }
        } else {
            0
        };

        if !self.recursive {
            ret
        } else {
            ret | self.dive_into(&root)
        }
    }

    fn dive_into<P: AsRef<Path>>(&self, root: P) -> i32 {
        let mut ret = 0;
        let root = root.as_ref();
        let follow = self.dereference || self.bit_flag & FTS_LOGICAL != 0;
        for entry in WalkDir::new(root).follow_links(follow).min_depth(1) {
            let entry = unwrap!(entry, e, {
                ret = 1;
                show_info!("{}", e);
                continue;
            });
            let path = entry.path();
            let meta = match self.obtain_meta(path, follow) {
                Some(m) => m,
                _ => {
                    ret = 1;
                    continue;
                }
            };

            if !self.matched(meta.uid(), meta.gid()) {
                continue;
            }

            ret = match wrap_chown(
                path,
                &meta,
                self.dest_uid,
                self.dest_gid,
                follow,
                self.verbosity.clone(),
            ) {
                Ok(n) => {
                    if n != "" {
                        show_info!("{}", n);
                    }
                    0
                }
                Err(e) => {
                    if self.verbosity != Verbosity::Silent {
                        show_info!("{}", e);
                    }
                    1
                }
            }
        }
        ret
    }

    fn obtain_meta<P: AsRef<Path>>(&self, path: P, follow: bool) -> Option<Metadata> {
        use self::Verbosity::*;
        let path = path.as_ref();
        let meta = if follow {
            unwrap!(path.metadata(), e, {
                match self.verbosity {
                    Silent => (),
                    _ => show_info!("cannot access '{}': {}", path.display(), e),
                }
                return None;
            })
        } else {
            unwrap!(path.symlink_metadata(), e, {
                match self.verbosity {
                    Silent => (),
                    _ => show_info!("cannot dereference '{}': {}", path.display(), e),
                }
                return None;
            })
        };
        Some(meta)
    }

    #[inline]
    fn matched(&self, uid: uid_t, gid: gid_t) -> bool {
        match self.filter {
            IfFrom::All => true,
            IfFrom::User(u) => u == uid,
            IfFrom::Group(g) => g == gid,
            IfFrom::UserGroup(u, g) => u == uid && g == gid,
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
use seccomp::{
    allow_syscall, allow_syscall_if, deserialize_binary, BpfProgram, BpfThreadMap,
    DeserializationError, Error, SeccompAction, SeccompCmpArgLen as ArgLen, SeccompCmpOp::Eq,
    SeccompCondition as Cond, SeccompFilter, SeccompRule,
};
use std::convert::TryInto;
use std::fs::File;
use utils::signal::sigrtmin;

/// The default filter containing the allowed syscall rules required by `Firecracker` to
/// function.
/// Any non-trivial modification to this allow list needs a proper comment to specify its source
/// or why the sycall/condition is needed.
pub fn get_default_filters() -> Result<BpfThreadMap, Error> {
    let mut filters = BpfThreadMap::new();

    let filter: BpfProgram = SeccompFilter::new(
        vec![
            // Called by the api thread to receive data on socket
            allow_syscall_if(
                libc::SYS_accept4,
                or![and![Cond::new(
                    3,
                    ArgLen::DWORD,
                    Eq,
                    libc::SOCK_CLOEXEC as u64
                )?],],
            ),
            // Called for expanding the heap
            allow_syscall(libc::SYS_brk),
            // Used for metrics and logging, via the helpers in utils/src/time.rs
            // It's not called on some platforms, because of vdso optimisations. In those cases,
            // musl falls back to the regular syscall.
            allow_syscall(libc::SYS_clock_gettime),
            allow_syscall(libc::SYS_close),
            // Needed for vsock
            allow_syscall(libc::SYS_connect),
            allow_syscall(libc::SYS_epoll_ctl),
            allow_syscall(libc::SYS_epoll_pwait),
            #[cfg(all(target_env = "gnu", target_arch = "x86_64"))]
            allow_syscall(libc::SYS_epoll_wait),
            allow_syscall(libc::SYS_exit),
            allow_syscall(libc::SYS_exit_group),
            // Used by snapshotting, drive patching and rescanning
            allow_syscall_if(
                libc::SYS_fcntl,
                or![and![
                    Cond::new(1, ArgLen::DWORD, Eq, super::FCNTL_F_SETFD)?,
                    Cond::new(2, ArgLen::DWORD, Eq, super::FCNTL_FD_CLOEXEC)?,
                ],],
            ),
            // Used for drive patching & rescanning, for reading the local timezone
            allow_syscall(libc::SYS_fstat),
            // Used for snapshotting
            allow_syscall(libc::SYS_ftruncate),
            // Used for synchronization
            allow_syscall_if(
                libc::SYS_futex,
                or![
                    and![Cond::new(1, ArgLen::DWORD, Eq, super::FUTEX_WAIT_PRIVATE)?],
                    and![Cond::new(1, ArgLen::DWORD, Eq, super::FUTEX_WAKE_PRIVATE)?],
                    #[cfg(target_env = "gnu")]
                    and![Cond::new(
                        1,
                        ArgLen::DWORD,
                        Eq,
                        super::FUTEX_CMP_REQUEUE_PRIVATE
                    )?],
                ],
            ),
            // Used by glibc's tgkill
            #[cfg(target_env = "gnu")]
            allow_syscall(libc::SYS_getpid),
            allow_syscall_if(libc::SYS_ioctl, super::create_ioctl_seccomp_rule()?),
            // Used by the block device
            allow_syscall(libc::SYS_lseek),
            // Triggered by musl for some customer workloads
            #[cfg(target_env = "musl")]
            allow_syscall_if(
                libc::SYS_madvise,
                or![and![Cond::new(
                    2,
                    ArgLen::DWORD,
                    Eq,
                    libc::MADV_DONTNEED as u64
                )?],],
            ),
            // Used for re-allocating large memory regions, for example vectors
            allow_syscall(libc::SYS_mremap),
            // Used for freeing memory
            allow_syscall(libc::SYS_munmap),
            allow_syscall_if(
                libc::SYS_mmap,
                or![
                    // Used for reading the timezone in LocalTime::now()
                    and![Cond::new(3, ArgLen::DWORD, Eq, libc::MAP_SHARED as u64)?],
                    // Used by the balloon device
                    and![Cond::new(
                        3,
                        ArgLen::DWORD,
                        Eq,
                        (libc::MAP_FIXED | libc::MAP_ANONYMOUS | libc::MAP_PRIVATE) as u64
                    )?],
                ],
            ),
            #[cfg(target_arch = "x86_64")]
            allow_syscall(libc::SYS_open),
            #[cfg(target_arch = "aarch64")]
            allow_syscall(libc::SYS_openat),
            allow_syscall(libc::SYS_read),
            // Used by the API thread and vsock
            allow_syscall(libc::SYS_recvfrom),
            // SYS_rt_sigreturn is needed in case a fault does occur, so that the signal handler
            // can return. Otherwise we get stuck in a fault loop.
            allow_syscall(libc::SYS_rt_sigreturn),
            // Used by the API thread and vsock
            allow_syscall_if(
                libc::SYS_socket,
                or![and![
                    Cond::new(0, ArgLen::DWORD, Eq, libc::AF_UNIX as u64)?,
                    Cond::new(
                        1,
                        ArgLen::DWORD,
                        Eq,
                        (libc::SOCK_STREAM as u64) | (libc::SOCK_CLOEXEC as u64)
                    )?,
                    Cond::new(2, ArgLen::DWORD, Eq, 0u64)?
                ],],
            ),
            // Used to kick vcpus
            allow_syscall_if(
                libc::SYS_tkill,
                or![and![Cond::new(
                    1,
                    ArgLen::DWORD,
                    Eq,
                    (sigrtmin() + super::super::vstate::vcpu::VCPU_RTSIG_OFFSET) as u64
                )?]],
            ),
            // Used to kick vcpus, on gnu
            #[cfg(target_env = "gnu")]
            allow_syscall(libc::SYS_tgkill),
            // Needed for rate limiting
            allow_syscall_if(
                libc::SYS_timerfd_create,
                or![and![
                    Cond::new(0, ArgLen::DWORD, Eq, libc::CLOCK_MONOTONIC as u64)?,
                    Cond::new(
                        1,
                        ArgLen::DWORD,
                        Eq,
                        (libc::TFD_CLOEXEC as u64) | (libc::TFD_NONBLOCK as u64)
                    )?,
                ],],
            ),
            // Needed for rate limiting
            allow_syscall_if(
                libc::SYS_timerfd_settime,
                or![and![Cond::new(1, ArgLen::DWORD, Eq, 0u64)?],],
            ),
            allow_syscall(libc::SYS_fsync),
            allow_syscall(libc::SYS_write),
        ]
        .into_iter()
        .collect(),
        SeccompAction::Trap,
        std::env::consts::ARCH,
    )?
    .try_into()?;

    filters.insert("api".to_string(), filter.clone());
    filters.insert("vmm".to_string(), filter.clone());
    filters.insert("vcpu".to_string(), filter);

    Ok(filters)
}

/// Retrieve empty seccomp filters
pub fn get_empty_filters() -> BpfThreadMap {
    let mut map = BpfThreadMap::new();
    map.insert("vmm".to_string(), vec![]);
    map.insert("api".to_string(), vec![]);
    map.insert("vcpu".to_string(), vec![]);
    map
}

/// Retrieve custom seccomp filters
pub fn get_custom_filters(mut file: File) -> Result<BpfThreadMap, DeserializationError> {
    deserialize_binary(&mut file)
}

#[cfg(test)]
mod tests {
    use super::{get_custom_filters, get_default_filters};
    use utils::tempfile::TempFile;

    #[test]
    fn get_filters() {
        assert!(get_default_filters().is_ok());

        let file = TempFile::new().unwrap();
        assert!(get_custom_filters(file.into_file()).is_ok());
    }
}
<commit_msg>Add get_custom_filters support<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
use seccomp::{
    allow_syscall, allow_syscall_if, deserialize_binary, BpfProgram, BpfThreadMap,
    DeserializationError, Error, SeccompAction, SeccompCmpArgLen as ArgLen, SeccompCmpOp::Eq,
    SeccompCondition as Cond, SeccompFilter, SeccompRule,
};
use std::convert::TryInto;
use std::fmt;
use std::fs::File;
use utils::signal::sigrtmin;

const THREAD_CATEGORIES: [&str; 3] = ["vmm", "api", "vcpu"];

/// Error retrieving seccomp filters.
#[derive(fmt::Debug)]
pub enum FilterError {
    /// Filter deserialitaion error.
    Deserialization(DeserializationError),
    /// Invalid thread categories.
    ThreadCategories(String),
    /// Missing Thread Category.
    MissingThreadCategory(String),
    /// Seccomp error occurred.
    Seccomp(Error),
}

impl fmt::Display for FilterError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::FilterError::*;

        match *self {
            Deserialization(ref err) => write!(f, "Filter (de)serialization failed: {}", err),
            ThreadCategories(ref categories) => {
                write!(f, "Invalid thread categories: {}", categories)
            }
            MissingThreadCategory(ref category) => {
                write!(f, "Missing thread category: {}", category)
            }
            Seccomp(ref err) => write!(f, "Seccomp error: {}", err),
        }
    }
}

/// Return an error if the BpfThreadMap contains invalid thread categories.
fn filter_thread_categories(map: BpfThreadMap) -> Result<BpfThreadMap, FilterError> {
    let (filters, invalid_filters): (BpfThreadMap, BpfThreadMap) = map
        .into_iter()
        .partition(|(k, _)| THREAD_CATEGORIES.contains(&k.as_str()));
    if !invalid_filters.is_empty() {
        // build the error message
        let mut thread_categories_string =
            invalid_filters
                .keys()
                .fold("".to_string(), |mut acc, elem| {
                    acc.push_str(elem);
                    acc.push_str(",");
                    acc
                });
        thread_categories_string.pop();
        return Err(FilterError::ThreadCategories(thread_categories_string));
    }

    for &category in THREAD_CATEGORIES.iter() {
        let category_string = category.to_string();
        if !filters.contains_key(&category_string) {
            return Err(FilterError::MissingThreadCategory(category_string));
        }
    }

    Ok(filters)
}

/// The default filter containing the allowed syscall rules required by `Firecracker` to
/// function.
/// Any non-trivial modification to this allow list needs a proper comment to specify its source
/// or why the sycall/condition is needed.
pub fn get_default_filters() -> Result<BpfThreadMap, Error> {
    let mut filters = BpfThreadMap::new();

    let filter: BpfProgram = SeccompFilter::new(
        vec![
            // Called by the api thread to receive data on socket
            allow_syscall_if(
                libc::SYS_accept4,
                or![and![Cond::new(
                    3,
                    ArgLen::DWORD,
                    Eq,
                    libc::SOCK_CLOEXEC as u64
                )?],],
            ),
            // Called for expanding the heap
            allow_syscall(libc::SYS_brk),
            // Used for metrics and logging, via the helpers in utils/src/time.rs
            // It's not called on some platforms, because of vdso optimisations. In those cases,
            // musl falls back to the regular syscall.
            allow_syscall(libc::SYS_clock_gettime),
            allow_syscall(libc::SYS_close),
            // Needed for vsock
            allow_syscall(libc::SYS_connect),
            allow_syscall(libc::SYS_epoll_ctl),
            allow_syscall(libc::SYS_epoll_pwait),
            #[cfg(all(target_env = "gnu", target_arch = "x86_64"))]
            allow_syscall(libc::SYS_epoll_wait),
            allow_syscall(libc::SYS_exit),
            allow_syscall(libc::SYS_exit_group),
            // Used by snapshotting, drive patching and rescanning
            allow_syscall_if(
                libc::SYS_fcntl,
                or![and![
                    Cond::new(1, ArgLen::DWORD, Eq, super::FCNTL_F_SETFD)?,
                    Cond::new(2, ArgLen::DWORD, Eq, super::FCNTL_FD_CLOEXEC)?,
                ],],
            ),
            // Used for drive patching & rescanning, for reading the local timezone
            allow_syscall(libc::SYS_fstat),
            // Used for snapshotting
            allow_syscall(libc::SYS_ftruncate),
            // Used for synchronization
            allow_syscall_if(
                libc::SYS_futex,
                or![
                    and![Cond::new(1, ArgLen::DWORD, Eq, super::FUTEX_WAIT_PRIVATE)?],
                    and![Cond::new(1, ArgLen::DWORD, Eq, super::FUTEX_WAKE_PRIVATE)?],
                    #[cfg(target_env = "gnu")]
                    and![Cond::new(
                        1,
                        ArgLen::DWORD,
                        Eq,
                        super::FUTEX_CMP_REQUEUE_PRIVATE
                    )?],
                ],
            ),
            // Used by glibc's tgkill
            #[cfg(target_env = "gnu")]
            allow_syscall(libc::SYS_getpid),
            allow_syscall_if(libc::SYS_ioctl, super::create_ioctl_seccomp_rule()?),
            // Used by the block device
            allow_syscall(libc::SYS_lseek),
            // Triggered by musl for some customer workloads
            #[cfg(target_env = "musl")]
            allow_syscall_if(
                libc::SYS_madvise,
                or![and![Cond::new(
                    2,
                    ArgLen::DWORD,
                    Eq,
                    libc::MADV_DONTNEED as u64
                )?],],
            ),
            // Used for re-allocating large memory regions, for example vectors
            allow_syscall(libc::SYS_mremap),
            // Used for freeing memory
            allow_syscall(libc::SYS_munmap),
            allow_syscall_if(
                libc::SYS_mmap,
                or![
                    // Used for reading the timezone in LocalTime::now()
                    and![Cond::new(3, ArgLen::DWORD, Eq, libc::MAP_SHARED as u64)?],
                    // Used by the balloon device
                    and![Cond::new(
                        3,
                        ArgLen::DWORD,
                        Eq,
                        (libc::MAP_FIXED | libc::MAP_ANONYMOUS | libc::MAP_PRIVATE) as u64
                    )?],
                ],
            ),
            #[cfg(target_arch = "x86_64")]
            allow_syscall(libc::SYS_open),
            #[cfg(target_arch = "aarch64")]
            allow_syscall(libc::SYS_openat),
            allow_syscall(libc::SYS_read),
            // Used by the API thread and vsock
            allow_syscall(libc::SYS_recvfrom),
            // SYS_rt_sigreturn is needed in case a fault does occur, so that the signal handler
            // can return. Otherwise we get stuck in a fault loop.
            allow_syscall(libc::SYS_rt_sigreturn),
            // Used by the API thread and vsock
            allow_syscall_if(
                libc::SYS_socket,
                or![and![
                    Cond::new(0, ArgLen::DWORD, Eq, libc::AF_UNIX as u64)?,
                    Cond::new(
                        1,
                        ArgLen::DWORD,
                        Eq,
                        (libc::SOCK_STREAM as u64) | (libc::SOCK_CLOEXEC as u64)
                    )?,
                    Cond::new(2, ArgLen::DWORD, Eq, 0u64)?
                ],],
            ),
            // Used to kick vcpus
            allow_syscall_if(
                libc::SYS_tkill,
                or![and![Cond::new(
                    1,
                    ArgLen::DWORD,
                    Eq,
                    (sigrtmin() + super::super::vstate::vcpu::VCPU_RTSIG_OFFSET) as u64
                )?]],
            ),
            // Used to kick vcpus, on gnu
            #[cfg(target_env = "gnu")]
            allow_syscall(libc::SYS_tgkill),
            // Needed for rate limiting
            allow_syscall_if(
                libc::SYS_timerfd_create,
                or![and![
                    Cond::new(0, ArgLen::DWORD, Eq, libc::CLOCK_MONOTONIC as u64)?,
                    Cond::new(
                        1,
                        ArgLen::DWORD,
                        Eq,
                        (libc::TFD_CLOEXEC as u64) | (libc::TFD_NONBLOCK as u64)
                    )?,
                ],],
            ),
            // Needed for rate limiting
            allow_syscall_if(
                libc::SYS_timerfd_settime,
                or![and![Cond::new(1, ArgLen::DWORD, Eq, 0u64)?],],
            ),
            allow_syscall(libc::SYS_fsync),
            allow_syscall(libc::SYS_write),
        ]
        .into_iter()
        .collect(),
        SeccompAction::Trap,
        std::env::consts::ARCH,
    )?
    .try_into()?;

    filters.insert("api".to_string(), filter.clone());
    filters.insert("vmm".to_string(), filter.clone());
    filters.insert("vcpu".to_string(), filter);

    Ok(filters)
}

/// Retrieve empty seccomp filters.
pub fn get_empty_filters() -> BpfThreadMap {
    let mut map = BpfThreadMap::new();
    map.insert("vmm".to_string(), vec![]);
    map.insert("api".to_string(), vec![]);
    map.insert("vcpu".to_string(), vec![]);
    map
}

/// Retrieve custom seccomp filters.
pub fn get_custom_filters(mut file: File) -> Result<BpfThreadMap, FilterError> {
    let map = deserialize_binary(&mut file).map_err(FilterError::Deserialization)?;
    filter_thread_categories(map)
}

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

    #[test]
    fn get_filters() {
        assert!(get_default_filters().is_ok());
    }

    #[test]
    fn test_filter_thread_categories() {
        // correct categories
        let mut map = BpfThreadMap::new();
        map.insert("vcpu".to_string(), vec![]);
        map.insert("vmm".to_string(), vec![]);
        map.insert("api".to_string(), vec![]);

        assert_eq!(filter_thread_categories(map).unwrap().len(), 3);

        // invalid categories
        let mut map = BpfThreadMap::new();
        map.insert("vcpu".to_string(), vec![]);
        map.insert("vmm".to_string(), vec![]);
        map.insert("thread1".to_string(), vec![]);
        map.insert("thread2".to_string(), vec![]);

        match filter_thread_categories(map).unwrap_err() {
            FilterError::ThreadCategories(err) => {
                assert!(err == "thread2,thread1" || err == "thread1,thread2")
            }
            _ => panic!("Expected ThreadCategories error."),
        }

        // missing category
        let mut map = BpfThreadMap::new();
        map.insert("vcpu".to_string(), vec![]);
        map.insert("vmm".to_string(), vec![]);

        match filter_thread_categories(map).unwrap_err() {
            FilterError::MissingThreadCategory(name) => assert_eq!(name, "api"),
            _ => panic!("Expected MissingThreadCategory error."),
        }
    }
}
<|endoftext|>


<commit_before>// Farbfeld filter that applies a cool blurring effect over images
// Usage: png2ff < greyscale.png | blur | ff2png > false-colored.png

use std::io;
use std::io::prelude::*;
use std::io::BufReader;
use std::io::BufWriter;

extern crate byteorder;
use byteorder::{ReadBytesExt, WriteBytesExt, BigEndian};

extern crate rayon;
use rayon::prelude::*;

#[derive (Clone, Copy)]
struct FfPixel {
	pub r: u16,
	pub g: u16,
	pub b: u16,
	pub a: u16,
}

#[derive (Clone, Copy, Debug, PartialEq)]
struct PlaneCoord {
	pub x: i32,
	pub y: i32,
}

impl PlaneCoord {
	pub fn contains_x (&self, x: i32) -> bool {
		x >= 0 && x < self.x
	}

	pub fn contains_y (&self, y: i32) -> bool {
		y >= 0 && y < self.y
	}

	pub fn contains (&self, x: i32, y: i32) -> bool {
		self.contains_x (x) &&
		self.contains_y (y)
	}

	pub fn clamp_x (&self, x: i32) -> i32 {
		return if x < 0 {
			0
		}
		else if x >= self.x {
			self.x - 1
		}
		else {
			x
		}
	}

	pub fn clamp_y (&self, y: i32) -> i32 {
		return if y < 0 {
			0
		}
		else if y >= self.y {
			self.y - 1
		}
		else {
			y
		}
	}

	pub fn clamp (&self, pc: &PlaneCoord) -> PlaneCoord {
		PlaneCoord {
			x: self.clamp_x (pc.x),
			y: self.clamp_y (pc.y),
		}
	}

	pub fn index (&self, pc: &PlaneCoord) -> i32 {
		let pc = self.clamp (pc);

		pc.y * self.x + pc.x
	}

	pub fn clamped_index (&self, pc: &PlaneCoord) -> i32 {
		self.index (&self.clamp (pc))
	}
}

struct HdrPlane {
	pub size: PlaneCoord,
	pub pixels: Vec <f64>,
}

impl HdrPlane {
	pub fn new (size: PlaneCoord) -> HdrPlane {
		HdrPlane {
			size: size,
			pixels: vec! [0.0; (size.x * size.y) as usize],
		}
	}
}

fn read_pixel <T> (reader: &mut T) -> FfPixel where T: ReadBytesExt {
	FfPixel {
		r: reader.read_u16::<BigEndian> ().unwrap (),
		g: reader.read_u16::<BigEndian> ().unwrap (),
		b: reader.read_u16::<BigEndian> ().unwrap (),
		a: reader.read_u16::<BigEndian> ().unwrap (),
	}
}

fn write_pixel <T> (writer: &mut T, pix: &FfPixel) where T: WriteBytesExt {
	writer.write_u16::<BigEndian> (pix.r).unwrap ();
	writer.write_u16::<BigEndian> (pix.g).unwrap ();
	writer.write_u16::<BigEndian> (pix.b).unwrap ();
	writer.write_u16::<BigEndian> (pix.a).unwrap ();
}

fn pascal (n: u32) -> Vec <u64> {
	match n {
		0 => {
			vec! []
		}
		1 => {
			vec! [1]
		},
		2 => {
			vec! [1, 1]
		},
		3 => {
			vec! [1, 2, 1]
		},
		4 => {
			vec! [1, 3, 3, 1]
		}
		5 => {
			vec! [1, 4, 6, 4, 1]
		},
		_ => {
			let step = 5;
			let prev_row = pascal (n - step + 1);
			let step_row = pascal (step);

			let mut output = vec! [];
			for _ in 0..n {
				output.push (0);
			}

			for a in 0..prev_row.len () {
				for b in 0..step_row.len () {
					output [a + b] += prev_row [a] * step_row [b];
				}
			}

			output
		}
	}
}

struct ScanlineJob <'a> {
	pub y: i32,
	pub row_chunk: &'a mut [f64],
}

fn blur_hor (input_plane: &HdrPlane, filter: &Vec <u64>) -> HdrPlane {
	let offset = -(filter.len () as i32 - 1) / 2;
	let filter_sum: u64 = filter.iter ().sum ();
	let filter_f: Vec <f64> = filter.iter ().map (|x| *x as f64 / filter_sum as f64).collect ();

	let sz = input_plane.size;

	let mut pixels = vec! [0.0f64; (sz.x * sz.y) as usize];

	(0..sz.y).zip (pixels.chunks_mut (sz.x as usize)).map (|(y, row_chunk)| ScanlineJob { y: y, row_chunk: row_chunk }).collect::<Vec <ScanlineJob>> ().par_iter_mut ().for_each (|job| {
		let scanline_index = job.y * sz.x;

		for x in 0..sz.x {
			let offset_2 = x + offset;

			job.row_chunk [x as usize] = (sz.clamp_x (x + 0 + offset) - offset_2..sz.clamp_x (x + filter_f.len () as i32 + offset) - offset_2).map (|i| filter_f [i as usize] * input_plane.pixels [(scanline_index + i + offset_2) as usize]).sum ();
		}
	});

	HdrPlane {
		size: sz,
		pixels: pixels,
	}
}

fn blur_vert (input_plane: &HdrPlane, filter: &Vec <u64>) -> HdrPlane {
	let offset = -(filter.len () as i32 - 1) / 2;
	let filter_sum: u64 = filter.iter ().sum ();
	// Rust plz
	let filter_sum = filter_sum as f64;
	let filter_f: Vec <f64> = filter.iter ().map (|x| *x as f64 / filter_sum).collect ();

	let sz = input_plane.size;

	let mut pixels = vec![];
	for y in 0..sz.y {
		for x in 0..sz.x {
			let dest_index = sz.clamped_index (&PlaneCoord {
				x: x,
				y: y
			});

			pixels.push (0.0f64);

			for src_y in sz.clamp_y (y + 0 + offset)..sz.clamp_y (y + filter_f.len () as i32 - 1 + offset)
			{
				let src_index = sz.clamped_index (&PlaneCoord {
					x: x,
					y: src_y,
				});

				pixels [dest_index as usize] += filter_f [(src_y - y - offset) as usize] * input_plane.pixels [src_index as usize];
			}
		}
	}

	HdrPlane {
		size: sz,
		pixels: pixels,
	}
}

fn read_farbfeld_grey <T> (reader: &mut T) -> HdrPlane where T: ReadBytesExt {
	let mut magic = [0; 8];
	reader.read (&mut magic [..]).unwrap ();

	// TODO: Assert that magic == 'farbfeld'
	//assert_eq! (magic, &"farbfeld");

	let width = reader.read_u32::<BigEndian>().unwrap() as i32;
	let height = reader.read_u32::<BigEndian>().unwrap() as i32;

	writeln!(&mut io::stderr(), "{} x {}", width, height).expect("failed printing to stderr");

	let mut pixels = vec![0.0; (width * height) as usize];

	for i in 0 .. width * height {
		let gamma = read_pixel (reader).g as f64 / 65536.0;
		pixels [i as usize] = gamma * gamma;
	}

	HdrPlane {
		size: PlaneCoord {
			x: width,
			y: height,
		},
		pixels: pixels,
	}
}

fn write_farbfeld_plane <T> (writer: &mut T, plane: &HdrPlane) where T: WriteBytesExt
{
	write_farbfeld (writer, &plane.size, plane.pixels.iter ().map (|x| {
		let val = (x.sqrt ().min (1.0) * 65535.0) as u16;

		FfPixel {
			r: val,
			g: val,
			b: val,
			a: 65535,
		}
	}));
}

fn double_to_ff (val: f64) -> u16 {
	(val.sqrt ().min (1.0) * 65535.0) as u16
}

fn write_farbfeld <T, I> (writer: &mut T, size: &PlaneCoord, pixels: I ) where T: WriteBytesExt, I: Iterator <Item = FfPixel>
{
	writer.write ("farbfeld".as_bytes ()).expect ("printing farbfeld");

	writer.write_u32::<BigEndian> (size.x as u32).unwrap ();
	writer.write_u32::<BigEndian> (size.y as u32).unwrap ();

	for pixel in pixels.take ((size.x * size.y) as usize) {
		write_pixel (writer, &pixel);
	}
}

fn blur (input_plane: &HdrPlane, filter: &Vec <u64>) -> HdrPlane {
	//blur_vert (&blur_hor (input_plane, filter), filter)
	blur_hor (input_plane, filter)
}

fn blur_n (input_plane: &HdrPlane, filter: &Vec <u64>, n: u32) -> HdrPlane {
	match n {
		0 => {
			// I know this is dumb
			panic! ("Can't blur an image 0 times")
		},
		1 => {
			blur (input_plane, filter)
		},
		_ => {
			blur (&blur_n (input_plane, filter, n - 1), filter)
		},
	}
}

fn main () {
	// Reader
	let mut reader = BufReader::new (io::stdin ());

	let filter = pascal (41);

	let input_plane = read_farbfeld_grey (&mut reader);
	let blurry_1 = blur_n (&input_plane, &filter, 1);
	let blurry_10 = blur_n (&blurry_1, &filter, 9);
	let blurry_20 = blur_n (&blurry_10, &filter, 10);

	let sum_iter = (0..(input_plane.size.x * input_plane.size.y) as usize).map (|i| {
		let r = input_plane.pixels [i];
		let g = r;
		let b = r;

		let r = r + blurry_1.pixels [i] * 512.0 / 256.0;
		let g = g + blurry_1.pixels [i] * 256.0 / 256.0;

		let r = r + blurry_20.pixels [i] * 59.0 / 256.0;
		let b = b + blurry_20.pixels [i] * 60.0 / 256.0;

		FfPixel {
			r: double_to_ff (r),
			g: double_to_ff (g),
			b: double_to_ff (b),
			a: 65535,
		}
	});

	// Write
	{
		let mut writer = BufWriter::new (io::stdout ());

		write_farbfeld (&mut writer, &input_plane.size, sum_iter);
	}

	writeln!(&mut io::stderr(), "filter {:?}", filter).expect("failed printing to stderr");
}
<commit_msg>Dicking with the colors<commit_after>// Farbfeld filter that applies a cool blurring effect over images
// Usage: png2ff < greyscale.png | blur | ff2png > false-colored.png

use std::io;
use std::io::prelude::*;
use std::io::BufReader;
use std::io::BufWriter;

extern crate byteorder;
use byteorder::{ReadBytesExt, WriteBytesExt, BigEndian};

extern crate rayon;
use rayon::prelude::*;

#[derive (Clone, Copy)]
struct FfPixel {
	pub r: u16,
	pub g: u16,
	pub b: u16,
	pub a: u16,
}

#[derive (Clone, Copy, Debug, PartialEq)]
struct PlaneCoord {
	pub x: i32,
	pub y: i32,
}

impl PlaneCoord {
	pub fn contains_x (&self, x: i32) -> bool {
		x >= 0 && x < self.x
	}

	pub fn contains_y (&self, y: i32) -> bool {
		y >= 0 && y < self.y
	}

	pub fn contains (&self, x: i32, y: i32) -> bool {
		self.contains_x (x) &&
		self.contains_y (y)
	}

	pub fn clamp_x (&self, x: i32) -> i32 {
		return if x < 0 {
			0
		}
		else if x >= self.x {
			self.x - 1
		}
		else {
			x
		}
	}

	pub fn clamp_y (&self, y: i32) -> i32 {
		return if y < 0 {
			0
		}
		else if y >= self.y {
			self.y - 1
		}
		else {
			y
		}
	}

	pub fn clamp (&self, pc: &PlaneCoord) -> PlaneCoord {
		PlaneCoord {
			x: self.clamp_x (pc.x),
			y: self.clamp_y (pc.y),
		}
	}

	pub fn index (&self, pc: &PlaneCoord) -> i32 {
		let pc = self.clamp (pc);

		pc.y * self.x + pc.x
	}

	pub fn clamped_index (&self, pc: &PlaneCoord) -> i32 {
		self.index (&self.clamp (pc))
	}
}

struct HdrPlane {
	pub size: PlaneCoord,
	pub pixels: Vec <f64>,
}

impl HdrPlane {
	pub fn new (size: PlaneCoord) -> HdrPlane {
		HdrPlane {
			size: size,
			pixels: vec! [0.0; (size.x * size.y) as usize],
		}
	}
}

fn read_pixel <T> (reader: &mut T) -> FfPixel where T: ReadBytesExt {
	FfPixel {
		r: reader.read_u16::<BigEndian> ().unwrap (),
		g: reader.read_u16::<BigEndian> ().unwrap (),
		b: reader.read_u16::<BigEndian> ().unwrap (),
		a: reader.read_u16::<BigEndian> ().unwrap (),
	}
}

fn write_pixel <T> (writer: &mut T, pix: &FfPixel) where T: WriteBytesExt {
	writer.write_u16::<BigEndian> (pix.r).unwrap ();
	writer.write_u16::<BigEndian> (pix.g).unwrap ();
	writer.write_u16::<BigEndian> (pix.b).unwrap ();
	writer.write_u16::<BigEndian> (pix.a).unwrap ();
}

fn pascal (n: u32) -> Vec <u64> {
	match n {
		0 => {
			vec! []
		}
		1 => {
			vec! [1]
		},
		2 => {
			vec! [1, 1]
		},
		3 => {
			vec! [1, 2, 1]
		},
		4 => {
			vec! [1, 3, 3, 1]
		}
		5 => {
			vec! [1, 4, 6, 4, 1]
		},
		_ => {
			let step = 5;
			let prev_row = pascal (n - step + 1);
			let step_row = pascal (step);

			let mut output = vec! [];
			for _ in 0..n {
				output.push (0);
			}

			for a in 0..prev_row.len () {
				for b in 0..step_row.len () {
					output [a + b] += prev_row [a] * step_row [b];
				}
			}

			output
		}
	}
}

struct ScanlineJob <'a> {
	pub y: i32,
	pub row_chunk: &'a mut [f64],
}

fn blur_hor (input_plane: &HdrPlane, filter: &Vec <u64>) -> HdrPlane {
	let offset = -(filter.len () as i32 - 1) / 2;
	let filter_sum: u64 = filter.iter ().sum ();
	let filter_f: Vec <f64> = filter.iter ().map (|x| *x as f64 / filter_sum as f64).collect ();

	let sz = input_plane.size;

	let mut pixels = vec! [0.0f64; (sz.x * sz.y) as usize];

	(0..sz.y).zip (pixels.chunks_mut (sz.x as usize)).map (|(y, row_chunk)| ScanlineJob { y: y, row_chunk: row_chunk }).collect::<Vec <ScanlineJob>> ().par_iter_mut ().for_each (|job| {
		let scanline_index = job.y * sz.x;

		for x in 0..sz.x {
			let offset_2 = x + offset;

			job.row_chunk [x as usize] = (sz.clamp_x (x + 0 + offset) - offset_2..sz.clamp_x (x + filter_f.len () as i32 + offset) - offset_2).map (|i| filter_f [i as usize] * input_plane.pixels [(scanline_index + i + offset_2) as usize]).sum ();
		}
	});

	HdrPlane {
		size: sz,
		pixels: pixels,
	}
}

fn blur_vert (input_plane: &HdrPlane, filter: &Vec <u64>) -> HdrPlane {
	let offset = -(filter.len () as i32 - 1) / 2;
	let filter_sum: u64 = filter.iter ().sum ();
	let filter_f: Vec <f64> = filter.iter ().map (|x| *x as f64 / filter_sum as f64).collect ();

	let sz = input_plane.size;

	let mut pixels = vec! [0.0f64; (sz.x * sz.y) as usize];

	(0..sz.y).zip (pixels.chunks_mut (sz.x as usize)).map (|(y, row_chunk)| ScanlineJob { y: y, row_chunk: row_chunk }).collect::<Vec <ScanlineJob>> ().par_iter_mut ().for_each (|job| {
		for x in 0..sz.x {
			let offset_2 = job.y + offset;

			job.row_chunk [x as usize] = (sz.clamp_y (job.y + 0 + offset) - offset_2..sz.clamp_y (job.y + filter_f.len () as i32 + offset) - offset_2).map (|i| filter_f [i as usize] * input_plane.pixels [( sz.index (&PlaneCoord {x: x, y: i + offset_2})) as usize]).sum ();
		}
	});

	HdrPlane {
		size: sz,
		pixels: pixels,
	}
}

fn read_farbfeld_grey <T> (reader: &mut T) -> HdrPlane where T: ReadBytesExt {
	let mut magic = [0; 8];
	reader.read (&mut magic [..]).unwrap ();

	// TODO: Assert that magic == 'farbfeld'
	//assert_eq! (magic, &"farbfeld");

	let width = reader.read_u32::<BigEndian>().unwrap() as i32;
	let height = reader.read_u32::<BigEndian>().unwrap() as i32;

	writeln!(&mut io::stderr(), "{} x {}", width, height).expect("failed printing to stderr");

	let mut pixels = vec![0.0; (width * height) as usize];

	for i in 0 .. width * height {
		let gamma = read_pixel (reader).g as f64 / 65536.0;
		pixels [i as usize] = gamma * gamma;
	}

	HdrPlane {
		size: PlaneCoord {
			x: width,
			y: height,
		},
		pixels: pixels,
	}
}

fn write_farbfeld_plane <T> (writer: &mut T, plane: &HdrPlane) where T: WriteBytesExt
{
	write_farbfeld (writer, &plane.size, plane.pixels.iter ().map (|x| {
		let val = (x.sqrt ().min (1.0) * 65535.0) as u16;

		FfPixel {
			r: val,
			g: val,
			b: val,
			a: 65535,
		}
	}));
}

fn double_to_ff (val: f64) -> u16 {
	(val.sqrt ().min (1.0) * 65535.0) as u16
}

fn write_farbfeld <T, I> (writer: &mut T, size: &PlaneCoord, pixels: I ) where T: WriteBytesExt, I: Iterator <Item = FfPixel>
{
	writer.write ("farbfeld".as_bytes ()).expect ("printing farbfeld");

	writer.write_u32::<BigEndian> (size.x as u32).unwrap ();
	writer.write_u32::<BigEndian> (size.y as u32).unwrap ();

	for pixel in pixels.take ((size.x * size.y) as usize) {
		write_pixel (writer, &pixel);
	}
}

fn blur (input_plane: &HdrPlane, filter: &Vec <u64>) -> HdrPlane {
	blur_vert (&blur_hor (input_plane, filter), filter)
	//blur_hor (input_plane, filter)
	//blur_vert (input_plane, filter)
}

fn blur_n (input_plane: &HdrPlane, filter: &Vec <u64>, n: u32) -> HdrPlane {
	match n {
		0 => {
			// I know this is dumb
			panic! ("Can't blur an image 0 times")
		},
		1 => {
			blur (input_plane, filter)
		},
		_ => {
			blur (&blur_n (input_plane, filter, n - 1), filter)
		},
	}
}

fn main () {
	// Reader
	let mut reader = BufReader::new (io::stdin ());

	let filter = pascal (41);

	let input_plane = read_farbfeld_grey (&mut reader);
	let blurry_1 = blur_n (&input_plane, &filter, 1);
	let blurry_10 = blur_n (&blurry_1, &filter, 9);
	let blurry_20 = blur_n (&blurry_10, &filter, 10);

	let sum_iter = (0..(input_plane.size.x * input_plane.size.y) as usize).map (|i| {
		let r = input_plane.pixels [i];
		let g = r;
		let b = r;

		let b = b + blurry_1.pixels [i] * 256.0 / 256.0;
		let g = g + blurry_1.pixels [i] * 128.0 / 256.0;

		let b = b + blurry_10.pixels [i] * 512.0 / 256.0;
		let g = g + blurry_10.pixels [i] * 128.0 / 256.0;

		let r = r + blurry_20.pixels [i] * 90.0 / 256.0;
		let b = b + blurry_20.pixels [i] * 30.0 / 256.0;

		FfPixel {
			r: double_to_ff (r),
			g: double_to_ff (g),
			b: double_to_ff (b),
			a: 65535,
		}
	});

	// Write
	{
		let mut writer = BufWriter::new (io::stdout ());

		write_farbfeld (&mut writer, &input_plane.size, sum_iter);
	}

	writeln!(&mut io::stderr(), "filter {:?}", filter).expect("failed printing to stderr");
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! This file defines ledger store APIs that are related to the main ledger accumulator, from the
//! root(LedgerInfo) to leaf(TransactionInfo).

use crate::utils::iterators::{EpochEndingLedgerInfoIter, ExpectContinuousVersions};
use crate::{
    errors::AptosDbError,
    schema::{
        epoch_by_version::EpochByVersionSchema, ledger_info::LedgerInfoSchema,
        transaction_accumulator::TransactionAccumulatorSchema,
        transaction_info::TransactionInfoSchema,
    },
};
use accumulator::{HashReader, MerkleAccumulator};
use anyhow::{ensure, format_err, Result};
use aptos_crypto::{
    hash::{CryptoHash, TransactionAccumulatorHasher},
    HashValue,
};
use aptos_types::{
    epoch_state::EpochState,
    ledger_info::LedgerInfoWithSignatures,
    proof::{
        definition::LeafCount, position::Position, AccumulatorConsistencyProof,
        TransactionAccumulatorProof, TransactionAccumulatorRangeProof, TransactionInfoWithProof,
    },
    transaction::{TransactionInfo, Version},
};
use arc_swap::ArcSwap;
use itertools::Itertools;
use schemadb::{ReadOptions, SchemaBatch, DB};
use std::{ops::Deref, sync::Arc};

#[derive(Debug)]
pub struct LedgerStore {
    db: Arc<DB>,

    /// We almost always need the latest ledger info and signatures to serve read requests, so we
    /// cache it in memory in order to avoid reading DB and deserializing the object frequently. It
    /// should be updated every time new ledger info and signatures are persisted.
    latest_ledger_info: ArcSwap<Option<LedgerInfoWithSignatures>>,
}

impl LedgerStore {
    pub fn new(db: Arc<DB>) -> Self {
        // Upon restart, read the latest ledger info and signatures and cache them in memory.
        let ledger_info = {
            let mut iter = db
                .iter::<LedgerInfoSchema>(ReadOptions::default())
                .expect("Constructing iterator should work.");
            iter.seek_to_last();
            iter.next()
                .transpose()
                .expect("Reading latest ledger info from DB should work.")
                .map(|kv| kv.1)
        };

        Self {
            db,
            latest_ledger_info: ArcSwap::from(Arc::new(ledger_info)),
        }
    }

    pub fn get_epoch(&self, version: Version) -> Result<u64> {
        let mut iter = self
            .db
            .iter::<EpochByVersionSchema>(ReadOptions::default())?;
        // Search for the end of the previous epoch.
        iter.seek_for_prev(&version)?;
        let (epoch_end_version, epoch) = match iter.next().transpose()? {
            Some(x) => x,
            None => {
                // There should be a genesis LedgerInfo at version 0 (genesis only consists of one
                // transaction), so this normally doesn't happen. However this part of
                // implementation doesn't need to rely on this assumption.
                return Ok(0);
            }
        };
        ensure!(
            epoch_end_version <= version,
            "DB corruption: looking for epoch for version {}, got epoch {} ends at version {}",
            version,
            epoch,
            epoch_end_version
        );
        // If the obtained epoch ended before the given version, return epoch+1, otherwise
        // the given version is exactly the last version of the found epoch.
        Ok(if epoch_end_version < version {
            epoch + 1
        } else {
            epoch
        })
    }

    /// Gets ledger info at specified version and ensures it's an epoch ending.
    pub fn get_epoch_ending_ledger_info(
        &self,
        version: Version,
    ) -> Result<LedgerInfoWithSignatures> {
        let epoch = self.get_epoch(version)?;
        let li = self
            .db
            .get::<LedgerInfoSchema>(&epoch)?
            .ok_or_else(|| AptosDbError::NotFound(format!("LedgerInfo for epoch {}.", epoch)))?;
        ensure!(
            li.ledger_info().version() == version,
            "Epoch {} didn't end at version {}",
            epoch,
            version,
        );
        li.ledger_info()
            .next_epoch_state()
            .ok_or_else(|| format_err!("Not an epoch change at version {}", version))?;

        Ok(li)
    }

    pub fn get_latest_ledger_info_option(&self) -> Option<LedgerInfoWithSignatures> {
        let ledger_info_ptr = self.latest_ledger_info.load();
        let ledger_info: &Option<_> = ledger_info_ptr.deref();
        ledger_info.clone()
    }

    pub fn get_latest_ledger_info(&self) -> Result<LedgerInfoWithSignatures> {
        self.get_latest_ledger_info_option()
            .ok_or_else(|| AptosDbError::NotFound(String::from("Genesis LedgerInfo")).into())
    }

    pub fn set_latest_ledger_info(&self, ledger_info_with_sigs: LedgerInfoWithSignatures) {
        self.latest_ledger_info
            .store(Arc::new(Some(ledger_info_with_sigs)));
    }

    pub fn get_latest_ledger_info_in_epoch(&self, epoch: u64) -> Result<LedgerInfoWithSignatures> {
        self.db.get::<LedgerInfoSchema>(&epoch)?.ok_or_else(|| {
            AptosDbError::NotFound(format!("Last LedgerInfo of epoch {}", epoch)).into()
        })
    }

    pub fn get_epoch_state(&self, epoch: u64) -> Result<EpochState> {
        ensure!(epoch > 0, "EpochState only queryable for epoch >= 1.",);

        let ledger_info_with_sigs =
            self.db
                .get::<LedgerInfoSchema>(&(epoch - 1))?
                .ok_or_else(|| {
                    AptosDbError::NotFound(format!("Last LedgerInfo of epoch {}", epoch - 1))
                })?;
        let latest_epoch_state = ledger_info_with_sigs
            .ledger_info()
            .next_epoch_state()
            .ok_or_else(|| format_err!("Last LedgerInfo in epoch must carry next_epoch_state."))?;

        Ok(latest_epoch_state.clone())
    }

    pub fn get_frozen_subtree_hashes(&self, num_transactions: LeafCount) -> Result<Vec<HashValue>> {
        Accumulator::get_frozen_subtree_hashes(self, num_transactions)
    }

    /// Get transaction info given `version`
    pub fn get_transaction_info(&self, version: Version) -> Result<TransactionInfo> {
        self.db
            .get::<TransactionInfoSchema>(&version)?
            .ok_or_else(|| format_err!("No TransactionInfo at version {}", version))
    }

    pub fn get_latest_transaction_info_option(&self) -> Result<Option<(Version, TransactionInfo)>> {
        let mut iter = self
            .db
            .iter::<TransactionInfoSchema>(ReadOptions::default())?;
        iter.seek_to_last();
        iter.next().transpose()
    }

    /// Get latest transaction info together with its version. Note that during node syncing, this
    /// version can be greater than what's in the latest LedgerInfo.
    pub fn get_latest_transaction_info(&self) -> Result<(Version, TransactionInfo)> {
        self.get_latest_transaction_info_option()?
            .ok_or_else(|| AptosDbError::NotFound(String::from("Genesis TransactionInfo.")).into())
    }

    /// Gets an iterator that yields `num_transaction_infos` transaction infos starting from
    /// `start_version`.
    pub(crate) fn get_transaction_info_iter(
        &self,
        start_version: Version,
        num_transaction_infos: usize,
    ) -> Result<impl Iterator<Item = Result<TransactionInfo>> + '_> {
        let mut iter = self
            .db
            .iter::<TransactionInfoSchema>(ReadOptions::default())?;
        iter.seek(&start_version)?;
        iter.expect_continuous_versions(start_version, num_transaction_infos)
    }

    /// Gets an iterator that yields epoch ending ledger infos, starting
    /// from `start_epoch`, and ends at the one before `end_epoch`
    pub fn get_epoch_ending_ledger_info_iter(
        &self,
        start_epoch: u64,
        end_epoch: u64,
    ) -> Result<EpochEndingLedgerInfoIter> {
        let mut iter = self.db.iter::<LedgerInfoSchema>(ReadOptions::default())?;
        iter.seek(&start_epoch)?;
        Ok(EpochEndingLedgerInfoIter::new(iter, start_epoch, end_epoch))
    }

    pub fn ensure_epoch_ending(&self, version: Version) -> Result<()> {
        self.db
            .get::<EpochByVersionSchema>(&version)?
            .ok_or_else(|| format_err!("Version {} is not epoch ending.", version))?;
        Ok(())
    }

    /// Get transaction info at `version` with proof towards root of ledger at `ledger_version`.
    pub fn get_transaction_info_with_proof(
        &self,
        version: Version,
        ledger_version: Version,
    ) -> Result<TransactionInfoWithProof> {
        Ok(TransactionInfoWithProof::new(
            self.get_transaction_proof(version, ledger_version)?,
            self.get_transaction_info(version)?,
        ))
    }

    /// Get proof for transaction at `version` towards root of ledger at `ledger_version`.
    pub fn get_transaction_proof(
        &self,
        version: Version,
        ledger_version: Version,
    ) -> Result<TransactionAccumulatorProof> {
        Accumulator::get_proof(self, ledger_version + 1 /* num_leaves */, version)
    }

    /// Get proof for `num_txns` consecutive transactions starting from `start_version` towards
    /// root of ledger at `ledger_version`.
    pub fn get_transaction_range_proof(
        &self,
        start_version: Option<Version>,
        num_txns: u64,
        ledger_version: Version,
    ) -> Result<TransactionAccumulatorRangeProof> {
        Accumulator::get_range_proof(
            self,
            ledger_version + 1, /* num_leaves */
            start_version,
            num_txns,
        )
    }

    /// Gets proof that shows the ledger at `ledger_version` is consistent with the ledger at
    /// `client_known_version`.
    pub fn get_consistency_proof(
        &self,
        client_known_version: Option<Version>,
        ledger_version: Version,
    ) -> Result<AccumulatorConsistencyProof> {
        let client_known_num_leaves = client_known_version
            .map(|v| v.saturating_add(1))
            .unwrap_or(0);
        let ledger_num_leaves = ledger_version.saturating_add(1);
        Accumulator::get_consistency_proof(self, ledger_num_leaves, client_known_num_leaves)
    }

    /// Write `txn_infos` to `batch`. Assigned `first_version` to the version number of the
    /// first transaction, and so on.
    pub fn put_transaction_infos(
        &self,
        first_version: u64,
        txn_infos: &[TransactionInfo],
        batch: &mut SchemaBatch,
    ) -> Result<HashValue> {
        // write txn_info
        (first_version..first_version + txn_infos.len() as u64)
            .zip_eq(txn_infos.iter())
            .try_for_each(|(version, txn_info)| {
                batch.put::<TransactionInfoSchema>(&version, txn_info)
            })?;

        // write hash of txn_info into the accumulator
        let txn_hashes: Vec<HashValue> = txn_infos.iter().map(TransactionInfo::hash).collect();
        let (root_hash, writes) = Accumulator::append(
            self,
            first_version, /* num_existing_leaves */
            &txn_hashes,
        )?;
        writes
            .iter()
            .try_for_each(|(pos, hash)| batch.put::<TransactionAccumulatorSchema>(pos, hash))?;
        Ok(root_hash)
    }

    /// Write `ledger_info` to `batch`.
    pub fn put_ledger_info(
        &self,
        ledger_info_with_sigs: &LedgerInfoWithSignatures,
        batch: &mut SchemaBatch,
    ) -> Result<()> {
        let ledger_info = ledger_info_with_sigs.ledger_info();

        if ledger_info.ends_epoch() {
            // This is the last version of the current epoch, update the epoch by version index.
            batch.put::<EpochByVersionSchema>(&ledger_info.version(), &ledger_info.epoch())?;
        }
        batch.put::<LedgerInfoSchema>(&ledger_info.epoch(), ledger_info_with_sigs)
    }

    pub fn get_root_hash(&self, version: Version) -> Result<HashValue> {
        Accumulator::get_root_hash(self, version + 1)
    }
}

pub(crate) type Accumulator = MerkleAccumulator<LedgerStore, TransactionAccumulatorHasher>;

impl HashReader for LedgerStore {
    fn get(&self, position: Position) -> Result<HashValue> {
        self.db
            .get::<TransactionAccumulatorSchema>(&position)?
            .ok_or_else(|| format_err!("{} does not exist.", position))
    }
}

#[cfg(test)]
mod ledger_info_test;
#[cfg(test)]
pub(crate) mod ledger_info_test_utils;
#[cfg(test)]
mod transaction_info_test;
<commit_msg>fix typo<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! This file defines ledger store APIs that are related to the main ledger accumulator, from the
//! root(LedgerInfo) to leaf(TransactionInfo).

use crate::utils::iterators::{EpochEndingLedgerInfoIter, ExpectContinuousVersions};
use crate::{
    errors::AptosDbError,
    schema::{
        epoch_by_version::EpochByVersionSchema, ledger_info::LedgerInfoSchema,
        transaction_accumulator::TransactionAccumulatorSchema,
        transaction_info::TransactionInfoSchema,
    },
};
use accumulator::{HashReader, MerkleAccumulator};
use anyhow::{ensure, format_err, Result};
use aptos_crypto::{
    hash::{CryptoHash, TransactionAccumulatorHasher},
    HashValue,
};
use aptos_types::{
    epoch_state::EpochState,
    ledger_info::LedgerInfoWithSignatures,
    proof::{
        definition::LeafCount, position::Position, AccumulatorConsistencyProof,
        TransactionAccumulatorProof, TransactionAccumulatorRangeProof, TransactionInfoWithProof,
    },
    transaction::{TransactionInfo, Version},
};
use arc_swap::ArcSwap;
use itertools::Itertools;
use schemadb::{ReadOptions, SchemaBatch, DB};
use std::{ops::Deref, sync::Arc};

#[derive(Debug)]
pub struct LedgerStore {
    db: Arc<DB>,

    /// We almost always need the latest ledger info and signatures to serve read requests, so we
    /// cache it in memory in order to avoid reading DB and deserializing the object frequently. It
    /// should be updated every time new ledger info and signatures are persisted.
    latest_ledger_info: ArcSwap<Option<LedgerInfoWithSignatures>>,
}

impl LedgerStore {
    pub fn new(db: Arc<DB>) -> Self {
        // Upon restart, read the latest ledger info and signatures and cache them in memory.
        let ledger_info = {
            let mut iter = db
                .iter::<LedgerInfoSchema>(ReadOptions::default())
                .expect("Constructing iterator should work.");
            iter.seek_to_last();
            iter.next()
                .transpose()
                .expect("Reading latest ledger info from DB should work.")
                .map(|kv| kv.1)
        };

        Self {
            db,
            latest_ledger_info: ArcSwap::from(Arc::new(ledger_info)),
        }
    }

    pub fn get_epoch(&self, version: Version) -> Result<u64> {
        let mut iter = self
            .db
            .iter::<EpochByVersionSchema>(ReadOptions::default())?;
        // Search for the end of the previous epoch.
        iter.seek_for_prev(&version)?;
        let (epoch_end_version, epoch) = match iter.next().transpose()? {
            Some(x) => x,
            None => {
                // There should be a genesis LedgerInfo at version 0 (genesis only consists of one
                // transaction), so this normally doesn't happen. However this part of
                // implementation doesn't need to rely on this assumption.
                return Ok(0);
            }
        };
        ensure!(
            epoch_end_version <= version,
            "DB corruption: looking for epoch for version {}, got epoch {} ends at version {}",
            version,
            epoch,
            epoch_end_version
        );
        // If the obtained epoch ended before the given version, return epoch+1, otherwise
        // the given version is exactly the last version of the found epoch.
        Ok(if epoch_end_version < version {
            epoch + 1
        } else {
            epoch
        })
    }

    /// Gets ledger info at specified version and ensures it's an epoch ending.
    pub fn get_epoch_ending_ledger_info(
        &self,
        version: Version,
    ) -> Result<LedgerInfoWithSignatures> {
        let epoch = self.get_epoch(version)?;
        let li = self
            .db
            .get::<LedgerInfoSchema>(&epoch)?
            .ok_or_else(|| AptosDbError::NotFound(format!("LedgerInfo for epoch {}.", epoch)))?;
        ensure!(
            li.ledger_info().version() == version,
            "Epoch {} didn't end at version {}",
            epoch,
            version,
        );
        li.ledger_info()
            .next_epoch_state()
            .ok_or_else(|| format_err!("Not an epoch change at version {}", version))?;

        Ok(li)
    }

    pub fn get_latest_ledger_info_option(&self) -> Option<LedgerInfoWithSignatures> {
        let ledger_info_ptr = self.latest_ledger_info.load();
        let ledger_info: &Option<_> = ledger_info_ptr.deref();
        ledger_info.clone()
    }

    pub fn get_latest_ledger_info(&self) -> Result<LedgerInfoWithSignatures> {
        self.get_latest_ledger_info_option()
            .ok_or_else(|| AptosDbError::NotFound(String::from("Genesis LedgerInfo")).into())
    }

    pub fn set_latest_ledger_info(&self, ledger_info_with_sigs: LedgerInfoWithSignatures) {
        self.latest_ledger_info
            .store(Arc::new(Some(ledger_info_with_sigs)));
    }

    pub fn get_latest_ledger_info_in_epoch(&self, epoch: u64) -> Result<LedgerInfoWithSignatures> {
        self.db.get::<LedgerInfoSchema>(&epoch)?.ok_or_else(|| {
            AptosDbError::NotFound(format!("Last LedgerInfo of epoch {}", epoch)).into()
        })
    }

    pub fn get_epoch_state(&self, epoch: u64) -> Result<EpochState> {
        ensure!(epoch > 0, "EpochState only queryable for epoch >= 1.",);

        let ledger_info_with_sigs =
            self.db
                .get::<LedgerInfoSchema>(&(epoch - 1))?
                .ok_or_else(|| {
                    AptosDbError::NotFound(format!("Last LedgerInfo of epoch {}", epoch - 1))
                })?;
        let latest_epoch_state = ledger_info_with_sigs
            .ledger_info()
            .next_epoch_state()
            .ok_or_else(|| format_err!("Last LedgerInfo in epoch must carry next_epoch_state."))?;

        Ok(latest_epoch_state.clone())
    }

    pub fn get_frozen_subtree_hashes(&self, num_transactions: LeafCount) -> Result<Vec<HashValue>> {
        Accumulator::get_frozen_subtree_hashes(self, num_transactions)
    }

    /// Get transaction info given `version`
    pub fn get_transaction_info(&self, version: Version) -> Result<TransactionInfo> {
        self.db
            .get::<TransactionInfoSchema>(&version)?
            .ok_or_else(|| format_err!("No TransactionInfo at version {}", version))
    }

    pub fn get_latest_transaction_info_option(&self) -> Result<Option<(Version, TransactionInfo)>> {
        let mut iter = self
            .db
            .iter::<TransactionInfoSchema>(ReadOptions::default())?;
        iter.seek_to_last();
        iter.next().transpose()
    }

    /// Get latest transaction info together with its version. Note that during node syncing, this
    /// version can be greater than what's in the latest LedgerInfo.
    pub fn get_latest_transaction_info(&self) -> Result<(Version, TransactionInfo)> {
        self.get_latest_transaction_info_option()?
            .ok_or_else(|| AptosDbError::NotFound(String::from("Genesis TransactionInfo.")).into())
    }

    /// Gets an iterator that yields `num_transaction_infos` transaction infos starting from
    /// `start_version`.
    pub(crate) fn get_transaction_info_iter(
        &self,
        start_version: Version,
        num_transaction_infos: usize,
    ) -> Result<impl Iterator<Item = Result<TransactionInfo>> + '_> {
        let mut iter = self
            .db
            .iter::<TransactionInfoSchema>(ReadOptions::default())?;
        iter.seek(&start_version)?;
        iter.expect_continuous_versions(start_version, num_transaction_infos)
    }

    /// Gets an iterator that yields epoch ending ledger infos, starting
    /// from `start_epoch`, and ends at the one before `end_epoch`
    pub fn get_epoch_ending_ledger_info_iter(
        &self,
        start_epoch: u64,
        end_epoch: u64,
    ) -> Result<EpochEndingLedgerInfoIter> {
        let mut iter = self.db.iter::<LedgerInfoSchema>(ReadOptions::default())?;
        iter.seek(&start_epoch)?;
        Ok(EpochEndingLedgerInfoIter::new(iter, start_epoch, end_epoch))
    }

    pub fn ensure_epoch_ending(&self, version: Version) -> Result<()> {
        self.db
            .get::<EpochByVersionSchema>(&version)?
            .ok_or_else(|| format_err!("Version {} is not epoch ending.", version))?;
        Ok(())
    }

    /// Get transaction info at `version` with proof towards root of ledger at `ledger_version`.
    pub fn get_transaction_info_with_proof(
        &self,
        version: Version,
        ledger_version: Version,
    ) -> Result<TransactionInfoWithProof> {
        Ok(TransactionInfoWithProof::new(
            self.get_transaction_proof(version, ledger_version)?,
            self.get_transaction_info(version)?,
        ))
    }

    /// Get proof for transaction at `version` towards root of ledger at `ledger_version`.
    pub fn get_transaction_proof(
        &self,
        version: Version,
        ledger_version: Version,
    ) -> Result<TransactionAccumulatorProof> {
        Accumulator::get_proof(self, ledger_version + 1 /* num_leaves */, version)
    }

    /// Get proof for `num_txns` consecutive transactions starting from `start_version` towards
    /// root of ledger at `ledger_version`.
    pub fn get_transaction_range_proof(
        &self,
        start_version: Option<Version>,
        num_txns: u64,
        ledger_version: Version,
    ) -> Result<TransactionAccumulatorRangeProof> {
        Accumulator::get_range_proof(
            self,
            ledger_version + 1, /* num_leaves */
            start_version,
            num_txns,
        )
    }

    /// Gets proof that shows the ledger at `ledger_version` is consistent with the ledger at
    /// `client_known_version`.
    pub fn get_consistency_proof(
        &self,
        client_known_version: Option<Version>,
        ledger_version: Version,
    ) -> Result<AccumulatorConsistencyProof> {
        let client_known_num_leaves = client_known_version
            .map(|v| v.saturating_add(1))
            .unwrap_or(0);
        let ledger_num_leaves = ledger_version.saturating_add(1);
        Accumulator::get_consistency_proof(self, ledger_num_leaves, client_known_num_leaves)
    }

    /// Write `txn_infos` to `batch`. Assigned `first_version` to the version number of the
    /// first transaction, and so on.
    pub fn put_transaction_infos(
        &self,
        first_version: u64,
        txn_infos: &[TransactionInfo],
        batch: &mut SchemaBatch,
    ) -> Result<HashValue> {
        // write txn_info
        (first_version..first_version + txn_infos.len() as u64)
            .zip_eq(txn_infos.iter())
            .try_for_each(|(version, txn_info)| {
                batch.put::<TransactionInfoSchema>(&version, txn_info)
            })?;

        // write hash of txn_info into the accumulator
        let txn_hashes: Vec<HashValue> = txn_infos.iter().map(TransactionInfo::hash).collect();
        let (root_hash, writes) = Accumulator::append(
            self,
            first_version, /* num_existing_leaves */
            &txn_hashes,
        )?;
        writes
            .iter()
            .try_for_each(|(pos, hash)| batch.put::<TransactionAccumulatorSchema>(pos, hash))?;
        Ok(root_hash)
    }

    /// Write `ledger_info_with_sigs` to `batch`.
    pub fn put_ledger_info(
        &self,
        ledger_info_with_sigs: &LedgerInfoWithSignatures,
        batch: &mut SchemaBatch,
    ) -> Result<()> {
        let ledger_info = ledger_info_with_sigs.ledger_info();

        if ledger_info.ends_epoch() {
            // This is the last version of the current epoch, update the epoch by version index.
            batch.put::<EpochByVersionSchema>(&ledger_info.version(), &ledger_info.epoch())?;
        }
        batch.put::<LedgerInfoSchema>(&ledger_info.epoch(), ledger_info_with_sigs)
    }

    pub fn get_root_hash(&self, version: Version) -> Result<HashValue> {
        Accumulator::get_root_hash(self, version + 1)
    }
}

pub(crate) type Accumulator = MerkleAccumulator<LedgerStore, TransactionAccumulatorHasher>;

impl HashReader for LedgerStore {
    fn get(&self, position: Position) -> Result<HashValue> {
        self.db
            .get::<TransactionAccumulatorSchema>(&position)?
            .ok_or_else(|| format_err!("{} does not exist.", position))
    }
}

#[cfg(test)]
mod ledger_info_test;
#[cfg(test)]
pub(crate) mod ledger_info_test_utils;
#[cfg(test)]
mod transaction_info_test;
<|endoftext|>