<commit_before>use std::env;
use std::io::{stdout, stderr, Write};
use std::path::{Path, PathBuf};
use std::fs::metadata;
use std::os::unix::fs::MetadataExt;

use libc::getuid;

use options::build_mode::{build_mode, BuildMode};
use config::{Config, Settings, find_config};
use config::read_settings::{read_settings, MergedSettings};
use argparse::{ArgumentParser, Store, List, Collect, Print, StoreFalse, StoreTrue};

mod list;
mod user;
mod pack;
mod push;
pub mod build;
mod storage;
mod underscore;
mod completion;
mod commands;

mod supervisor;
mod simple;
mod capsule;

pub mod wrap;
pub mod system;
pub mod network;
mod socket;
mod volumes;
mod options;
mod prerequisites;

pub struct Context {
    pub config: Config,
    pub settings: Settings,
    pub ext_settings: MergedSettings,
    pub workdir: PathBuf,
    pub config_dir: PathBuf,
    pub build_mode: BuildMode,
    pub prerequisites: bool,
    pub isolate_network: bool
}

fn check_export(cmd: &String) -> Option<String> {
    if cmd == "/proc/self/exe" {
        return None;
    }
    match cmd.rfind("/") {
        Some(slash) => {
            let exe = &cmd[slash+1..];
            if exe == "vagga" || exe.starts_with("vagga_") {
                None
            } else {
                Some(exe.to_owned())
            }
        }
        None => Some(cmd.clone()),
    }
}

pub fn run(input_args: Vec<String>) -> i32 {
    let mut err = stderr();
    let workdir = env::current_dir().unwrap();

    let (config, cfg_dir) = match find_config(&workdir, true) {
        Ok(tup) => tup,
        Err(e) => {
            writeln!(&mut err, "{}", e).ok();
            return 126;
        }
    };
    // Not sure this is a best place, but the variable is needed for correct
    // reading of settings
    if let Some(x) = env::var_os("HOME") {
        env::set_var("_VAGGA_HOME", x);
    }

    let (ext_settings, int_settings) = match read_settings(&cfg_dir)
    {
        Ok(tup) => tup,
        Err(e) => {
            writeln!(&mut err, "{}", e).ok();
            return 126;
        }
    };

    let mut commands = Vec::<String>::new();
    let mut cname = "".to_string();
    let mut args = vec!();
    let mut set_env = Vec::<String>::new();
    let mut propagate_env = Vec::<String>::new();
    let mut bmode = Default::default();
    let mut owner_check = true;
    let mut prerequisites = true;
    let mut isolate_network = false;

    let export_command = input_args.get(0).and_then(check_export);
    if !int_settings.run_symlinks_as_commands || export_command.is_none() {
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Runs a command in container, optionally builds container if that
            does not exists or outdated.

            Run `vagga` without arguments to see the list of commands.
            ");
        ap.add_option(&["-V", "--version"],
            Print(env!("VAGGA_VERSION").to_string()),
            "Show vagga version and exit");
        ap.refer(&mut set_env)
          .add_option(&["-E", "--env", "--environ"], Collect,
                "Set environment variable for running command")
          .metavar("NAME=VALUE");
        ap.refer(&mut propagate_env)
          .add_option(&["-e", "--use-env"], Collect,
                "Propagate variable VAR into command environment")
          .metavar("VAR");
        ap.refer(&mut owner_check)
          .add_option(&["--ignore-owner-check"], StoreFalse,
                "Ignore checking owner of the project directory");
        ap.refer(&mut prerequisites)
            .add_option(&["--no-prerequisites"], StoreFalse,
            "Run only specified command(s), don't run prerequisites");
        ap.refer(&mut isolate_network)
            .add_option(
                &["--isolate-network", "--no-network", "--no-net"],
                StoreTrue,
                "Run command(s) inside isolated network");
        build_mode(&mut ap, &mut bmode);
        ap.refer(&mut commands)
          .add_option(&["-m", "--run-multi"], List, "
            Run the following list of commands. Each without an arguments.
            When any of them fails, stop the chain. Basically it's the
            shortcut to `vagga cmd1 && vagga cmd2` except containers for
            `cmd2` are built beforehand, for your convenience. Also builtin
            commands (those starting with underscore) do not work with
            `vagga -m`");
        ap.refer(&mut cname)
          .add_argument("command", Store,
                "A vagga command to run");
        ap.refer(&mut args)
          .add_argument("args", List,
                "Arguments for the command");
        ap.stop_on_first_argument(true);
        ap.silence_double_dash(false);
        match ap.parse(input_args, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return 0,
            Err(_) => return 122,
        }
    } else {
        let export_command = export_command.unwrap();
        for (name, cmd) in &config.commands {
            match cmd.link() {
                Some(ref lnk) if lnk.name == export_command => {
                    cname = name.clone();
                }
                _ => {}
            }
        }
        if cname == "" {
            writeln!(&mut err, "Can't find command {:?}", export_command).ok();
            return 127;
        }
        args = input_args;
        args.remove(0);  // the command itself
    }

    if &cname[..] == "_network" {
        args.insert(0, "vagga _network".to_string());
        return ::network::run(args);
    }

    if owner_check {
        let uid = unsafe { getuid() };
        match metadata(&cfg_dir) {
            Ok(ref stat) if stat.uid() == uid  => {}
            Ok(_) => {
                if uid == 0 {
                    writeln!(&mut err, "You should not run vagga as root \
                        (see http://bit.ly/err_root)").ok();
                    return 122;
                } else {
                    warn!("You are running vagga as a user \
                        different from the owner of project directory. \
                        You may not have needed permissions \
                        (see http://bit.ly/err_root)");
                }
            }
            Err(e) => {
                writeln!(&mut err, "Can't stat {:?}: {}", cfg_dir, e).ok();
                return 126;
            }
        }
    }

    let int_workdir = workdir.strip_prefix(&cfg_dir)
                             .unwrap_or(&Path::new("."));

    for k in propagate_env.into_iter() {
        if k.chars().find(|&c| c == '=').is_some() {
            writeln!(&mut err, "Environment variable name \
                (for option `-e`/`--use-env`) \
                can't contain equals `=` character. \
                To set key-value pair use `-E`/`--environ` option").ok();
            return 126;
        } else {
            env::set_var(&("VAGGAENV_".to_string() + &k[..]),
                env::var_os(&k).unwrap_or(From::from("")));
        }
    }
    for pair in set_env.into_iter() {
        let mut pairiter = pair[..].splitn(2, '=');
        let key = "VAGGAENV_".to_string() + pairiter.next().unwrap();
        if let Some(value) = pairiter.next() {
            env::set_var(&key, value.to_string());
        } else {
            env::remove_var(&key);
        }
    }

    let context = Context {
        config: config,
        settings: int_settings,
        ext_settings: ext_settings,
        workdir: int_workdir.to_path_buf(),
        config_dir: cfg_dir.to_path_buf(),
        build_mode: bmode,
        prerequisites: prerequisites,
        isolate_network: isolate_network,
    };

    if commands.len() > 0 {
        let result = user::run_multiple_commands(&context, commands);
        match result {
            Ok(rc) => {
                return rc;
            }
            Err(text) =>  {
                writeln!(&mut err, "{}", text).ok();
                return 121;
            }
        }
    }

    let result:Result<i32, String> = match &cname[..] {
        "" => {
            list::print_help(&context.config)
        }
        "_create_netns" => {
            network::create_netns(&context.config, args)
        }
        "_destroy_netns" => {
            network::destroy_netns(&context.config, args)
        }
        "_list" => {
            list::print_list(&context.config, args)
        }
        "_version_hash" => {
            underscore::version_hash(&context, &cname, args)
        }
        "_build_shell" | "_clean" | "_check_overlayfs_support" => {
            underscore::passthrough(&context, &cname, args)
        }
        "_base_dir" => {
            println!("{}", cfg_dir.display());
            Ok(1)
        }
        "_relative_work_dir" => {
            println!("{}", int_workdir.display());
            Ok(1)
        }
        "_build" => {
            build::build_command(&context, args)
        }
        "_run" => {
            underscore::run_command(&context, args)
        }
        "_run_in_netns" => {
            underscore::run_in_netns(&context, cname, args)
        }
        "_pack_image" => {
            pack::pack_command(&context, args)
        }
        "_push_image" => {
            push::push_command(&context, args)
        }
        "_init_storage_dir" => {
            storage::init_dir(&context.ext_settings, &cfg_dir, args)
        }
        "_compgen" => {
            completion::generate_completions(&context.config, args)
        }
        "_help" => {
            completion::generate_command_help(&context, args)
        }
        "_update_symlinks" => {
            commands::update_symlinks(&context, args)
        }
        "_capsule" => {
            ::capsule::run_command(&context, args)
        }
        _ => {
            user::run_user_command(&context, cname, args)
        }
    };

    match result {
        Ok(rc) => {
            return rc;
        }
        Err(text) =>  {
            writeln!(&mut err, "{}", text).ok();
            return 121;
        }
    }
}
<commit_msg>Fixes exit codes of `_base_dir` and `_relative_base_dir` commands<commit_after>use std::env;
use std::io::{stdout, stderr, Write};
use std::path::{Path, PathBuf};
use std::fs::metadata;
use std::os::unix::fs::MetadataExt;

use libc::getuid;

use options::build_mode::{build_mode, BuildMode};
use config::{Config, Settings, find_config};
use config::read_settings::{read_settings, MergedSettings};
use argparse::{ArgumentParser, Store, List, Collect, Print, StoreFalse, StoreTrue};

mod list;
mod user;
mod pack;
mod push;
pub mod build;
mod storage;
mod underscore;
mod completion;
mod commands;

mod supervisor;
mod simple;
mod capsule;

pub mod wrap;
pub mod system;
pub mod network;
mod socket;
mod volumes;
mod options;
mod prerequisites;

pub struct Context {
    pub config: Config,
    pub settings: Settings,
    pub ext_settings: MergedSettings,
    pub workdir: PathBuf,
    pub config_dir: PathBuf,
    pub build_mode: BuildMode,
    pub prerequisites: bool,
    pub isolate_network: bool
}

fn check_export(cmd: &String) -> Option<String> {
    if cmd == "/proc/self/exe" {
        return None;
    }
    match cmd.rfind("/") {
        Some(slash) => {
            let exe = &cmd[slash+1..];
            if exe == "vagga" || exe.starts_with("vagga_") {
                None
            } else {
                Some(exe.to_owned())
            }
        }
        None => Some(cmd.clone()),
    }
}

pub fn run(input_args: Vec<String>) -> i32 {
    let mut err = stderr();
    let workdir = env::current_dir().unwrap();

    let (config, cfg_dir) = match find_config(&workdir, true) {
        Ok(tup) => tup,
        Err(e) => {
            writeln!(&mut err, "{}", e).ok();
            return 126;
        }
    };
    // Not sure this is a best place, but the variable is needed for correct
    // reading of settings
    if let Some(x) = env::var_os("HOME") {
        env::set_var("_VAGGA_HOME", x);
    }

    let (ext_settings, int_settings) = match read_settings(&cfg_dir)
    {
        Ok(tup) => tup,
        Err(e) => {
            writeln!(&mut err, "{}", e).ok();
            return 126;
        }
    };

    let mut commands = Vec::<String>::new();
    let mut cname = "".to_string();
    let mut args = vec!();
    let mut set_env = Vec::<String>::new();
    let mut propagate_env = Vec::<String>::new();
    let mut bmode = Default::default();
    let mut owner_check = true;
    let mut prerequisites = true;
    let mut isolate_network = false;

    let export_command = input_args.get(0).and_then(check_export);
    if !int_settings.run_symlinks_as_commands || export_command.is_none() {
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Runs a command in container, optionally builds container if that
            does not exists or outdated.

            Run `vagga` without arguments to see the list of commands.
            ");
        ap.add_option(&["-V", "--version"],
            Print(env!("VAGGA_VERSION").to_string()),
            "Show vagga version and exit");
        ap.refer(&mut set_env)
          .add_option(&["-E", "--env", "--environ"], Collect,
                "Set environment variable for running command")
          .metavar("NAME=VALUE");
        ap.refer(&mut propagate_env)
          .add_option(&["-e", "--use-env"], Collect,
                "Propagate variable VAR into command environment")
          .metavar("VAR");
        ap.refer(&mut owner_check)
          .add_option(&["--ignore-owner-check"], StoreFalse,
                "Ignore checking owner of the project directory");
        ap.refer(&mut prerequisites)
            .add_option(&["--no-prerequisites"], StoreFalse,
            "Run only specified command(s), don't run prerequisites");
        ap.refer(&mut isolate_network)
            .add_option(
                &["--isolate-network", "--no-network", "--no-net"],
                StoreTrue,
                "Run command(s) inside isolated network");
        build_mode(&mut ap, &mut bmode);
        ap.refer(&mut commands)
          .add_option(&["-m", "--run-multi"], List, "
            Run the following list of commands. Each without an arguments.
            When any of them fails, stop the chain. Basically it's the
            shortcut to `vagga cmd1 && vagga cmd2` except containers for
            `cmd2` are built beforehand, for your convenience. Also builtin
            commands (those starting with underscore) do not work with
            `vagga -m`");
        ap.refer(&mut cname)
          .add_argument("command", Store,
                "A vagga command to run");
        ap.refer(&mut args)
          .add_argument("args", List,
                "Arguments for the command");
        ap.stop_on_first_argument(true);
        ap.silence_double_dash(false);
        match ap.parse(input_args, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return 0,
            Err(_) => return 122,
        }
    } else {
        let export_command = export_command.unwrap();
        for (name, cmd) in &config.commands {
            match cmd.link() {
                Some(ref lnk) if lnk.name == export_command => {
                    cname = name.clone();
                }
                _ => {}
            }
        }
        if cname == "" {
            writeln!(&mut err, "Can't find command {:?}", export_command).ok();
            return 127;
        }
        args = input_args;
        args.remove(0);  // the command itself
    }

    if &cname[..] == "_network" {
        args.insert(0, "vagga _network".to_string());
        return ::network::run(args);
    }

    if owner_check {
        let uid = unsafe { getuid() };
        match metadata(&cfg_dir) {
            Ok(ref stat) if stat.uid() == uid  => {}
            Ok(_) => {
                if uid == 0 {
                    writeln!(&mut err, "You should not run vagga as root \
                        (see http://bit.ly/err_root)").ok();
                    return 122;
                } else {
                    warn!("You are running vagga as a user \
                        different from the owner of project directory. \
                        You may not have needed permissions \
                        (see http://bit.ly/err_root)");
                }
            }
            Err(e) => {
                writeln!(&mut err, "Can't stat {:?}: {}", cfg_dir, e).ok();
                return 126;
            }
        }
    }

    let int_workdir = workdir.strip_prefix(&cfg_dir)
                             .unwrap_or(&Path::new("."));

    for k in propagate_env.into_iter() {
        if k.chars().find(|&c| c == '=').is_some() {
            writeln!(&mut err, "Environment variable name \
                (for option `-e`/`--use-env`) \
                can't contain equals `=` character. \
                To set key-value pair use `-E`/`--environ` option").ok();
            return 126;
        } else {
            env::set_var(&("VAGGAENV_".to_string() + &k[..]),
                env::var_os(&k).unwrap_or(From::from("")));
        }
    }
    for pair in set_env.into_iter() {
        let mut pairiter = pair[..].splitn(2, '=');
        let key = "VAGGAENV_".to_string() + pairiter.next().unwrap();
        if let Some(value) = pairiter.next() {
            env::set_var(&key, value.to_string());
        } else {
            env::remove_var(&key);
        }
    }

    let context = Context {
        config: config,
        settings: int_settings,
        ext_settings: ext_settings,
        workdir: int_workdir.to_path_buf(),
        config_dir: cfg_dir.to_path_buf(),
        build_mode: bmode,
        prerequisites: prerequisites,
        isolate_network: isolate_network,
    };

    if commands.len() > 0 {
        let result = user::run_multiple_commands(&context, commands);
        match result {
            Ok(rc) => {
                return rc;
            }
            Err(text) =>  {
                writeln!(&mut err, "{}", text).ok();
                return 121;
            }
        }
    }

    let result:Result<i32, String> = match &cname[..] {
        "" => {
            list::print_help(&context.config)
        }
        "_create_netns" => {
            network::create_netns(&context.config, args)
        }
        "_destroy_netns" => {
            network::destroy_netns(&context.config, args)
        }
        "_list" => {
            list::print_list(&context.config, args)
        }
        "_version_hash" => {
            underscore::version_hash(&context, &cname, args)
        }
        "_build_shell" | "_clean" | "_check_overlayfs_support" => {
            underscore::passthrough(&context, &cname, args)
        }
        "_base_dir" => {
            println!("{}", cfg_dir.display());
            Ok(0)
        }
        "_relative_work_dir" => {
            println!("{}", int_workdir.display());
            Ok(0)
        }
        "_build" => {
            build::build_command(&context, args)
        }
        "_run" => {
            underscore::run_command(&context, args)
        }
        "_run_in_netns" => {
            underscore::run_in_netns(&context, cname, args)
        }
        "_pack_image" => {
            pack::pack_command(&context, args)
        }
        "_push_image" => {
            push::push_command(&context, args)
        }
        "_init_storage_dir" => {
            storage::init_dir(&context.ext_settings, &cfg_dir, args)
        }
        "_compgen" => {
            completion::generate_completions(&context.config, args)
        }
        "_help" => {
            completion::generate_command_help(&context, args)
        }
        "_update_symlinks" => {
            commands::update_symlinks(&context, args)
        }
        "_capsule" => {
            ::capsule::run_command(&context, args)
        }
        _ => {
            user::run_user_command(&context, cname, args)
        }
    };

    match result {
        Ok(rc) => {
            return rc;
        }
        Err(text) =>  {
            writeln!(&mut err, "{}", text).ok();
            return 121;
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Adam Perry <adam.n.perry@gmail.com>
//
// This software may be modified and distributed under the terms of the MIT license.  See the
// LICENSE file for details.

use libc::c_int;

use parasail_sys::{parasail_result_free, parasail_nw_striped_profile_sat,
                   parasail_sg_striped_profile_sat, parasail_sw_striped_profile_sat,
                   parasail_sg_stats_striped_sat};
use profile::Profile;

/// Provides a score for global pairwise alignment, using a vectorized version of [Needleman-Wunsch](https://en.wikipedia.org/wiki/Needleman%E2%80%93Wunsch_algorithm).
///
/// # Examples
///
/// ```
/// # use parasailors::*;
/// // create & lookup substitution matrices
/// let identity_matrix = Matrix::new(MatrixType::Identity);
///
/// let query = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// let profile_ident = Profile::new(query, &identity_matrix);
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// assert_eq!(50, global_alignment_score(&profile_ident, reference, 1, 1));
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTCCTTTTTTNNNNNNNNN";
/// assert_eq!(48, global_alignment_score(&profile_ident, reference, 1, 1));
/// ```
pub fn global_alignment_score(query_profile: &Profile,
                              database_sequence: &[u8],
                              open_cost: i32,
                              gap_extend_cost: i32)
                              -> i32 {

    unsafe {
        let result = parasail_nw_striped_profile_sat(**query_profile,
                                                     database_sequence.as_ptr(),
                                                     database_sequence.len() as c_int,
                                                     open_cost,
                                                     gap_extend_cost);
        let score = (*result).score;
        parasail_result_free(result);
        score
    }
}

/// Provides a score for semi-global pairwise alignment using a vectorized algorithm.
///
/// This results in a score that corresponds to a global alignment for the query sequence (i.e. the sequence in the `Profile`) and a local alignment for the reference sequence. This is particularly useful when checking for the presence of an NGS read in a much longer reference sequence. This behaves like a global alignment, except that gaps at the start or end of the reference sequence's alignment are ignored.
///
/// # Examples
///
/// ```
/// # use parasailors::*;
/// # let identity_matrix = Matrix::new(MatrixType::Identity);
/// let query = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// let profile_ident = Profile::new(query, &identity_matrix);
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// assert_eq!(50, semi_global_alignment_score(&profile_ident, reference, 1, 1));
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTCCTTTTTTNNNNNNNNN";
/// assert_eq!(48, semi_global_alignment_score(&profile_ident, reference, 1, 1));
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTT";
/// assert_eq!(35, semi_global_alignment_score(&profile_ident, reference, 1, 1));
/// ```
pub fn semi_global_alignment_score(query_profile: &Profile,
                                   database_sequence: &[u8],
                                   open_cost: i32,
                                   gap_extend_cost: i32)
                                   -> i32 {

    unsafe {
        let result = parasail_sg_striped_profile_sat(**query_profile,
                                                     database_sequence.as_ptr(),
                                                     database_sequence.len() as c_int,
                                                     open_cost,
                                                     gap_extend_cost);
        let score = (*result).score;
        parasail_result_free(result);
        score
    }
}

/// Returns a score for local pairwise alignment using a vectorized version of [Smith-Waterman](https://en.wikipedia.org/wiki/Smith%E2%80%93Waterman_algorithm).
///
/// # Examples
///
/// ```
/// # use parasailors::*;
/// # let identity_matrix = Matrix::new(MatrixType::Identity);
/// let query = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// let profile_ident = Profile::new(query, &identity_matrix);
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// assert_eq!(50, local_alignment_score(&profile_ident, reference, 1, 1));
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTCCTTTTTTNNNNNNNNN";
/// assert_eq!(48, local_alignment_score(&profile_ident, reference, 1, 1));
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTT";
/// assert_eq!(35, local_alignment_score(&profile_ident, reference, 1, 1));
/// ```
pub fn local_alignment_score(query_profile: &Profile,
                             database_sequence: &[u8],
                             open_cost: i32,
                             gap_extend_cost: i32)
                             -> i32 {

    unsafe {
        let result = parasail_sw_striped_profile_sat(**query_profile,
                                                     database_sequence.as_ptr(),
                                                     database_sequence.len() as c_int,
                                                     open_cost,
                                                     gap_extend_cost);
        let score = (*result).score;
        parasail_result_free(result);
        score
    }
}

/// Stores statistics from an alignment.
pub struct AlignmentStats {
    /// The score according to the substitution matrix and gap penalty scheme used.
    pub score: i64,
    /// Number of exactly matching characters.
    pub num_matches: u64,
    /// Number of positively scoring character substitutions (this is the same as num_matches when used an identity matrix).
    pub num_positive_subs: u64,
    /// The length of the found alignment.
    pub align_length: usize,
    /// The starting index (0-based) of the alignment in the query (usually 0).
    pub query_end: usize,
    /// The starting index (0-based) of the alignment in the reference.
    pub ref_end: usize,
}

/// Provides statistics for semi-global pairwise alignment using a vectorized algorithm.
///
/// This results in a series of statistics, including a score that corresponds to a global alignment for the query sequence (i.e. the sequence in the `Profile`) and a local alignment for the reference sequence. This is particularly useful when checking for the presence of an NGS read in a much longer reference sequence. This behaves like a global alignment, except that gaps at the start or end of the reference sequence's alignment are ignored.
///
/// Other statistics include the number of matching characters, the number of positive substitutions, the length of the found alignment, and the starting point in both sequences where the alignment starts.
///
/// # Examples
///
/// ```
/// # use parasailors::*;
/// let identity_matrix = Matrix::new(MatrixType::Identity);
/// let query = b"AAAACCCCCCCCCCGGG";
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// let stats = semi_global_alignment_stats(query, reference, 1, 1, &identity_matrix);
/// assert_eq!(17, stats.score);
/// assert_eq!(17, stats.num_matches);
/// assert_eq!(17, stats.num_positive_subs);
/// assert_eq!(17, stats.align_length);
/// assert_eq!(17, stats.query_end);
/// assert_eq!(23, stats.ref_end);
/// ```
pub fn semi_global_alignment_stats(query_sequence: &[u8],
                                   database_sequence: &[u8],
                                   open_cost: i32,
                                   gap_extend_cost: i32,
                                   substitution_matrix: &::matrix::Matrix)
                                   -> AlignmentStats {

    unsafe {
        let result = parasail_sg_stats_striped_sat(query_sequence.as_ptr(),
                                                   query_sequence.len() as c_int,
                                                   database_sequence.as_ptr(),
                                                   database_sequence.len() as c_int,
                                                   open_cost,
                                                   gap_extend_cost,
                                                   **substitution_matrix);

        let score = (*result).score as i64;
        let num_matches = (*result).matches as u64;
        let num_subs = (*result).similar as u64;
        let align_len = (*result).length as usize;

        // calculate start from end
        let query_end = (*result).end_query as usize + 1;
        let ref_end = (*result).end_ref as usize + 1;

        parasail_result_free(result);

        AlignmentStats {
            score: score,
            num_matches: num_matches,
            num_positive_subs: num_subs,
            align_length: align_len,
            query_end: query_end,
            ref_end: ref_end,
        }
    }
}

#[test]
fn test_semiglobal_stats() {
    use matrix::{Matrix, MatrixType};
    use std::str;
    let identity_matrix = Matrix::new(MatrixType::Identity);
    let query = b"AAAACCCCCCCCCCGGG";

    let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
    let stats = semi_global_alignment_stats(query, reference, 1, 1, &identity_matrix);
    assert_eq!(17, stats.score);
    assert_eq!(17, stats.num_matches);
    assert_eq!(17, stats.num_positive_subs);
    assert_eq!(17, stats.align_length);
    assert_eq!(17, stats.query_end);
    assert_eq!(23, stats.ref_end);

    assert_eq!(str::from_utf8(query).unwrap(),
               str::from_utf8(&query[stats.query_end - stats.align_length..stats.query_end])
                   .unwrap());
    assert_eq!(str::from_utf8(query).unwrap(),
               str::from_utf8(&reference[stats.ref_end - stats.align_length..stats.ref_end])
                   .unwrap());

    // these two test cases "borrowed" mutably from rust-bio
    let x = b"ACCGTGGAT";
    let y = b"AAAAACCGTTGAT";
    let ident_with_penalty = Matrix::new(MatrixType::IdentityWithPenalty);
    let alignment = semi_global_alignment_stats(x, y, 5, 1, &ident_with_penalty);
    assert_eq!(7, alignment.score);
    assert_eq!(13, alignment.ref_end);
    assert_eq!(9, alignment.query_end);

    let x = b"CCGGCA";
    let y = b"ACCGTTGACGC";
    let alignment = semi_global_alignment_stats(x, y, 5, 1, &ident_with_penalty);
    assert_eq!(1, alignment.score);
    assert_eq!(1, alignment.ref_end);
    assert_eq!(6, alignment.query_end);
}
<commit_msg>Adding local alignment statistics.<commit_after>// Copyright (c) 2016 Adam Perry <adam.n.perry@gmail.com>
//
// This software may be modified and distributed under the terms of the MIT license.  See the
// LICENSE file for details.

use libc::c_int;

use parasail_sys::{parasail_result_free, parasail_nw_striped_profile_sat,
                   parasail_sg_striped_profile_sat, parasail_sw_striped_profile_sat,
                   parasail_sg_stats_striped_sat, parasail_sw_stats_striped_sat};
use profile::Profile;

/// Provides a score for global pairwise alignment, using a vectorized version of [Needleman-Wunsch](https://en.wikipedia.org/wiki/Needleman%E2%80%93Wunsch_algorithm).
///
/// # Examples
///
/// ```
/// # use parasailors::*;
/// // create & lookup substitution matrices
/// let identity_matrix = Matrix::new(MatrixType::Identity);
///
/// let query = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// let profile_ident = Profile::new(query, &identity_matrix);
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// assert_eq!(50, global_alignment_score(&profile_ident, reference, 1, 1));
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTCCTTTTTTNNNNNNNNN";
/// assert_eq!(48, global_alignment_score(&profile_ident, reference, 1, 1));
/// ```
pub fn global_alignment_score(query_profile: &Profile,
                              database_sequence: &[u8],
                              open_cost: i32,
                              gap_extend_cost: i32)
                              -> i32 {

    unsafe {
        let result = parasail_nw_striped_profile_sat(**query_profile,
                                                     database_sequence.as_ptr(),
                                                     database_sequence.len() as c_int,
                                                     open_cost,
                                                     gap_extend_cost);
        let score = (*result).score;
        parasail_result_free(result);
        score
    }
}

/// Provides a score for semi-global pairwise alignment using a vectorized algorithm.
///
/// This results in a score that corresponds to a global alignment for the query sequence (i.e. the sequence in the `Profile`) and a local alignment for the reference sequence. This is particularly useful when checking for the presence of an NGS read in a much longer reference sequence. This behaves like a global alignment, except that gaps at the start or end of the reference sequence's alignment are ignored.
///
/// # Examples
///
/// ```
/// # use parasailors::*;
/// # let identity_matrix = Matrix::new(MatrixType::Identity);
/// let query = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// let profile_ident = Profile::new(query, &identity_matrix);
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// assert_eq!(50, semi_global_alignment_score(&profile_ident, reference, 1, 1));
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTCCTTTTTTNNNNNNNNN";
/// assert_eq!(48, semi_global_alignment_score(&profile_ident, reference, 1, 1));
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTT";
/// assert_eq!(35, semi_global_alignment_score(&profile_ident, reference, 1, 1));
/// ```
pub fn semi_global_alignment_score(query_profile: &Profile,
                                   database_sequence: &[u8],
                                   open_cost: i32,
                                   gap_extend_cost: i32)
                                   -> i32 {

    unsafe {
        let result = parasail_sg_striped_profile_sat(**query_profile,
                                                     database_sequence.as_ptr(),
                                                     database_sequence.len() as c_int,
                                                     open_cost,
                                                     gap_extend_cost);
        let score = (*result).score;
        parasail_result_free(result);
        score
    }
}

/// Returns a score for local pairwise alignment using a vectorized version of [Smith-Waterman](https://en.wikipedia.org/wiki/Smith%E2%80%93Waterman_algorithm).
///
/// # Examples
///
/// ```
/// # use parasailors::*;
/// # let identity_matrix = Matrix::new(MatrixType::Identity);
/// let query = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// let profile_ident = Profile::new(query, &identity_matrix);
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// assert_eq!(50, local_alignment_score(&profile_ident, reference, 1, 1));
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTCCTTTTTTNNNNNNNNN";
/// assert_eq!(48, local_alignment_score(&profile_ident, reference, 1, 1));
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTT";
/// assert_eq!(35, local_alignment_score(&profile_ident, reference, 1, 1));
/// ```
pub fn local_alignment_score(query_profile: &Profile,
                             database_sequence: &[u8],
                             open_cost: i32,
                             gap_extend_cost: i32)
                             -> i32 {

    unsafe {
        let result = parasail_sw_striped_profile_sat(**query_profile,
                                                     database_sequence.as_ptr(),
                                                     database_sequence.len() as c_int,
                                                     open_cost,
                                                     gap_extend_cost);
        let score = (*result).score;
        parasail_result_free(result);
        score
    }
}

/// Stores statistics from an alignment.
pub struct AlignmentStats {
    /// The score according to the substitution matrix and gap penalty scheme used.
    pub score: i64,
    /// Number of exactly matching characters.
    pub num_matches: u64,
    /// Number of positively scoring character substitutions (this is the same as num_matches when used an identity matrix).
    pub num_positive_subs: u64,
    /// The length of the found alignment.
    pub align_length: usize,
    /// The starting index (0-based) of the alignment in the query (usually 0).
    pub query_end: usize,
    /// The starting index (0-based) of the alignment in the reference.
    pub ref_end: usize,
}

/// Provides statistics for semi-global pairwise alignment using a vectorized algorithm.
///
/// This results in a series of statistics, including a score that corresponds to a global alignment for the query sequence and a local alignment for the reference sequence. This is particularly useful when checking for the presence of an NGS read in a much longer reference sequence. This behaves like a global alignment, except that gaps at the start or end of the reference sequence's alignment are ignored.
///
/// Other statistics include the number of matching characters, the number of positive substitutions, the length of the found alignment, and the starting point in both sequences where the alignment starts.
///
/// # Examples
///
/// ```
/// # use parasailors::*;
/// let identity_matrix = Matrix::new(MatrixType::Identity);
/// let query = b"AAAACCCCCCCCCCGGG";
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// let stats = semi_global_alignment_stats(query, reference, 1, 1, &identity_matrix);
/// assert_eq!(17, stats.score);
/// assert_eq!(17, stats.num_matches);
/// assert_eq!(17, stats.num_positive_subs);
/// assert_eq!(17, stats.align_length);
/// assert_eq!(17, stats.query_end);
/// assert_eq!(23, stats.ref_end);
/// ```
pub fn semi_global_alignment_stats(query_sequence: &[u8],
                                   database_sequence: &[u8],
                                   open_cost: i32,
                                   gap_extend_cost: i32,
                                   substitution_matrix: &::matrix::Matrix)
                                   -> AlignmentStats {

    unsafe {
        let result = parasail_sg_stats_striped_sat(query_sequence.as_ptr(),
                                                   query_sequence.len() as c_int,
                                                   database_sequence.as_ptr(),
                                                   database_sequence.len() as c_int,
                                                   open_cost,
                                                   gap_extend_cost,
                                                   **substitution_matrix);

        let score = (*result).score as i64;
        let num_matches = (*result).matches as u64;
        let num_subs = (*result).similar as u64;
        let align_len = (*result).length as usize;

        // calculate start from end
        let query_end = (*result).end_query as usize + 1;
        let ref_end = (*result).end_ref as usize + 1;

        parasail_result_free(result);

        AlignmentStats {
            score: score,
            num_matches: num_matches,
            num_positive_subs: num_subs,
            align_length: align_len,
            query_end: query_end,
            ref_end: ref_end,
        }
    }
}

/// Provides statistics for local pairwise alignment using a vectorized algorithm.
///
/// # Examples
///
/// ```
/// # use parasailors::*;
/// let identity_matrix = Matrix::new(MatrixType::Identity);
/// let query = b"AAAACCCCCCCCCCGGG";
///
/// let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
/// let stats = local_alignment_stats(query, reference, 1, 1, &identity_matrix);
/// assert_eq!(17, stats.score);
/// assert_eq!(17, stats.num_matches);
/// assert_eq!(17, stats.num_positive_subs);
/// assert_eq!(17, stats.align_length);
/// assert_eq!(17, stats.query_end);
/// assert_eq!(23, stats.ref_end);
/// ```
pub fn local_alignment_stats(query_sequence: &[u8],
                                   database_sequence: &[u8],
                                   open_cost: i32,
                                   gap_extend_cost: i32,
                                   substitution_matrix: &::matrix::Matrix)
                                   -> AlignmentStats {

    unsafe {
        let result = parasail_sw_stats_striped_sat(query_sequence.as_ptr(),
                                                   query_sequence.len() as c_int,
                                                   database_sequence.as_ptr(),
                                                   database_sequence.len() as c_int,
                                                   open_cost,
                                                   gap_extend_cost,
                                                   **substitution_matrix);

        let score = (*result).score as i64;
        let num_matches = (*result).matches as u64;
        let num_subs = (*result).similar as u64;
        let align_len = (*result).length as usize;

        // calculate start from end
        let query_end = (*result).end_query as usize + 1;
        let ref_end = (*result).end_ref as usize + 1;

        parasail_result_free(result);

        AlignmentStats {
            score: score,
            num_matches: num_matches,
            num_positive_subs: num_subs,
            align_length: align_len,
            query_end: query_end,
            ref_end: ref_end,
        }
    }
}

#[test]
fn test_semiglobal_stats() {
    use matrix::{Matrix, MatrixType};
    use std::str;
    let identity_matrix = Matrix::new(MatrixType::Identity);
    let query = b"AAAACCCCCCCCCCGGG";

    let reference = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTTTNNNNNNNNN";
    let stats = semi_global_alignment_stats(query, reference, 1, 1, &identity_matrix);
    assert_eq!(17, stats.score);
    assert_eq!(17, stats.num_matches);
    assert_eq!(17, stats.num_positive_subs);
    assert_eq!(17, stats.align_length);
    assert_eq!(17, stats.query_end);
    assert_eq!(23, stats.ref_end);

    assert_eq!(str::from_utf8(query).unwrap(),
               str::from_utf8(&query[stats.query_end - stats.align_length..stats.query_end])
                   .unwrap());
    assert_eq!(str::from_utf8(query).unwrap(),
               str::from_utf8(&reference[stats.ref_end - stats.align_length..stats.ref_end])
                   .unwrap());

    // these two test cases "borrowed" mutably from rust-bio
    let x = b"ACCGTGGAT";
    let y = b"AAAAACCGTTGAT";
    let ident_with_penalty = Matrix::new(MatrixType::IdentityWithPenalty);
    let alignment = semi_global_alignment_stats(x, y, 5, 1, &ident_with_penalty);
    assert_eq!(7, alignment.score);
    assert_eq!(13, alignment.ref_end);
    assert_eq!(9, alignment.query_end);

    let x = b"CCGGCA";
    let y = b"ACCGTTGACGC";
    let alignment = semi_global_alignment_stats(x, y, 5, 1, &ident_with_penalty);
    assert_eq!(1, alignment.score);
    assert_eq!(1, alignment.ref_end);
    assert_eq!(6, alignment.query_end);
}
<|endoftext|>


<commit_before>extern crate lru_cache;

use {SqliteResult, SqliteConnection, SqliteStatement};
use self::lru_cache::LruCache;

/// Prepared statements are cached for faster execution.
/// FIXME limitation: the same SQL can be cached only once...
#[derive(Debug)]
pub struct StatementCache<'conn> {
    pub conn: &'conn SqliteConnection,
    cache: LruCache<String, SqliteStatement<'conn>>,
}

impl<'conn> StatementCache<'conn> {
    pub fn new(conn: &'conn SqliteConnection, capacity: usize) -> StatementCache<'conn> {
        StatementCache{ conn: conn,  cache: LruCache::new(capacity) }
    }

    /// Search the cache for a prepared-statement object that implements `sql`.
    // If no such prepared-statement can be found, allocate and prepare a new one.
    pub fn get(&mut self, sql: &str) -> SqliteResult<SqliteStatement<'conn>> {
        let stmt = self.cache.remove(sql);
        match stmt {
            Some(stmt) => Ok(stmt),
            _ => self.conn.prepare(sql)
        }
    }

    /// If `discard` is true, then the statement is deleted immediately.
    /// Otherwise it is added to the LRU list and may be returned
    /// by a subsequent call to `get()`.
    pub fn release(&mut self, stmt: SqliteStatement<'conn>, discard: bool) -> SqliteResult<()> {
        if discard {
            return stmt.finalize();
        }
        // FIXME stmt.reset_if_needed();
        // clear bindings ???
        self.cache.insert(stmt.sql(), stmt).map_or(Ok(()), |stmt| stmt.finalize())
    }

    /// Flush the prepared statement cache
    pub fn flush(&mut self) {
        self.cache.clear();
    }

    /// Return (current, max) sizes.
    pub fn size(&self) -> (usize, usize) {
        (self.cache.len(), self.cache.capacity())
    }

    /// Set the maximum number of cached statements.
    pub fn set_size(&mut self, capacity: usize) {
        self.cache.set_capacity(capacity);
    }
}

#[cfg(test)]
mod test {
    use SqliteConnection;
    use super::StatementCache;

   #[test]
    fn test_cache() {
        let db = SqliteConnection::open_in_memory().unwrap();
        let mut cache = StatementCache::new(&db, 10);
        let sql = "PRAGMA schema_version";
        let mut stmt = cache.get(sql).unwrap();
        //println!("NEW {:?}", stmt);
        cache.release(stmt, false).unwrap();
        stmt = cache.get(sql).unwrap();
        //println!("CACHED {:?}", stmt);
        cache.release(stmt, true).unwrap();
        cache.flush();
    }
}
<commit_msg>Improve documentation<commit_after>//! Prepared statements cache for faster execution.
extern crate lru_cache;

use {SqliteResult, SqliteConnection, SqliteStatement};
use self::lru_cache::LruCache;

/// Prepared statements cache.
///
/// FIXME limitation: the same SQL can be cached only once...
#[derive(Debug)]
pub struct StatementCache<'conn> {
    pub conn: &'conn SqliteConnection,
    cache: LruCache<String, SqliteStatement<'conn>>,
}

impl<'conn> StatementCache<'conn> {
    /// Create a statement cache.
    pub fn new(conn: &'conn SqliteConnection, capacity: usize) -> StatementCache<'conn> {
        StatementCache{ conn: conn,  cache: LruCache::new(capacity) }
    }

    /// Search the cache for a prepared-statement object that implements `sql`.
    // If no such prepared-statement can be found, allocate and prepare a new one.
    pub fn get(&mut self, sql: &str) -> SqliteResult<SqliteStatement<'conn>> {
        let stmt = self.cache.remove(sql);
        match stmt {
            Some(stmt) => Ok(stmt),
            _ => self.conn.prepare(sql)
        }
    }

    /// If `discard` is true, then the statement is deleted immediately.
    /// Otherwise it is added to the LRU list and may be returned
    /// by a subsequent call to `get()`.
    pub fn release(&mut self, stmt: SqliteStatement<'conn>, discard: bool) -> SqliteResult<()> {
        if discard {
            return stmt.finalize();
        }
        // FIXME stmt.reset_if_needed();
        // clear bindings ???
        self.cache.insert(stmt.sql(), stmt).map_or(Ok(()), |stmt| stmt.finalize())
    }

    /// Flush the prepared statement cache
    pub fn flush(&mut self) {
        self.cache.clear();
    }

    /// Return (current, max) sizes.
    pub fn size(&self) -> (usize, usize) {
        (self.cache.len(), self.cache.capacity())
    }

    /// Set the maximum number of cached statements.
    pub fn set_size(&mut self, capacity: usize) {
        self.cache.set_capacity(capacity);
    }
}

#[cfg(test)]
mod test {
    use SqliteConnection;
    use super::StatementCache;

   #[test]
    fn test_cache() {
        let db = SqliteConnection::open_in_memory().unwrap();
        let mut cache = StatementCache::new(&db, 10);
        let sql = "PRAGMA schema_version";
        let mut stmt = cache.get(sql).unwrap();
        //println!("NEW {:?}", stmt);
        cache.release(stmt, false).unwrap();
        stmt = cache.get(sql).unwrap();
        //println!("CACHED {:?}", stmt);
        cache.release(stmt, true).unwrap();
        cache.flush();
    }
}
<|endoftext|>


<commit_before>use std::cmp::min;
use std::fmt;
use std::io;
use std::str;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

use bytes::{Bytes, BytesMut};
use bytes::{IntoBuf, Buf, BufMut, BigEndian};
use tokio_io::codec::{Encoder, Decoder};
use tokio_proto::streaming::multiplex::{Frame, RequestId};

use constants::*;
use packet::{PacketMagic, PTYPES};

pub struct PacketHeader {
    pub magic: PacketMagic,
    pub ptype: u32,
    pub psize: u32,
}

impl fmt::Debug for PacketHeader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "PacketHeader {{ magic: {:?}, ptype: {}, size: {} }}",
               match self.magic {
                   PacketMagic::REQ => "REQ",
                   PacketMagic::RES => "RES",
                   PacketMagic::TEXT => "TEXT",
                   _ => "UNKNOWN",
               },
               match self.ptype {
                   p @ 0...42 => PTYPES[p as usize].name,
                   _ => "__UNIMPLEMENTED__",
               },
               self.psize)
    }
}

pub struct PacketCodec {
    data_todo: Option<usize>,
    request_counter: Arc<AtomicUsize>,
    active_request_id: RequestId,
}

impl PacketCodec {
    pub fn new(request_counter: Arc<AtomicUsize>) -> PacketCodec {
        PacketCodec {
            data_todo: None,
            request_counter: request_counter,
            active_request_id: 0,
        }
    }
}

pub type PacketItem = Frame<PacketHeader, Bytes, io::Error>;

impl PacketHeader {
    pub fn admin_decode(buf: &mut BytesMut) -> Result<Option<PacketItem>, io::Error> {
        let newline = buf[..].iter().position(|b| *b == b'\n');
        if let Some(n) = newline {
            let line = buf.split_to(n);
            buf.split_to(1); // drop the newline itself
            let data_str = match str::from_utf8(&line[..]) {
                Ok(s) => s,
                Err(_) => return Err(io::Error::new(io::ErrorKind::Other, "invalid string")),
            };
            info!("admin command data: {:?}", data_str);
            let command = match data_str.trim() {
                "version" => ADMIN_VERSION,
                "status" => ADMIN_STATUS,
                _ => ADMIN_UNKNOWN,
            };
            return Ok(Some(Frame::Message {
                id: 0,
                message: PacketHeader {
                    magic: PacketMagic::TEXT,
                    ptype: command,
                    psize: 0,
                },
                body: false,
                solo: false,
            }));
        }
        Ok(None) // Wait for more data
    }

    pub fn decode(buf: &mut BytesMut,
                  request_counter: &Arc<AtomicUsize>)
                  -> Result<Option<PacketItem>, io::Error> {
        debug!("Decoding {:?}", buf);
        // Peek at first 4
        // Is this a req/res
        if buf.len() < 4 {
            return Ok(None);
        }
        let mut magic_buf: [u8; 4] = [0; 4];
        magic_buf.clone_from_slice(&buf[0..4]);
        let magic = match magic_buf {
            REQ => PacketMagic::REQ,
            RES => PacketMagic::RES,
            // TEXT/ADMIN protocol
            _ => PacketMagic::TEXT,
        };
        debug!("Magic is {:?}", magic);
        if magic == PacketMagic::TEXT {
            debug!("admin protocol detected");
            return PacketHeader::admin_decode(buf);
        }
        if buf.len() < 12 {
            return Ok(None);
        }
        buf.split_to(4);
        // Now get the type
        let ptype = buf.split_to(4).into_buf().get_u32::<BigEndian>();
        debug!("We got a {}", &PTYPES[ptype as usize].name);
        // Now the length
        let psize = buf.split_to(4).into_buf().get_u32::<BigEndian>();
        debug!("Data section is {} bytes", psize);
        let solo = match ptype {
            SUBMIT_JOB |
            SUBMIT_JOB_LOW |
            SUBMIT_JOB_HIGH |
            SUBMIT_REDUCE_JOB => false,
            GRAB_JOB | GRAB_JOB_UNIQ | GRAB_JOB_ALL => false,
            GET_STATUS |
            GET_STATUS_UNIQUE => false,
            _ => false,  // TODO: when we figure out how to Service::call on these, set to true
        };
        let req_id: RequestId = if solo {
            0
        } else {
            request_counter.fetch_add(1, Ordering::Relaxed) as RequestId
        };
        Ok(Some(Frame::Message {
            id: req_id,
            message: PacketHeader {
                magic: magic,
                ptype: ptype,
                psize: psize,
            },
            body: true, // TODO: false for 0 psize?
            solo: solo,
        }))
    }

    pub fn to_bytes(&self) -> Bytes {
        let magic = match self.magic {
            PacketMagic::UNKNOWN => panic!("Unknown packet magic cannot be sent"),
            PacketMagic::REQ => REQ,
            PacketMagic::RES => RES,
            PacketMagic::TEXT => {
                return Bytes::from_static(b"");
            }
        };
        let mut buf = BytesMut::with_capacity(12);
        buf.extend(magic.iter());
        buf.put_u32::<BigEndian>(self.ptype);
        buf.put_u32::<BigEndian>(self.psize);
        buf.freeze()
    }

    pub fn new_text_res(body: &Bytes) -> PacketHeader {
        PacketHeader {
            magic: PacketMagic::TEXT,
            ptype: ADMIN_RESPONSE,
            psize: body.len() as u32,
        }
    }
}

impl Decoder for PacketCodec {
    type Item = Frame<PacketHeader, Bytes, io::Error>;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, io::Error> {
        match self.data_todo {
            None => {
                match PacketHeader::decode(buf, &self.request_counter)? {
                    Some(Frame::Message { id, message, body, solo }) => {
                        self.active_request_id = id;
                        self.data_todo = Some(message.psize as usize);
                        Ok(Some(Frame::Message {
                            id: id,
                            message: message,
                            body: body,
                            solo: solo,
                        }))
                    }
                    Some(_) => panic!("Expecting Frame::Message, got something else"),
                    None => Ok(None),
                }
            }
            Some(0) => {
                self.data_todo = None;
                Ok(Some(Frame::Body {
                    id: self.active_request_id,
                    chunk: None,
                }))
            }
            Some(data_todo) => {
                let chunk_size = min(buf.len(), data_todo);
                self.data_todo = Some(data_todo - chunk_size);
                Ok(Some(Frame::Body {
                    id: self.active_request_id,
                    chunk: Some(buf.split_to(chunk_size).freeze()),
                }))
            }
        }
    }
}

impl Encoder for PacketCodec {
    type Item = Frame<PacketHeader, Bytes, io::Error>;
    type Error = io::Error;

    fn encode(&mut self, msg: Self::Item, buf: &mut BytesMut) -> io::Result<()> {
        debug!("Encoding {:?}", msg);
        match msg {
            Frame::Message { id, message, body, solo } => {
                trace!("Encoding header for id={} solo={}", id, solo);
                if body {
                    debug!("body follows")
                }
                buf.extend(message.to_bytes())
            }
            Frame::Body { id, chunk } => {
                trace!("Encoding body for id={}", id);
                match chunk {
                    Some(chunk) => buf.extend_from_slice(&chunk[..]),
                    None => {}
                }
            }
            Frame::Error { id, error } => {
                error!("Sending error frame for id={}. {}", id, error);
                buf.extend("ERR UNKNOWN_COMMAND Unknown+server+command".bytes())
            }
        }
        Ok(())
    }
}
<commit_msg>fix 0 length chunk bug<commit_after>use std::cmp::min;
use std::fmt;
use std::io;
use std::str;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

use bytes::{Bytes, BytesMut};
use bytes::{IntoBuf, Buf, BufMut, BigEndian};
use tokio_io::codec::{Encoder, Decoder};
use tokio_proto::streaming::multiplex::{Frame, RequestId};

use constants::*;
use packet::{PacketMagic, PTYPES};

pub struct PacketHeader {
    pub magic: PacketMagic,
    pub ptype: u32,
    pub psize: u32,
}

impl fmt::Debug for PacketHeader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "PacketHeader {{ magic: {:?}, ptype: {}, size: {} }}",
               match self.magic {
                   PacketMagic::REQ => "REQ",
                   PacketMagic::RES => "RES",
                   PacketMagic::TEXT => "TEXT",
                   _ => "UNKNOWN",
               },
               match self.ptype {
                   p @ 0...42 => PTYPES[p as usize].name,
                   _ => "__UNIMPLEMENTED__",
               },
               self.psize)
    }
}

pub struct PacketCodec {
    data_todo: Option<usize>,
    request_counter: Arc<AtomicUsize>,
    active_request_id: RequestId,
}

impl PacketCodec {
    pub fn new(request_counter: Arc<AtomicUsize>) -> PacketCodec {
        PacketCodec {
            data_todo: None,
            request_counter: request_counter,
            active_request_id: 0,
        }
    }
}

pub type PacketItem = Frame<PacketHeader, Bytes, io::Error>;

impl PacketHeader {
    pub fn admin_decode(buf: &mut BytesMut) -> Result<Option<PacketItem>, io::Error> {
        let newline = buf[..].iter().position(|b| *b == b'\n');
        if let Some(n) = newline {
            let line = buf.split_to(n);
            buf.split_to(1); // drop the newline itself
            let data_str = match str::from_utf8(&line[..]) {
                Ok(s) => s,
                Err(_) => return Err(io::Error::new(io::ErrorKind::Other, "invalid string")),
            };
            info!("admin command data: {:?}", data_str);
            let command = match data_str.trim() {
                "version" => ADMIN_VERSION,
                "status" => ADMIN_STATUS,
                _ => ADMIN_UNKNOWN,
            };
            return Ok(Some(Frame::Message {
                id: 0,
                message: PacketHeader {
                    magic: PacketMagic::TEXT,
                    ptype: command,
                    psize: 0,
                },
                body: false,
                solo: false,
            }));
        }
        Ok(None) // Wait for more data
    }

    pub fn decode(buf: &mut BytesMut,
                  request_counter: &Arc<AtomicUsize>)
                  -> Result<Option<PacketItem>, io::Error> {
        debug!("Decoding {:?}", buf);
        // Peek at first 4
        // Is this a req/res
        if buf.len() < 4 {
            return Ok(None);
        }
        let mut magic_buf: [u8; 4] = [0; 4];
        magic_buf.clone_from_slice(&buf[0..4]);
        let magic = match magic_buf {
            REQ => PacketMagic::REQ,
            RES => PacketMagic::RES,
            // TEXT/ADMIN protocol
            _ => PacketMagic::TEXT,
        };
        debug!("Magic is {:?}", magic);
        if magic == PacketMagic::TEXT {
            debug!("admin protocol detected");
            return PacketHeader::admin_decode(buf);
        }
        if buf.len() < 12 {
            return Ok(None);
        }
        buf.split_to(4);
        // Now get the type
        let ptype = buf.split_to(4).into_buf().get_u32::<BigEndian>();
        debug!("We got a {}", &PTYPES[ptype as usize].name);
        // Now the length
        let psize = buf.split_to(4).into_buf().get_u32::<BigEndian>();
        debug!("Data section is {} bytes", psize);
        let solo = match ptype {
            SUBMIT_JOB |
            SUBMIT_JOB_LOW |
            SUBMIT_JOB_HIGH |
            SUBMIT_REDUCE_JOB => false,
            GRAB_JOB | GRAB_JOB_UNIQ | GRAB_JOB_ALL => false,
            GET_STATUS |
            GET_STATUS_UNIQUE => false,
            _ => false,  // TODO: when we figure out how to Service::call on these, set to true
        };
        let req_id: RequestId = if solo {
            0
        } else {
            request_counter.fetch_add(1, Ordering::Relaxed) as RequestId
        };
        Ok(Some(Frame::Message {
            id: req_id,
            message: PacketHeader {
                magic: magic,
                ptype: ptype,
                psize: psize,
            },
            body: true, // TODO: false for 0 psize?
            solo: solo,
        }))
    }

    pub fn to_bytes(&self) -> Bytes {
        let magic = match self.magic {
            PacketMagic::UNKNOWN => panic!("Unknown packet magic cannot be sent"),
            PacketMagic::REQ => REQ,
            PacketMagic::RES => RES,
            PacketMagic::TEXT => {
                return Bytes::from_static(b"");
            }
        };
        let mut buf = BytesMut::with_capacity(12);
        buf.extend(magic.iter());
        buf.put_u32::<BigEndian>(self.ptype);
        buf.put_u32::<BigEndian>(self.psize);
        buf.freeze()
    }

    pub fn new_text_res(body: &Bytes) -> PacketHeader {
        PacketHeader {
            magic: PacketMagic::TEXT,
            ptype: ADMIN_RESPONSE,
            psize: body.len() as u32,
        }
    }
}

impl Decoder for PacketCodec {
    type Item = Frame<PacketHeader, Bytes, io::Error>;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, io::Error> {
        match self.data_todo {
            None => {
                match PacketHeader::decode(buf, &self.request_counter)? {
                    Some(Frame::Message { id, message, body, solo }) => {
                        self.active_request_id = id;
                        self.data_todo = Some(message.psize as usize);
                        Ok(Some(Frame::Message {
                            id: id,
                            message: message,
                            body: body,
                            solo: solo,
                        }))
                    }
                    Some(_) => panic!("Expecting Frame::Message, got something else"),
                    None => Ok(None),
                }
            }
            Some(0) => {
                self.data_todo = None;
                Ok(Some(Frame::Body {
                    id: self.active_request_id,
                    chunk: None,
                }))
            }
            Some(data_todo) => {
                let chunk_size = min(buf.len(), data_todo);
                if chunk_size == 0 {
                    debug!("0 length chunk?");
                    Ok(None)
                } else {
                    self.data_todo = Some(data_todo - chunk_size);
                    Ok(Some(Frame::Body {
                        id: self.active_request_id,
                        chunk: Some(buf.split_to(chunk_size).freeze()),
                    }))
                }
            }
        }
    }
}

impl Encoder for PacketCodec {
    type Item = Frame<PacketHeader, Bytes, io::Error>;
    type Error = io::Error;

    fn encode(&mut self, msg: Self::Item, buf: &mut BytesMut) -> io::Result<()> {
        debug!("Encoding {:?}", msg);
        match msg {
            Frame::Message { id, message, body, solo } => {
                trace!("Encoding header for id={} solo={}", id, solo);
                if body {
                    debug!("body follows")
                }
                buf.extend(message.to_bytes())
            }
            Frame::Body { id, chunk } => {
                trace!("Encoding body for id={}", id);
                match chunk {
                    Some(chunk) => buf.extend_from_slice(&chunk[..]),
                    None => {}
                }
            }
            Frame::Error { id, error } => {
                error!("Sending error frame for id={}. {}", id, error);
                buf.extend("ERR UNKNOWN_COMMAND Unknown+server+command".bytes())
            }
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::ffi::{CStr, CString};
use std::fs::{read_link, File};
use std::io::{self, Read, Seek, SeekFrom, Write};
use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};

use libc::{
    self, c_char, c_int, c_long, c_uint, close, fcntl, ftruncate64, off64_t, syscall,
    SYS_memfd_create, EINVAL, F_ADD_SEALS, F_GET_SEALS, F_SEAL_GROW, F_SEAL_SEAL, F_SEAL_SHRINK,
    F_SEAL_WRITE, MFD_ALLOW_SEALING,
};
use serde::{Deserialize, Serialize};

use crate::{errno, errno_result, Result};

/// A shared memory file descriptor and its size.
#[derive(Serialize, Deserialize)]
pub struct SharedMemory {
    #[serde(with = "crate::with_as_descriptor")]
    fd: File,
    size: u64,
}

// from <sys/memfd.h>
const MFD_CLOEXEC: c_uint = 0x0001;

unsafe fn memfd_create(name: *const c_char, flags: c_uint) -> c_int {
    syscall(SYS_memfd_create as c_long, name, flags) as c_int
}

/// A set of memfd seals.
///
/// An enumeration of each bit can be found at `fcntl(2)`.
#[derive(Copy, Clone, Default)]
pub struct MemfdSeals(i32);

impl MemfdSeals {
    /// Returns an empty set of memfd seals.
    #[inline]
    pub fn new() -> MemfdSeals {
        MemfdSeals(0)
    }

    /// Gets the raw bitmask of seals enumerated in `fcntl(2)`.
    #[inline]
    pub fn bitmask(self) -> i32 {
        self.0
    }

    /// True of the grow seal bit is present.
    #[inline]
    pub fn grow_seal(self) -> bool {
        self.0 & F_SEAL_GROW != 0
    }

    /// Sets the grow seal bit.
    #[inline]
    pub fn set_grow_seal(&mut self) {
        self.0 |= F_SEAL_GROW;
    }

    /// True of the shrink seal bit is present.
    #[inline]
    pub fn shrink_seal(self) -> bool {
        self.0 & F_SEAL_SHRINK != 0
    }

    /// Sets the shrink seal bit.
    #[inline]
    pub fn set_shrink_seal(&mut self) {
        self.0 |= F_SEAL_SHRINK;
    }

    /// True of the write seal bit is present.
    #[inline]
    pub fn write_seal(self) -> bool {
        self.0 & F_SEAL_WRITE != 0
    }

    /// Sets the write seal bit.
    #[inline]
    pub fn set_write_seal(&mut self) {
        self.0 |= F_SEAL_WRITE;
    }

    /// True of the seal seal bit is present.
    #[inline]
    pub fn seal_seal(self) -> bool {
        self.0 & F_SEAL_SEAL != 0
    }

    /// Sets the seal seal bit.
    #[inline]
    pub fn set_seal_seal(&mut self) {
        self.0 |= F_SEAL_SEAL;
    }
}

impl SharedMemory {
    /// Convenience function for `SharedMemory::new` that is always named and accepts a wide variety
    /// of string-like types.
    ///
    /// Note that the given name may not have NUL characters anywhere in it, or this will return an
    /// error.
    pub fn named<T: Into<Vec<u8>>>(name: T) -> Result<SharedMemory> {
        Self::new(Some(
            &CString::new(name).map_err(|_| errno::Error::new(EINVAL))?,
        ))
    }

    /// Convenience function for `SharedMemory::new` that has an arbitrary and unspecified name.
    pub fn anon() -> Result<SharedMemory> {
        Self::new(None)
    }

    /// Creates a new shared memory file descriptor with zero size.
    ///
    /// If a name is given, it will appear in `/proc/self/fd/<shm fd>` for the purposes of
    /// debugging. The name does not need to be unique.
    ///
    /// The file descriptor is opened with the close on exec flag and allows memfd sealing.
    pub fn new(name: Option<&CStr>) -> Result<SharedMemory> {
        let shm_name = name
            .map(|n| n.as_ptr())
            .unwrap_or(b"/crosvm_shm\0".as_ptr() as *const c_char);
        // The following are safe because we give a valid C string and check the
        // results of the memfd_create call.
        let fd = unsafe { memfd_create(shm_name, MFD_CLOEXEC | MFD_ALLOW_SEALING) };
        if fd < 0 {
            return errno_result();
        }

        let file = unsafe { File::from_raw_fd(fd) };

        Ok(SharedMemory { fd: file, size: 0 })
    }

    /// Constructs a `SharedMemory` instance from a `File` that represents shared memory.
    ///
    /// The size of the resulting shared memory will be determined using `File::seek`. If the given
    /// file's size can not be determined this way, this will return an error.
    pub fn from_file(mut file: File) -> Result<SharedMemory> {
        let file_size = file.seek(SeekFrom::End(0))?;
        Ok(SharedMemory {
            fd: file,
            size: file_size as u64,
        })
    }

    /// Gets the memfd seals that have already been added to this.
    ///
    /// This may fail if this instance was not constructed from a memfd.
    pub fn get_seals(&self) -> Result<MemfdSeals> {
        let ret = unsafe { fcntl(self.fd.as_raw_fd(), F_GET_SEALS) };
        if ret < 0 {
            return errno_result();
        }
        Ok(MemfdSeals(ret))
    }

    /// Adds the given set of memfd seals.
    ///
    /// This may fail if this instance was not constructed from a memfd with sealing allowed or if
    /// the seal seal (`F_SEAL_SEAL`) bit was already added.
    pub fn add_seals(&mut self, seals: MemfdSeals) -> Result<()> {
        let ret = unsafe { fcntl(self.fd.as_raw_fd(), F_ADD_SEALS, seals) };
        if ret < 0 {
            return errno_result();
        }
        Ok(())
    }

    /// Gets the size in bytes of the shared memory.
    ///
    /// The size returned here does not reflect changes by other interfaces or users of the shared
    /// memory file descriptor..
    pub fn size(&self) -> u64 {
        self.size
    }

    /// Sets the size in bytes of the shared memory.
    ///
    /// Note that if some process has already mapped this shared memory and the new size is smaller,
    /// that process may get signaled with SIGBUS if they access any page past the new size.
    pub fn set_size(&mut self, size: u64) -> Result<()> {
        let ret = unsafe { ftruncate64(self.fd.as_raw_fd(), size as off64_t) };
        if ret < 0 {
            return errno_result();
        }
        self.size = size;
        Ok(())
    }

    /// Reads the name from the underlying file as a `String`.
    ///
    /// If the underlying file was not created with `SharedMemory::new` or with `memfd_create`, the
    /// results are undefined. Because this returns a `String`, the name's bytes are interpreted as
    /// utf-8.
    pub fn read_name(&self) -> Result<String> {
        let fd_path = format!("/proc/self/fd/{}", self.as_raw_fd());
        let link_name = read_link(fd_path)?;
        link_name
            .to_str()
            .map(|s| {
                s.trim_start_matches("/memfd:")
                    .trim_end_matches(" (deleted)")
                    .to_owned()
            })
            .ok_or_else(|| errno::Error::new(EINVAL))
    }
}

impl Read for SharedMemory {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.fd.read(buf)
    }
}

impl Read for &SharedMemory {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        (&self.fd).read(buf)
    }
}

impl Write for SharedMemory {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.fd.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.fd.flush()
    }
}

impl Write for &SharedMemory {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        (&self.fd).write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        (&self.fd).flush()
    }
}

impl Seek for SharedMemory {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        self.fd.seek(pos)
    }
}

impl Seek for &SharedMemory {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        (&self.fd).seek(pos)
    }
}

impl AsRawFd for SharedMemory {
    fn as_raw_fd(&self) -> RawFd {
        self.fd.as_raw_fd()
    }
}

impl AsRawFd for &SharedMemory {
    fn as_raw_fd(&self) -> RawFd {
        self.fd.as_raw_fd()
    }
}

impl IntoRawFd for SharedMemory {
    fn into_raw_fd(self) -> RawFd {
        self.fd.into_raw_fd()
    }
}

impl From<SharedMemory> for File {
    fn from(s: SharedMemory) -> File {
        s.fd
    }
}

/// Checks if the kernel we are running on has memfd_create. It was introduced in 3.17.
/// Only to be used from tests to prevent running on ancient kernels that won't
/// support the functionality anyways.
pub fn kernel_has_memfd() -> bool {
    unsafe {
        let fd = memfd_create(b"/test_memfd_create\0".as_ptr() as *const c_char, 0);
        if fd < 0 {
            if errno::Error::last().errno() == libc::ENOSYS {
                return false;
            }
            return true;
        }
        close(fd);
    }
    true
}

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

    use std::ffi::CString;

    use data_model::VolatileMemory;

    use crate::MemoryMapping;

    #[test]
    fn named() {
        if !kernel_has_memfd() {
            return;
        }
        const TEST_NAME: &str = "Name McCool Person";
        let shm = SharedMemory::named(TEST_NAME).expect("failed to create shared memory");
        assert_eq!(shm.read_name(), Ok(TEST_NAME.to_owned()));
    }

    #[test]
    fn anon() {
        if !kernel_has_memfd() {
            return;
        }
        SharedMemory::anon().expect("failed to create shared memory");
    }

    #[test]
    fn new() {
        if !kernel_has_memfd() {
            return;
        }
        let shm = SharedMemory::anon().expect("failed to create shared memory");
        assert_eq!(shm.size(), 0);
    }

    #[test]
    fn new_sized() {
        if !kernel_has_memfd() {
            return;
        }
        let mut shm = SharedMemory::anon().expect("failed to create shared memory");
        shm.set_size(1024)
            .expect("failed to set shared memory size");
        assert_eq!(shm.size(), 1024);
    }

    #[test]
    fn new_huge() {
        if !kernel_has_memfd() {
            return;
        }
        let mut shm = SharedMemory::anon().expect("failed to create shared memory");
        shm.set_size(0x7fff_ffff_ffff_ffff)
            .expect("failed to set shared memory size");
        assert_eq!(shm.size(), 0x7fff_ffff_ffff_ffff);
    }

    #[test]
    fn new_too_huge() {
        if !kernel_has_memfd() {
            return;
        }
        let mut shm = SharedMemory::anon().expect("failed to create shared memory");
        shm.set_size(0x8000_0000_0000_0000).unwrap_err();
        assert_eq!(shm.size(), 0);
    }

    #[test]
    fn new_named() {
        if !kernel_has_memfd() {
            return;
        }
        let name = "very unique name";
        let cname = CString::new(name).unwrap();
        let shm = SharedMemory::new(Some(&cname)).expect("failed to create shared memory");
        assert_eq!(shm.read_name(), Ok(name.to_owned()));
    }

    #[test]
    fn new_sealed() {
        if !kernel_has_memfd() {
            return;
        }
        let mut shm = SharedMemory::anon().expect("failed to create shared memory");
        let mut seals = shm.get_seals().expect("failed to get seals");
        assert_eq!(seals.bitmask(), 0);
        seals.set_seal_seal();
        shm.add_seals(seals).expect("failed to add seals");
        seals = shm.get_seals().expect("failed to get seals");
        assert!(seals.seal_seal());
        // Adding more seals should be rejected by the kernel.
        shm.add_seals(seals).unwrap_err();
    }

    #[test]
    fn mmap_page() {
        if !kernel_has_memfd() {
            return;
        }
        let mut shm = SharedMemory::anon().expect("failed to create shared memory");
        shm.set_size(4096)
            .expect("failed to set shared memory size");

        let mmap1 =
            MemoryMapping::from_fd(&shm, shm.size() as usize).expect("failed to map shared memory");
        let mmap2 =
            MemoryMapping::from_fd(&shm, shm.size() as usize).expect("failed to map shared memory");

        assert_ne!(
            mmap1.get_slice(0, 1).unwrap().as_ptr(),
            mmap2.get_slice(0, 1).unwrap().as_ptr()
        );

        mmap1
            .get_slice(0, 4096)
            .expect("failed to get mmap slice")
            .write_bytes(0x45);

        for i in 0..4096 {
            assert_eq!(mmap2.get_ref::<u8>(i).unwrap().load(), 0x45u8);
        }
    }

    #[test]
    fn mmap_page_offset() {
        if !kernel_has_memfd() {
            return;
        }
        let mut shm = SharedMemory::anon().expect("failed to create shared memory");
        shm.set_size(8092)
            .expect("failed to set shared memory size");

        let mmap1 = MemoryMapping::from_fd_offset(&shm, shm.size() as usize, 4096)
            .expect("failed to map shared memory");
        let mmap2 =
            MemoryMapping::from_fd(&shm, shm.size() as usize).expect("failed to map shared memory");

        mmap1
            .get_slice(0, 4096)
            .expect("failed to get mmap slice")
            .write_bytes(0x45);

        for i in 0..4096 {
            assert_eq!(mmap2.get_ref::<u8>(i).unwrap().load(), 0);
        }
        for i in 4096..8092 {
            assert_eq!(mmap2.get_ref::<u8>(i).unwrap().load(), 0x45u8);
        }
    }
}
<commit_msg>sys_util: Add support for F_SEAL_FUTURE_WRITE.<commit_after>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::ffi::{CStr, CString};
use std::fs::{read_link, File};
use std::io::{self, Read, Seek, SeekFrom, Write};
use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};

use libc::{
    self, c_char, c_int, c_long, c_uint, close, fcntl, ftruncate64, off64_t, syscall,
    SYS_memfd_create, EINVAL, F_ADD_SEALS, F_GET_SEALS, F_SEAL_FUTURE_WRITE, F_SEAL_GROW,
    F_SEAL_SEAL, F_SEAL_SHRINK, F_SEAL_WRITE, MFD_ALLOW_SEALING,
};
use serde::{Deserialize, Serialize};

use crate::{errno, errno_result, Result};

/// A shared memory file descriptor and its size.
#[derive(Serialize, Deserialize)]
pub struct SharedMemory {
    #[serde(with = "crate::with_as_descriptor")]
    fd: File,
    size: u64,
}

// from <sys/memfd.h>
const MFD_CLOEXEC: c_uint = 0x0001;

unsafe fn memfd_create(name: *const c_char, flags: c_uint) -> c_int {
    syscall(SYS_memfd_create as c_long, name, flags) as c_int
}

/// A set of memfd seals.
///
/// An enumeration of each bit can be found at `fcntl(2)`.
#[derive(Copy, Clone, Default)]
pub struct MemfdSeals(i32);

impl MemfdSeals {
    /// Returns an empty set of memfd seals.
    #[inline]
    pub fn new() -> MemfdSeals {
        MemfdSeals(0)
    }

    /// Gets the raw bitmask of seals enumerated in `fcntl(2)`.
    #[inline]
    pub fn bitmask(self) -> i32 {
        self.0
    }

    /// True if the grow seal bit is present.
    #[inline]
    pub fn grow_seal(self) -> bool {
        self.0 & F_SEAL_GROW != 0
    }

    /// Sets the grow seal bit.
    #[inline]
    pub fn set_grow_seal(&mut self) {
        self.0 |= F_SEAL_GROW;
    }

    /// True if the shrink seal bit is present.
    #[inline]
    pub fn shrink_seal(self) -> bool {
        self.0 & F_SEAL_SHRINK != 0
    }

    /// Sets the shrink seal bit.
    #[inline]
    pub fn set_shrink_seal(&mut self) {
        self.0 |= F_SEAL_SHRINK;
    }

    /// True if the write seal bit is present.
    #[inline]
    pub fn write_seal(self) -> bool {
        self.0 & F_SEAL_WRITE != 0
    }

    /// Sets the write seal bit.
    #[inline]
    pub fn set_write_seal(&mut self) {
        self.0 |= F_SEAL_WRITE;
    }

    /// True if the future write seal bit is present.
    #[inline]
    pub fn future_write_seal(self) -> bool {
        self.0 & F_SEAL_FUTURE_WRITE != 0
    }

    /// Sets the future write seal bit.
    #[inline]
    pub fn set_future_write_seal(&mut self) {
        self.0 |= F_SEAL_FUTURE_WRITE;
    }

    /// True of the seal seal bit is present.
    #[inline]
    pub fn seal_seal(self) -> bool {
        self.0 & F_SEAL_SEAL != 0
    }

    /// Sets the seal seal bit.
    #[inline]
    pub fn set_seal_seal(&mut self) {
        self.0 |= F_SEAL_SEAL;
    }
}

impl SharedMemory {
    /// Convenience function for `SharedMemory::new` that is always named and accepts a wide variety
    /// of string-like types.
    ///
    /// Note that the given name may not have NUL characters anywhere in it, or this will return an
    /// error.
    pub fn named<T: Into<Vec<u8>>>(name: T) -> Result<SharedMemory> {
        Self::new(Some(
            &CString::new(name).map_err(|_| errno::Error::new(EINVAL))?,
        ))
    }

    /// Convenience function for `SharedMemory::new` that has an arbitrary and unspecified name.
    pub fn anon() -> Result<SharedMemory> {
        Self::new(None)
    }

    /// Creates a new shared memory file descriptor with zero size.
    ///
    /// If a name is given, it will appear in `/proc/self/fd/<shm fd>` for the purposes of
    /// debugging. The name does not need to be unique.
    ///
    /// The file descriptor is opened with the close on exec flag and allows memfd sealing.
    pub fn new(name: Option<&CStr>) -> Result<SharedMemory> {
        let shm_name = name
            .map(|n| n.as_ptr())
            .unwrap_or(b"/crosvm_shm\0".as_ptr() as *const c_char);
        // The following are safe because we give a valid C string and check the
        // results of the memfd_create call.
        let fd = unsafe { memfd_create(shm_name, MFD_CLOEXEC | MFD_ALLOW_SEALING) };
        if fd < 0 {
            return errno_result();
        }

        let file = unsafe { File::from_raw_fd(fd) };

        Ok(SharedMemory { fd: file, size: 0 })
    }

    /// Constructs a `SharedMemory` instance from a `File` that represents shared memory.
    ///
    /// The size of the resulting shared memory will be determined using `File::seek`. If the given
    /// file's size can not be determined this way, this will return an error.
    pub fn from_file(mut file: File) -> Result<SharedMemory> {
        let file_size = file.seek(SeekFrom::End(0))?;
        Ok(SharedMemory {
            fd: file,
            size: file_size as u64,
        })
    }

    /// Gets the memfd seals that have already been added to this.
    ///
    /// This may fail if this instance was not constructed from a memfd.
    pub fn get_seals(&self) -> Result<MemfdSeals> {
        let ret = unsafe { fcntl(self.fd.as_raw_fd(), F_GET_SEALS) };
        if ret < 0 {
            return errno_result();
        }
        Ok(MemfdSeals(ret))
    }

    /// Adds the given set of memfd seals.
    ///
    /// This may fail if this instance was not constructed from a memfd with sealing allowed or if
    /// the seal seal (`F_SEAL_SEAL`) bit was already added.
    pub fn add_seals(&mut self, seals: MemfdSeals) -> Result<()> {
        let ret = unsafe { fcntl(self.fd.as_raw_fd(), F_ADD_SEALS, seals) };
        if ret < 0 {
            return errno_result();
        }
        Ok(())
    }

    /// Gets the size in bytes of the shared memory.
    ///
    /// The size returned here does not reflect changes by other interfaces or users of the shared
    /// memory file descriptor..
    pub fn size(&self) -> u64 {
        self.size
    }

    /// Sets the size in bytes of the shared memory.
    ///
    /// Note that if some process has already mapped this shared memory and the new size is smaller,
    /// that process may get signaled with SIGBUS if they access any page past the new size.
    pub fn set_size(&mut self, size: u64) -> Result<()> {
        let ret = unsafe { ftruncate64(self.fd.as_raw_fd(), size as off64_t) };
        if ret < 0 {
            return errno_result();
        }
        self.size = size;
        Ok(())
    }

    /// Reads the name from the underlying file as a `String`.
    ///
    /// If the underlying file was not created with `SharedMemory::new` or with `memfd_create`, the
    /// results are undefined. Because this returns a `String`, the name's bytes are interpreted as
    /// utf-8.
    pub fn read_name(&self) -> Result<String> {
        let fd_path = format!("/proc/self/fd/{}", self.as_raw_fd());
        let link_name = read_link(fd_path)?;
        link_name
            .to_str()
            .map(|s| {
                s.trim_start_matches("/memfd:")
                    .trim_end_matches(" (deleted)")
                    .to_owned()
            })
            .ok_or_else(|| errno::Error::new(EINVAL))
    }
}

impl Read for SharedMemory {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.fd.read(buf)
    }
}

impl Read for &SharedMemory {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        (&self.fd).read(buf)
    }
}

impl Write for SharedMemory {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.fd.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.fd.flush()
    }
}

impl Write for &SharedMemory {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        (&self.fd).write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        (&self.fd).flush()
    }
}

impl Seek for SharedMemory {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        self.fd.seek(pos)
    }
}

impl Seek for &SharedMemory {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        (&self.fd).seek(pos)
    }
}

impl AsRawFd for SharedMemory {
    fn as_raw_fd(&self) -> RawFd {
        self.fd.as_raw_fd()
    }
}

impl AsRawFd for &SharedMemory {
    fn as_raw_fd(&self) -> RawFd {
        self.fd.as_raw_fd()
    }
}

impl IntoRawFd for SharedMemory {
    fn into_raw_fd(self) -> RawFd {
        self.fd.into_raw_fd()
    }
}

impl From<SharedMemory> for File {
    fn from(s: SharedMemory) -> File {
        s.fd
    }
}

/// Checks if the kernel we are running on has memfd_create. It was introduced in 3.17.
/// Only to be used from tests to prevent running on ancient kernels that won't
/// support the functionality anyways.
pub fn kernel_has_memfd() -> bool {
    unsafe {
        let fd = memfd_create(b"/test_memfd_create\0".as_ptr() as *const c_char, 0);
        if fd < 0 {
            if errno::Error::last().errno() == libc::ENOSYS {
                return false;
            }
            return true;
        }
        close(fd);
    }
    true
}

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

    use std::ffi::CString;

    use data_model::VolatileMemory;

    use crate::MemoryMapping;

    #[test]
    fn named() {
        if !kernel_has_memfd() {
            return;
        }
        const TEST_NAME: &str = "Name McCool Person";
        let shm = SharedMemory::named(TEST_NAME).expect("failed to create shared memory");
        assert_eq!(shm.read_name(), Ok(TEST_NAME.to_owned()));
    }

    #[test]
    fn anon() {
        if !kernel_has_memfd() {
            return;
        }
        SharedMemory::anon().expect("failed to create shared memory");
    }

    #[test]
    fn new() {
        if !kernel_has_memfd() {
            return;
        }
        let shm = SharedMemory::anon().expect("failed to create shared memory");
        assert_eq!(shm.size(), 0);
    }

    #[test]
    fn new_sized() {
        if !kernel_has_memfd() {
            return;
        }
        let mut shm = SharedMemory::anon().expect("failed to create shared memory");
        shm.set_size(1024)
            .expect("failed to set shared memory size");
        assert_eq!(shm.size(), 1024);
    }

    #[test]
    fn new_huge() {
        if !kernel_has_memfd() {
            return;
        }
        let mut shm = SharedMemory::anon().expect("failed to create shared memory");
        shm.set_size(0x7fff_ffff_ffff_ffff)
            .expect("failed to set shared memory size");
        assert_eq!(shm.size(), 0x7fff_ffff_ffff_ffff);
    }

    #[test]
    fn new_too_huge() {
        if !kernel_has_memfd() {
            return;
        }
        let mut shm = SharedMemory::anon().expect("failed to create shared memory");
        shm.set_size(0x8000_0000_0000_0000).unwrap_err();
        assert_eq!(shm.size(), 0);
    }

    #[test]
    fn new_named() {
        if !kernel_has_memfd() {
            return;
        }
        let name = "very unique name";
        let cname = CString::new(name).unwrap();
        let shm = SharedMemory::new(Some(&cname)).expect("failed to create shared memory");
        assert_eq!(shm.read_name(), Ok(name.to_owned()));
    }

    #[test]
    fn new_sealed() {
        if !kernel_has_memfd() {
            return;
        }
        let mut shm = SharedMemory::anon().expect("failed to create shared memory");
        let mut seals = shm.get_seals().expect("failed to get seals");
        assert_eq!(seals.bitmask(), 0);
        seals.set_seal_seal();
        shm.add_seals(seals).expect("failed to add seals");
        seals = shm.get_seals().expect("failed to get seals");
        assert!(seals.seal_seal());
        // Adding more seals should be rejected by the kernel.
        shm.add_seals(seals).unwrap_err();
    }

    #[test]
    fn mmap_page() {
        if !kernel_has_memfd() {
            return;
        }
        let mut shm = SharedMemory::anon().expect("failed to create shared memory");
        shm.set_size(4096)
            .expect("failed to set shared memory size");

        let mmap1 =
            MemoryMapping::from_fd(&shm, shm.size() as usize).expect("failed to map shared memory");
        let mmap2 =
            MemoryMapping::from_fd(&shm, shm.size() as usize).expect("failed to map shared memory");

        assert_ne!(
            mmap1.get_slice(0, 1).unwrap().as_ptr(),
            mmap2.get_slice(0, 1).unwrap().as_ptr()
        );

        mmap1
            .get_slice(0, 4096)
            .expect("failed to get mmap slice")
            .write_bytes(0x45);

        for i in 0..4096 {
            assert_eq!(mmap2.get_ref::<u8>(i).unwrap().load(), 0x45u8);
        }
    }

    #[test]
    fn mmap_page_offset() {
        if !kernel_has_memfd() {
            return;
        }
        let mut shm = SharedMemory::anon().expect("failed to create shared memory");
        shm.set_size(8092)
            .expect("failed to set shared memory size");

        let mmap1 = MemoryMapping::from_fd_offset(&shm, shm.size() as usize, 4096)
            .expect("failed to map shared memory");
        let mmap2 =
            MemoryMapping::from_fd(&shm, shm.size() as usize).expect("failed to map shared memory");

        mmap1
            .get_slice(0, 4096)
            .expect("failed to get mmap slice")
            .write_bytes(0x45);

        for i in 0..4096 {
            assert_eq!(mmap2.get_ref::<u8>(i).unwrap().load(), 0);
        }
        for i in 4096..8092 {
            assert_eq!(mmap2.get_ref::<u8>(i).unwrap().load(), 0x45u8);
        }
    }
}
<|endoftext|>