<commit_before>use std::io;
use std::io::prelude::*;
use std::net::TcpStream;
use std::fs::File;
use std::env;
use std::process;

extern crate chrono;
use chrono::*;

extern crate aggregated_stats;

#[macro_use]
extern crate clap;

extern crate env_logger;
#[macro_use]
extern crate log;

extern crate prometheus;
extern crate hyper;

mod analyzer;
mod args;
mod filter;
mod log_parser;
mod render;
mod request_response_matcher;
mod http_handler;
mod result;

fn main() {
    env_logger::init().expect("Failed to initialize logging.");

    let args = args::parse_args(env::args()).expect("Failed to parse arguments.");

    if args.prometheus_listen.is_some() {
        let binding_address = args.prometheus_listen.clone().unwrap();
        http_handler::listen_http(args, &binding_address);
    } else {
        let result = run(&args);

        let mut stream;
        let mut stdout;

        let mut renderers: Vec<Box<render::Renderer>>;
        renderers = vec![];

        if !args.quiet {
            stdout = io::stdout();
            renderers.push(Box::new(render::terminal::TerminalRenderer::new(&mut stdout)));
        }

        if args.graphite_server.is_some() {
            stream = TcpStream::connect((args.graphite_server.as_ref().unwrap().as_str(),
                                         args.graphite_port.unwrap()))
                .expect("Could not connect to the Graphite server");

            renderers.push(Box::new(render::graphite::GraphiteRenderer::new(Utc::now(),
                                                                            args.graphite_prefix
                                                                                .clone(),
                                                                            &mut stream)));
        }

        if args.influxdb_write_url.is_some() {
            renderers.push(Box::new(render::influxdb::InfluxDbRenderer::new(&args.influxdb_write_url
                                                                           .clone()
                                                                           .unwrap(),
                                                                       args.influxdb_tags
                                                                           .clone())));
        }

        for mut renderer in renderers {
            renderer.render(result.clone());
        }
    }
}

fn run(args: &args::RequestLogAnalyzerArgs) -> result::RequestLogAnalyzerResult {
    let input: Box<io::Read> = match args.filename.as_ref() {
        "-" => Box::new(io::stdin()),
        _ => {
            match File::open(&args.filename) {
                Ok(file) => Box::new(file),
                Err(err) => {
                    eprintln!("Failed to open file {}: {}", &args.filename, err);
                    process::exit(1);
                }
            }
        }
    };

    let reader = io::BufReader::new(input);

    let mut events_iterator = reader.lines()
        .map(log_parser::parse_line)
        .filter(|event| event.is_ok())
        .map(|event| event.unwrap());

    let pairs_iterator =
        request_response_matcher::RequestResponsePairIterator::new(&mut events_iterator)
            .filter(|pair| filter::matches_filter(pair, &args.conditions));

    analyzer::analyze_iterator(pairs_iterator)
}

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

    #[test]
    fn test_run() {
        let args = args::RequestLogAnalyzerArgs {
            filename: String::from("src/test/simple-1.log"),
            conditions: filter::FilterConditions {
                include_terms: None,
                exclude_terms: None,
                latest_time: None,
            },
            graphite_server: None,
            graphite_port: Some(2003),
            graphite_prefix: None,
            prometheus_listen: None,
            influxdb_write_url: None,
            influxdb_tags: None,
            quiet: false,
        };

        let result = run(&args);
        assert_eq!(result.count, 2);

        let timing = result.timing.unwrap();
        assert_eq!(timing.min, 7);
        assert_eq!(timing.max, 10);

        assert!(result.error.is_some());
    }

    #[test]
    fn test_run_ignore_broken_lines() {
        let args = args::RequestLogAnalyzerArgs {
            filename: String::from("src/test/broken.log"),
            conditions: filter::FilterConditions {
                include_terms: None,
                exclude_terms: None,
                latest_time: None,
            },
            graphite_server: None,
            graphite_port: Some(2003),
            graphite_prefix: None,
            prometheus_listen: None,
            influxdb_write_url: None,
            influxdb_tags: None,
            quiet: false,
        };

        let result = run(&args);
        assert_eq!(result.count, 1);
    }
}
<commit_msg>Handled failure instead of panic<commit_after>use std::io;
use std::io::prelude::*;
use std::net::TcpStream;
use std::fs::File;
use std::env;
use std::process;

extern crate chrono;
use chrono::*;

extern crate aggregated_stats;

#[macro_use]
extern crate clap;

extern crate env_logger;
#[macro_use]
extern crate log;

extern crate prometheus;
extern crate hyper;

mod analyzer;
mod args;
mod filter;
mod log_parser;
mod render;
mod request_response_matcher;
mod http_handler;
mod result;

fn main() {
    env_logger::init().expect("Failed to initialize logging.");

    let args = args::parse_args(env::args()).expect("Failed to parse arguments.");

    if args.prometheus_listen.is_some() {
        let binding_address = args.prometheus_listen.clone().unwrap();
        http_handler::listen_http(args, &binding_address);
    } else {
        let result = run(&args);

        let mut stream;
        let mut stdout;

        let mut renderers: Vec<Box<render::Renderer>>;
        renderers = vec![];

        if !args.quiet {
            stdout = io::stdout();
            renderers.push(Box::new(render::terminal::TerminalRenderer::new(&mut stdout)));
        }

        if args.graphite_server.is_some() {
            stream = match TcpStream::connect((args.graphite_server.as_ref().unwrap().as_str(),
                                               args.graphite_port.unwrap())) {
                Ok(stream) => stream,
                Err(err) => {
                    eprintln!("Could not connect to the Graphite server: {}", err);
                    process::exit(1);
                }
            };

            renderers.push(Box::new(render::graphite::GraphiteRenderer::new(Utc::now(),
                                                                            args.graphite_prefix
                                                                                .clone(),
                                                                            &mut stream)));
        }

        if args.influxdb_write_url.is_some() {
            renderers.push(Box::new(render::influxdb::InfluxDbRenderer::new(&args.influxdb_write_url
                                                                           .clone()
                                                                           .unwrap(),
                                                                       args.influxdb_tags
                                                                           .clone())));
        }

        for mut renderer in renderers {
            renderer.render(result.clone());
        }
    }
}

fn run(args: &args::RequestLogAnalyzerArgs) -> result::RequestLogAnalyzerResult {
    let input: Box<io::Read> = match args.filename.as_ref() {
        "-" => Box::new(io::stdin()),
        _ => {
            match File::open(&args.filename) {
                Ok(file) => Box::new(file),
                Err(err) => {
                    eprintln!("Failed to open file {}: {}", &args.filename, err);
                    process::exit(1);
                }
            }
        }
    };

    let reader = io::BufReader::new(input);

    let mut events_iterator = reader.lines()
        .map(log_parser::parse_line)
        .filter(|event| event.is_ok())
        .map(|event| event.unwrap());

    let pairs_iterator =
        request_response_matcher::RequestResponsePairIterator::new(&mut events_iterator)
            .filter(|pair| filter::matches_filter(pair, &args.conditions));

    analyzer::analyze_iterator(pairs_iterator)
}

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

    #[test]
    fn test_run() {
        let args = args::RequestLogAnalyzerArgs {
            filename: String::from("src/test/simple-1.log"),
            conditions: filter::FilterConditions {
                include_terms: None,
                exclude_terms: None,
                latest_time: None,
            },
            graphite_server: None,
            graphite_port: Some(2003),
            graphite_prefix: None,
            prometheus_listen: None,
            influxdb_write_url: None,
            influxdb_tags: None,
            quiet: false,
        };

        let result = run(&args);
        assert_eq!(result.count, 2);

        let timing = result.timing.unwrap();
        assert_eq!(timing.min, 7);
        assert_eq!(timing.max, 10);

        assert!(result.error.is_some());
    }

    #[test]
    fn test_run_ignore_broken_lines() {
        let args = args::RequestLogAnalyzerArgs {
            filename: String::from("src/test/broken.log"),
            conditions: filter::FilterConditions {
                include_terms: None,
                exclude_terms: None,
                latest_time: None,
            },
            graphite_server: None,
            graphite_port: Some(2003),
            graphite_prefix: None,
            prometheus_listen: None,
            influxdb_write_url: None,
            influxdb_tags: None,
            quiet: false,
        };

        let result = run(&args);
        assert_eq!(result.count, 1);
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// Needed to derive `Serialize` on ServiceProperties
#![feature(custom_derive, plugin)]
#![plugin(serde_macros)]
// For Docopt macro
#![plugin(docopt_macros)]

// Make linter fail for every warning
#![plugin(clippy)]
#![deny(clippy)]

#![cfg_attr(test, feature(const_fn))] // Dependency of stainless
#![cfg_attr(test, plugin(stainless))] // Test runner

#![feature(reflect_marker)]

#![feature(associated_consts)]

extern crate chrono;
extern crate core;
extern crate docopt;
extern crate env_logger;
extern crate foxbox_adapters;
extern crate foxbox_taxonomy;
extern crate foxbox_users;
#[macro_use]
extern crate hyper;
#[macro_use]
extern crate iron;
extern crate iron_cors;
#[cfg(test)]
extern crate iron_test;
extern crate libc;
#[macro_use]
extern crate log;
extern crate mio;
extern crate mount;
extern crate nix;
extern crate openssl;
extern crate openssl_sys;
extern crate rand;
extern crate router;
extern crate rustc_serialize;
extern crate rusqlite;
extern crate serde;
extern crate serde_json;
extern crate staticfile;
extern crate time;
extern crate timer;
extern crate transformable_channels;
extern crate unicase;
extern crate uuid;
extern crate ws;
extern crate multicast_dns;
extern crate xml;

#[cfg(test)]
extern crate regex;
#[cfg(test)]
extern crate tempdir;

// Need to be declared first so to let the macros be visible from other modules.
#[macro_use]
mod utils;
mod adapters;
mod config_store;
mod controller;
mod http_server;
mod managed_process;
mod profile_service;
mod registration;
mod service;
mod upnp;
mod service_router;
mod stable_uuid;
mod static_router;
mod taxonomy_router;
mod tls;
mod traits;
mod tunnel_controller;
mod ws_server;

#[cfg(test)]
mod stubs {
    #![allow(dead_code)]
    #![allow(unused_variables)]
    #![allow(boxed_local)]
    pub mod service;
    pub mod controller;
}

use controller::FoxBox;
use env_logger::LogBuilder;
use tunnel_controller:: { TunnelConfig, Tunnel };
use libc::SIGINT;
use log::{ LogRecord, LogLevelFilter };

use multicast_dns::host::HostManager;
use profile_service::ProfilePath;
use std::env;
use std::mem;
use std::sync::atomic::{ AtomicBool, Ordering, ATOMIC_BOOL_INIT };
use tls::TlsOption;
use traits::Controller;

docopt!(Args derive Debug, "
Usage: foxbox [-v] [-h] [-l <hostname>] [-p <port>] [-w <wsport>] [-d <profile_path>] [-r <url>] [-i <iface>] [-t <tunnel>] [-s <secret>] [--disable-tls] [--remote-name <hostname>] [-c <namespace;key;value>]...

Options:
    -v, --verbose            Toggle verbose output.
    -l, --local-name <hostname>    Set local hostname. Linux only. Requires to be a member of the netdev group.
    -p, --port <port>        Set port to listen on for http connections. [default: 3000]
    -w, --wsport <wsport>    Set port to listen on for websocket. [default: 4000]
    -d, --profile <path>     Set profile path to store user data.
    -r, --register <url>     Change the url of the registration endpoint. [default: http://localhost:4242/register]
    -i, --iface <iface>      Specify the local IP interface.
    -t, --tunnel <tunnel>    Set the tunnel endpoint's hostname. If omitted, the tunnel is disabled.
    -s, --tunnel-secret <secret>    Set the tunnel shared secret. [default: secret]
        --disable-tls               Run as a plain HTTP server, disabling encryption.
        --remote-name <hostname>    Set remote hostname. This the URL to access the box through the bridge. If omitted, the tunnel is disabled
    -c, --config <namespace;key;value>  Set configuration override
    -h, --help               Print this help menu.
",
        flag_local_name: Option<String>,
        flag_port: u16,
        flag_wsport: u16,
        flag_profile: Option<String>,
        flag_register: String,
        flag_iface: Option<String>,
        flag_tunnel: Option<String>,
        flag_tunnel_secret: String,
        flag_disable_tls: bool,
        flag_remote_name: Option<String>,
        flag_config: Option<Vec<String>>);

/// Updates local host name with the provided host name string. If requested host name
/// is not available (used by anyone else on the same network) then collision
/// resolution logic is triggered and alternative name is chosen automatically
/// (host name plus "-2", "-3" and etc. postfix). This function blocks until host name
/// is updated and returns actual host name.
///
/// # Panics
///
/// Panics if provided host name is not valid non-FQDN host name.
///
/// # Arguments
///
/// * `hostname` - host name name we'd like to set (should be a valid non-FQDN host name).
fn update_hostname(hostname: String) -> Option<String> {
    let host_manager = HostManager::new();

    if !host_manager.is_valid_name(&hostname) {
        panic!("Host name `{}` is not a valid host name!", &hostname);
    }

    Some(host_manager.set_name(&hostname))
}

// Handle SIGINT (Ctrl-C) for manual shutdown.
// Signal handlers must not do anything substantial. To trigger shutdown, we atomically
// flip this flag; the event loop checks the flag and exits accordingly.
static SHUTDOWN_FLAG: AtomicBool = ATOMIC_BOOL_INIT;
unsafe fn handle_sigint(_:i32) {
    SHUTDOWN_FLAG.store(true, Ordering::Release);
}

#[cfg(any(target_os = "linux", target_os = "android"))]
#[inline]
fn tid_str() -> String {
    // gettid only exists for the linux and android variants of nix
    format!("({}) ", nix::unistd::gettid())
}

#[cfg(not(any(target_os = "linux", target_os = "android")))]
#[inline]
fn tid_str() -> &'static str {
    ""
}

fn main() {
    unsafe {
        libc::signal(SIGINT, mem::transmute(handle_sigint));
    }

    let format = |record: &LogRecord| {
        let t = time::now();
        format!("{}.{:03}: {}{:5} {}",
            time::strftime("%Y-%m-%d %H:%M:%S", &t).unwrap(),
            t.tm_nsec / 1000_000,
            tid_str(),
            record.level(),
            record.args()
        )
    };
    let mut builder = LogBuilder::new();
    builder.format(format).filter(None, LogLevelFilter::Info);

    if env::var("RUST_LOG").is_ok() {
       builder.parse(&env::var("RUST_LOG").unwrap());
    }
    builder.init().unwrap();

    let args: Args = Args::docopt().decode().unwrap_or_else(|e| e.exit());

    // Start the tunnel.
    let mut tunnel: Option<Tunnel> = None;
    if let Some(tunnel_url) = args.flag_tunnel {
        if let Some(remote_name) = args.flag_remote_name {
            tunnel = Some(Tunnel::new(TunnelConfig::new(tunnel_url,
                                                        args.flag_tunnel_secret,
                                                        args.flag_port,
                                                        args.flag_wsport,
                                                        remote_name)));
            tunnel.as_mut().unwrap().start().unwrap();
        }
    }

    // Register with the nUPNP server.
    let registrar = registration::Registrar::new();
    registrar.start(args.flag_register, args.flag_iface, &tunnel);

    let mut controller = FoxBox::new(
        args.flag_verbose, args.flag_local_name.map_or(None, update_hostname), args.flag_port,
        args.flag_wsport,
        if args.flag_disable_tls { TlsOption::Disabled } else { TlsOption::Enabled },
        match args.flag_profile {
            Some(p) => ProfilePath::Custom(p),
            None => ProfilePath::Default
        });

    // Override config values
    {
        if let Some(flags) = args.flag_config {
            for flag in flags {
                let items: Vec<String> = utils::split_escaped(&flag, ';');
                if items.len() >= 3 {
                    let namespace = items[0].clone();
                    let key = items[1].clone();
                    let value = items[2..].join(";");
                    warn!("Setting config override: {}::{}->{}", namespace, key, value);
                    controller.config.set_override(&namespace, &key, &value);
                } else {
                    error!("Config override requires three fields: {}", flag)
                }
            }
        }
    }

    controller.run(&SHUTDOWN_FLAG);

    if let Some(mut tunnel) = tunnel {
        tunnel.stop().unwrap();
    }
}

#[cfg(test)]
describe! main {
    describe! args {
        it "should have default values" {
            let argv = || vec!["foxbox"];
            let args: super::super::Args = super::super::Args::docopt().argv(argv().into_iter())
                .decode().unwrap();

            assert_eq!(args.flag_verbose, false);
            assert_eq!(args.flag_local_name, None);
            assert_eq!(args.flag_port, 3000);
            assert_eq!(args.flag_wsport, 4000);
            assert_eq!(args.flag_register, "http://localhost:4242/register");
            assert_eq!(args.flag_iface, None);
            assert_eq!(args.flag_tunnel, None);
            assert_eq!(args.flag_config, None);
            assert_eq!(args.flag_help, false);
        }

        it "should support short form" {
            let argv = || vec!["foxbox",
                               "-v",
                               "-p", "1234",
                               "-l", "foobar",
                               "-w", "4567",
                               "-r", "http://foo.bar:6868/register",
                               "-i", "eth99",
                               "-t", "tunnel.host",
                               "-c", "ns;key;value"];

           let args: super::super::Args = super::super::Args::docopt().argv(argv().into_iter())
               .decode().unwrap();

            assert_eq!(args.flag_verbose, true);
            assert_eq!(args.flag_local_name.unwrap(), "foobar");
            assert_eq!(args.flag_port, 1234);
            assert_eq!(args.flag_wsport, 4567);
            assert_eq!(args.flag_register, "http://foo.bar:6868/register");
            assert_eq!(args.flag_iface.unwrap(), "eth99");
            assert_eq!(args.flag_tunnel.unwrap(), "tunnel.host");
            assert_eq!(args.flag_config.unwrap(), vec!["ns;key;value"]);
        }

        it "should support long form" {
            let argv = || vec!["foxbox",
                               "--verbose",
                               "--port", "1234",
                               "--local-name", "foobar",
                               "--wsport", "4567",
                               "--register", "http://foo.bar:6868/register",
                               "--iface", "eth99",
                               "--tunnel", "tunnel.host",
                               "--config", "ns;key;value"];

            let args: super::super::Args = super::super::Args::docopt().argv(argv().into_iter())
                .decode().unwrap();

            assert_eq!(args.flag_verbose, true);
            assert_eq!(args.flag_local_name.unwrap(), "foobar");
            assert_eq!(args.flag_port, 1234);
            assert_eq!(args.flag_wsport, 4567);
            assert_eq!(args.flag_register, "http://foo.bar:6868/register");
            assert_eq!(args.flag_iface.unwrap(), "eth99");
            assert_eq!(args.flag_tunnel.unwrap(), "tunnel.host");
            assert_eq!(args.flag_config.unwrap(), vec!["ns;key;value"]);
        }
    }
}
<commit_msg>Colorizing logger output on terminals<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// Needed to derive `Serialize` on ServiceProperties
#![feature(custom_derive, plugin)]
#![plugin(serde_macros)]
// For Docopt macro
#![plugin(docopt_macros)]

// Make linter fail for every warning
#![plugin(clippy)]
#![deny(clippy)]

#![cfg_attr(test, feature(const_fn))] // Dependency of stainless
#![cfg_attr(test, plugin(stainless))] // Test runner

#![feature(reflect_marker)]

#![feature(associated_consts)]

extern crate chrono;
extern crate core;
extern crate docopt;
extern crate env_logger;
extern crate foxbox_adapters;
extern crate foxbox_taxonomy;
extern crate foxbox_users;
#[macro_use]
extern crate hyper;
#[macro_use]
extern crate iron;
extern crate iron_cors;
#[cfg(test)]
extern crate iron_test;
extern crate libc;
#[macro_use]
extern crate log;
extern crate mio;
extern crate mount;
extern crate nix;
extern crate openssl;
extern crate openssl_sys;
extern crate rand;
extern crate router;
extern crate rustc_serialize;
extern crate rusqlite;
extern crate serde;
extern crate serde_json;
extern crate staticfile;
extern crate time;
extern crate timer;
extern crate transformable_channels;
extern crate unicase;
extern crate uuid;
extern crate ws;
extern crate multicast_dns;
extern crate xml;

#[cfg(test)]
extern crate regex;
#[cfg(test)]
extern crate tempdir;

// Need to be declared first so to let the macros be visible from other modules.
#[macro_use]
mod utils;
mod adapters;
mod config_store;
mod controller;
mod http_server;
mod managed_process;
mod profile_service;
mod registration;
mod service;
mod upnp;
mod service_router;
mod stable_uuid;
mod static_router;
mod taxonomy_router;
mod tls;
mod traits;
mod tunnel_controller;
mod ws_server;

#[cfg(test)]
mod stubs {
    #![allow(dead_code)]
    #![allow(unused_variables)]
    #![allow(boxed_local)]
    pub mod service;
    pub mod controller;
}

use controller::FoxBox;
use env_logger::LogBuilder;
use tunnel_controller:: { TunnelConfig, Tunnel };
use libc::SIGINT;
use log::{ LogRecord, LogLevelFilter };

use multicast_dns::host::HostManager;
use profile_service::ProfilePath;
use std::env;
use std::mem;
use std::sync::atomic::{ AtomicBool, Ordering, ATOMIC_BOOL_INIT };
use tls::TlsOption;
use traits::Controller;

docopt!(Args derive Debug, "
Usage: foxbox [-v] [-h] [-l <hostname>] [-p <port>] [-w <wsport>] [-d <profile_path>] [-r <url>] [-i <iface>] [-t <tunnel>] [-s <secret>] [--disable-tls] [--remote-name <hostname>] [-c <namespace;key;value>]...

Options:
    -v, --verbose            Toggle verbose output.
    -l, --local-name <hostname>    Set local hostname. Linux only. Requires to be a member of the netdev group.
    -p, --port <port>        Set port to listen on for http connections. [default: 3000]
    -w, --wsport <wsport>    Set port to listen on for websocket. [default: 4000]
    -d, --profile <path>     Set profile path to store user data.
    -r, --register <url>     Change the url of the registration endpoint. [default: http://localhost:4242/register]
    -i, --iface <iface>      Specify the local IP interface.
    -t, --tunnel <tunnel>    Set the tunnel endpoint's hostname. If omitted, the tunnel is disabled.
    -s, --tunnel-secret <secret>    Set the tunnel shared secret. [default: secret]
        --disable-tls               Run as a plain HTTP server, disabling encryption.
        --remote-name <hostname>    Set remote hostname. This the URL to access the box through the bridge. If omitted, the tunnel is disabled
    -c, --config <namespace;key;value>  Set configuration override
    -h, --help               Print this help menu.
",
        flag_local_name: Option<String>,
        flag_port: u16,
        flag_wsport: u16,
        flag_profile: Option<String>,
        flag_register: String,
        flag_iface: Option<String>,
        flag_tunnel: Option<String>,
        flag_tunnel_secret: String,
        flag_disable_tls: bool,
        flag_remote_name: Option<String>,
        flag_config: Option<Vec<String>>);

/// Updates local host name with the provided host name string. If requested host name
/// is not available (used by anyone else on the same network) then collision
/// resolution logic is triggered and alternative name is chosen automatically
/// (host name plus "-2", "-3" and etc. postfix). This function blocks until host name
/// is updated and returns actual host name.
///
/// # Panics
///
/// Panics if provided host name is not valid non-FQDN host name.
///
/// # Arguments
///
/// * `hostname` - host name name we'd like to set (should be a valid non-FQDN host name).
fn update_hostname(hostname: String) -> Option<String> {
    let host_manager = HostManager::new();

    if !host_manager.is_valid_name(&hostname) {
        panic!("Host name `{}` is not a valid host name!", &hostname);
    }

    Some(host_manager.set_name(&hostname))
}

// Handle SIGINT (Ctrl-C) for manual shutdown.
// Signal handlers must not do anything substantial. To trigger shutdown, we atomically
// flip this flag; the event loop checks the flag and exits accordingly.
static SHUTDOWN_FLAG: AtomicBool = ATOMIC_BOOL_INIT;
unsafe fn handle_sigint(_:i32) {
    SHUTDOWN_FLAG.store(true, Ordering::Release);
}

#[cfg(any(target_os = "linux", target_os = "android"))]
#[inline]
fn tid_str() -> String {
    // gettid only exists for the linux and android variants of nix
    format!("({}) ", nix::unistd::gettid())
}

#[cfg(not(any(target_os = "linux", target_os = "android")))]
#[inline]
fn tid_str() -> &'static str {
    ""
}

fn main() {
    unsafe {
        libc::signal(SIGINT, mem::transmute(handle_sigint));
    }

    let mut builder = LogBuilder::new();
    let istty = unsafe { libc::isatty(libc::STDERR_FILENO as i32) } != 0;
    if istty {
        // Colorized output formatter
        let format = |record: &LogRecord| {
            let t = time::now();
            let level_color = match record.level() {
                log::LogLevel::Error => "\x1b[1;31m",  // bold red
                log::LogLevel::Warn  => "\x1b[1;33m",  // bold yellow
                log::LogLevel::Info  => "\x1b[1;32m",  // bold green
                log::LogLevel::Debug => "\x1b[1;34m",  // bold blue
                log::LogLevel::Trace => "\x1b[1;35m"   // bold magenta
            };
            format!("[\x1b[90m{}.{:03}\x1b[0m] {}{}{:5}\x1b[0m {}",
                time::strftime("%Y-%m-%d %H:%M:%S", &t).unwrap(),
                t.tm_nsec / 1_000_000,
                tid_str(),
                level_color,
                record.level(),
                record.args()
            )
        };
        builder.format(format).filter(None, LogLevelFilter::Info);
    } else {
        // Plain output formatter
        let format = |record: &LogRecord| {
            let t = time::now();
            format!("{}.{:03} {}{:5} {}",
                time::strftime("%Y-%m-%d %H:%M:%S", &t).unwrap(),
                t.tm_nsec / 1_000_000,
                tid_str(),
                record.level(),
                record.args()
            )
        };
        builder.format(format).filter(None, LogLevelFilter::Info);
    }

    if env::var("RUST_LOG").is_ok() {
       builder.parse(&env::var("RUST_LOG").unwrap());
    }
    builder.init().unwrap();

    let args: Args = Args::docopt().decode().unwrap_or_else(|e| e.exit());

    // Start the tunnel.
    let mut tunnel: Option<Tunnel> = None;
    if let Some(tunnel_url) = args.flag_tunnel {
        if let Some(remote_name) = args.flag_remote_name {
            tunnel = Some(Tunnel::new(TunnelConfig::new(tunnel_url,
                                                        args.flag_tunnel_secret,
                                                        args.flag_port,
                                                        args.flag_wsport,
                                                        remote_name)));
            tunnel.as_mut().unwrap().start().unwrap();
        }
    }

    // Register with the nUPNP server.
    let registrar = registration::Registrar::new();
    registrar.start(args.flag_register, args.flag_iface, &tunnel);

    let mut controller = FoxBox::new(
        args.flag_verbose, args.flag_local_name.map_or(None, update_hostname), args.flag_port,
        args.flag_wsport,
        if args.flag_disable_tls { TlsOption::Disabled } else { TlsOption::Enabled },
        match args.flag_profile {
            Some(p) => ProfilePath::Custom(p),
            None => ProfilePath::Default
        });

    // Override config values
    {
        if let Some(flags) = args.flag_config {
            for flag in flags {
                let items: Vec<String> = utils::split_escaped(&flag, ';');
                if items.len() >= 3 {
                    let namespace = items[0].clone();
                    let key = items[1].clone();
                    let value = items[2..].join(";");
                    warn!("Setting config override: {}::{}->{}", namespace, key, value);
                    controller.config.set_override(&namespace, &key, &value);
                } else {
                    error!("Config override requires three fields: {}", flag)
                }
            }
        }
    }

    controller.run(&SHUTDOWN_FLAG);

    if let Some(mut tunnel) = tunnel {
        tunnel.stop().unwrap();
    }
}

#[cfg(test)]
describe! main {
    describe! args {
        it "should have default values" {
            let argv = || vec!["foxbox"];
            let args: super::super::Args = super::super::Args::docopt().argv(argv().into_iter())
                .decode().unwrap();

            assert_eq!(args.flag_verbose, false);
            assert_eq!(args.flag_local_name, None);
            assert_eq!(args.flag_port, 3000);
            assert_eq!(args.flag_wsport, 4000);
            assert_eq!(args.flag_register, "http://localhost:4242/register");
            assert_eq!(args.flag_iface, None);
            assert_eq!(args.flag_tunnel, None);
            assert_eq!(args.flag_config, None);
            assert_eq!(args.flag_help, false);
        }

        it "should support short form" {
            let argv = || vec!["foxbox",
                               "-v",
                               "-p", "1234",
                               "-l", "foobar",
                               "-w", "4567",
                               "-r", "http://foo.bar:6868/register",
                               "-i", "eth99",
                               "-t", "tunnel.host",
                               "-c", "ns;key;value"];

           let args: super::super::Args = super::super::Args::docopt().argv(argv().into_iter())
               .decode().unwrap();

            assert_eq!(args.flag_verbose, true);
            assert_eq!(args.flag_local_name.unwrap(), "foobar");
            assert_eq!(args.flag_port, 1234);
            assert_eq!(args.flag_wsport, 4567);
            assert_eq!(args.flag_register, "http://foo.bar:6868/register");
            assert_eq!(args.flag_iface.unwrap(), "eth99");
            assert_eq!(args.flag_tunnel.unwrap(), "tunnel.host");
            assert_eq!(args.flag_config.unwrap(), vec!["ns;key;value"]);
        }

        it "should support long form" {
            let argv = || vec!["foxbox",
                               "--verbose",
                               "--port", "1234",
                               "--local-name", "foobar",
                               "--wsport", "4567",
                               "--register", "http://foo.bar:6868/register",
                               "--iface", "eth99",
                               "--tunnel", "tunnel.host",
                               "--config", "ns;key;value"];

            let args: super::super::Args = super::super::Args::docopt().argv(argv().into_iter())
                .decode().unwrap();

            assert_eq!(args.flag_verbose, true);
            assert_eq!(args.flag_local_name.unwrap(), "foobar");
            assert_eq!(args.flag_port, 1234);
            assert_eq!(args.flag_wsport, 4567);
            assert_eq!(args.flag_register, "http://foo.bar:6868/register");
            assert_eq!(args.flag_iface.unwrap(), "eth99");
            assert_eq!(args.flag_tunnel.unwrap(), "tunnel.host");
            assert_eq!(args.flag_config.unwrap(), vec!["ns;key;value"]);
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use] extern crate clap;
extern crate ini;
extern crate sha1;

use std::ffi::{OsStr, OsString};
use std::fmt;
use std::fs::{self, File};
use std::io::{self, Read, Write, BufWriter};
use std::iter::FromIterator;
use std::path::{Path, PathBuf};

use ini::Ini;
use ini::ini::Error as IniError;

type Sha1DigestBytes = [u8; 20];
type ModList = Vec<String>;
type MediaSet = Vec<Asset>;

struct Asset {
	path: PathBuf,
	hash: Sha1DigestBytes,
}

impl Asset {
	pub fn new(pb: PathBuf, h: Sha1DigestBytes) -> Self {
		Asset {
			path: pb,
			hash: h,
		}
	}
}

impl PartialEq for Asset {
	fn eq(&self, other: &Self) -> bool {
		self.hash == other.hash
	}
}

enum Error {
	Io(io::Error),
	Ini(IniError),
}

impl fmt::Display for Error {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		match *self {
			Error::Io(ref e)  => write!(f, "IO error: {}", e),
			Error::Ini(ref e) => write!(f, "Settings file error: {}", e),
		}
	}
}

impl From<io::Error> for Error {
	fn from(e: io::Error) -> Self { Error::Io(e) }
}

impl From<IniError> for Error {
	fn from(e: IniError) -> Self { Error::Ini(e) }
}

enum AssetCopyMode {
	Symlink,
	Hardlink,
	Copy,
	None,
}

fn to_hex(input: &[u8]) -> String {
	String::from_iter(input.iter().map(|b| format!("{:02x}", b)))
}

fn make_absolute(path: &Path) -> PathBuf {
	if path.is_absolute() {
		path.to_path_buf()
	} else {
		std::env::current_dir()
			.and_then(|cd| Ok(cd.join(path)))
			.or_else(|_err| -> io::Result<_> {
				Ok(path.to_path_buf())
			})
			.unwrap()
	}
}

fn hash_file(path: &Path) -> io::Result<Sha1DigestBytes> {
	let mut buf = [0u8; 8192];
	let mut hash = sha1::Sha1::new();
	let mut file = File::open(&path)?;
	loop {
		match file.read(&mut buf) {
			Ok(0) => break,
			Ok(len) => hash.update(&buf[..len]),
			Err(e) => return Err(e),
		}
	}
	Ok(hash.digest().bytes())
}

fn search_media_dir(ms: &mut MediaSet, path: &Path) -> io::Result<()> {
	for entry in path.read_dir()? {
		let pb = entry?.path();
		if pb.is_file() {
			let h = hash_file(pb.as_path())?;
			ms.push(Asset::new(pb, h));
		}
	}
	Ok(())
}

fn search_mod_dir(ms: &mut MediaSet, path: &Path) -> io::Result<()> {
	static MEDIA_DIRS: &'static [&'static str] = &["textures", "models", "sounds"];
	for media_dir in MEDIA_DIRS {
		let media_pb = path.join(media_dir);
		if media_pb.is_dir() {
			search_media_dir(ms, media_pb.as_path())?;
		}
	}
	Ok(())
}

fn search_modpack_dir(ms: &mut MediaSet, path: &Path, mods: Option<&ModList>) -> io::Result<()> {
	for entry in path.read_dir()? {
		let entry_path = entry?.path();
		if !entry_path.is_dir() {
			continue;
		} else if entry_path.join("modpack.txt").exists() {
			search_modpack_dir(ms, entry_path.as_path(), mods)?;
		} else if entry_path.join("init.lua").exists() {
			if let Some(mod_list) = mods {
				let mod_name = &entry_path.file_name()
					.expect("Mod directory has no name!")
					.to_str()
					.expect("Mod directory name is not valid Unicode")
					.to_string();
				if !mod_list.contains(mod_name) {
					continue;
				}
			}
			search_mod_dir(ms, entry_path.as_path())?;
		}
		// Otherwise it's probably a VCS directory or something similar
	}
	Ok(())
}

fn write_index(ms: &MediaSet, path: &Path) -> io::Result<()> {
	let file = File::create(&path)?;
	let mut writer = BufWriter::new(file);
	writer.write_all(b"MTHS\x00\x01")?;
	for asset in ms {
		writer.write_all(&asset.hash)?;
	}
	Ok(())
}

fn copy_assets(ms: &MediaSet, path: &Path, mode: AssetCopyMode) -> io::Result<()> {
	fn copy_no_result<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> io::Result<()> {
		fs::copy(src, dst).map(|_| ())
	}

	#[cfg(unix)]
	fn symlink_file<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> io::Result<()> {
		std::os::unix::fs::symlink(src, dst)
	}

	#[cfg(windows)]
	fn symlink_file<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> io::Result<()> {
		std::os::windows::fs::symlink_file(src, dst)
	}

	#[cfg(not(any(unix, windows)))]
	fn symlink_file<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> io::Result<()> {
		Err(io::Error::new(io::ErrorKind::Other,
				"Symlinking not supported on this platform!"))
	}

	let copy_func = match mode {
		AssetCopyMode::Symlink => symlink_file,
		AssetCopyMode::Hardlink => fs::hard_link,
		AssetCopyMode::Copy => copy_no_result,
		AssetCopyMode::None => return Ok(()),
	};

	for asset in ms {
		copy_func(&asset.path, path.join(to_hex(&asset.hash)))?;
	}
	Ok(())
}

fn get_mod_list(path: &Path) -> Result<ModList, IniError> {
	let world_mt = Ini::load_from_file(path.join("world.mt"))?;
	let main_sec = world_mt.general_section();

	let mut list: ModList = vec![];
	for (key, value) in main_sec {
		if !key.starts_with("load_mod_") || value != "true" {
			continue;
		}
		let (_, mod_name) = key.split_at(9);
		list.push(mod_name.to_string());
	}
	Ok(list)
}

fn get_args<'a>() -> clap::ArgMatches<'a> {
	use clap::{App, Arg, ArgGroup};

	fn check_new_dir(s: &OsStr) -> Result<(), OsString> {
		let p = make_absolute(Path::new(&s));
		if p.is_dir() {
			return Ok(())
		}
		if let Some(parent) = p.parent() {
			if parent.is_dir() {
				return Ok(())
			}
		}
		return Err("Invalid path.".into())
	}

	fn check_existing_dir(s: &OsStr) -> Result<(), OsString> {
		if make_absolute(Path::new(&s)).is_dir() {
			Ok(())
		} else {
			Err("Invalid path.".into())
		}
	}

	let app = clap_app! { @app (app_from_crate!())
		(version_short: "v")

		(@arg mod_paths: [PATHS] ... validator_os(check_existing_dir) "Additional mod paths to search.")

		(@arg out:   -o --out   <PATH> validator_os(check_new_dir) "Path to the output directory.")
		(@arg world: -w --world <PATH> validator_os(check_existing_dir) "Path to the world directory.")
		(@arg game:  -g --game  <PATH> validator_os(check_existing_dir) "Path to the game directory.")

		// Group these together with display_order
		(@arg copy: -c --copy display_order(1000) "Copy assets to output folder.")
		// Symlink added below if applicable
		(@arg hardlink: -l --hardlink display_order(1000) "Hard link assets to output folder.")
	};

	// Add symlink option if supported
	#[cfg(not(any(unix, windows)))]
	let add_symlink_arg = |app| app;

	#[cfg(any(unix, windows))]
	fn add_symlink_arg<'a>(app: App<'a, 'a>) -> App<'a, 'a> {
		app.arg(Arg::with_name("symlink")
			.short("s")
			.long("symlink")
			.display_order(1000)
			.help("Symbolically link assets to output folder."))
	}

	add_symlink_arg(app)
		// Link group has to be added manually because the
		// symlink argument is added conditionally.
		.group(ArgGroup::with_name("link")
			.args(&["copy", "symlink", "hardlink"]))

		.get_matches()
}

fn run(args: clap::ArgMatches) -> Result<(), Error> {
	// These unwraps are safe since the values are required
	// and clap will exit if the value is missing.
	let out_opt = args.value_of_os("out").unwrap();
	let out_path = Path::new(&out_opt);
	let world_opt = args.value_of_os("world").unwrap();
	let world_path = Path::new(&world_opt);
	let game_opt = args.value_of_os("game").unwrap();
	let game_path = Path::new(&game_opt);

	let copy_type = if args.is_present("copy") {
			AssetCopyMode::Copy
		} else if args.is_present("symlink") {
			AssetCopyMode::Symlink
		} else if args.is_present("hardlink") {
			AssetCopyMode::Hardlink
		} else {
			AssetCopyMode::None
		};

	let mut ms = MediaSet::new();
	let mods = get_mod_list(world_path)?;

	// Search world mods.
	let worldmods_path = world_path.join("worldmods");
	if worldmods_path.exists() {
		search_modpack_dir(&mut ms, worldmods_path.as_path(), Some(&mods))?;
	}

	// Search game mods.
	// Note: Game mods can not currently be disabled.
	search_modpack_dir(&mut ms, game_path.join("mods").as_path(), None)?;

	if let Some(mod_paths) = args.values_of_os("mod_paths") {
		for mod_path in mod_paths {
			search_modpack_dir(&mut ms,
					Path::new(&mod_path),
					Some(&mods))?;
		}
	}

	// Deduplicate list.  Otherwise linking will fail and the index will
	// be unnecessarily large.
	ms.sort_by(|a, b| a.hash.cmp(&b.hash));
	ms.dedup();

	if !out_path.exists() {
		fs::create_dir(out_path)?;
	}

	write_index(&ms, out_path.join("index.mth").as_path())?;

	copy_assets(&ms, out_path, copy_type)?;

	Ok(())
}

fn main() {
	match run(get_args()) {
		Ok(()) => return,
		Err(e) => {
			println!("{}", e);
			std::process::exit(1)
		}
	}
}
<commit_msg>Don't overwrite assets in output directory<commit_after>#[macro_use] extern crate clap;
extern crate ini;
extern crate sha1;

use std::ffi::{OsStr, OsString};
use std::fmt;
use std::fs::{self, File};
use std::io::{self, Read, Write, BufWriter};
use std::iter::FromIterator;
use std::path::{Path, PathBuf};

use ini::Ini;
use ini::ini::Error as IniError;

type Sha1DigestBytes = [u8; 20];
type ModList = Vec<String>;
type MediaSet = Vec<Asset>;

struct Asset {
	path: PathBuf,
	hash: Sha1DigestBytes,
}

impl Asset {
	pub fn new(pb: PathBuf, h: Sha1DigestBytes) -> Self {
		Asset {
			path: pb,
			hash: h,
		}
	}
}

impl PartialEq for Asset {
	fn eq(&self, other: &Self) -> bool {
		self.hash == other.hash
	}
}

enum Error {
	Io(io::Error),
	Ini(IniError),
}

impl fmt::Display for Error {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		match *self {
			Error::Io(ref e)  => write!(f, "IO error: {}", e),
			Error::Ini(ref e) => write!(f, "Settings file error: {}", e),
		}
	}
}

impl From<io::Error> for Error {
	fn from(e: io::Error) -> Self { Error::Io(e) }
}

impl From<IniError> for Error {
	fn from(e: IniError) -> Self { Error::Ini(e) }
}

enum AssetCopyMode {
	Symlink,
	Hardlink,
	Copy,
	None,
}

fn to_hex(input: &[u8]) -> String {
	String::from_iter(input.iter().map(|b| format!("{:02x}", b)))
}

fn make_absolute(path: &Path) -> PathBuf {
	if path.is_absolute() {
		path.to_path_buf()
	} else {
		std::env::current_dir()
			.and_then(|cd| Ok(cd.join(path)))
			.or_else(|_err| -> io::Result<_> {
				Ok(path.to_path_buf())
			})
			.unwrap()
	}
}

fn hash_file(path: &Path) -> io::Result<Sha1DigestBytes> {
	let mut buf = [0u8; 8192];
	let mut hash = sha1::Sha1::new();
	let mut file = File::open(&path)?;
	loop {
		match file.read(&mut buf) {
			Ok(0) => break,
			Ok(len) => hash.update(&buf[..len]),
			Err(e) => return Err(e),
		}
	}
	Ok(hash.digest().bytes())
}

fn search_media_dir(ms: &mut MediaSet, path: &Path) -> io::Result<()> {
	for entry in path.read_dir()? {
		let pb = entry?.path();
		if pb.is_file() {
			let h = hash_file(pb.as_path())?;
			ms.push(Asset::new(pb, h));
		}
	}
	Ok(())
}

fn search_mod_dir(ms: &mut MediaSet, path: &Path) -> io::Result<()> {
	static MEDIA_DIRS: &'static [&'static str] = &["textures", "models", "sounds"];
	for media_dir in MEDIA_DIRS {
		let media_pb = path.join(media_dir);
		if media_pb.is_dir() {
			search_media_dir(ms, media_pb.as_path())?;
		}
	}
	Ok(())
}

fn search_modpack_dir(ms: &mut MediaSet, path: &Path, mods: Option<&ModList>) -> io::Result<()> {
	for entry in path.read_dir()? {
		let entry_path = entry?.path();
		if !entry_path.is_dir() {
			continue;
		} else if entry_path.join("modpack.txt").exists() {
			search_modpack_dir(ms, entry_path.as_path(), mods)?;
		} else if entry_path.join("init.lua").exists() {
			if let Some(mod_list) = mods {
				let mod_name = &entry_path.file_name()
					.expect("Mod directory has no name!")
					.to_str()
					.expect("Mod directory name is not valid Unicode")
					.to_string();
				if !mod_list.contains(mod_name) {
					continue;
				}
			}
			search_mod_dir(ms, entry_path.as_path())?;
		}
		// Otherwise it's probably a VCS directory or something similar
	}
	Ok(())
}

fn write_index(ms: &MediaSet, path: &Path) -> io::Result<()> {
	let file = File::create(&path)?;
	let mut writer = BufWriter::new(file);
	writer.write_all(b"MTHS\x00\x01")?;
	for asset in ms {
		writer.write_all(&asset.hash)?;
	}
	Ok(())
}

fn copy_assets(ms: &MediaSet, path: &Path, mode: AssetCopyMode) -> io::Result<()> {
	fn copy_no_result<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> io::Result<()> {
		fs::copy(src, dst).map(|_| ())
	}

	#[cfg(unix)]
	fn symlink_file<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> io::Result<()> {
		std::os::unix::fs::symlink(src, dst)
	}

	#[cfg(windows)]
	fn symlink_file<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> io::Result<()> {
		std::os::windows::fs::symlink_file(src, dst)
	}

	#[cfg(not(any(unix, windows)))]
	fn symlink_file<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> io::Result<()> {
		Err(io::Error::new(io::ErrorKind::Other,
				"Symlinking not supported on this platform!"))
	}

	let copy_func = match mode {
		AssetCopyMode::Symlink => symlink_file,
		AssetCopyMode::Hardlink => fs::hard_link,
		AssetCopyMode::Copy => copy_no_result,
		AssetCopyMode::None => return Ok(()),
	};

	for asset in ms {
		let to_path = path.join(to_hex(&asset.hash));
		if !to_path.exists() {
			copy_func(&asset.path, to_path)?;
		}
	}
	Ok(())
}

fn get_mod_list(path: &Path) -> Result<ModList, IniError> {
	let world_mt = Ini::load_from_file(path.join("world.mt"))?;
	let main_sec = world_mt.general_section();

	let mut list: ModList = vec![];
	for (key, value) in main_sec {
		if !key.starts_with("load_mod_") || value != "true" {
			continue;
		}
		let (_, mod_name) = key.split_at(9);
		list.push(mod_name.to_string());
	}
	Ok(list)
}

fn get_args<'a>() -> clap::ArgMatches<'a> {
	use clap::{App, Arg, ArgGroup};

	fn check_new_dir(s: &OsStr) -> Result<(), OsString> {
		let p = make_absolute(Path::new(&s));
		if p.is_dir() {
			return Ok(())
		}
		if let Some(parent) = p.parent() {
			if parent.is_dir() {
				return Ok(())
			}
		}
		return Err("Invalid path.".into())
	}

	fn check_existing_dir(s: &OsStr) -> Result<(), OsString> {
		if make_absolute(Path::new(&s)).is_dir() {
			Ok(())
		} else {
			Err("Invalid path.".into())
		}
	}

	let app = clap_app! { @app (app_from_crate!())
		(version_short: "v")

		(@arg mod_paths: [PATHS] ... validator_os(check_existing_dir) "Additional mod paths to search.")

		(@arg out:   -o --out   <PATH> validator_os(check_new_dir) "Path to the output directory.")
		(@arg world: -w --world <PATH> validator_os(check_existing_dir) "Path to the world directory.")
		(@arg game:  -g --game  <PATH> validator_os(check_existing_dir) "Path to the game directory.")

		// Group these together with display_order
		(@arg copy: -c --copy display_order(1000) "Copy assets to output folder.")
		// Symlink added below if applicable
		(@arg hardlink: -l --hardlink display_order(1000) "Hard link assets to output folder.")
	};

	// Add symlink option if supported
	#[cfg(not(any(unix, windows)))]
	let add_symlink_arg = |app| app;

	#[cfg(any(unix, windows))]
	fn add_symlink_arg<'a>(app: App<'a, 'a>) -> App<'a, 'a> {
		app.arg(Arg::with_name("symlink")
			.short("s")
			.long("symlink")
			.display_order(1000)
			.help("Symbolically link assets to output folder."))
	}

	add_symlink_arg(app)
		// Link group has to be added manually because the
		// symlink argument is added conditionally.
		.group(ArgGroup::with_name("link")
			.args(&["copy", "symlink", "hardlink"]))

		.get_matches()
}

fn run(args: clap::ArgMatches) -> Result<(), Error> {
	// These unwraps are safe since the values are required
	// and clap will exit if the value is missing.
	let out_opt = args.value_of_os("out").unwrap();
	let out_path = Path::new(&out_opt);
	let world_opt = args.value_of_os("world").unwrap();
	let world_path = Path::new(&world_opt);
	let game_opt = args.value_of_os("game").unwrap();
	let game_path = Path::new(&game_opt);

	let copy_type = if args.is_present("copy") {
			AssetCopyMode::Copy
		} else if args.is_present("symlink") {
			AssetCopyMode::Symlink
		} else if args.is_present("hardlink") {
			AssetCopyMode::Hardlink
		} else {
			AssetCopyMode::None
		};

	let mut ms = MediaSet::new();
	let mods = get_mod_list(world_path)?;

	// Search world mods.
	let worldmods_path = world_path.join("worldmods");
	if worldmods_path.exists() {
		search_modpack_dir(&mut ms, worldmods_path.as_path(), Some(&mods))?;
	}

	// Search game mods.
	// Note: Game mods can not currently be disabled.
	search_modpack_dir(&mut ms, game_path.join("mods").as_path(), None)?;

	if let Some(mod_paths) = args.values_of_os("mod_paths") {
		for mod_path in mod_paths {
			search_modpack_dir(&mut ms,
					Path::new(&mod_path),
					Some(&mods))?;
		}
	}

	// Deduplicate list.  Otherwise linking will fail and the index will
	// be unnecessarily large.
	ms.sort_by(|a, b| a.hash.cmp(&b.hash));
	ms.dedup();

	if !out_path.exists() {
		fs::create_dir(out_path)?;
	}

	write_index(&ms, out_path.join("index.mth").as_path())?;

	copy_assets(&ms, out_path, copy_type)?;

	Ok(())
}

fn main() {
	match run(get_args()) {
		Ok(()) => return,
		Err(e) => {
			println!("{}", e);
			std::process::exit(1)
		}
	}
}
<|endoftext|>


<commit_before>use std::env;
use std::fmt;
use std::fs::{File, rename, remove_file};
use std::io::{BufRead, Write};
use std::io;
use std::path::PathBuf;
use std::process::Command;

extern crate tempfile;

use tempfile::NamedTempFile;

#[derive(Debug)]
enum Operation {
    NoOp,
    Remove,
    Rename { new_path: PathBuf }
}

#[derive(Debug)]
struct Change {
    path: PathBuf,
    operation: Operation
}

trait Operate {
    fn operate(&self) -> io::Result<()>;
}

impl Operate for Change {
    fn operate(&self) -> io::Result<()> {
        match self.operation {
            Operation::NoOp => Ok(()),
            Operation::Remove => remove_file(&self.path),
            Operation::Rename { ref new_path } => rename(&self.path, new_path)
        }
    }
}

impl fmt::Display for Change {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.operation {
            Operation::NoOp => write!(f, ""),
            Operation::Remove => write!(f, "{}", &self.path.to_str().unwrap()),
            Operation::Rename { ref new_path } => write!(f, "{} → {}", &self.path.to_str().unwrap(), new_path.to_str().unwrap())
        }
    }
}

fn lookup_program() -> String {
    match env::var_os("VISUAL") {
        Some(val) => val,
        None => env::var_os("EDITOR").expect("Set $VISUAL or $EDITOR to use edfp")
    }.into_string().unwrap()
}

fn parse_lines(lines: (String, String)) -> Change {
    let (given, new) = lines;
    let path = PathBuf::from(&given);
    let empty_string = String::new();
    if new == given {
        Change { operation: Operation::NoOp, path }
    } else if new == String::new() {
        Change { operation: Operation::Remove, path }
    } else {
        Change {
            operation: Operation::Rename {
                new_path: PathBuf::from(new)
            },
            path
        }
    }
}

fn describe_modifying_changes<O: io::Write>(header: &str, changes: &Vec<&Change>, mut output: O) -> O {
    if !changes.is_empty() {
        writeln!(output, "{}", header);
        for change in changes {
            writeln!(output, "\t{}", change);
        }
    }
    output
}

fn describe_changes<O: io::Write>(changes: &Vec<Change>, mut output: O) {
    let removals: Vec<&Change> = changes.iter().filter(|c| match c.operation {
        Operation::Remove => true,
        _ => false
    }).collect();
    let renamings: Vec<&Change> = changes.iter().filter(|c| match c.operation {
        Operation::Rename { ref new_path } => true,
        _ => false
    }).collect();

    if renamings.is_empty() && removals.is_empty() {
        writeln!(output, "No changes.");
    }

    let output = describe_modifying_changes("Remove/delete the following files:", &removals, output);
    describe_modifying_changes("Rename/move the following files:", &renamings, output);
}

fn user_approves_changes<I: io::Read>(mut input: I) -> bool {
    true
}

fn edfp<I: io::Read, O: io::Write>(mut input: I, mut output: O) {
    let mut buf = String::new();
    input.read_to_string(&mut buf).expect("Failed to read input");

    let mut tmp: NamedTempFile = NamedTempFile::new().expect("failed to create temporary file");
    tmp.write_all(&buf.as_bytes()).expect("Could not write to tempfile");

    let editor_program = lookup_program();
    let tmp_path_str = tmp.path().to_str().expect("failed to parse tmp path");
    let editor = Command::new(editor_program.as_str()).arg(tmp_path_str).spawn();

    let success = editor
        .ok()
        .expect(format!("failed to launch {}", editor_program).as_str())
        .wait()
        .expect(format!("failed to wait for {}", editor_program).as_str())
        .success();

    if !success {
        panic!(format!("{} exited non-zero", editor_program));
    }

    let commands_file = tmp.reopen().ok().expect("failed to re-open");
    let commands = io::BufReader::new(commands_file).lines()
        .map(|line| line.ok().expect("failed to read line from commands file"));
    let changes: Vec<Change> = buf.lines().map(String::from).zip(commands).map(parse_lines).collect();

    describe_changes(&changes, output);

    if user_approves_changes(input) {
        for change in changes {
            change.operate().ok().expect("failed to change")
        }
    }
}

fn main() {
    let input = io::stdin();
    let mut output = io::stdout();
    edfp(input, output)
}
<commit_msg>it pretty much works<commit_after>use std::env;
use std::fmt;
use std::fs::{OpenOptions, File, rename, remove_file};
use std::io::{BufRead, Write};
use std::io;
use std::path::PathBuf;
use std::process::Command;

extern crate tempfile;

use tempfile::NamedTempFile;

#[derive(Debug)]
enum Operation {
    NoOp,
    Remove,
    Rename { new_path: PathBuf }
}

#[derive(Debug)]
struct Change {
    path: PathBuf,
    operation: Operation
}

trait Operate {
    fn operate(&self) -> io::Result<()>;
}

impl Operate for Change {
    fn operate(&self) -> io::Result<()> {
        match self.operation {
            Operation::NoOp => Ok(()),
            Operation::Remove => remove_file(&self.path),
            Operation::Rename { ref new_path } => rename(&self.path, new_path)
        }
    }
}

impl fmt::Display for Change {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.operation {
            Operation::NoOp => write!(f, ""),
            Operation::Remove => write!(f, "{}", &self.path.to_str().unwrap()),
            Operation::Rename { ref new_path } => write!(f, "{} → {}", &self.path.to_str().unwrap(), new_path.to_str().unwrap())
        }
    }
}

pub fn get_tty() -> io::Result<File> {
    OpenOptions::new().read(true).write(true).open("/dev/tty")
}

fn lookup_program() -> String {
    match env::var_os("VISUAL") {
        Some(val) => val,
        None => env::var_os("EDITOR").expect("Set $VISUAL or $EDITOR to use edfp")
    }.into_string().unwrap()
}

fn parse_lines(lines: (String, String)) -> Change {
    let (given, new) = lines;
    let path = PathBuf::from(&given);
    let empty_string = String::new();
    if new == given {
        Change { operation: Operation::NoOp, path }
    } else if new == String::new() {
        Change { operation: Operation::Remove, path }
    } else {
        Change {
            operation: Operation::Rename {
                new_path: PathBuf::from(new)
            },
            path
        }
    }
}

fn describe_modifying_changes<O: io::Write>(header: &str, changes: &Vec<&Change>, mut output: O) -> O {
    if !changes.is_empty() {
        writeln!(output, "{}", header);
        for change in changes {
            writeln!(output, "\t{}", change);
        }
        writeln!(output);
    }
    output
}

fn describe_changes<O: io::Write>(changes: &Vec<Change>, mut output: O) -> bool {
    let removals: Vec<&Change> = changes.iter().filter(|c| match c.operation {
        Operation::Remove => true,
        _ => false
    }).collect();
    let renamings: Vec<&Change> = changes.iter().filter(|c| match c.operation {
        Operation::Rename { ref new_path } => true,
        _ => false
    }).collect();

    if renamings.is_empty() && removals.is_empty() {
        writeln!(output, "No changes");
        false
    } else {
        let output = describe_modifying_changes("DELETE the following files:\n", &removals, output);
        describe_modifying_changes("RENAME the following files:\n", &renamings, output);
        true
    }
}

fn user_approves_changes() -> bool {
    let mut input = String::new();
    let tty: File = get_tty().ok().expect("failed to open tty");
    write!(&tty, "Would you like to make the changes described above (y/n)? ");
    io::BufReader::new(tty).read_line(&mut input).ok().expect("failed to read user input");
    input.to_lowercase().starts_with("y")
}

fn edfp<I: io::Read, O: io::Write>(mut input: I, mut output: O) {
    let mut buf = String::new();
    input.read_to_string(&mut buf).expect("Failed to read input");

    let mut tmp: NamedTempFile = NamedTempFile::new().expect("failed to create temporary file");
    tmp.write_all(&buf.as_bytes()).expect("Could not write to tempfile");

    let editor_program = lookup_program();
    let tmp_path_str = tmp.path().to_str().expect("failed to parse tmp path");
    let editor = Command::new(editor_program.as_str()).arg(tmp_path_str).spawn();

    let success = editor
        .ok()
        .expect(format!("failed to launch {}", editor_program).as_str())
        .wait()
        .expect(format!("failed to wait for {}", editor_program).as_str())
        .success();

    if !success {
        panic!(format!("{} exited non-zero", editor_program));
    }

    let commands_file = tmp.reopen().ok().expect("failed to re-open");
    let commands = io::BufReader::new(commands_file).lines()
        .map(|line| line.ok().expect("failed to read line from commands file"));
    let changes: Vec<Change> = buf.lines().map(String::from).zip(commands).map(parse_lines).collect();

    let any_changes: bool = describe_changes(&changes, output);

    if any_changes && user_approves_changes() {
        for change in changes {
            change.operate().ok().expect("failed to change")
        }
    }
}

fn main() {
    let input = io::stdin();
    let mut output = io::stdout();
    edfp(input, output)
}
<|endoftext|>


<commit_before>extern crate byteorder;
use std::io::Cursor;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

macro_rules! define_instructions {
    ($($variant:ident, $value:expr, $name:expr, $num_operands:expr;)*) => (
        #[derive(Clone, Copy, Debug, Eq, PartialEq)]
        #[repr(u8)]
        enum Inst {
            $($variant = $value),*
        }

        impl Inst {
            fn from_u8(inst: u8) -> Option<Inst> {
                match inst {
                    $($value => Some(Inst::$variant),)*
                    _ => None
                }
            }

            fn from_str(inst: &str) -> Option<Inst> {
                match inst {
                    $($name => Some(Inst::$variant),)*
                    _ => None
                }
            }

            fn num_operands(&self) -> u8 {
                match *self {
                    $(Inst::$variant => $num_operands),*
                }
            }
        }
    )
}

// Bytecode instruction opcodes. The values of these opcodes should never change, to remain
// compatible with existing bytecode programs.
define_instructions! {
    Nop,   0,  "nop",   0;
    Print, 1,  "print", 0;
    Halt,  2,  "halt",  0;
    Push,  3,  "push",  1;
    Dup,   4,  "dup",   0;
    Pop,   5,  "pop",   0;
    Swap,  6,  "swap",  0;
    Add,   7,  "add",   0;
    Sub,   8,  "sub",   0;
    Mul,   9,  "mul",   0;
    Div,   10, "div",   0;
    Mod,   11, "mod",   0;
    Eq,    12, "eq",    0;
    Lt,    13, "lt",    0;
    Lte,   14, "lte",   0;
    Gt,    15, "gt",    0;
    Gte,   16, "gte",   0;
    Jz,    17, "jz",    1;
    Jnz,   18, "jnz",   1;
    Jump,  19, "jump",  1;
    Call,  20, "call",  1;
    Ret,   21, "ret",   0;
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum VmError {
    InvalidOpcode,
    UnexpectedProgramEnd, // Hit end of program while reading operand.
    StackOverflow,
    StackUnderflow,
    ControlStackOverflow,
    ControlStackUnderflow,
}

#[derive(Debug, Eq, PartialEq)]
struct Vm {
    stack: Box<[i32]>,
    stack_idx: usize,
    control_stack: Box<[i32]>,
    control_stack_idx: usize,
}

impl Vm {
    fn execute(&mut self, program: &[u8]) -> Result<(), VmError> {
        use VmError::*;

        let stack = &mut self.stack;
        let mut stack_idx = self.stack_idx;
        let mut opcodes = Cursor::new(program);

        while let Ok(opcode) = opcodes.read_u8() {
            let inst = try!(Inst::from_u8(opcode).ok_or(InvalidOpcode));
            match inst {
                Inst::Nop => {},

                Inst::Print => {
                    let val = *try!(stack.get(stack_idx - 1).ok_or(StackUnderflow));
                    println!("{}", val);
                    stack_idx -= 1;
                },

                Inst::Halt => {
                    break;
                },

                Inst::Push => {
                    let val = try!(opcodes.read_i32::<LittleEndian>().or(Err(UnexpectedProgramEnd)));
                    let stack_top = try!(stack.get_mut(stack_idx).ok_or(StackOverflow));
                    *stack_top = val;
                    stack_idx += 1;
                },

                Inst::Dup => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    if stack_idx >= stack.len() { return Err(StackOverflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx) = *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx += 1;
                },

                Inst::Pop => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    stack_idx -= 1;
                },

                Inst::Swap => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let tmp = *stack.get_unchecked(stack_idx - 1);
                        *stack.get_unchecked_mut(stack_idx - 1) = *stack.get_unchecked(stack_idx - 2);
                        *stack.get_unchecked_mut(stack_idx - 2) = tmp;
                    }
                },

                Inst::Add => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) += *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Sub => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) -= *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Mul => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) *= *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Div => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) /= *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Mod => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) %= *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Eq => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 == val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Lt => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 < val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Lte => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 <= val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Gt => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 > val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Gte => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 >= val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Jz => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    let condition = unsafe { *stack.get_unchecked(stack_idx - 1) };
                    try!(jump(&mut opcodes, condition == 0));
                    stack_idx -= 1;
                },

                Inst::Jnz => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    let condition = unsafe { *stack.get_unchecked(stack_idx - 1) };
                    try!(jump(&mut opcodes, condition != 0));
                    stack_idx -= 1;
                },

                Inst::Jump => {
                    try!(jump(&mut opcodes, true));
                },

                Inst::Call => {
                    let control_stack_top = try!(self.control_stack.get_mut(self.control_stack_idx).ok_or(ControlStackOverflow));
                    *control_stack_top = opcodes.position() as i32 + 4;
                    try!(jump(&mut opcodes, true));
                    self.control_stack_idx += 1;
                },

                Inst::Ret => {
                    let addr = *try!(self.control_stack.get_mut(self.control_stack_idx - 1).ok_or(ControlStackOverflow));
                    opcodes.set_position(addr as u64);
                    self.control_stack_idx -= 1;
                },
            }
        }

        #[inline(always)]
        fn jump(opcodes: &mut Cursor<&[u8]>, condition: bool) -> Result<(), VmError> {
            let delta = try!(opcodes.read_i32::<LittleEndian>().or(Err(UnexpectedProgramEnd)));
            if condition {
                let operand_size = std::mem::size_of::<i32>() as i64;
                let addr = (opcodes.position() as i64 + delta as i64 - operand_size) as u64;
                opcodes.set_position(addr);
            }
            Ok(())
        }

        self.stack_idx = stack_idx;
        Ok(())
    }
}

fn assemble(source: &str) -> Vec<u8> {
    use std::collections::HashMap;

    let mut program: Vec<u8> = Vec::new();
    let mut label_definitions: HashMap<&str, usize> = HashMap::new();
    let mut label_uses: Vec<(&str, usize)> = Vec::new();

    for line in source.split(|c| c == '\n' || c == ';') {
        let mut tokens = line.split(char::is_whitespace).filter(|s| !s.is_empty());
        let mut first_token = tokens.next();

        // Parse an optional label at the start of the line.
        if let Some(label) = first_token {
            if label.chars().next_back() == Some(':') {
                let label_name = &label[..label.len() - 1];
                if label_definitions.insert(label_name, program.len()).is_some() {
                    panic!("Attempted to redefine label '{}'", label_name);
                }
                first_token = tokens.next()
            }
        }

        // Parse the rest of the line if it's not blank.
        if let Some(opcode) = first_token {
            // Parse the instruction name.
            let inst = Inst::from_str(opcode).unwrap_or_else(|| {
                panic!("Unrecognized instruction '{}'", opcode)
            });
            program.push(inst as u8);

            // Parse the operands.
            for _ in 0..inst.num_operands() {
                let operand = tokens.next().unwrap_or_else(|| {
                    panic!("Missing one or more operands after '{}'", opcode)
                });

                if operand.chars().next() == Some('@') {
                    let label_name = &operand[1..];
                    label_uses.push((label_name, program.len()));

                    // Push four zero bytes to be overwritten by the label location later.
                    for _ in 0..4 { program.push(0); }
                } else if let Ok(number) = operand.parse::<i32>() {
                    let operand_index = program.len();
                    for _ in 0..4 { program.push(0); }
                    (&mut program[operand_index..]).write_i32::<LittleEndian>(number).unwrap();
                } else {
                    panic!("Expected label or valid 32-bit signed integer after '{}', not '{}'",
                           opcode, operand);
                }
            }
        }
    }

    // Resolve label references and fill in their values in the bytecode.
    for (label_name, use_index) in label_uses {
        let target_index = *label_definitions.get(label_name).unwrap_or_else(|| {
            panic!("Reference to undefined label '{}'", label_name);
        });
        let delta = target_index as i32 - use_index as i32;
        (&mut program[use_index..]).write_i32::<LittleEndian>(delta).unwrap();
    }

    program
}

fn main() {
    let source = r"
        push 1
loop:   call @plus2
        dup; print
        jump @loop

plus2:  push 2
        add
        ret
    ";

    // let source = r"
    //     push 1
    //     push 10
// loop:   dup; push 0; eq; jnz @done
    //     jump @loop
// done:   print
    // ";

    let program = assemble(source);

    let mut vm = Vm {
        stack: Box::new([0; 256]),
        stack_idx: 0,
        control_stack: Box::new([0; 256]),
        control_stack_idx: 0,
    };
    vm.execute(&program).unwrap();
}
<commit_msg>Add control stack cpush/cpop/cdup instructions.<commit_after>extern crate byteorder;
use std::io::Cursor;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

macro_rules! define_instructions {
    ($($variant:ident, $value:expr, $name:expr, $num_operands:expr;)*) => (
        #[derive(Clone, Copy, Debug, Eq, PartialEq)]
        #[repr(u8)]
        enum Inst {
            $($variant = $value),*
        }

        impl Inst {
            fn from_u8(inst: u8) -> Option<Inst> {
                match inst {
                    $($value => Some(Inst::$variant),)*
                    _ => None
                }
            }

            fn from_str(inst: &str) -> Option<Inst> {
                match inst {
                    $($name => Some(Inst::$variant),)*
                    _ => None
                }
            }

            fn num_operands(&self) -> u8 {
                match *self {
                    $(Inst::$variant => $num_operands),*
                }
            }
        }
    )
}

// Bytecode instruction opcodes. The values of these opcodes should never change, to remain
// compatible with existing bytecode programs.
define_instructions! {
    Nop,   0,  "nop",   0;
    Print, 1,  "print", 0;
    Halt,  2,  "halt",  0;
    Push,  3,  "push",  1;
    Dup,   4,  "dup",   0;
    Pop,   5,  "pop",   0;
    Swap,  6,  "swap",  0;
    Add,   7,  "add",   0;
    Sub,   8,  "sub",   0;
    Mul,   9,  "mul",   0;
    Div,   10, "div",   0;
    Mod,   11, "mod",   0;
    Eq,    12, "eq",    0;
    Lt,    13, "lt",    0;
    Lte,   14, "lte",   0;
    Gt,    15, "gt",    0;
    Gte,   16, "gte",   0;
    Jz,    17, "jz",    1;
    Jnz,   18, "jnz",   1;
    Jump,  19, "jump",  1;
    Call,  20, "call",  1;
    Ret,   21, "ret",   0;
    CPush, 22, "cpush", 0;
    CPop,  23, "cpop",  0;
    CDup,  24, "cdup",  0;
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum VmError {
    InvalidOpcode,
    UnexpectedProgramEnd, // Hit end of program while reading operand.
    StackOverflow,
    StackUnderflow,
    ControlStackOverflow,
    ControlStackUnderflow,
}

#[derive(Debug, Eq, PartialEq)]
struct Vm {
    stack: Box<[i32]>,
    stack_idx: usize,
    control_stack: Box<[i32]>,
    control_stack_idx: usize,
}

impl Vm {
    fn execute(&mut self, program: &[u8]) -> Result<(), VmError> {
        use VmError::*;

        let stack = &mut self.stack;
        let mut stack_idx = self.stack_idx;
        let mut opcodes = Cursor::new(program);

        while let Ok(opcode) = opcodes.read_u8() {
            let inst = try!(Inst::from_u8(opcode).ok_or(InvalidOpcode));
            match inst {
                Inst::Nop => {},

                Inst::Print => {
                    let val = *try!(stack.get(stack_idx - 1).ok_or(StackUnderflow));
                    println!("{}", val);
                    stack_idx -= 1;
                },

                Inst::Halt => {
                    break;
                },

                Inst::Push => {
                    let val = try!(opcodes.read_i32::<LittleEndian>().or(Err(UnexpectedProgramEnd)));
                    let stack_top = try!(stack.get_mut(stack_idx).ok_or(StackOverflow));
                    *stack_top = val;
                    stack_idx += 1;
                },

                Inst::Dup => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    if stack_idx >= stack.len() { return Err(StackOverflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx) = *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx += 1;
                },

                Inst::Pop => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    stack_idx -= 1;
                },

                Inst::Swap => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let tmp = *stack.get_unchecked(stack_idx - 1);
                        *stack.get_unchecked_mut(stack_idx - 1) = *stack.get_unchecked(stack_idx - 2);
                        *stack.get_unchecked_mut(stack_idx - 2) = tmp;
                    }
                },

                Inst::Add => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) += *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Sub => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) -= *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Mul => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) *= *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Div => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) /= *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Mod => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) %= *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Eq => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 == val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Lt => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 < val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Lte => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 <= val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Gt => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 > val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Gte => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 >= val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Jz => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    let condition = unsafe { *stack.get_unchecked(stack_idx - 1) };
                    try!(jump(&mut opcodes, condition == 0));
                    stack_idx -= 1;
                },

                Inst::Jnz => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    let condition = unsafe { *stack.get_unchecked(stack_idx - 1) };
                    try!(jump(&mut opcodes, condition != 0));
                    stack_idx -= 1;
                },

                Inst::Jump => {
                    try!(jump(&mut opcodes, true));
                },

                Inst::Call => {
                    let control_stack_top = try!(self.control_stack.get_mut(self.control_stack_idx).ok_or(ControlStackOverflow));
                    *control_stack_top = opcodes.position() as i32 + 4;
                    try!(jump(&mut opcodes, true));
                    self.control_stack_idx += 1;
                },

                Inst::Ret => {
                    let addr = *try!(self.control_stack.get_mut(self.control_stack_idx - 1).ok_or(ControlStackOverflow));
                    opcodes.set_position(addr as u64);
                    self.control_stack_idx -= 1;
                },

                Inst::CPush => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    if self.control_stack_idx >= self.control_stack.len() { return Err(ControlStackOverflow); }
                    unsafe {
                        *self.control_stack.get_unchecked_mut(self.control_stack_idx) = *stack.get_unchecked(stack_idx - 1);
                    }
                    self.control_stack_idx += 1;
                    stack_idx -= 1;
                },

                Inst::CPop => {
                    if self.control_stack_idx < 1 { return Err(ControlStackUnderflow); }
                    if stack_idx >= stack.len() { return Err(StackOverflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx) = *self.control_stack.get_unchecked(self.control_stack_idx - 1);
                    }
                    self.control_stack_idx -= 1;
                    stack_idx += 1;
                },

                Inst::CDup => {
                    if self.control_stack_idx < 1 { return Err(ControlStackUnderflow); }
                    if stack_idx >= stack.len() { return Err(StackOverflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx) = *self.control_stack.get_unchecked(self.control_stack_idx - 1);
                    }
                    stack_idx += 1;
                },
            }
        }

        #[inline(always)]
        fn jump(opcodes: &mut Cursor<&[u8]>, condition: bool) -> Result<(), VmError> {
            let delta = try!(opcodes.read_i32::<LittleEndian>().or(Err(UnexpectedProgramEnd)));
            if condition {
                let operand_size = std::mem::size_of::<i32>() as i64;
                let addr = (opcodes.position() as i64 + delta as i64 - operand_size) as u64;
                opcodes.set_position(addr);
            }
            Ok(())
        }

        self.stack_idx = stack_idx;
        Ok(())
    }
}

fn assemble(source: &str) -> Vec<u8> {
    use std::collections::HashMap;

    let mut program: Vec<u8> = Vec::new();
    let mut label_definitions: HashMap<&str, usize> = HashMap::new();
    let mut label_uses: Vec<(&str, usize)> = Vec::new();

    for line in source.split(|c| c == '\n' || c == ';') {
        let mut tokens = line.split(char::is_whitespace).filter(|s| !s.is_empty());
        let mut first_token = tokens.next();

        // Parse an optional label at the start of the line.
        if let Some(label) = first_token {
            if label.chars().next_back() == Some(':') {
                let label_name = &label[..label.len() - 1];
                if label_definitions.insert(label_name, program.len()).is_some() {
                    panic!("Attempted to redefine label '{}'", label_name);
                }
                first_token = tokens.next()
            }
        }

        // Parse the rest of the line if it's not blank.
        if let Some(opcode) = first_token {
            // Parse the instruction name.
            let inst = Inst::from_str(opcode).unwrap_or_else(|| {
                panic!("Unrecognized instruction '{}'", opcode)
            });
            program.push(inst as u8);

            // Parse the operands.
            for _ in 0..inst.num_operands() {
                let operand = tokens.next().unwrap_or_else(|| {
                    panic!("Missing one or more operands after '{}'", opcode)
                });

                if operand.chars().next() == Some('@') {
                    let label_name = &operand[1..];
                    label_uses.push((label_name, program.len()));

                    // Push four zero bytes to be overwritten by the label location later.
                    for _ in 0..4 { program.push(0); }
                } else if let Ok(number) = operand.parse::<i32>() {
                    let operand_index = program.len();
                    for _ in 0..4 { program.push(0); }
                    (&mut program[operand_index..]).write_i32::<LittleEndian>(number).unwrap();
                } else {
                    panic!("Expected label or valid 32-bit signed integer after '{}', not '{}'",
                           opcode, operand);
                }
            }
        }
    }

    // Resolve label references and fill in their values in the bytecode.
    for (label_name, use_index) in label_uses {
        let target_index = *label_definitions.get(label_name).unwrap_or_else(|| {
            panic!("Reference to undefined label '{}'", label_name);
        });
        let delta = target_index as i32 - use_index as i32;
        (&mut program[use_index..]).write_i32::<LittleEndian>(delta).unwrap();
    }

    program
}

fn main() {
    // let source = r"
    //     push 1
// loop:   call @plus2
    //     dup; print
    //     jump @loop

// plus2:  push 2
    //     add
    //     ret
    // ";

    let source = r"
        push 10
        call @fact
        print
        halt

fact:   push 1
        swap
loop:   dup; jz @done
        dup; cpush
        mul
        cpop; push 1; sub
        jump @loop
done:   pop
        ret
    ";

    let program = assemble(source);

    let mut vm = Vm {
        stack: Box::new([0; 256]),
        stack_idx: 0,
        control_stack: Box::new([0; 256]),
        control_stack_idx: 0,
    };
    vm.execute(&program).unwrap();
}
<|endoftext|>