<commit_before>//! Configurator for the platform.  Initializes config files for the modules and sets up
//! the initial environment for platform runtime.
//!
//! This requires that the package libncurses-dev is installed!

#![feature(plugin)]
#![plugin(indoc)]

extern crate cursive;

use std::fs::File;
use std::process::{Child, Command, Stdio};
use std::path::Path;

use cursive::Cursive;
use cursive::views::{Dialog, TextView};

mod theme;
use theme::THEME;

fn main() {
    // Check if this is the first run of the configurator
    let path = Path::new(".platform_conf");
    if !path.exists() {
        // File::create(path).expect("Unable to create config placeholder file; do you have write permissions to this directory?");
        first_time();
    }

    // clear console + reset colored background before exiting
    print!(".{}[0m{}[2J", 27 as char, 27 as char);
}

/// Displays welcome and walks the user through first time configuration of the platform.
fn first_time() {
    let mut siv = Cursive::new();
    siv.set_theme(THEME);

    siv.add_layer(Dialog::around(TextView::new(
        indoc!("Welcome to the Bot4 Algorithmic Trading Platform!

                This tool will set up the environment for the trading platform.  It will walk you through the process of \
                installing all prerequisite software and initializing all necessary configuration settings for the platform's \
                modules.") ))
        .title("Welcome")
        .button("Continue", |s| redis_config(s) )
    );

    // Start the event loop
    siv.run();
}

fn redis_config(s: &mut Cursive) {
    let mut message_text =
        String::from(
            indoc!("The first thing that needs to be configured is Redis.  Redis PubSub is used as a messaging service \
                    that allows for communication between the platform's different modules.  It is recommended that you \
                    use a local instance of Redis since it's easy to create a security vulnerability if it is not properly \
                    configured.\n\n"
            )
        );

    let installed = redis_installed();

    if installed {
        message_text +=
            indoc!("I detected that you currently have Redis installed.  To use this local Redis install for the platform, \
                    select the \"Local\" option below."
            );
    } else {
        message_text +=
            indoc!("I was unable to detect a local Redis installation.  If Redis is currently installed and you \
                    want to use a local installation, please add the `redis-server` executable to the system PATH and \
                    restart the configuration process."
            );
    }

    s.pop_layer();
    s.add_layer(Dialog::text(
        message_text)
        .title("Information")
        .button("Local", |s| redis_local(s) )
        .button("Remote", |s| redis_remote(s) )
        .button("Exit", |s| s.quit() )
    )
}

fn redis_local(s: &Cursive) {
    unimplemented!();
}

fn redis_remote(s: &Cursive) {
    unimplemented!();
}

fn postgres_config(s: &mut Cursive) {
    unimplemented!();
}

pub fn redis_installed() -> bool {
    let child = Command::new("which")
        .arg("redis-server")
        .stdout(Stdio::piped())
        .spawn()
        .expect("Unable to spawn `which redis_server`");

    let output = child.wait_with_output()
        .expect("Unable to get output from `which redis_server`");

    output.stdout.len() > 0
}

#[test]
fn redis_installed_test() {
    redis_installed();
}
<commit_msg>Fleshed out much more of the Configurator  - Now reads in configuration for both Redis and PostgreSQL for remote and local locations.<commit_after>//! Configurator for the platform.  Initializes config files for the modules and sets up
//! the initial environment for platform runtime.
//!
//! This requires that the package libncurses-dev is installed!

#![feature(plugin)]
#![plugin(indoc)]

extern crate cursive;

use std::fs::File;
use std::process::{Child, Command, Stdio};
use std::collections::HashMap;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::rc::Rc;

use cursive::Cursive;
use cursive::views::{Dialog, TextView, EditView, ListView, BoxView, IdView};
use cursive::view::{SizeConstraint, ViewWrapper};

mod theme;
use theme::THEME;

// type Settings = HashMap<String, String>;
// type SettingsRef = Arc<Mutex<HashMap<String, String>>>;

#[derive(Clone)]
struct Settings {
    inner: Arc<Mutex<HashMap<String, String>>>,
}

impl Settings {
    pub fn new() -> Settings {
        Settings {
            inner: Arc::new(Mutex::new(HashMap::new()))
        }
    }

    pub fn set(&self, key: &str, val: &str) {
        let mut ul = self.inner.lock().unwrap();
        ul.insert(key.to_string(), val.to_string());
    }

    pub fn get(&self, key: String) -> Option<String> {
        let ul = self.inner.lock().unwrap();
        match ul.get(&key) {
            Some(val) => Some(val.clone()),
            None => None
        }
    }

    /// Dumps the Settings object to a JSON file that can be used to populate the Settings object from scratch
    pub fn write_json(filename: &str) {
        unimplemented!();
    }

    /// Reads the supplied JSON file and generates a Settings object from its contents.
    pub fn read_json(filename: &str) -> Settings {
        unimplemented!();
    }
}

const MIN15: SizeConstraint = SizeConstraint::AtLeast(10);
const FREE: SizeConstraint = SizeConstraint::Free;

fn main() {
    // Check if this is the first run of the configurator
    let path = Path::new(".platform_conf");
    if !path.exists() {
        // File::create(path).expect("Unable to create config placeholder file; do you have write permissions to this directory?");
        first_time();
    }

    // clear console + reset colored background before exiting
    print!(".{}[0m{}[2J", 27 as char, 27 as char);
}

fn get_by_id(id: &str, s: &mut Cursive) -> Option<Rc<String>> {
    match s.find_id::<BoxView<EditView>>(id) {
        Some(container) => Some(container.get_view().get_content()),
        None => None
    }
}

/// Displays welcome and walks the user through first time configuration of the platform.
fn first_time() {
    let mut siv = Cursive::new();
    siv.set_theme(THEME);

    // Main Key:Value settings for the application
    let settings = Settings::new();

    siv.add_layer(
        Dialog::around(TextView::new(
            indoc!(
                "Welcome to the Bot4 Algorithmic Trading Platform!

                This tool will set up the environment for the trading platform.  It will walk you through the process of \
                installing all prerequisite software and initializing all necessary configuration settings for the platform's \
                modules."
            )
        )).title("Welcome")
            .button("Continue", move |s| redis_config(s, settings.clone()) )
    );

    // Start the event loop
    siv.run();
}

fn redis_config(s: &mut Cursive, settings: Settings) {
    let settings_clone = settings.clone();

    let mut message_text =
        String::from(
            indoc!(
                "The first thing that needs to be configured is Redis.  Redis PubSub is used as a messaging service \
                that allows for communication between the platform's different modules.  It is recommended that you \
                use a local instance of Redis since it's easy to create a security vulnerability if it is not properly \
                configured.\n\n"
            )
        );

    let installed = is_installed("redis-server");

    if installed {
        message_text +=
            indoc!(
                "I detected that you currently have Redis installed.  To use this local Redis install for the platform, \
                select the \"Local\" option below."
            );
    } else {
        message_text +=
            indoc!(
                "I was unable to detect a local Redis installation.  If Redis is currently installed and you \
                want to use a local installation, please add the `redis-server` executable to the system PATH and \
                restart the configuration process."
            );
    }

    s.pop_layer();
    s.add_layer(Dialog::text(message_text)
        .title("Information")
        .button("Local", move |s| redis_local(s, is_installed("redis-server"), settings.clone()) )
        .button("Remote", move |s| redis_remote(s, settings_clone.clone()) )
        .button("Exit", |s| s.quit() )
    )
}

fn redis_local(s: &mut Cursive, installed: bool, settings: Settings) {
    settings.set("redis_host", "localhost");

    if !installed {
        s.add_layer(Dialog::text(
            indoc!(
                "You must install redis in order to use it locally.  Install Redis, add the `redis-server` \
                binary to to the system PATH, and select local again."
            )
        ).dismiss_button("Ok"))
    } else {
        s.pop_layer();
        let mut port_box = EditView::new()
            .on_submit(move |s, port| {
                settings.set("redis_port", port);
                postgres_config(s, settings.clone())
            });
        port_box.set_content("6379");

        s.add_layer(Dialog::around(port_box)
            .title("Redis Port")
        );
    }
}

fn redis_remote(s: &mut Cursive, settings: Settings) {
    s.pop_layer();
    s.add_layer(Dialog::new()
        .content(ListView::new()
            .child("Redis Host", IdView::new("redis-host", BoxView::new(MIN15, FREE, EditView::new())))
            .child("Redis Port", IdView::new("redis-port", BoxView::new(MIN15, FREE, EditView::new().content("6379"))))
        ).title("Remote Redis Server Settings")
        .button("Ok", move |s| {
            settings.set("redis-host", &*get_by_id("redis-host", s).unwrap());
            settings.set("redis-port", &*get_by_id("redis-port", s).unwrap());
            postgres_config(s, settings.clone())
        })
    );
}

fn postgres_config(s: &mut Cursive, settings: Settings) {
    let settings_clone = settings.clone();
    let installed = is_installed("psql");

    let mut message_text =
        String::from(
            indoc!("The platform also relies on PostgreSQL to store historical Tick data, persistant platform \
                    data, and other information for the platform.\n\n"
            )
        );

    if !installed {
        message_text += indoc!(
            "I was unable do detect an active PostgreSQL installation on this host.  In order to use this \
            host for the platform, you must first install PostgreSQL and add the `psql` binary to the system \
            path.  Once you've installed it, select \"Local\" again.create"
        )
    } else {
        message_text += indoc!(
            "I detected that you have PostgreSQL installed locally.  To configure the platform to use the \
            local PostgreSQL installation, select the \"Local\" option below."
        );
    }

    s.pop_layer();
    s.add_layer(Dialog::text(message_text)
        .title("PostgreSQL Configuration")
        .button("Local", move |s| postgres_local(s, is_installed("psql"), settings.clone()) )
        .button("Remote", move |s| postgres_remote(s, settings_clone.clone()) )
        .button("Exit", |s| s.quit() )
    );
}

fn postgres_remote(s: &mut Cursive, settings: Settings) {
    s.pop_layer();
    s.add_layer(Dialog::new()
        .content(ListView::new()
            .child("Postgres Host", IdView::new("postgres-host", BoxView::new(MIN15, FREE, EditView::new())))
            .child("Postgres User", IdView::new("postgres-user", BoxView::new(MIN15, FREE, EditView::new())))
            .child("Postgres Password", IdView::new("postgres-password", BoxView::new(MIN15, FREE, EditView::new().secret())))
            .child("Postgres Port", IdView::new("postgres-port", BoxView::new(MIN15, FREE, EditView::new().content("5432"))))
            .child("Postgres Database", IdView::new("postgres-db", BoxView::new(MIN15, FREE, EditView::new())))
        ).title("Remote PostgreSQL Server Settings")
            .button("Ok", move |s| save_settings(s, settings.clone()) )
    )
}

fn postgres_local(s: &mut Cursive, installed: bool, settings: Settings) {
    settings.set("postgres_host", "localhost");

    if !installed {
        s.add_layer(Dialog::text(
            indoc!(
                "You must install PostgreSQL in order to use it locally.  Install PostgreSQL, add the `psql` \
                binary to to the system PATH, and select local again."
            )
        ).dismiss_button("Ok"))
    } else {
        s.pop_layer();
        s.add_layer(Dialog::new()
            .content(ListView::new()
                .child("Postgres User", IdView::new("postgres-user", BoxView::new(MIN15, FREE, EditView::new())))
                .child("Postgres Password", IdView::new("postgres-password", BoxView::new(MIN15, FREE, EditView::new().secret())))
                .child("Postgres Port", IdView::new("postgres-port", BoxView::new(MIN15, FREE, EditView::new())))
                .child("Postgres Database", IdView::new("postgres-db", BoxView::new(MIN15, FREE, EditView::new())))
            ).title("Local PostgreSQL Server Settings")
                .button("Ok", move |s| save_settings(s, settings.clone()) )
        )
    }
}

fn is_installed(command: &str) -> bool {
    let child = Command::new("which")
        .arg(command)
        .stdout(Stdio::piped())
        .spawn()
        .expect("Unable to spawn `which redis_server`");

    let output = child.wait_with_output()
        .expect("Unable to get output from `which redis_server`");

    output.stdout.len() > 0
}

fn save_settings(s: &mut Cursive, settings: Settings) {
    s.pop_layer();
    for id in ["postgres-host", "postgres-user", "postgres-password", "postgres-port", "postgres-db"].iter() {
        let val = get_by_id(id, s);
        if val.is_some() && *id == "postgres-host" {
            settings.set(id, &(*val.unwrap()) );
        }
    }

    s.add_layer(Dialog::text(
        indoc!(
            "The trading platform has been successfully configured.  Run `run.sh` and visit `localhost:8002` in \
            your web browser to start using the platform."
        )
    ).button("Ok", |s| s.quit() ))
}

#[test]
fn redis_installed_test() {
    is_isntalled("redis-server");
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate clap;
extern crate env_logger;
extern crate habitat_core as hcore;
extern crate habitat_common as common;
extern crate habitat_pkg_export_docker as export_docker;
extern crate handlebars;
#[macro_use]
extern crate serde_json;
#[macro_use]
extern crate log;

use clap::{App, Arg};
use handlebars::Handlebars;
use std::env;
use std::result;
use std::str::FromStr;
use std::io::prelude::*;
use std::io;
use std::fs::File;

use hcore::channel;
use hcore::PROGRAM_NAME;
use hcore::url as hurl;
use hcore::env as henv;
use hcore::package::PackageIdent;
use common::ui::{Coloring, UI, NOCOLORING_ENVVAR, NONINTERACTIVE_ENVVAR};

use export_docker::{Cli, BuildSpec, Error, Naming};

// Synced with the version of the Habitat operator.
pub const VERSION: &'static str = "0.1.0";

// Kubernetes manifest template
const MANIFESTFILE: &'static str = include_str!("../defaults/KubernetesManifest.hbs");

fn main() {
    env_logger::init().unwrap();
    let mut ui = get_ui();
    if let Err(e) = start(&mut ui) {
        let _ = ui.fatal(e);
        std::process::exit(1)
    }
}

fn get_ui() -> UI {
    let isatty = if henv::var(NONINTERACTIVE_ENVVAR)
        .map(|val| val == "true")
        .unwrap_or(false)
    {
        Some(false)
    } else {
        None
    };
    let coloring = if henv::var(NOCOLORING_ENVVAR)
        .map(|val| val == "true")
        .unwrap_or(false)
    {
        Coloring::Never
    } else {
        Coloring::Auto
    };
    UI::default_with(coloring, isatty)
}

fn start(ui: &mut UI) -> result::Result<(), Error> {
    let m = cli().get_matches();
    debug!("clap cli args: {:?}", m);

    if !m.is_present("NO_DOCKER_IMAGE") {
        gen_docker_img(ui, &m)?;
    }
    gen_k8s_manifest(ui, &m)
}

fn gen_docker_img(ui: &mut UI, matches: &clap::ArgMatches) -> result::Result<(), Error> {
    let default_channel = channel::default();
    let default_url = hurl::default_bldr_url();
    let spec = BuildSpec::new_from_cli_matches(&matches, &default_channel, &default_url);
    let naming = Naming::new_from_cli_matches(&matches);

    let docker_image = export_docker::export(ui, spec, &naming)?;
    docker_image.create_report(
        ui,
        env::current_dir()?.join("results"),
    )?;

    Ok(())
}

fn gen_k8s_manifest(_ui: &mut UI, matches: &clap::ArgMatches) -> result::Result<(), Error> {
    let count = matches.value_of("COUNT").unwrap_or("1");
    let topology = matches.value_of("TOPOLOGY").unwrap_or("standalone");
    let group = matches.value_of("GROUP");
    let config_secret_name = matches.value_of("CONFIG_SECRET_NAME");
    let ring_secret_name = matches.value_of("RING_SECRET_NAME");
    // clap ensures that we do have the mandatory args so unwrap() is fine here
    let pkg_ident_str = matches.value_of("PKG_IDENT_OR_ARTIFACT").unwrap();
    let pkg_ident = PackageIdent::from_str(pkg_ident_str)?;
    let image = matches.value_of("IMAGE_NAME").unwrap_or(pkg_ident_str);

    let json = json!({
        "metadata_name": pkg_ident.name,
        "image": image,
        "count": count,
        "service_topology": topology,
        "service_group": group,
        "config_secret_name": config_secret_name,
        "ring_secret_name": ring_secret_name,
    });

    let mut write: Box<Write> = match matches.value_of("OUTPUT") {
        Some(o) if o != "-" => Box::new(File::create(o)?),
        _ => Box::new(io::stdout()),
    };

    let r = Handlebars::new().template_render(MANIFESTFILE, &json)?;
    let out = r.lines().filter(|l| *l != "").collect::<Vec<_>>().join(
        "\n",
    ) + "\n";

    write.write(out.as_bytes())?;

    Ok(())
}

fn cli<'a, 'b>() -> App<'a, 'b> {
    let name: &str = &*PROGRAM_NAME;
    let about = "Creates a docker image and Kubernetes manifest for a Habitat package. Habitat \
                 operator must be deployed within the Kubernetes cluster to intercept the created \
                 objects.";

    let app = Cli::new(name, about)
        .add_base_packages_args()
        .add_builder_args()
        .add_tagging_args()
        .app;

    app.arg(
        Arg::with_name("OUTPUT")
            .value_name("OUTPUT")
            .long("output")
            .short("o")
            .help(
                "Name of manifest file to create. Pass '-' for stdout (default: -)",
            ),
    ).arg(
            Arg::with_name("COUNT")
                .value_name("COUNT")
                .long("count")
                .validator(valid_natural_number)
                .help("Count is the number of desired instances"),
        )
        .arg(
            Arg::with_name("TOPOLOGY")
                .value_name("TOPOLOGY")
                .long("topology")
                .short("t")
                .possible_values(&["standalone", "leader"])
                .help(
                    "A topology describes the intended relationship between peers \
                    within a service group. Specify either standalone or leader \
                    topology (default: standalone)",
                ),
        )
        .arg(
            Arg::with_name("GROUP")
                .value_name("GROUP")
                .long("service-group")
                .short("g")
                .help(
                    "group is a logical grouping of services with the same package and \
                    topology type connected together in a ring (default: default)",
                ),
        )
        .arg(
            Arg::with_name("CONFIG_SECRET_NAME")
                .value_name("CONFIG_SECRET_NAME")
                .long("config-secret-name")
                .short("n")
                .help(
                    "name of the Kubernetes Secret containing the config file - \
                    user.toml - that the user has previously created. Habitat will \
                    use it for initial configuration of the service",
                ),
        )
        .arg(
            Arg::with_name("RING_SECRET_NAME")
                .value_name("RING_SECRET_NAME")
                .long("ring-secret-name")
                .short("r")
                .help(
                    "name of the Kubernetes Secret that contains the ring key, which \
                    encrypts the communication between Habitat supervisors",
                ),
        )
        .arg(
            Arg::with_name("NO_DOCKER_IMAGE")
                .long("no-docker-image")
                .short("d")
                .help(
                    "Disable creation of docker image and only create Kubernetes manifest",
                ),
        )
        .arg(
            Arg::with_name("PKG_IDENT_OR_ARTIFACT")
                .value_name("PKG_IDENT_OR_ARTIFACT")
                .required(true)
                .help("Habitat package identifier (ex: acme/redis)"),
        )
}

fn valid_natural_number(val: String) -> result::Result<(), String> {
    match val.parse::<u32>() {
        Ok(_) => Ok(()),
        Err(_) => Err(format!("{} is not a natural number", val)),
    }
}
<commit_msg>pkg-export-kubernetes: Support package files<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate clap;
extern crate env_logger;
extern crate habitat_core as hcore;
extern crate habitat_common as common;
extern crate habitat_pkg_export_docker as export_docker;
extern crate handlebars;
#[macro_use]
extern crate serde_json;
#[macro_use]
extern crate log;

use clap::{App, Arg};
use handlebars::Handlebars;
use std::env;
use std::result;
use std::str::FromStr;
use std::io::prelude::*;
use std::io;
use std::fs::File;
use std::path::Path;

use hcore::channel;
use hcore::PROGRAM_NAME;
use hcore::url as hurl;
use hcore::env as henv;
use hcore::package::{PackageArchive, PackageIdent};
use common::ui::{Coloring, UI, NOCOLORING_ENVVAR, NONINTERACTIVE_ENVVAR};

use export_docker::{Cli, BuildSpec, Error, Naming};

// Synced with the version of the Habitat operator.
pub const VERSION: &'static str = "0.1.0";

// Kubernetes manifest template
const MANIFESTFILE: &'static str = include_str!("../defaults/KubernetesManifest.hbs");

fn main() {
    env_logger::init().unwrap();
    let mut ui = get_ui();
    if let Err(e) = start(&mut ui) {
        let _ = ui.fatal(e);
        std::process::exit(1)
    }
}

fn get_ui() -> UI {
    let isatty = if henv::var(NONINTERACTIVE_ENVVAR)
        .map(|val| val == "true")
        .unwrap_or(false)
    {
        Some(false)
    } else {
        None
    };
    let coloring = if henv::var(NOCOLORING_ENVVAR)
        .map(|val| val == "true")
        .unwrap_or(false)
    {
        Coloring::Never
    } else {
        Coloring::Auto
    };
    UI::default_with(coloring, isatty)
}

fn start(ui: &mut UI) -> result::Result<(), Error> {
    let m = cli().get_matches();
    debug!("clap cli args: {:?}", m);

    if !m.is_present("NO_DOCKER_IMAGE") {
        gen_docker_img(ui, &m)?;
    }
    gen_k8s_manifest(ui, &m)
}

fn gen_docker_img(ui: &mut UI, matches: &clap::ArgMatches) -> result::Result<(), Error> {
    let default_channel = channel::default();
    let default_url = hurl::default_bldr_url();
    let spec = BuildSpec::new_from_cli_matches(&matches, &default_channel, &default_url);
    let naming = Naming::new_from_cli_matches(&matches);

    let docker_image = export_docker::export(ui, spec, &naming)?;
    docker_image.create_report(
        ui,
        env::current_dir()?.join("results"),
    )?;

    Ok(())
}

fn gen_k8s_manifest(_ui: &mut UI, matches: &clap::ArgMatches) -> result::Result<(), Error> {
    let count = matches.value_of("COUNT").unwrap_or("1");
    let topology = matches.value_of("TOPOLOGY").unwrap_or("standalone");
    let group = matches.value_of("GROUP");
    let config_secret_name = matches.value_of("CONFIG_SECRET_NAME");
    let ring_secret_name = matches.value_of("RING_SECRET_NAME");
    // clap ensures that we do have the mandatory args so unwrap() is fine here
    let pkg_ident_str = matches.value_of("PKG_IDENT_OR_ARTIFACT").unwrap();
    let pkg_ident = if Path::new(pkg_ident_str).is_file() {
        // We're going to use the `$pkg_origin/$pkg_name`, fuzzy form of a package
        // identifier to ensure that update strategies will work if desired
        PackageArchive::new(pkg_ident_str).ident()?
    } else {
        PackageIdent::from_str(pkg_ident_str)?
    };
    let image = match matches.value_of("IMAGE_NAME") {
        Some(i) => i.to_string(),
        None => pkg_ident.origin + "/" + &pkg_ident.name,
    };

    let json = json!({
        "metadata_name": pkg_ident.name,
        "image": image,
        "count": count,
        "service_topology": topology,
        "service_group": group,
        "config_secret_name": config_secret_name,
        "ring_secret_name": ring_secret_name,
    });

    let mut write: Box<Write> = match matches.value_of("OUTPUT") {
        Some(o) if o != "-" => Box::new(File::create(o)?),
        _ => Box::new(io::stdout()),
    };

    let r = Handlebars::new().template_render(MANIFESTFILE, &json)?;
    let out = r.lines().filter(|l| *l != "").collect::<Vec<_>>().join(
        "\n",
    ) + "\n";

    write.write(out.as_bytes())?;

    Ok(())
}

fn cli<'a, 'b>() -> App<'a, 'b> {
    let name: &str = &*PROGRAM_NAME;
    let about = "Creates a docker image and Kubernetes manifest for a Habitat package. Habitat \
                 operator must be deployed within the Kubernetes cluster to intercept the created \
                 objects.";

    let app = Cli::new(name, about)
        .add_base_packages_args()
        .add_builder_args()
        .add_tagging_args()
        .app;

    app.arg(
        Arg::with_name("OUTPUT")
            .value_name("OUTPUT")
            .long("output")
            .short("o")
            .help(
                "Name of manifest file to create. Pass '-' for stdout (default: -)",
            ),
    ).arg(
            Arg::with_name("COUNT")
                .value_name("COUNT")
                .long("count")
                .validator(valid_natural_number)
                .help("Count is the number of desired instances"),
        )
        .arg(
            Arg::with_name("TOPOLOGY")
                .value_name("TOPOLOGY")
                .long("topology")
                .short("t")
                .possible_values(&["standalone", "leader"])
                .help(
                    "A topology describes the intended relationship between peers \
                    within a service group. Specify either standalone or leader \
                    topology (default: standalone)",
                ),
        )
        .arg(
            Arg::with_name("GROUP")
                .value_name("GROUP")
                .long("service-group")
                .short("g")
                .help(
                    "group is a logical grouping of services with the same package and \
                    topology type connected together in a ring (default: default)",
                ),
        )
        .arg(
            Arg::with_name("CONFIG_SECRET_NAME")
                .value_name("CONFIG_SECRET_NAME")
                .long("config-secret-name")
                .short("n")
                .help(
                    "name of the Kubernetes Secret containing the config file - \
                    user.toml - that the user has previously created. Habitat will \
                    use it for initial configuration of the service",
                ),
        )
        .arg(
            Arg::with_name("RING_SECRET_NAME")
                .value_name("RING_SECRET_NAME")
                .long("ring-secret-name")
                .short("r")
                .help(
                    "name of the Kubernetes Secret that contains the ring key, which \
                    encrypts the communication between Habitat supervisors",
                ),
        )
        .arg(
            Arg::with_name("NO_DOCKER_IMAGE")
                .long("no-docker-image")
                .short("d")
                .help(
                    "Disable creation of docker image and only create Kubernetes manifest",
                ),
        )
        .arg(
            Arg::with_name("PKG_IDENT_OR_ARTIFACT")
                .value_name("PKG_IDENT_OR_ARTIFACT")
                .required(true)
                .help("Habitat package identifier (ex: acme/redis)"),
        )
}

fn valid_natural_number(val: String) -> result::Result<(), String> {
    match val.parse::<u32>() {
        Ok(_) => Ok(()),
        Err(_) => Err(format!("{} is not a natural number", val)),
    }
}
<|endoftext|>


<commit_before>//! Simple module to store files in database.
//!
//! cratesfyi is generating more than 5 million files, they are small and mostly html files.
//! They are using so many inodes and it is better to store them in database instead of
//! filesystem. This module is adding files into database and retrieving them.

use std::path::{PathBuf, Path};
use postgres::Connection;
use rustc_serialize::json::{Json, ToJson};
use std::fs;
use std::io::Read;
use error::Result;
use failure::err_msg;
use rusoto_s3::{S3, PutObjectRequest, GetObjectRequest, S3Client};
use rusoto_core::region::Region;
use rusoto_credential::EnvironmentProvider;

fn get_file_list_from_dir<P: AsRef<Path>>(path: P,
                                          files: &mut Vec<PathBuf>)
                                          -> Result<()> {
    let path = path.as_ref();

    for file in try!(path.read_dir()) {
        let file = try!(file);

        if try!(file.file_type()).is_file() {
            files.push(file.path());
        } else if try!(file.file_type()).is_dir() {
            try!(get_file_list_from_dir(file.path(), files));
        }
    }

    Ok(())
}

pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>> {
    let path = path.as_ref();
    let mut files = Vec::new();

    if !path.exists() {
        return Err(err_msg("File not found"));
    } else if path.is_file() {
        files.push(PathBuf::from(path.file_name().unwrap()));
    } else if path.is_dir() {
        try!(get_file_list_from_dir(path, &mut files));
        for file_path in &mut files {
            // We want the paths in this list to not be {path}/bar.txt but just bar.txt
            *file_path = PathBuf::from(file_path.strip_prefix(path).unwrap());
        }
    }

    Ok(files)
}

pub struct Blob {
    pub path: String,
    pub mime: String,
    pub date_updated: time::Timespec,
    pub content: Vec<u8>,
}

pub fn get_path(conn: &Connection, path: &str) -> Option<Blob> {
    let rows = conn.query("SELECT path, mime, date_updated, content
                           FROM files
                           WHERE path = $1", &[&path]).unwrap();

    if rows.len() == 0 {
        None
    } else {
        let row = rows.get(0);
        let mut content = row.get(3);
        if content == b"in-s3" {
            let client = s3_client();
            content = client.and_then(|c| c.get_object(GetObjectRequest {
                bucket: "rust-docs-rs".into(),
                key: path.into(),
                ..Default::default()
            }).sync().ok()).and_then(|r| r.body).map(|b| {
                let mut b = b.into_blocking_read();
                let mut content = Vec::new();
                b.read_to_end(&mut content).unwrap();
                content
            }).unwrap();
        };

        Some(Blob {
            path: row.get(0),
            mime: row.get(1),
            date_updated: row.get(2),
            content,
        })
    }
}

fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() {
        return None;
    }
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        EnvironmentProvider::default(),
        std::env::var("S3_ENDPOINT").ok().map(|e| Region::Custom {
            name: "us-west-1".to_owned(),
            endpoint: e,
        }).unwrap_or(Region::UsWest1),
    ))
}

/// Adds files into database and returns list of files with their mime type in Json
pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
                                              prefix: &str,
                                              path: P)
                                              -> Result<Json> {
    use magic::{Cookie, flags};
    let cookie = try!(Cookie::open(flags::MIME_TYPE));
    try!(cookie.load::<&str>(&[]));

    let trans = try!(conn.transaction());
    let client = s3_client();
    let mut file_list_with_mimes: Vec<(String, PathBuf)> = Vec::new();

    for file_path in try!(get_file_list(&path)) {
        let (path, content, mime) = {
            let path = Path::new(path.as_ref()).join(&file_path);
            // Some files have insufficient permissions (like .lock file created by cargo in
            // documentation directory). We are skipping this files.
            let mut file = match fs::File::open(path) {
                Ok(f) => f,
                Err(_) => continue,
            };
            let mut content: Vec<u8> = Vec::new();
            try!(file.read_to_end(&mut content));
            let bucket_path = Path::new(prefix).join(&file_path)
                .into_os_string().into_string().unwrap();

            let mime = {
                let mime = try!(cookie.buffer(&content));
                // css's are causing some problem in browsers
                // magic will return text/plain for css file types
                // convert them to text/css
                // do the same for javascript files
                if mime == "text/plain" {
                    let e = file_path.extension().unwrap_or_default();
                    if e == "css" {
                        "text/css".to_owned()
                    } else if e == "js" {
                        "application/javascript".to_owned()
                    } else {
                        mime.to_owned()
                    }
                } else {
                    mime.to_owned()
                }
            };

            let content: Option<Vec<u8>> = if let Some(client) = &client {
                let s3_res = client.put_object(PutObjectRequest {
                    bucket: "rust-docs-rs".into(),
                    key: bucket_path.clone(),
                    body: Some(content.clone().into()),
                    content_type: Some(mime.clone()),
                    ..Default::default()
                }).sync();
                match s3_res {
                    // we've successfully uploaded the content, so steal it;
                    // we don't want to put it in the DB
                    Ok(_) => None,
                    // Since s3 was configured, we want to panic on failure to upload.
                    Err(e) => {
                        panic!("failed to upload to {}: {:?}", bucket_path, e)
                    },
                }
            } else {
                Some(content.clone().into())
            };

            file_list_with_mimes.push((mime.clone(), file_path.clone()));

            (
                bucket_path,
                content,
                mime,
            )
        };

        // check if file already exists in database
        let rows = try!(conn.query("SELECT COUNT(*) FROM files WHERE path = $1", &[&path]));

        let content = content.unwrap_or_else(|| "in-s3".to_owned().into());

        if rows.get(0).get::<usize, i64>(0) == 0 {
            try!(trans.query("INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                             &[&path, &mime, &content]));
        } else {
            try!(trans.query("UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                              WHERE path = $1",
                             &[&path, &mime, &content]));
        }
    }

    try!(trans.commit());

    file_list_to_json(file_list_with_mimes)
}

fn file_list_to_json(file_list: Vec<(String, PathBuf)>) -> Result<Json> {

    let mut file_list_json: Vec<Json> = Vec::new();

    for file in file_list {
        let mut v: Vec<String> = Vec::new();
        v.push(file.0.clone());
        v.push(file.1.into_os_string().into_string().unwrap());
        file_list_json.push(v.to_json());
    }

    Ok(file_list_json.to_json())
}

pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize> {
    let trans = try!(conn.transaction());
    let client = s3_client().expect("configured s3");

    let rows = try!(trans.query(
            &format!("SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}", n),
            &[]));
    let count = rows.len();

    let mut rt = ::tokio::runtime::current_thread::Runtime::new().unwrap();
    let mut futures = Vec::new();
    for row in &rows {
        let path: String = row.get(0);
        let mime: String = row.get(1);
        let content: Vec<u8> = row.get(2);
        let path_1 = path.clone();
        futures.push(client.put_object(PutObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.clone(),
            body: Some(content.into()),
            content_type: Some(mime),
            ..Default::default()
        }).map(move |_| {
            path_1
        }).map_err(move |e| {
            panic!("failed to upload to {}: {:?}", path, e)
        }));
    }

    use ::futures::future::Future;
    match rt.block_on(::futures::future::join_all(futures)) {
        Ok(paths) => {
            let statement = trans.prepare("UPDATE files SET content = E'in-s3' WHERE path = $1")
                .unwrap();
            for path in paths {
                statement.execute(&[&path]).unwrap();
            }
        }
        Err(e) => {
            panic!("results err: {:?}", e);
        }
    }

    try!(trans.commit());

    Ok(count)
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use std::env;
    use super::get_file_list;

    #[test]
    fn test_get_file_list() {
        let _ = env_logger::try_init();

        let files = get_file_list(env::current_dir().unwrap());
        assert!(files.is_ok());
        assert!(files.unwrap().len() > 0);

        let files = get_file_list(env::current_dir().unwrap().join("Cargo.toml")).unwrap();
        assert_eq!(files[0], std::path::Path::new("Cargo.toml"));
    }
}
<commit_msg>Utilize multi-threaded tokio runtime for move-to-s3<commit_after>//! Simple module to store files in database.
//!
//! cratesfyi is generating more than 5 million files, they are small and mostly html files.
//! They are using so many inodes and it is better to store them in database instead of
//! filesystem. This module is adding files into database and retrieving them.

use std::path::{PathBuf, Path};
use postgres::Connection;
use rustc_serialize::json::{Json, ToJson};
use std::fs;
use std::io::Read;
use error::Result;
use failure::err_msg;
use rusoto_s3::{S3, PutObjectRequest, GetObjectRequest, S3Client};
use rusoto_core::region::Region;
use rusoto_credential::EnvironmentProvider;

fn get_file_list_from_dir<P: AsRef<Path>>(path: P,
                                          files: &mut Vec<PathBuf>)
                                          -> Result<()> {
    let path = path.as_ref();

    for file in try!(path.read_dir()) {
        let file = try!(file);

        if try!(file.file_type()).is_file() {
            files.push(file.path());
        } else if try!(file.file_type()).is_dir() {
            try!(get_file_list_from_dir(file.path(), files));
        }
    }

    Ok(())
}

pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>> {
    let path = path.as_ref();
    let mut files = Vec::new();

    if !path.exists() {
        return Err(err_msg("File not found"));
    } else if path.is_file() {
        files.push(PathBuf::from(path.file_name().unwrap()));
    } else if path.is_dir() {
        try!(get_file_list_from_dir(path, &mut files));
        for file_path in &mut files {
            // We want the paths in this list to not be {path}/bar.txt but just bar.txt
            *file_path = PathBuf::from(file_path.strip_prefix(path).unwrap());
        }
    }

    Ok(files)
}

pub struct Blob {
    pub path: String,
    pub mime: String,
    pub date_updated: time::Timespec,
    pub content: Vec<u8>,
}

pub fn get_path(conn: &Connection, path: &str) -> Option<Blob> {
    let rows = conn.query("SELECT path, mime, date_updated, content
                           FROM files
                           WHERE path = $1", &[&path]).unwrap();

    if rows.len() == 0 {
        None
    } else {
        let row = rows.get(0);
        let mut content = row.get(3);
        if content == b"in-s3" {
            let client = s3_client();
            content = client.and_then(|c| c.get_object(GetObjectRequest {
                bucket: "rust-docs-rs".into(),
                key: path.into(),
                ..Default::default()
            }).sync().ok()).and_then(|r| r.body).map(|b| {
                let mut b = b.into_blocking_read();
                let mut content = Vec::new();
                b.read_to_end(&mut content).unwrap();
                content
            }).unwrap();
        };

        Some(Blob {
            path: row.get(0),
            mime: row.get(1),
            date_updated: row.get(2),
            content,
        })
    }
}

fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() {
        return None;
    }
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        EnvironmentProvider::default(),
        std::env::var("S3_ENDPOINT").ok().map(|e| Region::Custom {
            name: "us-west-1".to_owned(),
            endpoint: e,
        }).unwrap_or(Region::UsWest1),
    ))
}

/// Adds files into database and returns list of files with their mime type in Json
pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
                                              prefix: &str,
                                              path: P)
                                              -> Result<Json> {
    use magic::{Cookie, flags};
    let cookie = try!(Cookie::open(flags::MIME_TYPE));
    try!(cookie.load::<&str>(&[]));

    let trans = try!(conn.transaction());
    let client = s3_client();
    let mut file_list_with_mimes: Vec<(String, PathBuf)> = Vec::new();

    for file_path in try!(get_file_list(&path)) {
        let (path, content, mime) = {
            let path = Path::new(path.as_ref()).join(&file_path);
            // Some files have insufficient permissions (like .lock file created by cargo in
            // documentation directory). We are skipping this files.
            let mut file = match fs::File::open(path) {
                Ok(f) => f,
                Err(_) => continue,
            };
            let mut content: Vec<u8> = Vec::new();
            try!(file.read_to_end(&mut content));
            let bucket_path = Path::new(prefix).join(&file_path)
                .into_os_string().into_string().unwrap();

            let mime = {
                let mime = try!(cookie.buffer(&content));
                // css's are causing some problem in browsers
                // magic will return text/plain for css file types
                // convert them to text/css
                // do the same for javascript files
                if mime == "text/plain" {
                    let e = file_path.extension().unwrap_or_default();
                    if e == "css" {
                        "text/css".to_owned()
                    } else if e == "js" {
                        "application/javascript".to_owned()
                    } else {
                        mime.to_owned()
                    }
                } else {
                    mime.to_owned()
                }
            };

            let content: Option<Vec<u8>> = if let Some(client) = &client {
                let s3_res = client.put_object(PutObjectRequest {
                    bucket: "rust-docs-rs".into(),
                    key: bucket_path.clone(),
                    body: Some(content.clone().into()),
                    content_type: Some(mime.clone()),
                    ..Default::default()
                }).sync();
                match s3_res {
                    // we've successfully uploaded the content, so steal it;
                    // we don't want to put it in the DB
                    Ok(_) => None,
                    // Since s3 was configured, we want to panic on failure to upload.
                    Err(e) => {
                        panic!("failed to upload to {}: {:?}", bucket_path, e)
                    },
                }
            } else {
                Some(content.clone().into())
            };

            file_list_with_mimes.push((mime.clone(), file_path.clone()));

            (
                bucket_path,
                content,
                mime,
            )
        };

        // check if file already exists in database
        let rows = try!(conn.query("SELECT COUNT(*) FROM files WHERE path = $1", &[&path]));

        let content = content.unwrap_or_else(|| "in-s3".to_owned().into());

        if rows.get(0).get::<usize, i64>(0) == 0 {
            try!(trans.query("INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                             &[&path, &mime, &content]));
        } else {
            try!(trans.query("UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                              WHERE path = $1",
                             &[&path, &mime, &content]));
        }
    }

    try!(trans.commit());

    file_list_to_json(file_list_with_mimes)
}

fn file_list_to_json(file_list: Vec<(String, PathBuf)>) -> Result<Json> {

    let mut file_list_json: Vec<Json> = Vec::new();

    for file in file_list {
        let mut v: Vec<String> = Vec::new();
        v.push(file.0.clone());
        v.push(file.1.into_os_string().into_string().unwrap());
        file_list_json.push(v.to_json());
    }

    Ok(file_list_json.to_json())
}

pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize> {
    let trans = try!(conn.transaction());
    let client = s3_client().expect("configured s3");

    let rows = try!(trans.query(
            &format!("SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}", n),
            &[]));
    let count = rows.len();

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();
    let mut futures = Vec::new();
    for row in &rows {
        let path: String = row.get(0);
        let mime: String = row.get(1);
        let content: Vec<u8> = row.get(2);
        let path_1 = path.clone();
        futures.push(client.put_object(PutObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.clone(),
            body: Some(content.into()),
            content_type: Some(mime),
            ..Default::default()
        }).map(move |_| {
            path_1
        }).map_err(move |e| {
            panic!("failed to upload to {}: {:?}", path, e)
        }));
    }

    use ::futures::future::Future;
    match rt.block_on(::futures::future::join_all(futures)) {
        Ok(paths) => {
            let statement = trans.prepare("UPDATE files SET content = E'in-s3' WHERE path = $1")
                .unwrap();
            for path in paths {
                statement.execute(&[&path]).unwrap();
            }
        }
        Err(e) => {
            panic!("results err: {:?}", e);
        }
    }

    try!(trans.commit());

    Ok(count)
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use std::env;
    use super::get_file_list;

    #[test]
    fn test_get_file_list() {
        let _ = env_logger::try_init();

        let files = get_file_list(env::current_dir().unwrap());
        assert!(files.is_ok());
        assert!(files.unwrap().len() > 0);

        let files = get_file_list(env::current_dir().unwrap().join("Cargo.toml")).unwrap();
        assert_eq!(files[0], std::path::Path::new("Cargo.toml"));
    }
}
<|endoftext|>


<commit_before>//! Simple module to store files in database.
//!
//! cratesfyi is generating more than 5 million files, they are small and mostly html files.
//! They are using so many inodes and it is better to store them in database instead of
//! filesystem. This module is adding files into database and retrieving them.

use std::path::{PathBuf, Path};
use postgres::Connection;
use rustc_serialize::json::{Json, ToJson};
use std::cmp;
use std::fs;
use std::io::Read;
use error::Result;
use failure::err_msg;
use rusoto_s3::{S3, PutObjectRequest, GetObjectRequest, S3Client};
use rusoto_core::region::Region;
use rusoto_credential::DefaultCredentialsProvider;

const MAX_CONCURRENT_UPLOADS: usize = 1000;

fn get_file_list_from_dir<P: AsRef<Path>>(path: P,
                                          files: &mut Vec<PathBuf>)
                                          -> Result<()> {
    let path = path.as_ref();

    for file in path.read_dir()? {
        let file = file?;

        if file.file_type()?.is_file() {
            files.push(file.path());
        } else if file.file_type()?.is_dir() {
            get_file_list_from_dir(file.path(), files)?;
        }
    }

    Ok(())
}

pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>> {
    let path = path.as_ref();
    let mut files = Vec::new();

    if !path.exists() {
        return Err(err_msg("File not found"));
    } else if path.is_file() {
        files.push(PathBuf::from(path.file_name().unwrap()));
    } else if path.is_dir() {
        get_file_list_from_dir(path, &mut files)?;
        for file_path in &mut files {
            // We want the paths in this list to not be {path}/bar.txt but just bar.txt
            *file_path = PathBuf::from(file_path.strip_prefix(path).unwrap());
        }
    }

    Ok(files)
}

pub struct Blob {
    pub path: String,
    pub mime: String,
    pub date_updated: time::Timespec,
    pub content: Vec<u8>,
}

pub fn get_path(conn: &Connection, path: &str) -> Option<Blob> {
    if let Some(client) = s3_client() {
        let res = client.get_object(GetObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.into(),
            ..Default::default()
        }).sync();

        let res = match res {
            Ok(r) => r,
            Err(_) => {
                return None;
            }
        };

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::new();
        b.read_to_end(&mut content).unwrap();

        let last_modified = res.last_modified.unwrap();
        let last_modified = time::strptime(&last_modified, "%a, %d %b %Y %H:%M:%S %Z")
            .unwrap_or_else(|e| panic!("failed to parse {:?} as timespec: {:?}", last_modified, e))
            .to_timespec();

        Some(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated: last_modified,
            content,
        })
    } else {
        let rows = conn.query("SELECT path, mime, date_updated, content
                            FROM files
                            WHERE path = $1", &[&path]).unwrap();

        if rows.len() == 0 {
            None
        } else {
            let row = rows.get(0);

            Some(Blob {
                path: row.get(0),
                mime: row.get(1),
                date_updated: row.get(2),
                content: row.get(3),
            })
        }
    }
}

fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
        return None;
    }
    let creds = match DefaultCredentialsProvider::new() {
        Ok(creds) => creds,
        Err(err) => {
            warn!("failed to retrieve AWS credentials: {}", err);
            return None;
        }
    };
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        creds,
        std::env::var("S3_ENDPOINT").ok().map(|e| Region::Custom {
            name: "us-west-1".to_owned(),
            endpoint: e,
        }).unwrap_or(Region::UsWest1),
    ))
}

/// Adds files into database and returns list of files with their mime type in Json
pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
                                              prefix: &str,
                                              path: P)
                                              -> Result<Json> {
    use magic::{Cookie, flags};
    let cookie = Cookie::open(flags::MIME_TYPE)?;
    cookie.load::<&str>(&[])?;

    let trans = conn.transaction()?;

    use std::collections::HashMap;
    let mut file_paths_and_mimes: HashMap<PathBuf, String> = HashMap::new();

    use futures::future::Future;

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();

    let mut to_upload = get_file_list(&path)?;
    let mut batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
    let mut currently_uploading: Vec<_> = to_upload.drain(..batch_size).collect();
    let mut attempts = 0;

    while !to_upload.is_empty() || !currently_uploading.is_empty() {

        let mut futures = Vec::new();
        let client = s3_client();

        for file_path in &currently_uploading {
            let path = Path::new(path.as_ref()).join(&file_path);
            // Some files have insufficient permissions (like .lock file created by cargo in
            // documentation directory). We are skipping this files.
            let mut file = match fs::File::open(path) {
                Ok(f) => f,
                Err(_) => continue,
            };
            let mut content: Vec<u8> = Vec::new();
            file.read_to_end(&mut content)?;
            let bucket_path = Path::new(prefix).join(&file_path)
                .into_os_string().into_string().unwrap();

            let mime = {
                let mime = cookie.buffer(&content)?;
                // css's are causing some problem in browsers
                // magic will return text/plain for css file types
                // convert them to text/css
                // do the same for javascript files
                if mime == "text/plain" {
                    let e = file_path.extension().unwrap_or_default();
                    if e == "css" {
                        "text/css".to_owned()
                    } else if e == "js" {
                        "application/javascript".to_owned()
                    } else {
                        mime.to_owned()
                    }
                } else {
                    mime.to_owned()
                }
            };

            if let Some(client) = &client {
                futures.push(client.put_object(PutObjectRequest {
                    bucket: "rust-docs-rs".into(),
                    key: bucket_path.clone(),
                    body: Some(content.clone().into()),
                    content_type: Some(mime.clone()),
                    ..Default::default()
                }).inspect(|_| {
                    crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(1);
                }));
            } else {
                // If AWS credentials are configured, don't insert/update the database
                // check if file already exists in database
                let rows = conn.query("SELECT COUNT(*) FROM files WHERE path = $1", &[&bucket_path])?;

                if rows.get(0).get::<usize, i64>(0) == 0 {
                    trans.query("INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                                    &[&bucket_path, &mime, &content])?;
                } else {
                    trans.query("UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                                    WHERE path = $1",
                                    &[&bucket_path, &mime, &content])?;
                }
            }

            file_paths_and_mimes.insert(file_path.clone(), mime.clone());
        }

        if !futures.is_empty() {
            attempts += 1;

            if rt.block_on(::futures::future::join_all(futures)).is_ok() {
                // this batch was successful, start another batch if there are still more files
                batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
                currently_uploading = to_upload.drain(..batch_size).collect();
                attempts = 0;
            } else {
                // if any futures error, leave `currently_uploading` in tact so that we can retry the batch
                if attempts > 2 {
                    panic!("failed to upload 3 times, exiting");
                }
            }
        } else {
            batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
            currently_uploading = to_upload.drain(..batch_size).collect();
        }
    }

    trans.commit()?;

    let file_list_with_mimes: Vec<(String, PathBuf)> = file_paths_and_mimes
        .into_iter()
        .map(|(file_path, mime)| (mime, file_path))
        .collect();
    file_list_to_json(file_list_with_mimes)
}

fn file_list_to_json(file_list: Vec<(String, PathBuf)>) -> Result<Json> {

    let mut file_list_json: Vec<Json> = Vec::new();

    for file in file_list {
        let mut v: Vec<String> = Vec::new();
        v.push(file.0.clone());
        v.push(file.1.into_os_string().into_string().unwrap());
        file_list_json.push(v.to_json());
    }

    Ok(file_list_json.to_json())
}

pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize> {
    let trans = conn.transaction()?;
    let client = s3_client().expect("configured s3");

    let rows = trans.query(
            &format!("SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}", n),
            &[])?;
    let count = rows.len();

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();
    let mut futures = Vec::new();
    for row in &rows {
        let path: String = row.get(0);
        let mime: String = row.get(1);
        let content: Vec<u8> = row.get(2);
        let path_1 = path.clone();
        futures.push(client.put_object(PutObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.clone(),
            body: Some(content.into()),
            content_type: Some(mime),
            ..Default::default()
        }).map(move |_| {
            path_1
        }).map_err(move |e| {
            panic!("failed to upload to {}: {:?}", path, e)
        }));
    }

    use ::futures::future::Future;
    match rt.block_on(::futures::future::join_all(futures)) {
        Ok(paths) => {
            let statement = trans.prepare("DELETE FROM files WHERE path = $1").unwrap();
            for path in paths {
                statement.execute(&[&path]).unwrap();
            }
        }
        Err(e) => {
            panic!("results err: {:?}", e);
        }
    }

    trans.commit()?;

    Ok(count)
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use std::env;
    use super::get_file_list;

    #[test]
    fn test_get_file_list() {
        let _ = env_logger::try_init();

        let files = get_file_list(env::current_dir().unwrap());
        assert!(files.is_ok());
        assert!(files.unwrap().len() > 0);

        let files = get_file_list(env::current_dir().unwrap().join("Cargo.toml")).unwrap();
        assert_eq!(files[0], std::path::Path::new("Cargo.toml"));
    }
}
<commit_msg>db.file: allow overriding the region<commit_after>//! Simple module to store files in database.
//!
//! cratesfyi is generating more than 5 million files, they are small and mostly html files.
//! They are using so many inodes and it is better to store them in database instead of
//! filesystem. This module is adding files into database and retrieving them.

use std::path::{PathBuf, Path};
use postgres::Connection;
use rustc_serialize::json::{Json, ToJson};
use std::cmp;
use std::fs;
use std::io::Read;
use error::Result;
use failure::err_msg;
use rusoto_s3::{S3, PutObjectRequest, GetObjectRequest, S3Client};
use rusoto_core::region::Region;
use rusoto_credential::DefaultCredentialsProvider;

const MAX_CONCURRENT_UPLOADS: usize = 1000;

fn get_file_list_from_dir<P: AsRef<Path>>(path: P,
                                          files: &mut Vec<PathBuf>)
                                          -> Result<()> {
    let path = path.as_ref();

    for file in path.read_dir()? {
        let file = file?;

        if file.file_type()?.is_file() {
            files.push(file.path());
        } else if file.file_type()?.is_dir() {
            get_file_list_from_dir(file.path(), files)?;
        }
    }

    Ok(())
}

pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>> {
    let path = path.as_ref();
    let mut files = Vec::new();

    if !path.exists() {
        return Err(err_msg("File not found"));
    } else if path.is_file() {
        files.push(PathBuf::from(path.file_name().unwrap()));
    } else if path.is_dir() {
        get_file_list_from_dir(path, &mut files)?;
        for file_path in &mut files {
            // We want the paths in this list to not be {path}/bar.txt but just bar.txt
            *file_path = PathBuf::from(file_path.strip_prefix(path).unwrap());
        }
    }

    Ok(files)
}

pub struct Blob {
    pub path: String,
    pub mime: String,
    pub date_updated: time::Timespec,
    pub content: Vec<u8>,
}

pub fn get_path(conn: &Connection, path: &str) -> Option<Blob> {
    if let Some(client) = s3_client() {
        let res = client.get_object(GetObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.into(),
            ..Default::default()
        }).sync();

        let res = match res {
            Ok(r) => r,
            Err(_) => {
                return None;
            }
        };

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::new();
        b.read_to_end(&mut content).unwrap();

        let last_modified = res.last_modified.unwrap();
        let last_modified = time::strptime(&last_modified, "%a, %d %b %Y %H:%M:%S %Z")
            .unwrap_or_else(|e| panic!("failed to parse {:?} as timespec: {:?}", last_modified, e))
            .to_timespec();

        Some(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated: last_modified,
            content,
        })
    } else {
        let rows = conn.query("SELECT path, mime, date_updated, content
                            FROM files
                            WHERE path = $1", &[&path]).unwrap();

        if rows.len() == 0 {
            None
        } else {
            let row = rows.get(0);

            Some(Blob {
                path: row.get(0),
                mime: row.get(1),
                date_updated: row.get(2),
                content: row.get(3),
            })
        }
    }
}

fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
        return None;
    }
    let creds = match DefaultCredentialsProvider::new() {
        Ok(creds) => creds,
        Err(err) => {
            warn!("failed to retrieve AWS credentials: {}", err);
            return None;
        }
    };
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        creds,
        std::env::var("S3_ENDPOINT").ok().map(|e| Region::Custom {
            name: std::env::var("S3_REGION")
                .unwrap_or_else(|| "us-west-1".to_owned()),
            endpoint: e,
        }).unwrap_or(Region::UsWest1),
    ))
}

/// Adds files into database and returns list of files with their mime type in Json
pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
                                              prefix: &str,
                                              path: P)
                                              -> Result<Json> {
    use magic::{Cookie, flags};
    let cookie = Cookie::open(flags::MIME_TYPE)?;
    cookie.load::<&str>(&[])?;

    let trans = conn.transaction()?;

    use std::collections::HashMap;
    let mut file_paths_and_mimes: HashMap<PathBuf, String> = HashMap::new();

    use futures::future::Future;

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();

    let mut to_upload = get_file_list(&path)?;
    let mut batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
    let mut currently_uploading: Vec<_> = to_upload.drain(..batch_size).collect();
    let mut attempts = 0;

    while !to_upload.is_empty() || !currently_uploading.is_empty() {

        let mut futures = Vec::new();
        let client = s3_client();

        for file_path in &currently_uploading {
            let path = Path::new(path.as_ref()).join(&file_path);
            // Some files have insufficient permissions (like .lock file created by cargo in
            // documentation directory). We are skipping this files.
            let mut file = match fs::File::open(path) {
                Ok(f) => f,
                Err(_) => continue,
            };
            let mut content: Vec<u8> = Vec::new();
            file.read_to_end(&mut content)?;
            let bucket_path = Path::new(prefix).join(&file_path)
                .into_os_string().into_string().unwrap();

            let mime = {
                let mime = cookie.buffer(&content)?;
                // css's are causing some problem in browsers
                // magic will return text/plain for css file types
                // convert them to text/css
                // do the same for javascript files
                if mime == "text/plain" {
                    let e = file_path.extension().unwrap_or_default();
                    if e == "css" {
                        "text/css".to_owned()
                    } else if e == "js" {
                        "application/javascript".to_owned()
                    } else {
                        mime.to_owned()
                    }
                } else {
                    mime.to_owned()
                }
            };

            if let Some(client) = &client {
                futures.push(client.put_object(PutObjectRequest {
                    bucket: "rust-docs-rs".into(),
                    key: bucket_path.clone(),
                    body: Some(content.clone().into()),
                    content_type: Some(mime.clone()),
                    ..Default::default()
                }).inspect(|_| {
                    crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(1);
                }));
            } else {
                // If AWS credentials are configured, don't insert/update the database
                // check if file already exists in database
                let rows = conn.query("SELECT COUNT(*) FROM files WHERE path = $1", &[&bucket_path])?;

                if rows.get(0).get::<usize, i64>(0) == 0 {
                    trans.query("INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                                    &[&bucket_path, &mime, &content])?;
                } else {
                    trans.query("UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                                    WHERE path = $1",
                                    &[&bucket_path, &mime, &content])?;
                }
            }

            file_paths_and_mimes.insert(file_path.clone(), mime.clone());
        }

        if !futures.is_empty() {
            attempts += 1;

            if rt.block_on(::futures::future::join_all(futures)).is_ok() {
                // this batch was successful, start another batch if there are still more files
                batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
                currently_uploading = to_upload.drain(..batch_size).collect();
                attempts = 0;
            } else {
                // if any futures error, leave `currently_uploading` in tact so that we can retry the batch
                if attempts > 2 {
                    panic!("failed to upload 3 times, exiting");
                }
            }
        } else {
            batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
            currently_uploading = to_upload.drain(..batch_size).collect();
        }
    }

    trans.commit()?;

    let file_list_with_mimes: Vec<(String, PathBuf)> = file_paths_and_mimes
        .into_iter()
        .map(|(file_path, mime)| (mime, file_path))
        .collect();
    file_list_to_json(file_list_with_mimes)
}

fn file_list_to_json(file_list: Vec<(String, PathBuf)>) -> Result<Json> {

    let mut file_list_json: Vec<Json> = Vec::new();

    for file in file_list {
        let mut v: Vec<String> = Vec::new();
        v.push(file.0.clone());
        v.push(file.1.into_os_string().into_string().unwrap());
        file_list_json.push(v.to_json());
    }

    Ok(file_list_json.to_json())
}

pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize> {
    let trans = conn.transaction()?;
    let client = s3_client().expect("configured s3");

    let rows = trans.query(
            &format!("SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}", n),
            &[])?;
    let count = rows.len();

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();
    let mut futures = Vec::new();
    for row in &rows {
        let path: String = row.get(0);
        let mime: String = row.get(1);
        let content: Vec<u8> = row.get(2);
        let path_1 = path.clone();
        futures.push(client.put_object(PutObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.clone(),
            body: Some(content.into()),
            content_type: Some(mime),
            ..Default::default()
        }).map(move |_| {
            path_1
        }).map_err(move |e| {
            panic!("failed to upload to {}: {:?}", path, e)
        }));
    }

    use ::futures::future::Future;
    match rt.block_on(::futures::future::join_all(futures)) {
        Ok(paths) => {
            let statement = trans.prepare("DELETE FROM files WHERE path = $1").unwrap();
            for path in paths {
                statement.execute(&[&path]).unwrap();
            }
        }
        Err(e) => {
            panic!("results err: {:?}", e);
        }
    }

    trans.commit()?;

    Ok(count)
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use std::env;
    use super::get_file_list;

    #[test]
    fn test_get_file_list() {
        let _ = env_logger::try_init();

        let files = get_file_list(env::current_dir().unwrap());
        assert!(files.is_ok());
        assert!(files.unwrap().len() > 0);

        let files = get_file_list(env::current_dir().unwrap().join("Cargo.toml")).unwrap();
        assert_eq!(files[0], std::path::Path::new("Cargo.toml"));
    }
}
<|endoftext|>


<commit_before>use types::{OscPacket, OscType, OscResult, OscMessage, OscBundle};
use errors::OscError;
use utils;

use std::io;
use std::io::{Read, BufRead};

use byteorder::{BigEndian, ReadBytesExt};

/// Common MTP size for ethernet
pub const MTP: usize = 1536;

pub fn decode(msg: &[u8]) -> OscResult<OscPacket> {
    match msg[0] as char {
        '/' => {
            decode_message(msg)
        }
        '#' => {
            decode_bundle(msg)
        }
        _ => Err(OscError::BadPacket("Unknown message format.".to_string())),
    }
}

fn decode_message(msg: &[u8]) -> OscResult<OscPacket> {
    let mut cursor: io::Cursor<&[u8]> = io::Cursor::new(msg);

    let addr: String = try!(read_osc_string(&mut cursor));
    let type_tags: String = try!(read_osc_string(&mut cursor));

    if type_tags.len() > 1 {
        let args: Vec<OscType> = try!(read_osc_args(&mut cursor, type_tags));

        Ok(OscPacket::Message(OscMessage {
            addr: addr,
            args: Some(args),
        }))
    } else {
        Ok(OscPacket::Message(OscMessage {
            addr: addr,
            args: None,
        }))
    }
}

fn decode_bundle(msg: &[u8]) -> OscResult<OscPacket> {
    let mut cursor: io::Cursor<&[u8]> = io::Cursor::new(msg);

    let b = try!(read_osc_string(&mut cursor));
    if b != "bundle" {
        return Err(OscError::BadBundle);
    }

    let time_tag = try!(read_time_tag(&mut cursor));
    let mut bundle: Vec<OscPacket> = Vec::new();
    let size: usize = try!(cursor.read_u32::<BigEndian>()
                                 .map_err(OscError::ByteOrderError)) as usize;
    let mut buf: Vec<u8> = Vec::new();
    let cnt: usize = try!(cursor.take(size as u64)
                                .read_to_end(&mut buf)
                                .map_err(OscError::ReadError));
    if cnt == size {
        try!(decode(&mut buf).map(|p| bundle.push(p)));
    } else {
        return Err(OscError::BadBundle);
    }

    Ok(OscPacket::Bundle(OscBundle {
        timetag: time_tag,
        content: bundle,
    }))
}

fn read_osc_string(cursor: &mut io::Cursor<&[u8]>) -> OscResult<String> {
    let mut str_buf: Vec<u8> = Vec::new();
    try!(cursor.read_until(0, &mut str_buf)
               .map_err(OscError::ReadError)); // ignore returned byte count
    pad_cursor(cursor);
    // convert to String and remove nul bytes
    String::from_utf8(str_buf)
        .map_err(|e| OscError::StringError(e))
        .map(|s| {
            s.trim_matches(0u8 as char)
             .to_string()
        })
}

fn read_osc_args(cursor: &mut io::Cursor<&[u8]>, raw_type_tags: String) -> OscResult<Vec<OscType>> {
    let type_tags: Vec<char> = raw_type_tags.chars()
                                            .skip(1)
                                            .map(|c| c as char)
                                            .collect();
    let mut args: Vec<OscType> = Vec::with_capacity(type_tags.len());
    for tag in type_tags {
        match read_osc_arg(cursor, tag) {
            Ok(arg) => {
                args.push(arg);
            }
            Err(e) => return Err(e),
        }
    }
    Ok(args)
}

fn read_osc_arg(cursor: &mut io::Cursor<&[u8]>, tag: char) -> OscResult<OscType> {
    match tag {
        'f' => {
            cursor.read_f32::<BigEndian>()
                  .map(|f| OscType::Float(f))
                  .map_err(|e| OscError::ByteOrderError(e))
        }
        'd' => {
            cursor.read_f64::<BigEndian>()
                  .map(|d| OscType::Double(d))
                  .map_err(|e| OscError::ByteOrderError(e))
        }
        'i' => {
            cursor.read_i32::<BigEndian>()
                  .map(|i| OscType::Int(i))
                  .map_err(|e| OscError::ByteOrderError(e))
        }
        'h' => {
            cursor.read_i64::<BigEndian>()
                  .map(|l| OscType::Long(l))
                  .map_err(|e| OscError::ByteOrderError(e))
        }
        's' => {
            read_osc_string(cursor).map(|s| OscType::String(s))
        }
        't' => {
            // http://opensoundcontrol.org/node/3/#timetags
            read_time_tag(cursor)
        }
        'b' => {
            match cursor.read_u32::<BigEndian>() {
                Ok(size) => {
                    let mut buf: Vec<u8> = Vec::with_capacity(size as usize);
                    match cursor.take(size as u64).read(&mut buf) {
                        Ok(blob_size) => Ok(OscType::Blob(buf)),
                        Err(e) => Err(OscError::ReadError(e)),
                    }
                }
                // TODO: use correct error type
                Err(e) => return Err(OscError::BadBundle),
            }
        }
        _ => Err(OscError::BadArg(format!("Type tag \"{}\" is not implemented!", tag))),
    }
}

fn read_time_tag(cursor: &mut io::Cursor<&[u8]>) -> OscResult<OscType> {
    match cursor.read_u32::<BigEndian>() {
        Ok(date) => {
            match cursor.read_u32::<BigEndian>() {
                Ok(frac) => Ok(OscType::Time(date, frac)),
                Err(e) => Err(OscError::ByteOrderError(e)),
            }
        }
        Err(e) => Err(OscError::ByteOrderError(e)),
    }
}

fn pad_cursor(cursor: &mut io::Cursor<&[u8]>) {
    let pos = cursor.position();
    cursor.set_position(utils::pad(pos));
}
<commit_msg>Refactor `read_osc_args`<commit_after>use types::{OscPacket, OscType, OscResult, OscMessage, OscBundle};
use errors::OscError;
use utils;

use std::io;
use std::io::{Read, BufRead};

use byteorder::{BigEndian, ReadBytesExt};

/// Common MTP size for ethernet
pub const MTP: usize = 1536;

pub fn decode(msg: &[u8]) -> OscResult<OscPacket> {
    match msg[0] as char {
        '/' => {
            decode_message(msg)
        }
        '#' => {
            decode_bundle(msg)
        }
        _ => Err(OscError::BadPacket("Unknown message format.".to_string())),
    }
}

fn decode_message(msg: &[u8]) -> OscResult<OscPacket> {
    let mut cursor: io::Cursor<&[u8]> = io::Cursor::new(msg);

    let addr: String = try!(read_osc_string(&mut cursor));
    let type_tags: String = try!(read_osc_string(&mut cursor));

    if type_tags.len() > 1 {
        let args: Vec<OscType> = try!(read_osc_args(&mut cursor, type_tags));

        Ok(OscPacket::Message(OscMessage {
            addr: addr,
            args: Some(args),
        }))
    } else {
        Ok(OscPacket::Message(OscMessage {
            addr: addr,
            args: None,
        }))
    }
}

fn decode_bundle(msg: &[u8]) -> OscResult<OscPacket> {
    let mut cursor: io::Cursor<&[u8]> = io::Cursor::new(msg);

    let b = try!(read_osc_string(&mut cursor));
    if b != "bundle" {
        return Err(OscError::BadBundle);
    }

    let time_tag = try!(read_time_tag(&mut cursor));
    let mut bundle: Vec<OscPacket> = Vec::new();
    let size: usize = try!(cursor.read_u32::<BigEndian>()
                                 .map_err(OscError::ByteOrderError)) as usize;
    let mut buf: Vec<u8> = Vec::new();
    let cnt: usize = try!(cursor.take(size as u64)
                                .read_to_end(&mut buf)
                                .map_err(OscError::ReadError));
    if cnt == size {
        try!(decode(&mut buf).map(|p| bundle.push(p)));
    } else {
        return Err(OscError::BadBundle);
    }

    Ok(OscPacket::Bundle(OscBundle {
        timetag: time_tag,
        content: bundle,
    }))
}

fn read_osc_string(cursor: &mut io::Cursor<&[u8]>) -> OscResult<String> {
    let mut str_buf: Vec<u8> = Vec::new();
    try!(cursor.read_until(0, &mut str_buf)
               .map_err(OscError::ReadError)); // ignore returned byte count
    pad_cursor(cursor);
    // convert to String and remove nul bytes
    String::from_utf8(str_buf)
        .map_err(|e| OscError::StringError(e))
        .map(|s| {
            s.trim_matches(0u8 as char)
             .to_string()
        })
}

fn read_osc_args(cursor: &mut io::Cursor<&[u8]>, raw_type_tags: String) -> OscResult<Vec<OscType>> {
    let type_tags: Vec<char> = raw_type_tags.chars()
                                            .skip(1)
                                            .map(|c| c as char)
                                            .collect();
    let mut args: Vec<OscType> = Vec::with_capacity(type_tags.len());
    for tag in type_tags {
        let arg: OscType = try!(read_osc_arg(cursor, tag));
        args.push(arg);
    }
    Ok(args)
}

fn read_osc_arg(cursor: &mut io::Cursor<&[u8]>, tag: char) -> OscResult<OscType> {
    match tag {
        'f' => {
            cursor.read_f32::<BigEndian>()
                  .map(|f| OscType::Float(f))
                  .map_err(|e| OscError::ByteOrderError(e))
        }
        'd' => {
            cursor.read_f64::<BigEndian>()
                  .map(|d| OscType::Double(d))
                  .map_err(|e| OscError::ByteOrderError(e))
        }
        'i' => {
            cursor.read_i32::<BigEndian>()
                  .map(|i| OscType::Int(i))
                  .map_err(|e| OscError::ByteOrderError(e))
        }
        'h' => {
            cursor.read_i64::<BigEndian>()
                  .map(|l| OscType::Long(l))
                  .map_err(|e| OscError::ByteOrderError(e))
        }
        's' => {
            read_osc_string(cursor).map(|s| OscType::String(s))
        }
        't' => {
            // http://opensoundcontrol.org/node/3/#timetags
            read_time_tag(cursor)
        }
        'b' => {
            match cursor.read_u32::<BigEndian>() {
                Ok(size) => {
                    let mut buf: Vec<u8> = Vec::with_capacity(size as usize);
                    match cursor.take(size as u64).read(&mut buf) {
                        Ok(blob_size) => Ok(OscType::Blob(buf)),
                        Err(e) => Err(OscError::ReadError(e)),
                    }
                }
                // TODO: use correct error type
                Err(e) => return Err(OscError::BadBundle),
            }
        }
        _ => Err(OscError::BadArg(format!("Type tag \"{}\" is not implemented!", tag))),
    }
}

fn read_time_tag(cursor: &mut io::Cursor<&[u8]>) -> OscResult<OscType> {
    match cursor.read_u32::<BigEndian>() {
        Ok(date) => {
            match cursor.read_u32::<BigEndian>() {
                Ok(frac) => Ok(OscType::Time(date, frac)),
                Err(e) => Err(OscError::ByteOrderError(e)),
            }
        }
        Err(e) => Err(OscError::ByteOrderError(e)),
    }
}

fn pad_cursor(cursor: &mut io::Cursor<&[u8]>) {
    let pos = cursor.position();
    cursor.set_position(utils::pad(pos));
}
<|endoftext|>