<commit_before>

extern crate cratesfyi;
extern crate clap;
extern crate log;
extern crate env_logger;
extern crate time;
extern crate rustwide;

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

use clap::{Arg, App, SubCommand};
use cratesfyi::{DocBuilder, RustwideBuilder, DocBuilderOptions, db};
use cratesfyi::utils::add_crate_to_queue;
use cratesfyi::start_web_server;
use cratesfyi::db::{add_path_into_database, connect_db};

pub fn main() {
    logger_init();

    let matches = App::new("cratesfyi")
        .version(cratesfyi::BUILD_VERSION)
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .subcommand(SubCommand::with_name("build")
            .about("Builds documentation in a chroot environment")
            .arg(Arg::with_name("PREFIX")
                .short("P")
                .long("prefix")
                .takes_value(true))
            .arg(Arg::with_name("CRATES_IO_INDEX_PATH")
                .long("crates-io-index-path")
                .help("Sets crates.io-index path")
                .takes_value(true))
            .arg(Arg::with_name("SKIP_IF_EXISTS")
                .short("s")
                .long("skip")
                .help("Skips building documentation if \
                                                      documentation exists"))
            .arg(Arg::with_name("SKIP_IF_LOG_EXISTS")
                .long("skip-if-log-exists")
                .help("Skips building documentation if build \
                                                      log exists"))
            .arg(Arg::with_name("KEEP_BUILD_DIRECTORY")
                .short("-k")
                .long("keep-build-directory")
                .help("Keeps build directory after build."))
            .subcommand(SubCommand::with_name("world").about("Builds documentation of every \
                                                              crate"))
            .subcommand(SubCommand::with_name("crate")
                .about("Builds documentation for a crate")
                .arg(Arg::with_name("CRATE_NAME")
                    .index(1)
                    .required_unless("LOCAL")
                    .requires("CRATE_VERSION")
                    .help("Crate name"))
                .arg(Arg::with_name("CRATE_VERSION")
                    .index(2)
                    .help("Version of crate"))
                .arg(Arg::with_name("LOCAL")
                    .short("-l")
                    .long("--local")
                    .takes_value(true)
                    .conflicts_with_all(&["CRATE_NAME", "CRATE_VERSION"])
                    .help("Build a crate at a specific path")))
            .subcommand(SubCommand::with_name("update-toolchain")
                .about("update the curretntly installed rustup toolchain")
                .arg(Arg::with_name("ONLY_FIRST_TIME")
                     .long("only-first-time")
                     .help("update toolchain only if no toolchain is currently installed")))
            .subcommand(SubCommand::with_name("add-essential-files")
                .about("Adds essential files for the installed version of rustc"))
            .subcommand(SubCommand::with_name("lock").about("Locks cratesfyi daemon to stop \
                                                              building new crates"))
            .subcommand(SubCommand::with_name("unlock")
                .about("Unlocks cratesfyi daemon to continue \
                                                              building new crates"))
            .subcommand(SubCommand::with_name("print-options")))
        .subcommand(SubCommand::with_name("start-web-server")
            .about("Starts web server")
            .arg(Arg::with_name("SOCKET_ADDR")
                .index(1)
                .required(false)
                .help("Socket address to listen to")))
        .subcommand(SubCommand::with_name("daemon").about("Starts cratesfyi daemon")
            .arg(Arg::with_name("FOREGROUND")
                .short("-f")
                .long("foreground")
                .help("run the server in the foreground instead of detaching a child")))
        .subcommand(SubCommand::with_name("database")
            .about("Database operations")
            .subcommand(SubCommand::with_name("move-to-s3"))
            .subcommand(SubCommand::with_name("migrate")
                .about("Run database migrations")
                .arg(Arg::with_name("VERSION")))
            .subcommand(SubCommand::with_name("update-github-fields")
                .about("Updates github stats for crates."))
            .subcommand(SubCommand::with_name("add-directory")
                .about("Adds a directory into database")
                .arg(Arg::with_name("DIRECTORY")
                    .index(1)
                    .required(true)
                    .help("Path of file or \
                                                                      directory"))
                .arg(Arg::with_name("PREFIX")
                    .index(2)
                    .help("Prefix of files in \
                                                                      database")))
            .subcommand(SubCommand::with_name("update-release-activity"))
            .about("Updates montly release activity \
                                                              chart")
            .subcommand(SubCommand::with_name("update-search-index")
            .about("Updates search index"))
            .subcommand(SubCommand::with_name("delete-crate")
                .about("Removes a whole crate from the database")
                .arg(Arg::with_name("CRATE_NAME").help("Name of the crate to delete"))))
        .subcommand(SubCommand::with_name("queue")
            .about("Interactions with the build queue")
            .subcommand(SubCommand::with_name("add")
                .about("Add a crate to the build queue")
                .arg(Arg::with_name("CRATE_NAME")
                    .index(1)
                    .required(true)
                    .help("Name of crate to build"))
                .arg(Arg::with_name("CRATE_VERSION")
                    .index(2)
                    .required(true)
                    .help("Version of crate to build"))
                .arg(Arg::with_name("BUILD_PRIORITY")
                    .short("p")
                    .long("priority")
                    .help("Priority of build (default: 5) (new crate builds get priority 0)")
                    .takes_value(true))))
        .get_matches();

    if let Some(matches) = matches.subcommand_matches("build") {
        let docbuilder_opts = {
            let mut docbuilder_opts = if let Some(prefix) = matches.value_of("PREFIX") {
                DocBuilderOptions::from_prefix(PathBuf::from(prefix))
            } else if let Ok(prefix) = env::var("CRATESFYI_PREFIX") {
                DocBuilderOptions::from_prefix(PathBuf::from(prefix))
            } else {
                DocBuilderOptions::default()
            };

            if let Some(crates_io_index_path) = matches.value_of("CRATES_IO_INDEX_PATH") {
                docbuilder_opts.crates_io_index_path = PathBuf::from(crates_io_index_path);
            }

            docbuilder_opts.skip_if_exists = matches.is_present("SKIP_IF_EXISTS");
            docbuilder_opts.skip_if_log_exists = matches.is_present("SKIP_IF_LOG_EXISTS");
            docbuilder_opts.keep_build_directory = matches.is_present("KEEP_BUILD_DIRECTORY");

            docbuilder_opts.check_paths().unwrap();

            docbuilder_opts
        };

        let mut docbuilder = DocBuilder::new(docbuilder_opts);

        if let Some(_) = matches.subcommand_matches("world") {
            docbuilder.load_cache().expect("Failed to load cache");
            let mut builder = RustwideBuilder::init().unwrap();
            builder.build_world(&mut docbuilder).expect("Failed to build world");
            docbuilder.save_cache().expect("Failed to save cache");
        } else if let Some(matches) = matches.subcommand_matches("crate") {
            docbuilder.load_cache().expect("Failed to load cache");
            let mut builder = RustwideBuilder::init().expect("failed to initialize rustwide");
            match matches.value_of("LOCAL") {
                Some(path) => builder.build_local_package(&mut docbuilder, Path::new(path)),
                None => builder.build_package(&mut docbuilder,
                    matches.value_of("CRATE_NAME").unwrap(), matches.value_of("CRATE_VERSION").unwrap(), None),
            }.expect("Building documentation failed");
            docbuilder.save_cache().expect("Failed to save cache");
        } else if let Some(m) = matches.subcommand_matches("update-toolchain") {
            if m.is_present("ONLY_FIRST_TIME") {
                let conn = db::connect_db().unwrap();
                let res = conn.query("SELECT * FROM config WHERE name = 'rustc_version';", &[]).unwrap();
                if !res.is_empty() {
                    println!("update-toolchain was already called in the past, exiting");
                    return;
                }
            }
            let mut builder = RustwideBuilder::init().unwrap();
            builder.update_toolchain().expect("failed to update toolchain");
        } else if matches.subcommand_matches("add-essential-files").is_some() {
            let mut builder = RustwideBuilder::init().unwrap();
            builder.add_essential_files().expect("failed to add essential files");
        } else if let Some(_) = matches.subcommand_matches("lock") {
            docbuilder.lock().expect("Failed to lock");
        } else if let Some(_) = matches.subcommand_matches("unlock") {
            docbuilder.unlock().expect("Failed to unlock");
        } else if let Some(_) = matches.subcommand_matches("print-options") {
            println!("{:?}", docbuilder.options());
        }

    } else if let Some(matches) = matches.subcommand_matches("database") {
        if let Some(matches) = matches.subcommand_matches("migrate") {
            let version = matches.value_of("VERSION").map(|v| v.parse::<i64>()
                                                          .expect("Version should be an integer"));
            db::migrate(version, &connect_db().expect("failed to connect to the database"))
                .expect("Failed to run database migrations");
        } else if let Some(_) = matches.subcommand_matches("update-github-fields") {
            cratesfyi::utils::github_updater().expect("Failed to update github fields");
        } else if let Some(matches) = matches.subcommand_matches("add-directory") {
            add_path_into_database(&db::connect_db().unwrap(),
                                   matches.value_of("PREFIX").unwrap_or(""),
                                   matches.value_of("DIRECTORY").unwrap())
                .expect("Failed to add directory into database");
        } else if let Some(_) = matches.subcommand_matches("update-release-activity") {
            // FIXME: This is actually util command not database
            cratesfyi::utils::update_release_activity().expect("Failed to update release activity");
        } else if let Some(_) = matches.subcommand_matches("update-search-index") {
            let conn = db::connect_db().unwrap();
            db::update_search_index(&conn).expect("Failed to update search index");
        } else if let Some(_) = matches.subcommand_matches("move-to-s3") {
            let conn = db::connect_db().unwrap();
            let mut count = 1;
            let mut total = 0;
            while count != 0 {
                count = db::file::move_to_s3(&conn, 5_000).expect("Failed to upload batch to S3");
                total += count;
                eprintln!(
                    "moved {} rows to s3 in this batch, total moved so far: {}",
                    count, total
                );
            }
        } else if let Some(matches) = matches.subcommand_matches("delete-crate") {
            let name = matches.value_of("CRATE_NAME").expect("missing crate name");
            let conn = db::connect_db().expect("failed to connect to the database");
            db::delete_crate(&conn, &name).expect("failed to delete the crate");
        }
    } else if let Some(matches) = matches.subcommand_matches("start-web-server") {
        start_web_server(Some(matches.value_of("SOCKET_ADDR").unwrap_or("0.0.0.0:3000")));
    } else if let Some(_) = matches.subcommand_matches("daemon") {
        let foreground = matches.subcommand_matches("daemon")
            .map_or(false, |opts| opts.is_present("FOREGROUND"));
        cratesfyi::utils::start_daemon(!foreground);
    } else if let Some(matches) = matches.subcommand_matches("queue") {
        if let Some(matches) = matches.subcommand_matches("add") {
            let priority = matches.value_of("BUILD_PRIORITY").unwrap_or("5");
            let priority: i32 = priority.parse().expect("--priority was not a number");
            let conn = connect_db().expect("Could not connect to database");

            add_crate_to_queue(&conn,
                               matches.value_of("CRATE_NAME").unwrap(),
                               matches.value_of("CRATE_VERSION").unwrap(),
                               priority).expect("Could not add crate to queue");
        }
    } else {
        println!("{}", matches.usage());
    }
}

fn logger_init() {
    use std::io::Write;

    let mut builder = env_logger::Builder::new();
    builder.format(|buf, record| {
        writeln!(buf, "{} [{}] {}: {}",
                time::now().strftime("%Y/%m/%d %H:%M:%S").unwrap(),
                record.level(),
                record.target(),
                record.args())
    });
    builder.parse(&env::var("RUST_LOG").unwrap_or("cratesfyi=info".to_owned()));

    rustwide::logging::init_with(builder.build());
}
<commit_msg>Update src/bin/cratesfyi.rs<commit_after>

extern crate cratesfyi;
extern crate clap;
extern crate log;
extern crate env_logger;
extern crate time;
extern crate rustwide;

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

use clap::{Arg, App, SubCommand};
use cratesfyi::{DocBuilder, RustwideBuilder, DocBuilderOptions, db};
use cratesfyi::utils::add_crate_to_queue;
use cratesfyi::start_web_server;
use cratesfyi::db::{add_path_into_database, connect_db};

pub fn main() {
    logger_init();

    let matches = App::new("cratesfyi")
        .version(cratesfyi::BUILD_VERSION)
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .subcommand(SubCommand::with_name("build")
            .about("Builds documentation in a chroot environment")
            .arg(Arg::with_name("PREFIX")
                .short("P")
                .long("prefix")
                .takes_value(true))
            .arg(Arg::with_name("CRATES_IO_INDEX_PATH")
                .long("crates-io-index-path")
                .help("Sets crates.io-index path")
                .takes_value(true))
            .arg(Arg::with_name("SKIP_IF_EXISTS")
                .short("s")
                .long("skip")
                .help("Skips building documentation if \
                                                      documentation exists"))
            .arg(Arg::with_name("SKIP_IF_LOG_EXISTS")
                .long("skip-if-log-exists")
                .help("Skips building documentation if build \
                                                      log exists"))
            .arg(Arg::with_name("KEEP_BUILD_DIRECTORY")
                .short("-k")
                .long("keep-build-directory")
                .help("Keeps build directory after build."))
            .subcommand(SubCommand::with_name("world").about("Builds documentation of every \
                                                              crate"))
            .subcommand(SubCommand::with_name("crate")
                .about("Builds documentation for a crate")
                .arg(Arg::with_name("CRATE_NAME")
                    .index(1)
                    .required_unless("LOCAL")
                    .requires("CRATE_VERSION")
                    .help("Crate name"))
                .arg(Arg::with_name("CRATE_VERSION")
                    .index(2)
                    .help("Version of crate"))
                .arg(Arg::with_name("LOCAL")
                    .short("-l")
                    .long("--local")
                    .takes_value(true)
                    .conflicts_with_all(&["CRATE_NAME", "CRATE_VERSION"])
                    .help("Build a crate at a specific path")))
            .subcommand(SubCommand::with_name("update-toolchain")
                .about("update the curretntly installed rustup toolchain")
                .arg(Arg::with_name("ONLY_FIRST_TIME")
                     .long("only-first-time")
                     .help("update toolchain only if no toolchain is currently installed")))
            .subcommand(SubCommand::with_name("add-essential-files")
                .about("Adds essential files for the installed version of rustc"))
            .subcommand(SubCommand::with_name("lock").about("Locks cratesfyi daemon to stop \
                                                              building new crates"))
            .subcommand(SubCommand::with_name("unlock")
                .about("Unlocks cratesfyi daemon to continue \
                                                              building new crates"))
            .subcommand(SubCommand::with_name("print-options")))
        .subcommand(SubCommand::with_name("start-web-server")
            .about("Starts web server")
            .arg(Arg::with_name("SOCKET_ADDR")
                .index(1)
                .required(false)
                .help("Socket address to listen to")))
        .subcommand(SubCommand::with_name("daemon").about("Starts cratesfyi daemon")
            .arg(Arg::with_name("FOREGROUND")
                .short("-f")
                .long("foreground")
                .help("run the server in the foreground instead of detaching a child")))
        .subcommand(SubCommand::with_name("database")
            .about("Database operations")
            .subcommand(SubCommand::with_name("move-to-s3"))
            .subcommand(SubCommand::with_name("migrate")
                .about("Run database migrations")
                .arg(Arg::with_name("VERSION")))
            .subcommand(SubCommand::with_name("update-github-fields")
                .about("Updates github stats for crates."))
            .subcommand(SubCommand::with_name("add-directory")
                .about("Adds a directory into database")
                .arg(Arg::with_name("DIRECTORY")
                    .index(1)
                    .required(true)
                    .help("Path of file or \
                                                                      directory"))
                .arg(Arg::with_name("PREFIX")
                    .index(2)
                    .help("Prefix of files in \
                                                                      database")))
            .subcommand(SubCommand::with_name("update-release-activity"))
            .about("Updates montly release activity \
                                                              chart")
            .subcommand(SubCommand::with_name("update-search-index")
            .about("Updates search index"))
            .subcommand(SubCommand::with_name("delete-crate")
                .about("Removes a whole crate from the database")
                .arg(Arg::with_name("CRATE_NAME")
                    .takes_value(true)
                    .help("Name of the crate to delete"))))
        .subcommand(SubCommand::with_name("queue")
            .about("Interactions with the build queue")
            .subcommand(SubCommand::with_name("add")
                .about("Add a crate to the build queue")
                .arg(Arg::with_name("CRATE_NAME")
                    .index(1)
                    .required(true)
                    .help("Name of crate to build"))
                .arg(Arg::with_name("CRATE_VERSION")
                    .index(2)
                    .required(true)
                    .help("Version of crate to build"))
                .arg(Arg::with_name("BUILD_PRIORITY")
                    .short("p")
                    .long("priority")
                    .help("Priority of build (default: 5) (new crate builds get priority 0)")
                    .takes_value(true))))
        .get_matches();

    if let Some(matches) = matches.subcommand_matches("build") {
        let docbuilder_opts = {
            let mut docbuilder_opts = if let Some(prefix) = matches.value_of("PREFIX") {
                DocBuilderOptions::from_prefix(PathBuf::from(prefix))
            } else if let Ok(prefix) = env::var("CRATESFYI_PREFIX") {
                DocBuilderOptions::from_prefix(PathBuf::from(prefix))
            } else {
                DocBuilderOptions::default()
            };

            if let Some(crates_io_index_path) = matches.value_of("CRATES_IO_INDEX_PATH") {
                docbuilder_opts.crates_io_index_path = PathBuf::from(crates_io_index_path);
            }

            docbuilder_opts.skip_if_exists = matches.is_present("SKIP_IF_EXISTS");
            docbuilder_opts.skip_if_log_exists = matches.is_present("SKIP_IF_LOG_EXISTS");
            docbuilder_opts.keep_build_directory = matches.is_present("KEEP_BUILD_DIRECTORY");

            docbuilder_opts.check_paths().unwrap();

            docbuilder_opts
        };

        let mut docbuilder = DocBuilder::new(docbuilder_opts);

        if let Some(_) = matches.subcommand_matches("world") {
            docbuilder.load_cache().expect("Failed to load cache");
            let mut builder = RustwideBuilder::init().unwrap();
            builder.build_world(&mut docbuilder).expect("Failed to build world");
            docbuilder.save_cache().expect("Failed to save cache");
        } else if let Some(matches) = matches.subcommand_matches("crate") {
            docbuilder.load_cache().expect("Failed to load cache");
            let mut builder = RustwideBuilder::init().expect("failed to initialize rustwide");
            match matches.value_of("LOCAL") {
                Some(path) => builder.build_local_package(&mut docbuilder, Path::new(path)),
                None => builder.build_package(&mut docbuilder,
                    matches.value_of("CRATE_NAME").unwrap(), matches.value_of("CRATE_VERSION").unwrap(), None),
            }.expect("Building documentation failed");
            docbuilder.save_cache().expect("Failed to save cache");
        } else if let Some(m) = matches.subcommand_matches("update-toolchain") {
            if m.is_present("ONLY_FIRST_TIME") {
                let conn = db::connect_db().unwrap();
                let res = conn.query("SELECT * FROM config WHERE name = 'rustc_version';", &[]).unwrap();
                if !res.is_empty() {
                    println!("update-toolchain was already called in the past, exiting");
                    return;
                }
            }
            let mut builder = RustwideBuilder::init().unwrap();
            builder.update_toolchain().expect("failed to update toolchain");
        } else if matches.subcommand_matches("add-essential-files").is_some() {
            let mut builder = RustwideBuilder::init().unwrap();
            builder.add_essential_files().expect("failed to add essential files");
        } else if let Some(_) = matches.subcommand_matches("lock") {
            docbuilder.lock().expect("Failed to lock");
        } else if let Some(_) = matches.subcommand_matches("unlock") {
            docbuilder.unlock().expect("Failed to unlock");
        } else if let Some(_) = matches.subcommand_matches("print-options") {
            println!("{:?}", docbuilder.options());
        }

    } else if let Some(matches) = matches.subcommand_matches("database") {
        if let Some(matches) = matches.subcommand_matches("migrate") {
            let version = matches.value_of("VERSION").map(|v| v.parse::<i64>()
                                                          .expect("Version should be an integer"));
            db::migrate(version, &connect_db().expect("failed to connect to the database"))
                .expect("Failed to run database migrations");
        } else if let Some(_) = matches.subcommand_matches("update-github-fields") {
            cratesfyi::utils::github_updater().expect("Failed to update github fields");
        } else if let Some(matches) = matches.subcommand_matches("add-directory") {
            add_path_into_database(&db::connect_db().unwrap(),
                                   matches.value_of("PREFIX").unwrap_or(""),
                                   matches.value_of("DIRECTORY").unwrap())
                .expect("Failed to add directory into database");
        } else if let Some(_) = matches.subcommand_matches("update-release-activity") {
            // FIXME: This is actually util command not database
            cratesfyi::utils::update_release_activity().expect("Failed to update release activity");
        } else if let Some(_) = matches.subcommand_matches("update-search-index") {
            let conn = db::connect_db().unwrap();
            db::update_search_index(&conn).expect("Failed to update search index");
        } else if let Some(_) = matches.subcommand_matches("move-to-s3") {
            let conn = db::connect_db().unwrap();
            let mut count = 1;
            let mut total = 0;
            while count != 0 {
                count = db::file::move_to_s3(&conn, 5_000).expect("Failed to upload batch to S3");
                total += count;
                eprintln!(
                    "moved {} rows to s3 in this batch, total moved so far: {}",
                    count, total
                );
            }
        } else if let Some(matches) = matches.subcommand_matches("delete-crate") {
            let name = matches.value_of("CRATE_NAME").expect("missing crate name");
            let conn = db::connect_db().expect("failed to connect to the database");
            db::delete_crate(&conn, &name).expect("failed to delete the crate");
        }
    } else if let Some(matches) = matches.subcommand_matches("start-web-server") {
        start_web_server(Some(matches.value_of("SOCKET_ADDR").unwrap_or("0.0.0.0:3000")));
    } else if let Some(_) = matches.subcommand_matches("daemon") {
        let foreground = matches.subcommand_matches("daemon")
            .map_or(false, |opts| opts.is_present("FOREGROUND"));
        cratesfyi::utils::start_daemon(!foreground);
    } else if let Some(matches) = matches.subcommand_matches("queue") {
        if let Some(matches) = matches.subcommand_matches("add") {
            let priority = matches.value_of("BUILD_PRIORITY").unwrap_or("5");
            let priority: i32 = priority.parse().expect("--priority was not a number");
            let conn = connect_db().expect("Could not connect to database");

            add_crate_to_queue(&conn,
                               matches.value_of("CRATE_NAME").unwrap(),
                               matches.value_of("CRATE_VERSION").unwrap(),
                               priority).expect("Could not add crate to queue");
        }
    } else {
        println!("{}", matches.usage());
    }
}

fn logger_init() {
    use std::io::Write;

    let mut builder = env_logger::Builder::new();
    builder.format(|buf, record| {
        writeln!(buf, "{} [{}] {}: {}",
                time::now().strftime("%Y/%m/%d %H:%M:%S").unwrap(),
                record.level(),
                record.target(),
                record.args())
    });
    builder.parse(&env::var("RUST_LOG").unwrap_or("cratesfyi=info".to_owned()));

    rustwide::logging::init_with(builder.build());
}
<|endoftext|>


<commit_before>// run-pass
#![feature(const_discriminant)]

use std::mem::{discriminant, Discriminant};

fn identity<T>(x: T) -> T { x }

enum Test {
    A(u8),
    B,
    C { a: u8, b: u8 },
}

const TEST_A: Discriminant<Test> = discriminant(&Test::A(5));
const TEST_A_OTHER: Discriminant<Test> = discriminant(&Test::A(17));
const TEST_B: Discriminant<Test> = discriminant(&Test::B);

fn main() {
    assert_eq!(TEST_A, TEST_A_OTHER);
    assert_eq!(TEST_A, discriminant(identity(&Test::A(17))));
    assert_eq!(TEST_B, discriminant(identity(&Test::B)));
    assert_ne!(TEST_A, TEST_B);
    assert_ne!(TEST_B, discriminant(identity(&Test::C { a: 42, b: 7 })));
}
<commit_msg>test discriminant of enum with uninhabited variant<commit_after>// run-pass
#![feature(const_discriminant)]

use std::mem::{discriminant, Discriminant};

fn identity<T>(x: T) -> T { x }

enum Test {
    A(u8),
    B,
    C { a: u8, b: u8 },
}

const TEST_A: Discriminant<Test> = discriminant(&Test::A(5));
const TEST_A_OTHER: Discriminant<Test> = discriminant(&Test::A(17));
const TEST_B: Discriminant<Test> = discriminant(&Test::B);

enum Void {}

enum SingleVariant {
    V,
    Never(Void),
}

const TEST_V: Discriminant<SingleVariant> = discriminant(&SingleVariant::V);

fn main() {
    assert_eq!(TEST_A, TEST_A_OTHER);
    assert_eq!(TEST_A, discriminant(identity(&Test::A(17))));
    assert_eq!(TEST_B, discriminant(identity(&Test::B)));
    assert_ne!(TEST_A, TEST_B);
    assert_ne!(TEST_B, discriminant(identity(&Test::C { a: 42, b: 7 })));

    assert_eq!(TEST_V, discriminant(identity(&SingleVariant::V)));
}
<|endoftext|>


<commit_before>use byteorder::{ByteOrder, NetworkEndian, ReadBytesExt, WriteBytesExt};
use log::*;
use std::borrow::Cow;
use std::default::Default;
use std::fmt;
use std::io::{Cursor, ErrorKind, Read, Write};
use std::result::Result as StdResult;
use std::string::{FromUtf8Error, String};

use super::coding::{CloseCode, Control, Data, OpCode};
use super::mask::{apply_mask, generate_mask};
use crate::error::{Error, Result};

/// A struct representing the close command.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct CloseFrame<'t> {
    /// The reason as a code.
    pub code: CloseCode,
    /// The reason as text string.
    pub reason: Cow<'t, str>,
}

impl<'t> CloseFrame<'t> {
    /// Convert into a owned string.
    pub fn into_owned(self) -> CloseFrame<'static> {
        CloseFrame {
            code: self.code,
            reason: self.reason.into_owned().into(),
        }
    }
}

impl<'t> fmt::Display for CloseFrame<'t> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} ({})", self.reason, self.code)
    }
}

/// A struct representing a WebSocket frame header.
#[allow(missing_copy_implementations)]
#[derive(Debug, Clone)]
pub struct FrameHeader {
    /// Indicates that the frame is the last one of a possibly fragmented message.
    pub is_final: bool,
    /// Reserved for protocol extensions.
    pub rsv1: bool,
    /// Reserved for protocol extensions.
    pub rsv2: bool,
    /// Reserved for protocol extensions.
    pub rsv3: bool,
    /// WebSocket protocol opcode.
    pub opcode: OpCode,
    /// A frame mask, if any.
    mask: Option<[u8; 4]>,
}

impl Default for FrameHeader {
    fn default() -> Self {
        FrameHeader {
            is_final: true,
            rsv1: false,
            rsv2: false,
            rsv3: false,
            opcode: OpCode::Control(Control::Close),
            mask: None,
        }
    }
}

impl FrameHeader {
    /// Parse a header from an input stream.
    /// Returns `None` if insufficient data and does not consume anything in this case.
    /// Payload size is returned along with the header.
    pub fn parse(cursor: &mut Cursor<impl AsRef<[u8]>>) -> Result<Option<(Self, u64)>> {
        let initial = cursor.position();
        match Self::parse_internal(cursor) {
            ret @ Ok(None) => {
                cursor.set_position(initial);
                ret
            }
            ret => ret,
        }
    }

    /// Get the size of the header formatted with given payload length.
    pub fn len(&self, length: u64) -> usize {
        2 + LengthFormat::for_length(length).extra_bytes() + if self.mask.is_some() { 4 } else { 0 }
    }

    /// Format a header for given payload size.
    pub fn format(&self, length: u64, output: &mut impl Write) -> Result<()> {
        let code: u8 = self.opcode.into();

        let one = {
            code | if self.is_final { 0x80 } else { 0 }
                | if self.rsv1 { 0x40 } else { 0 }
                | if self.rsv2 { 0x20 } else { 0 }
                | if self.rsv3 { 0x10 } else { 0 }
        };

        let lenfmt = LengthFormat::for_length(length);

        let two = { lenfmt.length_byte() | if self.mask.is_some() { 0x80 } else { 0 } };

        output.write_all(&[one, two])?;
        match lenfmt {
            LengthFormat::U8(_) => (),
            LengthFormat::U16 => output.write_u16::<NetworkEndian>(length as u16)?,
            LengthFormat::U64 => output.write_u64::<NetworkEndian>(length)?,
        }

        if let Some(ref mask) = self.mask {
            output.write_all(mask)?
        }

        Ok(())
    }

    /// Generate a random frame mask and store this in the header.
    ///
    /// Of course this does not change frame contents. It just generates a mask.
    pub(crate) fn set_random_mask(&mut self) {
        self.mask = Some(generate_mask())
    }
}

impl FrameHeader {
    /// Internal parse engine.
    /// Returns `None` if insufficient data.
    /// Payload size is returned along with the header.
    fn parse_internal(cursor: &mut impl Read) -> Result<Option<(Self, u64)>> {
        let (first, second) = {
            let mut head = [0u8; 2];
            if cursor.read(&mut head)? != 2 {
                return Ok(None);
            }
            trace!("Parsed headers {:?}", head);
            (head[0], head[1])
        };

        trace!("First: {:b}", first);
        trace!("Second: {:b}", second);

        let is_final = first & 0x80 != 0;

        let rsv1 = first & 0x40 != 0;
        let rsv2 = first & 0x20 != 0;
        let rsv3 = first & 0x10 != 0;

        let opcode = OpCode::from(first & 0x0F);
        trace!("Opcode: {:?}", opcode);

        let masked = second & 0x80 != 0;
        trace!("Masked: {:?}", masked);

        let length = {
            let length_byte = second & 0x7F;
            let length_length = LengthFormat::for_byte(length_byte).extra_bytes();
            if length_length > 0 {
                match cursor.read_uint::<NetworkEndian>(length_length) {
                    Err(ref err) if err.kind() == ErrorKind::UnexpectedEof => {
                        return Ok(None);
                    }
                    Err(err) => {
                        return Err(err.into());
                    }
                    Ok(read) => read,
                }
            } else {
                u64::from(length_byte)
            }
        };

        let mask = if masked {
            let mut mask_bytes = [0u8; 4];
            if cursor.read(&mut mask_bytes)? != 4 {
                return Ok(None);
            } else {
                Some(mask_bytes)
            }
        } else {
            None
        };

        // Disallow bad opcode
        match opcode {
            OpCode::Control(Control::Reserved(_)) | OpCode::Data(Data::Reserved(_)) => {
                return Err(Error::Protocol(
                    format!("Encountered invalid opcode: {}", first & 0x0F).into(),
                ))
            }
            _ => (),
        }

        let hdr = FrameHeader {
            is_final,
            rsv1,
            rsv2,
            rsv3,
            opcode,
            mask,
        };

        Ok(Some((hdr, length)))
    }
}

/// A struct representing a WebSocket frame.
#[derive(Debug, Clone)]
pub struct Frame {
    header: FrameHeader,
    payload: Vec<u8>,
}

impl Frame {
    /// Get the length of the frame.
    /// This is the length of the header + the length of the payload.
    #[inline]
    pub fn len(&self) -> usize {
        let length = self.payload.len();
        self.header.len(length as u64) + length
    }

    /// Check if the frame is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Get a reference to the frame's header.
    #[inline]
    pub fn header(&self) -> &FrameHeader {
        &self.header
    }

    /// Get a mutable reference to the frame's header.
    #[inline]
    pub fn header_mut(&mut self) -> &mut FrameHeader {
        &mut self.header
    }

    /// Get a reference to the frame's payload.
    #[inline]
    pub fn payload(&self) -> &Vec<u8> {
        &self.payload
    }

    /// Get a mutable reference to the frame's payload.
    #[inline]
    pub fn payload_mut(&mut self) -> &mut Vec<u8> {
        &mut self.payload
    }

    /// Test whether the frame is masked.
    #[inline]
    pub(crate) fn is_masked(&self) -> bool {
        self.header.mask.is_some()
    }

    /// Generate a random mask for the frame.
    ///
    /// This just generates a mask, payload is not changed. The actual masking is performed
    /// either on `format()` or on `apply_mask()` call.
    #[inline]
    pub(crate) fn set_random_mask(&mut self) {
        self.header.set_random_mask()
    }

    /// This method unmasks the payload and should only be called on frames that are actually
    /// masked. In other words, those frames that have just been received from a client endpoint.
    #[inline]
    pub(crate) fn apply_mask(&mut self) {
        if let Some(mask) = self.header.mask.take() {
            apply_mask(&mut self.payload, mask)
        }
    }

    /// Consume the frame into its payload as binary.
    #[inline]
    pub fn into_data(self) -> Vec<u8> {
        self.payload
    }

    /// Consume the frame into its payload as string.
    #[inline]
    pub fn into_string(self) -> StdResult<String, FromUtf8Error> {
        String::from_utf8(self.payload)
    }

    /// Consume the frame into a closing frame.
    #[inline]
    pub(crate) fn into_close(self) -> Result<Option<CloseFrame<'static>>> {
        match self.payload.len() {
            0 => Ok(None),
            1 => Err(Error::Protocol("Invalid close sequence".into())),
            _ => {
                let mut data = self.payload;
                let code = NetworkEndian::read_u16(&data[0..2]).into();
                data.drain(0..2);
                let text = String::from_utf8(data)?;
                Ok(Some(CloseFrame {
                    code,
                    reason: text.into(),
                }))
            }
        }
    }

    /// Create a new data frame.
    #[inline]
    pub fn message(data: Vec<u8>, opcode: OpCode, is_final: bool) -> Frame {
        debug_assert!(
            match opcode {
                OpCode::Data(_) => true,
                _ => false,
            },
            "Invalid opcode for data frame."
        );

        Frame {
            header: FrameHeader {
                is_final,
                opcode,
                ..FrameHeader::default()
            },
            payload: data,
        }
    }

    /// Create a new Pong control frame.
    #[inline]
    pub fn pong(data: Vec<u8>) -> Frame {
        Frame {
            header: FrameHeader {
                opcode: OpCode::Control(Control::Pong),
                ..FrameHeader::default()
            },
            payload: data,
        }
    }

    /// Create a new Ping control frame.
    #[inline]
    pub fn ping(data: Vec<u8>) -> Frame {
        Frame {
            header: FrameHeader {
                opcode: OpCode::Control(Control::Ping),
                ..FrameHeader::default()
            },
            payload: data,
        }
    }

    /// Create a new Close control frame.
    #[inline]
    pub fn close(msg: Option<CloseFrame>) -> Frame {
        let payload = if let Some(CloseFrame { code, reason }) = msg {
            let mut p = Vec::with_capacity(reason.as_bytes().len() + 2);
            p.write_u16::<NetworkEndian>(code.into()).unwrap(); // can't fail
            p.extend_from_slice(reason.as_bytes());
            p
        } else {
            Vec::new()
        };

        Frame {
            header: FrameHeader::default(),
            payload,
        }
    }

    /// Create a frame from given header and data.
    pub fn from_payload(header: FrameHeader, payload: Vec<u8>) -> Self {
        Frame { header, payload }
    }

    /// Write a frame out to a buffer
    pub fn format(mut self, output: &mut impl Write) -> Result<()> {
        self.header.format(self.payload.len() as u64, output)?;
        self.apply_mask();
        output.write_all(self.payload())?;
        Ok(())
    }
}

impl fmt::Display for Frame {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "
<FRAME>
final: {}
reserved: {} {} {}
opcode: {}
length: {}
payload length: {}
payload: 0x{}
            ",
            self.header.is_final,
            self.header.rsv1,
            self.header.rsv2,
            self.header.rsv3,
            self.header.opcode,
            // self.mask.map(|mask| format!("{:?}", mask)).unwrap_or("NONE".into()),
            self.len(),
            self.payload.len(),
            self.payload
                .iter()
                .map(|byte| format!("{:x}", byte))
                .collect::<String>()
        )
    }
}

/// Handling of the length format.
enum LengthFormat {
    U8(u8),
    U16,
    U64,
}

impl LengthFormat {
    /// Get length format for a given data size.
    #[inline]
    fn for_length(length: u64) -> Self {
        if length < 126 {
            LengthFormat::U8(length as u8)
        } else if length < 65536 {
            LengthFormat::U16
        } else {
            LengthFormat::U64
        }
    }

    /// Get the size of length encoding.
    #[inline]
    fn extra_bytes(&self) -> usize {
        match *self {
            LengthFormat::U8(_) => 0,
            LengthFormat::U16 => 2,
            LengthFormat::U64 => 8,
        }
    }

    /// Encode the givem length.
    #[inline]
    fn length_byte(&self) -> u8 {
        match *self {
            LengthFormat::U8(b) => b,
            LengthFormat::U16 => 126,
            LengthFormat::U64 => 127,
        }
    }

    /// Get length format for a given length byte.
    #[inline]
    fn for_byte(byte: u8) -> Self {
        match byte & 0x7F {
            126 => LengthFormat::U16,
            127 => LengthFormat::U64,
            b => LengthFormat::U8(b),
        }
    }
}

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

    use super::super::coding::{Data, OpCode};
    use std::io::Cursor;

    #[test]
    fn parse() {
        let mut raw: Cursor<Vec<u8>> =
            Cursor::new(vec![0x82, 0x07, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]);
        let (header, length) = FrameHeader::parse(&mut raw).unwrap().unwrap();
        assert_eq!(length, 7);
        let mut payload = Vec::new();
        raw.read_to_end(&mut payload).unwrap();
        let frame = Frame::from_payload(header, payload);
        assert_eq!(
            frame.into_data(),
            vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]
        );
    }

    #[test]
    fn format() {
        let frame = Frame::ping(vec![0x01, 0x02]);
        let mut buf = Vec::with_capacity(frame.len());
        frame.format(&mut buf).unwrap();
        assert_eq!(buf, vec![0x89, 0x02, 0x01, 0x02]);
    }

    #[test]
    fn display() {
        let f = Frame::message("hi there".into(), OpCode::Data(Data::Text), true);
        let view = format!("{}", f);
        assert!(view.contains("payload:"));
    }
}
<commit_msg>Make mask public on FrameHeader<commit_after>use byteorder::{ByteOrder, NetworkEndian, ReadBytesExt, WriteBytesExt};
use log::*;
use std::borrow::Cow;
use std::default::Default;
use std::fmt;
use std::io::{Cursor, ErrorKind, Read, Write};
use std::result::Result as StdResult;
use std::string::{FromUtf8Error, String};

use super::coding::{CloseCode, Control, Data, OpCode};
use super::mask::{apply_mask, generate_mask};
use crate::error::{Error, Result};

/// A struct representing the close command.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct CloseFrame<'t> {
    /// The reason as a code.
    pub code: CloseCode,
    /// The reason as text string.
    pub reason: Cow<'t, str>,
}

impl<'t> CloseFrame<'t> {
    /// Convert into a owned string.
    pub fn into_owned(self) -> CloseFrame<'static> {
        CloseFrame {
            code: self.code,
            reason: self.reason.into_owned().into(),
        }
    }
}

impl<'t> fmt::Display for CloseFrame<'t> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} ({})", self.reason, self.code)
    }
}

/// A struct representing a WebSocket frame header.
#[allow(missing_copy_implementations)]
#[derive(Debug, Clone)]
pub struct FrameHeader {
    /// Indicates that the frame is the last one of a possibly fragmented message.
    pub is_final: bool,
    /// Reserved for protocol extensions.
    pub rsv1: bool,
    /// Reserved for protocol extensions.
    pub rsv2: bool,
    /// Reserved for protocol extensions.
    pub rsv3: bool,
    /// WebSocket protocol opcode.
    pub opcode: OpCode,
    /// A frame mask, if any.
    pub mask: Option<[u8; 4]>,
}

impl Default for FrameHeader {
    fn default() -> Self {
        FrameHeader {
            is_final: true,
            rsv1: false,
            rsv2: false,
            rsv3: false,
            opcode: OpCode::Control(Control::Close),
            mask: None,
        }
    }
}

impl FrameHeader {
    /// Parse a header from an input stream.
    /// Returns `None` if insufficient data and does not consume anything in this case.
    /// Payload size is returned along with the header.
    pub fn parse(cursor: &mut Cursor<impl AsRef<[u8]>>) -> Result<Option<(Self, u64)>> {
        let initial = cursor.position();
        match Self::parse_internal(cursor) {
            ret @ Ok(None) => {
                cursor.set_position(initial);
                ret
            }
            ret => ret,
        }
    }

    /// Get the size of the header formatted with given payload length.
    pub fn len(&self, length: u64) -> usize {
        2 + LengthFormat::for_length(length).extra_bytes() + if self.mask.is_some() { 4 } else { 0 }
    }

    /// Format a header for given payload size.
    pub fn format(&self, length: u64, output: &mut impl Write) -> Result<()> {
        let code: u8 = self.opcode.into();

        let one = {
            code | if self.is_final { 0x80 } else { 0 }
                | if self.rsv1 { 0x40 } else { 0 }
                | if self.rsv2 { 0x20 } else { 0 }
                | if self.rsv3 { 0x10 } else { 0 }
        };

        let lenfmt = LengthFormat::for_length(length);

        let two = { lenfmt.length_byte() | if self.mask.is_some() { 0x80 } else { 0 } };

        output.write_all(&[one, two])?;
        match lenfmt {
            LengthFormat::U8(_) => (),
            LengthFormat::U16 => output.write_u16::<NetworkEndian>(length as u16)?,
            LengthFormat::U64 => output.write_u64::<NetworkEndian>(length)?,
        }

        if let Some(ref mask) = self.mask {
            output.write_all(mask)?
        }

        Ok(())
    }

    /// Generate a random frame mask and store this in the header.
    ///
    /// Of course this does not change frame contents. It just generates a mask.
    pub(crate) fn set_random_mask(&mut self) {
        self.mask = Some(generate_mask())
    }
}

impl FrameHeader {
    /// Internal parse engine.
    /// Returns `None` if insufficient data.
    /// Payload size is returned along with the header.
    fn parse_internal(cursor: &mut impl Read) -> Result<Option<(Self, u64)>> {
        let (first, second) = {
            let mut head = [0u8; 2];
            if cursor.read(&mut head)? != 2 {
                return Ok(None);
            }
            trace!("Parsed headers {:?}", head);
            (head[0], head[1])
        };

        trace!("First: {:b}", first);
        trace!("Second: {:b}", second);

        let is_final = first & 0x80 != 0;

        let rsv1 = first & 0x40 != 0;
        let rsv2 = first & 0x20 != 0;
        let rsv3 = first & 0x10 != 0;

        let opcode = OpCode::from(first & 0x0F);
        trace!("Opcode: {:?}", opcode);

        let masked = second & 0x80 != 0;
        trace!("Masked: {:?}", masked);

        let length = {
            let length_byte = second & 0x7F;
            let length_length = LengthFormat::for_byte(length_byte).extra_bytes();
            if length_length > 0 {
                match cursor.read_uint::<NetworkEndian>(length_length) {
                    Err(ref err) if err.kind() == ErrorKind::UnexpectedEof => {
                        return Ok(None);
                    }
                    Err(err) => {
                        return Err(err.into());
                    }
                    Ok(read) => read,
                }
            } else {
                u64::from(length_byte)
            }
        };

        let mask = if masked {
            let mut mask_bytes = [0u8; 4];
            if cursor.read(&mut mask_bytes)? != 4 {
                return Ok(None);
            } else {
                Some(mask_bytes)
            }
        } else {
            None
        };

        // Disallow bad opcode
        match opcode {
            OpCode::Control(Control::Reserved(_)) | OpCode::Data(Data::Reserved(_)) => {
                return Err(Error::Protocol(
                    format!("Encountered invalid opcode: {}", first & 0x0F).into(),
                ))
            }
            _ => (),
        }

        let hdr = FrameHeader {
            is_final,
            rsv1,
            rsv2,
            rsv3,
            opcode,
            mask,
        };

        Ok(Some((hdr, length)))
    }
}

/// A struct representing a WebSocket frame.
#[derive(Debug, Clone)]
pub struct Frame {
    header: FrameHeader,
    payload: Vec<u8>,
}

impl Frame {
    /// Get the length of the frame.
    /// This is the length of the header + the length of the payload.
    #[inline]
    pub fn len(&self) -> usize {
        let length = self.payload.len();
        self.header.len(length as u64) + length
    }

    /// Check if the frame is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Get a reference to the frame's header.
    #[inline]
    pub fn header(&self) -> &FrameHeader {
        &self.header
    }

    /// Get a mutable reference to the frame's header.
    #[inline]
    pub fn header_mut(&mut self) -> &mut FrameHeader {
        &mut self.header
    }

    /// Get a reference to the frame's payload.
    #[inline]
    pub fn payload(&self) -> &Vec<u8> {
        &self.payload
    }

    /// Get a mutable reference to the frame's payload.
    #[inline]
    pub fn payload_mut(&mut self) -> &mut Vec<u8> {
        &mut self.payload
    }

    /// Test whether the frame is masked.
    #[inline]
    pub(crate) fn is_masked(&self) -> bool {
        self.header.mask.is_some()
    }

    /// Generate a random mask for the frame.
    ///
    /// This just generates a mask, payload is not changed. The actual masking is performed
    /// either on `format()` or on `apply_mask()` call.
    #[inline]
    pub(crate) fn set_random_mask(&mut self) {
        self.header.set_random_mask()
    }

    /// This method unmasks the payload and should only be called on frames that are actually
    /// masked. In other words, those frames that have just been received from a client endpoint.
    #[inline]
    pub(crate) fn apply_mask(&mut self) {
        if let Some(mask) = self.header.mask.take() {
            apply_mask(&mut self.payload, mask)
        }
    }

    /// Consume the frame into its payload as binary.
    #[inline]
    pub fn into_data(self) -> Vec<u8> {
        self.payload
    }

    /// Consume the frame into its payload as string.
    #[inline]
    pub fn into_string(self) -> StdResult<String, FromUtf8Error> {
        String::from_utf8(self.payload)
    }

    /// Consume the frame into a closing frame.
    #[inline]
    pub(crate) fn into_close(self) -> Result<Option<CloseFrame<'static>>> {
        match self.payload.len() {
            0 => Ok(None),
            1 => Err(Error::Protocol("Invalid close sequence".into())),
            _ => {
                let mut data = self.payload;
                let code = NetworkEndian::read_u16(&data[0..2]).into();
                data.drain(0..2);
                let text = String::from_utf8(data)?;
                Ok(Some(CloseFrame {
                    code,
                    reason: text.into(),
                }))
            }
        }
    }

    /// Create a new data frame.
    #[inline]
    pub fn message(data: Vec<u8>, opcode: OpCode, is_final: bool) -> Frame {
        debug_assert!(
            match opcode {
                OpCode::Data(_) => true,
                _ => false,
            },
            "Invalid opcode for data frame."
        );

        Frame {
            header: FrameHeader {
                is_final,
                opcode,
                ..FrameHeader::default()
            },
            payload: data,
        }
    }

    /// Create a new Pong control frame.
    #[inline]
    pub fn pong(data: Vec<u8>) -> Frame {
        Frame {
            header: FrameHeader {
                opcode: OpCode::Control(Control::Pong),
                ..FrameHeader::default()
            },
            payload: data,
        }
    }

    /// Create a new Ping control frame.
    #[inline]
    pub fn ping(data: Vec<u8>) -> Frame {
        Frame {
            header: FrameHeader {
                opcode: OpCode::Control(Control::Ping),
                ..FrameHeader::default()
            },
            payload: data,
        }
    }

    /// Create a new Close control frame.
    #[inline]
    pub fn close(msg: Option<CloseFrame>) -> Frame {
        let payload = if let Some(CloseFrame { code, reason }) = msg {
            let mut p = Vec::with_capacity(reason.as_bytes().len() + 2);
            p.write_u16::<NetworkEndian>(code.into()).unwrap(); // can't fail
            p.extend_from_slice(reason.as_bytes());
            p
        } else {
            Vec::new()
        };

        Frame {
            header: FrameHeader::default(),
            payload,
        }
    }

    /// Create a frame from given header and data.
    pub fn from_payload(header: FrameHeader, payload: Vec<u8>) -> Self {
        Frame { header, payload }
    }

    /// Write a frame out to a buffer
    pub fn format(mut self, output: &mut impl Write) -> Result<()> {
        self.header.format(self.payload.len() as u64, output)?;
        self.apply_mask();
        output.write_all(self.payload())?;
        Ok(())
    }
}

impl fmt::Display for Frame {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "
<FRAME>
final: {}
reserved: {} {} {}
opcode: {}
length: {}
payload length: {}
payload: 0x{}
            ",
            self.header.is_final,
            self.header.rsv1,
            self.header.rsv2,
            self.header.rsv3,
            self.header.opcode,
            // self.mask.map(|mask| format!("{:?}", mask)).unwrap_or("NONE".into()),
            self.len(),
            self.payload.len(),
            self.payload
                .iter()
                .map(|byte| format!("{:x}", byte))
                .collect::<String>()
        )
    }
}

/// Handling of the length format.
enum LengthFormat {
    U8(u8),
    U16,
    U64,
}

impl LengthFormat {
    /// Get length format for a given data size.
    #[inline]
    fn for_length(length: u64) -> Self {
        if length < 126 {
            LengthFormat::U8(length as u8)
        } else if length < 65536 {
            LengthFormat::U16
        } else {
            LengthFormat::U64
        }
    }

    /// Get the size of length encoding.
    #[inline]
    fn extra_bytes(&self) -> usize {
        match *self {
            LengthFormat::U8(_) => 0,
            LengthFormat::U16 => 2,
            LengthFormat::U64 => 8,
        }
    }

    /// Encode the givem length.
    #[inline]
    fn length_byte(&self) -> u8 {
        match *self {
            LengthFormat::U8(b) => b,
            LengthFormat::U16 => 126,
            LengthFormat::U64 => 127,
        }
    }

    /// Get length format for a given length byte.
    #[inline]
    fn for_byte(byte: u8) -> Self {
        match byte & 0x7F {
            126 => LengthFormat::U16,
            127 => LengthFormat::U64,
            b => LengthFormat::U8(b),
        }
    }
}

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

    use super::super::coding::{Data, OpCode};
    use std::io::Cursor;

    #[test]
    fn parse() {
        let mut raw: Cursor<Vec<u8>> =
            Cursor::new(vec![0x82, 0x07, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]);
        let (header, length) = FrameHeader::parse(&mut raw).unwrap().unwrap();
        assert_eq!(length, 7);
        let mut payload = Vec::new();
        raw.read_to_end(&mut payload).unwrap();
        let frame = Frame::from_payload(header, payload);
        assert_eq!(
            frame.into_data(),
            vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]
        );
    }

    #[test]
    fn format() {
        let frame = Frame::ping(vec![0x01, 0x02]);
        let mut buf = Vec::with_capacity(frame.len());
        frame.format(&mut buf).unwrap();
        assert_eq!(buf, vec![0x89, 0x02, 0x01, 0x02]);
    }

    #[test]
    fn display() {
        let f = Frame::message("hi there".into(), OpCode::Data(Data::Text), true);
        let view = format!("{}", f);
        assert!(view.contains("payload:"));
    }
}
<|endoftext|>


<commit_before>import std._io.stdio_reader;

fn in_range(char c, char lo, char hi) -> bool {
    ret lo <= c && c <= hi;
}

fn is_alpha(char c) -> bool {
    ret in_range(c, 'a', 'z') ||
        in_range(c, 'A', 'Z');
}

fn is_dec_digit(char c) -> bool {
    ret in_range(c, '0', '9');
}

fn is_hex_digit(char c) -> bool {
    ret in_range(c, '0', '9') ||
        in_range(c, 'a', 'f') ||
        in_range(c, 'A', 'F');
}

fn is_bin_digit(char c) -> bool {
    ret c == '0' || c == '1';
}

fn is_whitespace(char c) -> bool {
    ret c == ' ' || c == '\t' || c == '\r' || c == '\n';
}

fn consume_any_whitespace(stdio_reader rdr, char c) -> char {
    auto c1 = c;
    while (is_whitespace(c1)) {
        c1 = rdr.getc() as char;
    }
    be consume_any_line_comment(rdr, c1);
}

fn consume_any_line_comment(stdio_reader rdr, char c) -> char {
    auto c1 = c;
    if (c1 == '/') {
        auto c2 = rdr.getc() as char;
        if (c2 == '/') {
            while (c1 != '\n') {
                c1 = rdr.getc() as char;
            }
            // Restart whitespace munch.
            be consume_any_whitespace(rdr, c1);
        }
    }
    ret c;
}

fn next_token(stdio_reader rdr) -> token.token {
    auto eof = (-1) as char;
    auto c = rdr.getc() as char;
    auto accum_str = "";
    auto accum_int = 0;

    c = consume_any_whitespace(rdr, c);

    if (c == eof) { ret token.EOF(); }

    if (is_alpha(c)) {
        while (is_alpha(c)) {
            accum_str += (c as u8);
            c = rdr.getc() as char;
        }
        rdr.ungetc(c as int);
        ret token.IDENT(accum_str);
    }

    if (is_dec_digit(c)) {
        if (c == '0') {
        } else {
            while (is_dec_digit(c)) {
                accum_int *= 10;
                accum_int += (c as int) - ('0' as int);
                c = rdr.getc() as char;
            }
            rdr.ungetc(c as int);
            ret token.LIT_INT(accum_int);
        }
    }

    fn op_or_opeq(stdio_reader rdr, char c2,
                  token.op op) -> token.token {
        if (c2 == '=') {
            ret token.OPEQ(op);
        } else {
            rdr.ungetc(c2 as int);
            ret token.OP(op);
        }
    }

    alt (c) {
        // One-byte tokens.
        case (';') { ret token.SEMI(); }
        case (',') { ret token.COMMA(); }
        case ('.') { ret token.DOT(); }
        case ('(') { ret token.LPAREN(); }
        case (')') { ret token.RPAREN(); }
        case ('{') { ret token.LBRACE(); }
        case ('}') { ret token.RBRACE(); }
        case ('[') { ret token.LBRACKET(); }
        case (']') { ret token.RBRACKET(); }
        case ('@') { ret token.AT(); }
        case ('#') { ret token.POUND(); }

        // Multi-byte tokens.
        case ('=') {
            auto c2 = rdr.getc() as char;
            if (c2 == '=') {
                ret token.OP(token.EQEQ());
            } else {
                rdr.ungetc(c2 as int);
                ret token.OP(token.EQ());
            }
        }

        case ('-') {
            auto c2 = rdr.getc() as char;
            if (c2 == '>') {
                ret token.RARROW();
            } else {
                ret op_or_opeq(rdr, c2, token.MINUS());
            }
        }

        case ('&') {
            auto c2 = rdr.getc() as char;
            if (c2 == '&') {
                ret token.OP(token.ANDAND());
            } else {
                ret op_or_opeq(rdr, c2, token.AND());
            }
        }

        case ('+') {
            ret op_or_opeq(rdr, rdr.getc() as char, token.PLUS());
        }

        case ('*') {
            ret op_or_opeq(rdr, rdr.getc() as char, token.STAR());
        }

        case ('/') {
            ret op_or_opeq(rdr, rdr.getc() as char, token.STAR());
        }

        case ('!') {
            ret op_or_opeq(rdr, rdr.getc() as char, token.NOT());
        }

        case ('^') {
            ret op_or_opeq(rdr, rdr.getc() as char, token.CARET());
        }

        case ('%') {
            ret op_or_opeq(rdr, rdr.getc() as char, token.PERCENT());
        }

    }

    log "lexer stopping at ";
    log c;
    ret token.EOF();
}

//
// Local Variables:
// mode: rust
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C ../.. 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
//
<commit_msg>Tidy lexer.<commit_after>import std._io.stdio_reader;

fn in_range(char c, char lo, char hi) -> bool {
    ret lo <= c && c <= hi;
}

fn is_alpha(char c) -> bool {
    ret in_range(c, 'a', 'z') ||
        in_range(c, 'A', 'Z');
}

fn is_dec_digit(char c) -> bool {
    ret in_range(c, '0', '9');
}

fn is_hex_digit(char c) -> bool {
    ret in_range(c, '0', '9') ||
        in_range(c, 'a', 'f') ||
        in_range(c, 'A', 'F');
}

fn is_bin_digit(char c) -> bool {
    ret c == '0' || c == '1';
}

fn is_whitespace(char c) -> bool {
    ret c == ' ' || c == '\t' || c == '\r' || c == '\n';
}

fn consume_any_whitespace(stdio_reader rdr, char c) -> char {
    auto c1 = c;
    while (is_whitespace(c1)) {
        c1 = rdr.getc() as char;
    }
    be consume_any_line_comment(rdr, c1);
}

fn consume_any_line_comment(stdio_reader rdr, char c) -> char {
    auto c1 = c;
    if (c1 == '/') {
        auto c2 = rdr.getc() as char;
        if (c2 == '/') {
            while (c1 != '\n') {
                c1 = rdr.getc() as char;
            }
            // Restart whitespace munch.
            be consume_any_whitespace(rdr, c1);
        }
    }
    ret c;
}

fn next_token(stdio_reader rdr) -> token.token {
    auto eof = (-1) as char;
    auto c = rdr.getc() as char;
    auto accum_str = "";
    auto accum_int = 0;

    fn next(stdio_reader rdr) -> char {
        ret rdr.getc() as char;
    }

    fn forget(stdio_reader rdr, char c) {
        rdr.ungetc(c as int);
    }

    c = consume_any_whitespace(rdr, c);

    if (c == eof) { ret token.EOF(); }

    if (is_alpha(c)) {
        while (is_alpha(c)) {
            accum_str += (c as u8);
            c = next(rdr);
        }
        forget(rdr, c);
        ret token.IDENT(accum_str);
    }

    if (is_dec_digit(c)) {
        if (c == '0') {
        } else {
            while (is_dec_digit(c)) {
                accum_int *= 10;
                accum_int += (c as int) - ('0' as int);
                c = next(rdr);
            }
            forget(rdr, c);
            ret token.LIT_INT(accum_int);
        }
    }

    fn op_or_opeq(stdio_reader rdr, char c2,
                  token.op op) -> token.token {
        if (c2 == '=') {
            ret token.OPEQ(op);
        } else {
            forget(rdr, c2);
            ret token.OP(op);
        }
    }

    alt (c) {
        // One-byte tokens.
        case (';') { ret token.SEMI(); }
        case (',') { ret token.COMMA(); }
        case ('.') { ret token.DOT(); }
        case ('(') { ret token.LPAREN(); }
        case (')') { ret token.RPAREN(); }
        case ('{') { ret token.LBRACE(); }
        case ('}') { ret token.RBRACE(); }
        case ('[') { ret token.LBRACKET(); }
        case (']') { ret token.RBRACKET(); }
        case ('@') { ret token.AT(); }
        case ('#') { ret token.POUND(); }

        // Multi-byte tokens.
        case ('=') {
            auto c2 = next(rdr);
            if (c2 == '=') {
                ret token.OP(token.EQEQ());
            } else {
                forget(rdr, c2);
                ret token.OP(token.EQ());
            }
        }

        case ('-') {
            auto c2 = next(rdr);
            if (c2 == '>') {
                ret token.RARROW();
            } else {
                ret op_or_opeq(rdr, c2, token.MINUS());
            }
        }

        case ('&') {
            auto c2 = next(rdr);
            if (c2 == '&') {
                ret token.OP(token.ANDAND());
            } else {
                ret op_or_opeq(rdr, c2, token.AND());
            }
        }

        case ('+') {
            ret op_or_opeq(rdr, next(rdr), token.PLUS());
        }

        case ('*') {
            ret op_or_opeq(rdr, next(rdr), token.STAR());
        }

        case ('/') {
            ret op_or_opeq(rdr, next(rdr), token.STAR());
        }

        case ('!') {
            ret op_or_opeq(rdr, next(rdr), token.NOT());
        }

        case ('^') {
            ret op_or_opeq(rdr, next(rdr), token.CARET());
        }

        case ('%') {
            ret op_or_opeq(rdr, next(rdr), token.PERCENT());
        }

    }

    log "lexer stopping at ";
    log c;
    ret token.EOF();
}

//
// Local Variables:
// mode: rust
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C ../.. 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
//
<|endoftext|>


<commit_before>use std::slice::SliceExt;

#[derive(Show)]
pub struct Configuration {
    pub choices: Vec<String>,
    pub visible_limit: usize,
    pub initial_search: String,
}

impl Configuration {
    pub fn from_inputs(choice: Vec<String>,
                       initial_search: Option<String>,
                       visible_limit: Option<usize>) -> Configuration {

        let cleaned = choice.iter().map(clean as fn(&String) -> String).collect::<Vec<String>>();
        let limit = visible_limit.unwrap_or(choice.len());
        let search = initial_search.unwrap_or("".to_string());

        Configuration { choices: cleaned,
                        initial_search: search,
                        visible_limit: limit }
    }

    fn parse_options(input: Vec<String>) -> Option<String> {
        match input.position_elem(&"-s".to_string()) {
            Some(ref idx) => Some(input[*idx + 1].clone()),
            None => None,
        }
    }
}

fn clean(input: &String) -> String {
    let slice = input.as_slice();
    slice.trim().to_string()
}

#[cfg(test)]

#[test]
fn removes_leading_and_trailing_whitespace() {
    let input = vec![" a choice ".to_string()];
    let config = Configuration::from_inputs(input, None, None);

    assert_eq!(config.choices, vec!("a choice"));
}

#[test]
fn can_specify_initial_search() {
    let input = vec!["foo".to_string()];
    let options = vec!["-s".to_string(),
                       "some search".to_string()];
    let config = Configuration::from_inputs(input, Configuration::parse_options(options), None);

    assert_eq!(config.initial_search.as_slice(), "some search");
}

#[test]
fn initial_search_is_optional() {
    let input = vec!["foo".to_string()];
    let options = vec![];
    let config = Configuration::from_inputs(input, Configuration::parse_options(options), None);

    assert_eq!(config.initial_search, "");
}
<commit_msg>Remove method that was just setting up tests<commit_after>use std::slice::SliceExt;

#[derive(Show)]
pub struct Configuration {
    pub choices: Vec<String>,
    pub visible_limit: usize,
    pub initial_search: String,
}

impl Configuration {
    pub fn from_inputs(choice: Vec<String>,
                       initial_search: Option<String>,
                       visible_limit: Option<usize>) -> Configuration {

        let cleaned = choice.iter().map(clean as fn(&String) -> String).collect::<Vec<String>>();
        let limit = visible_limit.unwrap_or(choice.len());
        let search = initial_search.unwrap_or("".to_string());

        Configuration { choices: cleaned,
                        initial_search: search,
                        visible_limit: limit }
    }
}

fn clean(input: &String) -> String {
    let slice = input.as_slice();
    slice.trim().to_string()
}

#[cfg(test)]

#[test]
fn removes_leading_and_trailing_whitespace() {
    let input = vec![" a choice ".to_string()];
    let config = Configuration::from_inputs(input, None, None);

    assert_eq!(config.choices, vec!("a choice"));
}

#[test]
fn can_specify_initial_search() {
    let input = vec!["foo".to_string()];
    let config = Configuration::from_inputs(input, Some("some search".to_string()), None);

    assert_eq!(config.initial_search.as_slice(), "some search");
}

#[test]
fn initial_search_is_optional() {
    let input = vec!["foo".to_string()];
    let config = Configuration::from_inputs(input, None, None);

    assert_eq!(config.initial_search, "");
}
<|endoftext|>


<commit_before>// Copyright 2021 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::io;
use std::os::unix::io::AsRawFd;
use std::os::unix::net::UnixListener;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use anyhow::{anyhow, bail, Context};
use argh::FromArgs;
use base::{error, get_max_open_files, warn, Event, RawDescriptor, Tube, UnlinkUnixListener};
use cros_async::{EventAsync, Executor};
use data_model::{DataInit, Le32};
use fuse::Server;
use futures::future::{AbortHandle, Abortable};
use hypervisor::ProtectionType;
use minijail::{self, Minijail};
use once_cell::sync::OnceCell;
use sync::Mutex;
use vm_memory::GuestMemory;
use vmm_vhost::message::{VhostUserProtocolFeatures, VhostUserVirtioFeatures};

use crate::virtio;
use crate::virtio::copy_config;
use crate::virtio::fs::passthrough::PassthroughFs;
use crate::virtio::fs::{process_fs_queue, virtio_fs_config, FS_MAX_TAG_LEN};
use crate::virtio::vhost::user::device::handler::{
    DeviceRequestHandler, Doorbell, VhostUserBackend,
};

static FS_EXECUTOR: OnceCell<Executor> = OnceCell::new();

async fn handle_fs_queue(
    mut queue: virtio::Queue,
    mem: GuestMemory,
    doorbell: Arc<Mutex<Doorbell>>,
    kick_evt: EventAsync,
    server: Arc<fuse::Server<PassthroughFs>>,
    tube: Arc<Mutex<Tube>>,
) {
    // Slot is always going to be 0 because we do not support DAX
    let slot: u32 = 0;

    loop {
        if let Err(e) = kick_evt.next_val().await {
            error!("Failed to read kick event for fs queue: {}", e);
            break;
        }
        if let Err(e) = process_fs_queue(&mem, &doorbell, &mut queue, &server, &tube, slot) {
            error!("Process FS queue failed: {}", e);
            break;
        }
    }
}

fn default_uidmap() -> String {
    let euid = unsafe { libc::geteuid() };
    format!("{} {} 1", euid, euid)
}

fn default_gidmap() -> String {
    let egid = unsafe { libc::getegid() };
    format!("{} {} 1", egid, egid)
}

fn jail_and_fork(
    mut keep_rds: Vec<RawDescriptor>,
    dir_path: PathBuf,
    uid_map: Option<String>,
    gid_map: Option<String>,
) -> anyhow::Result<i32> {
    // Create new minijail sandbox
    let mut j = Minijail::new()?;

    j.namespace_pids();
    j.namespace_user();
    j.namespace_user_disable_setgroups();
    j.uidmap(&uid_map.unwrap_or_else(default_uidmap))?;
    j.gidmap(&gid_map.unwrap_or_else(default_gidmap))?;
    j.run_as_init();

    j.namespace_vfs();
    j.namespace_net();
    j.no_new_privs();

    // Only pivot_root if we are not re-using the current root directory.
    if dir_path != Path::new("/") {
        // It's safe to call `namespace_vfs` multiple times.
        j.namespace_vfs();
        j.enter_pivot_root(&dir_path)?;
    }
    j.set_remount_mode(libc::MS_SLAVE);

    let limit = get_max_open_files().context("failed to get max open files")?;
    j.set_rlimit(libc::RLIMIT_NOFILE as i32, limit, limit)?;

    // Make sure there are no duplicates in keep_rds
    keep_rds.dedup();

    // fork on the jail here
    let pid = unsafe { j.fork(Some(&keep_rds))? };

    if pid > 0 {
        unsafe { libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM) };
    }

    if pid < 0 {
        bail!("Fork error! {}", std::io::Error::last_os_error());
    }

    Ok(pid)
}

struct FsBackend {
    server: Arc<fuse::Server<PassthroughFs>>,
    tag: [u8; FS_MAX_TAG_LEN],
    avail_features: u64,
    acked_features: u64,
    acked_protocol_features: VhostUserProtocolFeatures,
    workers: [Option<AbortHandle>; Self::MAX_QUEUE_NUM],
    keep_rds: Vec<RawDescriptor>,
}

impl FsBackend {
    pub fn new(tag: &str) -> anyhow::Result<Self> {
        if tag.len() > FS_MAX_TAG_LEN {
            bail!(
                "fs tag is too long: {} (max supported: {})",
                tag.len(),
                FS_MAX_TAG_LEN
            );
        }
        let mut fs_tag = [0u8; FS_MAX_TAG_LEN];
        fs_tag[..tag.len()].copy_from_slice(tag.as_bytes());

        let avail_features = virtio::base_features(ProtectionType::Unprotected)
            | VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits();

        // Use default passthroughfs config
        let fs = PassthroughFs::new(Default::default())?;

        let mut keep_rds: Vec<RawDescriptor> = [0, 1, 2].to_vec();
        keep_rds.append(&mut fs.keep_rds());

        let server = Arc::new(Server::new(fs));

        Ok(FsBackend {
            server,
            tag: fs_tag,
            avail_features,
            acked_features: 0,
            acked_protocol_features: VhostUserProtocolFeatures::empty(),
            workers: Default::default(),
            keep_rds,
        })
    }
}

impl VhostUserBackend for FsBackend {
    const MAX_QUEUE_NUM: usize = 2; /* worker queue and high priority queue */
    const MAX_VRING_LEN: u16 = 1024;

    type Error = anyhow::Error;

    fn features(&self) -> u64 {
        self.avail_features
    }

    fn ack_features(&mut self, value: u64) -> anyhow::Result<()> {
        let unrequested_features = value & !self.avail_features;
        if unrequested_features != 0 {
            bail!("invalid features are given: {:#x}", unrequested_features);
        }

        self.acked_features |= value;

        Ok(())
    }

    fn acked_features(&self) -> u64 {
        self.acked_features
    }

    fn protocol_features(&self) -> VhostUserProtocolFeatures {
        VhostUserProtocolFeatures::CONFIG | VhostUserProtocolFeatures::MQ
    }

    fn ack_protocol_features(&mut self, features: u64) -> anyhow::Result<()> {
        let features = VhostUserProtocolFeatures::from_bits(features)
            .ok_or_else(|| anyhow!("invalid protocol features are given: {:#x}", features))?;
        let supported = self.protocol_features();
        self.acked_protocol_features = features & supported;
        Ok(())
    }

    fn acked_protocol_features(&self) -> u64 {
        self.acked_protocol_features.bits()
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        let config = virtio_fs_config {
            tag: self.tag,
            num_request_queues: Le32::from(1),
        };
        copy_config(data, 0, config.as_slice(), offset);
    }

    fn reset(&mut self) {
        for handle in self.workers.iter_mut().filter_map(Option::take) {
            handle.abort();
        }
    }

    fn start_queue(
        &mut self,
        idx: usize,
        mut queue: virtio::Queue,
        mem: GuestMemory,
        doorbell: Arc<Mutex<Doorbell>>,
        kick_evt: Event,
    ) -> anyhow::Result<()> {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            warn!("Starting new queue handler without stopping old handler");
            handle.abort();
        }

        // Safe because the executor is initialized in main() below.
        let ex = FS_EXECUTOR.get().expect("Executor not initialized");

        // Enable any virtqueue features that were negotiated (like VIRTIO_RING_F_EVENT_IDX).
        queue.ack_features(self.acked_features);

        let kick_evt =
            EventAsync::new(kick_evt.0, ex).context("failed to create EventAsync for kick_evt")?;
        let (handle, registration) = AbortHandle::new_pair();
        let (_, fs_device_tube) = Tube::pair()?;

        ex.spawn_local(Abortable::new(
            handle_fs_queue(
                queue,
                mem,
                doorbell,
                kick_evt,
                self.server.clone(),
                Arc::new(Mutex::new(fs_device_tube)),
            ),
            registration,
        ))
        .detach();

        self.workers[idx] = Some(handle);
        Ok(())
    }

    fn stop_queue(&mut self, idx: usize) {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            handle.abort();
        }
    }
}

#[derive(FromArgs)]
#[argh(description = "")]
struct Options {
    #[argh(option, description = "path to a socket", arg_name = "PATH")]
    socket: String,
    #[argh(option, description = "the virtio-fs tag", arg_name = "TAG")]
    tag: String,
    #[argh(option, description = "path to a directory to share", arg_name = "DIR")]
    shared_dir: PathBuf,
    #[argh(option, description = "uid map to use", arg_name = "UIDMAP")]
    uid_map: Option<String>,
    #[argh(option, description = "gid map to use", arg_name = "GIDMAP")]
    gid_map: Option<String>,
}

/// Starts a vhost-user fs device.
/// Returns an error if the given `args` is invalid or the device fails to run.
pub fn run_fs_device(program_name: &str, args: &[&str]) -> anyhow::Result<()> {
    let opts = match Options::from_args(&[program_name], args) {
        Ok(opts) => opts,
        Err(e) => {
            if e.status.is_err() {
                bail!(e.output);
            } else {
                println!("{}", e.output);
            }
            return Ok(());
        }
    };

    base::syslog::init().context("Failed to initialize syslog")?;

    let fs_device = FsBackend::new(&opts.tag)?;

    // Create and bind unix socket
    let listener = UnixListener::bind(opts.socket).map(UnlinkUnixListener)?;
    let mut keep_rds = fs_device.keep_rds.clone();
    keep_rds.push(listener.as_raw_fd());
    base::syslog::push_descriptors(&mut keep_rds);

    let handler = DeviceRequestHandler::new(fs_device);

    let pid = jail_and_fork(keep_rds, opts.shared_dir, opts.uid_map, opts.gid_map)?;

    // Parent, nothing to do but wait and then exit
    if pid != 0 {
        unsafe { libc::waitpid(pid, std::ptr::null_mut(), 0) };
        return Ok(());
    }

    // We need to set the no setuid fixup secure bit so that we don't drop capabilities when
    // changing the thread uid/gid. Without this, creating new entries can fail in some corner
    // cases.
    const SECBIT_NO_SETUID_FIXUP: i32 = 1 << 2;
    // TODO(crbug.com/1199487): Remove this once libc provides the wrapper for all targets.
    #[cfg(target_os = "linux")]
    {
        // Safe because this doesn't modify any memory and we check the return value.
        let mut securebits = unsafe { libc::prctl(libc::PR_GET_SECUREBITS) };
        if securebits < 0 {
            bail!(io::Error::last_os_error());
        }
        securebits |= SECBIT_NO_SETUID_FIXUP;
        // Safe because this doesn't modify any memory and we check the return value.
        let ret = unsafe { libc::prctl(libc::PR_SET_SECUREBITS, securebits) };
        if ret < 0 {
            bail!(io::Error::last_os_error());
        }
    }

    // Child, we can continue by spawning the executor and set up the device
    let ex = Executor::new().context("Failed to create executor")?;

    let _ = FS_EXECUTOR.set(ex.clone());

    // run_until() returns an Result<Result<..>> which the ? operator lets us flatten.
    ex.run_until(handler.run_with_listener(listener, &ex))?
}
<commit_msg>devices: vhost-user: fs: remove global executor variable<commit_after>// Copyright 2021 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::io;
use std::os::unix::io::AsRawFd;
use std::os::unix::net::UnixListener;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use anyhow::{anyhow, bail, Context};
use argh::FromArgs;
use base::{error, get_max_open_files, warn, Event, RawDescriptor, Tube, UnlinkUnixListener};
use cros_async::{EventAsync, Executor};
use data_model::{DataInit, Le32};
use fuse::Server;
use futures::future::{AbortHandle, Abortable};
use hypervisor::ProtectionType;
use minijail::{self, Minijail};
use sync::Mutex;
use vm_memory::GuestMemory;
use vmm_vhost::message::{VhostUserProtocolFeatures, VhostUserVirtioFeatures};

use crate::virtio;
use crate::virtio::copy_config;
use crate::virtio::fs::passthrough::PassthroughFs;
use crate::virtio::fs::{process_fs_queue, virtio_fs_config, FS_MAX_TAG_LEN};
use crate::virtio::vhost::user::device::handler::{
    DeviceRequestHandler, Doorbell, VhostUserBackend,
};

async fn handle_fs_queue(
    mut queue: virtio::Queue,
    mem: GuestMemory,
    doorbell: Arc<Mutex<Doorbell>>,
    kick_evt: EventAsync,
    server: Arc<fuse::Server<PassthroughFs>>,
    tube: Arc<Mutex<Tube>>,
) {
    // Slot is always going to be 0 because we do not support DAX
    let slot: u32 = 0;

    loop {
        if let Err(e) = kick_evt.next_val().await {
            error!("Failed to read kick event for fs queue: {}", e);
            break;
        }
        if let Err(e) = process_fs_queue(&mem, &doorbell, &mut queue, &server, &tube, slot) {
            error!("Process FS queue failed: {}", e);
            break;
        }
    }
}

fn default_uidmap() -> String {
    let euid = unsafe { libc::geteuid() };
    format!("{} {} 1", euid, euid)
}

fn default_gidmap() -> String {
    let egid = unsafe { libc::getegid() };
    format!("{} {} 1", egid, egid)
}

fn jail_and_fork(
    mut keep_rds: Vec<RawDescriptor>,
    dir_path: PathBuf,
    uid_map: Option<String>,
    gid_map: Option<String>,
) -> anyhow::Result<i32> {
    // Create new minijail sandbox
    let mut j = Minijail::new()?;

    j.namespace_pids();
    j.namespace_user();
    j.namespace_user_disable_setgroups();
    j.uidmap(&uid_map.unwrap_or_else(default_uidmap))?;
    j.gidmap(&gid_map.unwrap_or_else(default_gidmap))?;
    j.run_as_init();

    j.namespace_vfs();
    j.namespace_net();
    j.no_new_privs();

    // Only pivot_root if we are not re-using the current root directory.
    if dir_path != Path::new("/") {
        // It's safe to call `namespace_vfs` multiple times.
        j.namespace_vfs();
        j.enter_pivot_root(&dir_path)?;
    }
    j.set_remount_mode(libc::MS_SLAVE);

    let limit = get_max_open_files().context("failed to get max open files")?;
    j.set_rlimit(libc::RLIMIT_NOFILE as i32, limit, limit)?;

    // Make sure there are no duplicates in keep_rds
    keep_rds.dedup();

    // fork on the jail here
    let pid = unsafe { j.fork(Some(&keep_rds))? };

    if pid > 0 {
        unsafe { libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM) };
    }

    if pid < 0 {
        bail!("Fork error! {}", std::io::Error::last_os_error());
    }

    Ok(pid)
}

struct FsBackend {
    ex: Executor,
    server: Arc<fuse::Server<PassthroughFs>>,
    tag: [u8; FS_MAX_TAG_LEN],
    avail_features: u64,
    acked_features: u64,
    acked_protocol_features: VhostUserProtocolFeatures,
    workers: [Option<AbortHandle>; Self::MAX_QUEUE_NUM],
    keep_rds: Vec<RawDescriptor>,
}

impl FsBackend {
    pub fn new(ex: &Executor, tag: &str) -> anyhow::Result<Self> {
        if tag.len() > FS_MAX_TAG_LEN {
            bail!(
                "fs tag is too long: {} (max supported: {})",
                tag.len(),
                FS_MAX_TAG_LEN
            );
        }
        let mut fs_tag = [0u8; FS_MAX_TAG_LEN];
        fs_tag[..tag.len()].copy_from_slice(tag.as_bytes());

        let avail_features = virtio::base_features(ProtectionType::Unprotected)
            | VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits();

        // Use default passthroughfs config
        let fs = PassthroughFs::new(Default::default())?;

        let mut keep_rds: Vec<RawDescriptor> = [0, 1, 2].to_vec();
        keep_rds.append(&mut fs.keep_rds());

        let server = Arc::new(Server::new(fs));

        Ok(FsBackend {
            ex: ex.clone(),
            server,
            tag: fs_tag,
            avail_features,
            acked_features: 0,
            acked_protocol_features: VhostUserProtocolFeatures::empty(),
            workers: Default::default(),
            keep_rds,
        })
    }
}

impl VhostUserBackend for FsBackend {
    const MAX_QUEUE_NUM: usize = 2; /* worker queue and high priority queue */
    const MAX_VRING_LEN: u16 = 1024;

    type Error = anyhow::Error;

    fn features(&self) -> u64 {
        self.avail_features
    }

    fn ack_features(&mut self, value: u64) -> anyhow::Result<()> {
        let unrequested_features = value & !self.avail_features;
        if unrequested_features != 0 {
            bail!("invalid features are given: {:#x}", unrequested_features);
        }

        self.acked_features |= value;

        Ok(())
    }

    fn acked_features(&self) -> u64 {
        self.acked_features
    }

    fn protocol_features(&self) -> VhostUserProtocolFeatures {
        VhostUserProtocolFeatures::CONFIG | VhostUserProtocolFeatures::MQ
    }

    fn ack_protocol_features(&mut self, features: u64) -> anyhow::Result<()> {
        let features = VhostUserProtocolFeatures::from_bits(features)
            .ok_or_else(|| anyhow!("invalid protocol features are given: {:#x}", features))?;
        let supported = self.protocol_features();
        self.acked_protocol_features = features & supported;
        Ok(())
    }

    fn acked_protocol_features(&self) -> u64 {
        self.acked_protocol_features.bits()
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        let config = virtio_fs_config {
            tag: self.tag,
            num_request_queues: Le32::from(1),
        };
        copy_config(data, 0, config.as_slice(), offset);
    }

    fn reset(&mut self) {
        for handle in self.workers.iter_mut().filter_map(Option::take) {
            handle.abort();
        }
    }

    fn start_queue(
        &mut self,
        idx: usize,
        mut queue: virtio::Queue,
        mem: GuestMemory,
        doorbell: Arc<Mutex<Doorbell>>,
        kick_evt: Event,
    ) -> anyhow::Result<()> {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            warn!("Starting new queue handler without stopping old handler");
            handle.abort();
        }

        // Enable any virtqueue features that were negotiated (like VIRTIO_RING_F_EVENT_IDX).
        queue.ack_features(self.acked_features);

        let kick_evt = EventAsync::new(kick_evt.0, &self.ex)
            .context("failed to create EventAsync for kick_evt")?;
        let (handle, registration) = AbortHandle::new_pair();
        let (_, fs_device_tube) = Tube::pair()?;

        self.ex
            .spawn_local(Abortable::new(
                handle_fs_queue(
                    queue,
                    mem,
                    doorbell,
                    kick_evt,
                    self.server.clone(),
                    Arc::new(Mutex::new(fs_device_tube)),
                ),
                registration,
            ))
            .detach();

        self.workers[idx] = Some(handle);
        Ok(())
    }

    fn stop_queue(&mut self, idx: usize) {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            handle.abort();
        }
    }
}

#[derive(FromArgs)]
#[argh(description = "")]
struct Options {
    #[argh(option, description = "path to a socket", arg_name = "PATH")]
    socket: String,
    #[argh(option, description = "the virtio-fs tag", arg_name = "TAG")]
    tag: String,
    #[argh(option, description = "path to a directory to share", arg_name = "DIR")]
    shared_dir: PathBuf,
    #[argh(option, description = "uid map to use", arg_name = "UIDMAP")]
    uid_map: Option<String>,
    #[argh(option, description = "gid map to use", arg_name = "GIDMAP")]
    gid_map: Option<String>,
}

/// Starts a vhost-user fs device.
/// Returns an error if the given `args` is invalid or the device fails to run.
pub fn run_fs_device(program_name: &str, args: &[&str]) -> anyhow::Result<()> {
    let opts = match Options::from_args(&[program_name], args) {
        Ok(opts) => opts,
        Err(e) => {
            if e.status.is_err() {
                bail!(e.output);
            } else {
                println!("{}", e.output);
            }
            return Ok(());
        }
    };

    base::syslog::init().context("Failed to initialize syslog")?;

    let ex = Executor::new().context("Failed to create executor")?;
    let fs_device = FsBackend::new(&ex, &opts.tag)?;

    // Create and bind unix socket
    let listener = UnixListener::bind(opts.socket).map(UnlinkUnixListener)?;
    let mut keep_rds = fs_device.keep_rds.clone();
    keep_rds.push(listener.as_raw_fd());
    base::syslog::push_descriptors(&mut keep_rds);

    let handler = DeviceRequestHandler::new(fs_device);

    let pid = jail_and_fork(keep_rds, opts.shared_dir, opts.uid_map, opts.gid_map)?;

    // Parent, nothing to do but wait and then exit
    if pid != 0 {
        unsafe { libc::waitpid(pid, std::ptr::null_mut(), 0) };
        return Ok(());
    }

    // We need to set the no setuid fixup secure bit so that we don't drop capabilities when
    // changing the thread uid/gid. Without this, creating new entries can fail in some corner
    // cases.
    const SECBIT_NO_SETUID_FIXUP: i32 = 1 << 2;
    // TODO(crbug.com/1199487): Remove this once libc provides the wrapper for all targets.
    #[cfg(target_os = "linux")]
    {
        // Safe because this doesn't modify any memory and we check the return value.
        let mut securebits = unsafe { libc::prctl(libc::PR_GET_SECUREBITS) };
        if securebits < 0 {
            bail!(io::Error::last_os_error());
        }
        securebits |= SECBIT_NO_SETUID_FIXUP;
        // Safe because this doesn't modify any memory and we check the return value.
        let ret = unsafe { libc::prctl(libc::PR_SET_SECUREBITS, securebits) };
        if ret < 0 {
            bail!(io::Error::last_os_error());
        }
    }

    // run_until() returns an Result<Result<..>> which the ? operator lets us flatten.
    ex.run_until(handler.run_with_listener(listener, &ex))?
}
<|endoftext|>