<commit_before>#![cfg(windows)]

extern crate libc;
extern crate winapi;
extern crate "kernel32-sys" as kernel32;

use self::kernel32::*;
use self::winapi::*;
use std::ptr;
use std::iter::repeat;
use std::cmp::min;
use screen::Key;
use screen::Key::*;

macro_rules! win32 {
    ($funcall:expr) => (
        if unsafe { $funcall } == 0 {
            panic!("Win32 call failed");
        }
    );
}

pub struct Screen {
    pub height: u16,
    pub width: u16,
    pub visible_choices: u16,
    pub start_line: u16,
    original_console_mode: DWORD,
    original_colors: WORD,
    conin: HANDLE,
    conout: HANDLE,
}

impl Drop for Screen {
    fn drop(&mut self) {
        win32!(SetConsoleMode(self.conout, self.original_console_mode));
    }
}

impl Screen {
    pub fn open_screen() -> Screen {
        let mut orig_mode;
        let conin: HANDLE;
        let conout: HANDLE;
        unsafe {
            // Unlike in Linux, we *could* use the default stdin/stdout file handles to talk
            // directly to the console. However, we need the default stdin handle for the list of
            // choices, and we need to write just the final selection to stdout. Therefore we have
            // to explicitly create new handles to talk to the console. The gory details are
            // available at:
            //
            // https://msdn.microsoft.com/en-us/library/windows/desktop/ms682075%28v=vs.85%29.aspx
            // http://stackoverflow.com/questions/377152/what-does-createfileconin-do
            const OPEN_EXISTING: DWORD = 3;
            let rw_access = GENERIC_READ | GENERIC_WRITE;
            conin = CreateFileA("CONIN$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            conout = CreateFileA("CONOUT$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            orig_mode = ::std::mem::uninitialized();
        }

        if conin == INVALID_HANDLE_VALUE || conout == INVALID_HANDLE_VALUE {
            panic!("Unable to open console");
        }

        let (cols, rows) = Screen::winsize(conout).unwrap();

        win32!(GetConsoleMode(conout, &mut orig_mode));
        let new_mode = orig_mode & !(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT);
        win32!(SetConsoleMode(conin, new_mode));

        let visible_choices = min(20, rows - 1);
        let start_line = rows - visible_choices - 1;
        let original_colors = Screen::get_original_colors(conout);
        Screen {
            height: rows,
            width: cols,
            visible_choices: visible_choices,
            start_line: start_line,
            original_console_mode: orig_mode,
            original_colors: original_colors,
            conin: conin,
            conout: conout,
        }
    }

    pub fn move_cursor(&mut self, line: u16, column: u16) {
        win32!(SetConsoleCursorPosition(self.conout, COORD { X: column as i16, Y: line as i16}));
    }

    pub fn move_cursor_to_bottom(&mut self) {
        let end_line = self.start_line + self.visible_choices;
        self.move_cursor(end_line, 0);
        self.write("\n");
    }

    pub fn blank_screen(&mut self) {
        let start_line = self.start_line;
        self.move_cursor(start_line, 0);
        let blank_line = repeat(' ').take(self.width as usize).collect::<String>();
        for _ in 0..self.height {
            self.write(&blank_line);
        }
        self.move_cursor(start_line, 0);
    }

    pub fn show_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: TRUE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn hide_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: FALSE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn write(&mut self, s: &str) {
        let mut copy = s.to_string();
        let len: DWORD = copy.len() as DWORD;
        copy.push('\0');
        let mut bytes_written: DWORD = 0;
        win32!(WriteFile(self.conout, copy.as_ptr() as PVOID, len, &mut bytes_written as LPDWORD, ptr::null_mut()));
    }

    pub fn write_inverted(&mut self, s: &str) {
        let orig = self.original_colors;
        const BLACK_ON_WHITE: WORD = (BACKGROUND_RED | BACKGROUND_GREEN | BACKGROUND_BLUE) as WORD;
        self.set_colors(BLACK_ON_WHITE);
        self.write(s);
        self.set_colors(orig);
    }

    fn set_colors(&mut self, colors: WORD) {
        win32!(SetConsoleTextAttribute(self.conout, colors));
    }

    // Currently the Windows implementation does not buffer, so this function just performs a
    // blocking read of a single key.
    pub fn get_buffered_keys(&mut self) -> Vec<Key> {
        let mut buf: Vec<u16> = repeat(0u16).take(0x1000).collect();
        let mut chars_read: DWORD = 0;
        win32!(ReadFile(self.conin, buf.as_mut_ptr() as LPVOID, 1, &mut chars_read as LPDWORD, ptr::null_mut()));

        let mut ret = Vec::new();
        for i in 0..chars_read {
            ret.push(Screen::translate_byte(buf[i as usize]));
        }
        ret
    }

    fn translate_byte(byte: u16) -> Key {
        if byte == '\r' as u16 {
            Enter
        } else if byte == 127 {
            Backspace
        } else if byte & 96 == 0 {
            Control(((byte + 96u16) as u8) as char)
        } else {
            Char((byte as u8) as char)
        }
    }

    fn get_original_colors(handle: HANDLE) -> WORD {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(handle, &mut buffer_info));
        buffer_info.wAttributes
    }

    fn winsize(conout: HANDLE) -> Option<(u16, u16)> {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        let result = unsafe { GetConsoleScreenBufferInfo(conout, &mut buffer_info) };
        if result != 0 {
            // This code specifically computes the size of the window,
            // *not* the size of the buffer (which is easily available
            // from dwSize). I got the algorithm from:
            //
            // http://stackoverflow.com/a/12642749
            let left = buffer_info.srWindow.Left;
            let top = buffer_info.srWindow.Top;
            let right = buffer_info.srWindow.Right;
            let bottom = buffer_info.srWindow.Bottom;
            let cols = right - left + 1;
            let rows = bottom - top + 1;
            Some((cols as u16, rows as u16))
        } else {
            None
        }
    }
}

#[test]
fn winsize_test() {
    let conout = unsafe { kernel32::GetStdHandle(-11) };
    let (cols, rows) = Screen::winsize(conout).expect("Failed to get window size!");
    // We don't know the window size a priori, but we can at least
    // assert that it is within some kind of sensible range.
    assert!(cols > 40);
    assert!(rows > 40);
    assert!(cols < 1000);
    assert!(rows < 1000);
}
<commit_msg>Use multithreaded input handling on Windows<commit_after>#![cfg(windows)]

extern crate libc;
extern crate winapi;
extern crate "kernel32-sys" as kernel32;

use self::kernel32::*;
use self::winapi::*;
use std::ptr;
use std::iter::repeat;
use std::cmp::min;
use screen::Key;
use screen::Key::*;

use std::thread;
use std::sync::mpsc::Receiver;
use std::sync::mpsc;

macro_rules! win32 {
    ($funcall:expr) => (
        if unsafe { $funcall } == 0 {
            panic!("Win32 call failed");
        }
    );
}

pub struct Screen {
    pub height: u16,
    pub width: u16,
    pub visible_choices: u16,
    pub start_line: u16,
    original_console_mode: DWORD,
    original_colors: WORD,
    input: Receiver<u16>,
    conout: HANDLE,
}

impl Drop for Screen {
    fn drop(&mut self) {
        win32!(SetConsoleMode(self.conout, self.original_console_mode));
    }
}

impl Screen {
    pub fn open_screen() -> Screen {
        let mut orig_mode;
        let conin: HANDLE;
        let conout: HANDLE;
        unsafe {
            // Unlike in Linux, we *could* use the default stdin/stdout file handles to talk
            // directly to the console. However, we need the default stdin handle for the list of
            // choices, and we need to write just the final selection to stdout. Therefore we have
            // to explicitly create new handles to talk to the console. The gory details are
            // available at:
            //
            // https://msdn.microsoft.com/en-us/library/windows/desktop/ms682075%28v=vs.85%29.aspx
            // http://stackoverflow.com/questions/377152/what-does-createfileconin-do
            const OPEN_EXISTING: DWORD = 3;
            let rw_access = GENERIC_READ | GENERIC_WRITE;
            conin = CreateFileA("CONIN$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            conout = CreateFileA("CONOUT$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            orig_mode = ::std::mem::uninitialized();
        }

        if conin == INVALID_HANDLE_VALUE || conout == INVALID_HANDLE_VALUE {
            panic!("Unable to open console");
        }

        let (cols, rows) = Screen::winsize(conout).unwrap();

        win32!(GetConsoleMode(conout, &mut orig_mode));
        let new_mode = orig_mode & !(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT);
        win32!(SetConsoleMode(conin, new_mode));

        let rx = Screen::spawn_input_thread(conin as usize);
        let visible_choices = min(20, rows - 1);
        let start_line = rows - visible_choices - 1;
        let original_colors = Screen::get_original_colors(conout);
        Screen {
            height: rows,
            width: cols,
            visible_choices: visible_choices,
            start_line: start_line,
            original_console_mode: orig_mode,
            original_colors: original_colors,
            input: rx,
            conout: conout,
        }
    }

    // We have to take the conin handle as a usize instead of a *mut c_void in order to avoid a
    // lecture from the compiler about how the latter type cannot be safely sent between threads.
    // I'm not sure if a better solution exists at this time.
    fn spawn_input_thread(conin: usize) -> Receiver<u16> {
        let (tx, rx) = mpsc::channel();

        thread::spawn(move || {
            loop {
                let conin = conin as *mut c_void;
                let mut buf: Vec<u16> = repeat(0u16).take(0x1000).collect();
                let mut chars_read: DWORD = 0;
                win32!(ReadFile(conin, buf.as_mut_ptr() as LPVOID, 1, &mut chars_read as LPDWORD, ptr::null_mut()));
                for i in 0..chars_read as usize {
                    tx.send(buf[i]).unwrap();
                }
            }
        });

        rx
    }

    pub fn move_cursor(&mut self, line: u16, column: u16) {
        win32!(SetConsoleCursorPosition(self.conout, COORD { X: column as i16, Y: line as i16}));
    }

    pub fn move_cursor_to_bottom(&mut self) {
        let end_line = self.start_line + self.visible_choices;
        self.move_cursor(end_line, 0);
        self.write("\n");
    }

    pub fn blank_screen(&mut self) {
        let start_line = self.start_line;
        self.move_cursor(start_line, 0);
        let blank_line = repeat(' ').take(self.width as usize).collect::<String>();
        for _ in 0..self.height {
            self.write(&blank_line);
        }
        self.move_cursor(start_line, 0);
    }

    pub fn show_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: TRUE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn hide_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: FALSE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn write(&mut self, s: &str) {
        let mut copy = s.to_string();
        let len: DWORD = copy.len() as DWORD;
        copy.push('\0');
        let mut bytes_written: DWORD = 0;
        win32!(WriteFile(self.conout, copy.as_ptr() as PVOID, len, &mut bytes_written as LPDWORD, ptr::null_mut()));
    }

    pub fn write_inverted(&mut self, s: &str) {
        let orig = self.original_colors;
        const BLACK_ON_WHITE: WORD = (BACKGROUND_RED | BACKGROUND_GREEN | BACKGROUND_BLUE) as WORD;
        self.set_colors(BLACK_ON_WHITE);
        self.write(s);
        self.set_colors(orig);
    }

    fn set_colors(&mut self, colors: WORD) {
        win32!(SetConsoleTextAttribute(self.conout, colors));
    }

    pub fn get_buffered_keys(&mut self) -> Vec<Key> {
        let mut ret = Vec::new();
        while let Ok(byte) = self.input.try_recv() {
            ret.push(Screen::translate_byte(byte));
        }
        if ret.is_empty() {
            let byte = self.input.recv().unwrap();
            ret.push(Screen::translate_byte(byte));
        }
        ret
    }

    fn translate_byte(byte: u16) -> Key {
        if byte == '\r' as u16 {
            Enter
        } else if byte == 127 {
            Backspace
        } else if byte & 96 == 0 {
            Control(((byte + 96u16) as u8) as char)
        } else {
            Char((byte as u8) as char)
        }
    }

    fn get_original_colors(handle: HANDLE) -> WORD {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(handle, &mut buffer_info));
        buffer_info.wAttributes
    }

    fn winsize(conout: HANDLE) -> Option<(u16, u16)> {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        let result = unsafe { GetConsoleScreenBufferInfo(conout, &mut buffer_info) };
        if result != 0 {
            // This code specifically computes the size of the window,
            // *not* the size of the buffer (which is easily available
            // from dwSize). I got the algorithm from:
            //
            // http://stackoverflow.com/a/12642749
            let left = buffer_info.srWindow.Left;
            let top = buffer_info.srWindow.Top;
            let right = buffer_info.srWindow.Right;
            let bottom = buffer_info.srWindow.Bottom;
            let cols = right - left + 1;
            let rows = bottom - top + 1;
            Some((cols as u16, rows as u16))
        } else {
            None
        }
    }
}

#[test]
fn winsize_test() {
    let conout = unsafe { kernel32::GetStdHandle(-11) };
    let (cols, rows) = Screen::winsize(conout).expect("Failed to get window size!");
    // We don't know the window size a priori, but we can at least
    // assert that it is within some kind of sensible range.
    assert!(cols > 40);
    assert!(rows > 40);
    assert!(cols < 1000);
    assert!(rows < 1000);
}
<|endoftext|>


<commit_before>//! SEGIMAP is an IMAP server implementation.
#![deny(non_camel_case_types)]
#![cfg_attr(feature = "unstable", feature(test))]
#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]

extern crate bufstream;
extern crate crypto;
#[macro_use] extern crate log;
#[macro_use]
extern crate nom;
extern crate num;
extern crate rand;
extern crate regex;
extern crate rustc_serialize;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate test;
extern crate time;
extern crate toml;
extern crate walkdir;

use server::Server;
use server::lmtp_serve;
use user::Session;
use user::User;
use user::Email;

use std::sync::Arc;
use std::thread::spawn;
use bufstream::BufStream;

mod command;
mod error;
mod folder;
mod message;
mod mime;
mod parser;
mod server;
mod user;
mod util;

fn main() {
    // This function only needs to be run once, really.
    // create_default_users(config.users.clone());

    // Create the server. We wrap it so that it is atomically reference
    // counted. This allows us to safely share it across threads
    let serv = Arc::new(Server::new());

    // Spawn a separate thread for listening for LMTP connections
    match serv.lmtp_listener() {
        Err(e) => {
            error!("Error listening on LMTP port: {}", e);
        }
        Ok(v) => {
            let lmtp_serv = serv.clone();
            spawn(move || {
                // let mut acceptor = v.listen();
                // We spawn a separate thread for each LMTP session
                for stream in v.incoming() {
                    match stream {
                        Err(e) => {
                            error!("Error accepting incoming LMTP connection: {}",
                                   e);
                        }
                        Ok(stream) => {
                            let session_serv = lmtp_serv.clone();
                            spawn(move || { lmtp_serve(session_serv, BufStream::new(stream)) });
                        }
                    }
                }
            });
        }
    }

    // The main thread handles listening for IMAP connections
    match serv.imap_listener() {
        Err(e) => {
            error!("Error listening on IMAP port: {}", e);
        }
        Ok(v) => {
            // let mut acceptor = v.listen();
            // For each IMAP session, we spawn a separate thread.
            for stream in v.incoming() {
                match stream {
                    Err(e) => {
                        error!("Error accepting incoming IMAP connection: {}",
                               e)
                    }
                    Ok(stream) => {
                        let session_serv = serv.clone();
                        spawn(move || {
                            let mut session = Session::new(
                                               BufStream::new(stream),
                                               session_serv);
                            session.handle();
                        });
                    }
                }
            }
            drop(v);
        }
    }
}

/// Function to create our default users.json for testing
#[allow(dead_code)]
fn create_default_users(filename: &str) {
    let mut users = Vec::new();
    users.push(User::new(Email::new("will".to_string(), "xqz.ca".to_string()),
                         "54321".to_string(), "./maildir".to_string()));
    users.push(User::new(Email::new("nikitapekin".to_string(),
                                    "gmail.com".to_string()),
                         "12345".to_string(), "./maildir".to_string()));
    user::save_users(filename, &users);
}
<commit_msg>fix: gate the test crate behind the "unstable" feature<commit_after>//! SEGIMAP is an IMAP server implementation.
#![deny(non_camel_case_types)]
#![cfg_attr(feature = "unstable", feature(test))]
#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]

extern crate bufstream;
extern crate crypto;
#[macro_use] extern crate log;
#[macro_use]
extern crate nom;
extern crate num;
extern crate rand;
extern crate regex;
extern crate rustc_serialize;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[cfg(feature = "unstable")]
extern crate test;
extern crate time;
extern crate toml;
extern crate walkdir;

use server::Server;
use server::lmtp_serve;
use user::Session;
use user::User;
use user::Email;

use std::sync::Arc;
use std::thread::spawn;
use bufstream::BufStream;

mod command;
mod error;
mod folder;
mod message;
mod mime;
mod parser;
mod server;
mod user;
mod util;

fn main() {
    // This function only needs to be run once, really.
    // create_default_users(config.users.clone());

    // Create the server. We wrap it so that it is atomically reference
    // counted. This allows us to safely share it across threads
    let serv = Arc::new(Server::new());

    // Spawn a separate thread for listening for LMTP connections
    match serv.lmtp_listener() {
        Err(e) => {
            error!("Error listening on LMTP port: {}", e);
        }
        Ok(v) => {
            let lmtp_serv = serv.clone();
            spawn(move || {
                // let mut acceptor = v.listen();
                // We spawn a separate thread for each LMTP session
                for stream in v.incoming() {
                    match stream {
                        Err(e) => {
                            error!("Error accepting incoming LMTP connection: {}",
                                   e);
                        }
                        Ok(stream) => {
                            let session_serv = lmtp_serv.clone();
                            spawn(move || { lmtp_serve(session_serv, BufStream::new(stream)) });
                        }
                    }
                }
            });
        }
    }

    // The main thread handles listening for IMAP connections
    match serv.imap_listener() {
        Err(e) => {
            error!("Error listening on IMAP port: {}", e);
        }
        Ok(v) => {
            // let mut acceptor = v.listen();
            // For each IMAP session, we spawn a separate thread.
            for stream in v.incoming() {
                match stream {
                    Err(e) => {
                        error!("Error accepting incoming IMAP connection: {}",
                               e)
                    }
                    Ok(stream) => {
                        let session_serv = serv.clone();
                        spawn(move || {
                            let mut session = Session::new(
                                               BufStream::new(stream),
                                               session_serv);
                            session.handle();
                        });
                    }
                }
            }
            drop(v);
        }
    }
}

/// Function to create our default users.json for testing
#[allow(dead_code)]
fn create_default_users(filename: &str) {
    let mut users = Vec::new();
    users.push(User::new(Email::new("will".to_string(), "xqz.ca".to_string()),
                         "54321".to_string(), "./maildir".to_string()));
    users.push(User::new(Email::new("nikitapekin".to_string(),
                                    "gmail.com".to_string()),
                         "12345".to_string(), "./maildir".to_string()));
    user::save_users(filename, &users);
}
<|endoftext|>


<commit_before>extern crate ansi_term;
#[macro_use]
extern crate clap;
extern crate globset;
extern crate ignore;
#[macro_use]
extern crate nom;
extern crate rayon;

// Import helper macros before `parsers`
#[macro_use]
mod helper_parsers;

mod error;
mod file_utils;
mod parsers;
mod transpile;

fn main() {
    use ansi_term::Colour::{Red, Green};
    use clap::{App, Arg};
    use file_utils::walk_indentex_files;
    use std::path::{Path, PathBuf};
    use std::process;
    use transpile::transpile_file;
    use rayon::prelude::*;
    use std::cmp;

    let m = App::new("indentex")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Transpiler for an indentation-based superset of LaTeX")
        .arg(Arg::with_name("path")
            .help("Path to a single indentex file or a directory (recursively transpile all \
                   indentex files)")
            .index(1)
            .required(true))
        .arg(Arg::with_name("verbose")
            .help("Show transpilation progress")
            .short("v")
            .long("verbose"))
        .get_matches();

    let path = Path::new(m.value_of("path").unwrap());
    let verbose = m.is_present("verbose");

    let mut ret_val: i32 = 0;

    let batch: Vec<PathBuf> = if path.is_file() {
        vec![path.to_path_buf()]
    } else if path.is_dir() {
        match walk_indentex_files(&path) {
            Ok(b) => b,
            Err(e) => {
                ret_val = 2;
                println!("{}", Red.bold().paint(format!("{}", e)));
                Vec::new()
            }
        }
    } else {
        ret_val = 4;
        println!("{}",
                 Red.bold().paint(format!("Error: path '{}' is neither a file nor a directory",
                                          path.display())));
        Vec::new()
    };

    let ret_val_transpilation = batch.par_iter()
        .map(|p| match transpile_file(&p) {
            Ok(_) => {
                if verbose {
                    println!("Transpiling file '{}'... {}",
                             p.display(),
                             Green.paint("ok"));
                }
                return 0 as i32;
            }
            Err(e) => {
                if verbose {
                    println!("Transpiling file '{}'... {}",
                             p.display(),
                             Red.paint("failed"));
                }
                println!("{}",
                         Red.bold().paint(format!("Could not transpile '{}': {}", p.display(), e)));
                return 8 as i32;
            }
        })
        .max()
        .unwrap();

    process::exit(cmp::max(ret_val, ret_val_transpilation));
}
<commit_msg>Add ErrorCode enum<commit_after>extern crate ansi_term;
#[macro_use]
extern crate clap;
extern crate globset;
extern crate ignore;
#[macro_use]
extern crate nom;
extern crate rayon;

// Import helper macros before `parsers`
#[macro_use]
mod helper_parsers;

mod error;
mod file_utils;
mod parsers;
mod transpile;

enum ErrorCode {
    Ok = 0,
    WalkError = 2,
    FileTypeError = 4,
    TranspilationError = 8,
}

fn main() {
    use ansi_term::Colour::{Red, Green};
    use clap::{App, Arg};
    use file_utils::walk_indentex_files;
    use std::path::{Path, PathBuf};
    use std::process;
    use transpile::transpile_file;
    use rayon::prelude::*;
    use std::cmp;

    let m = App::new("indentex")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Transpiler for an indentation-based superset of LaTeX")
        .arg(Arg::with_name("path")
            .help("Path to a single indentex file or a directory (recursively transpile all \
                   indentex files)")
            .index(1)
            .required(true))
        .arg(Arg::with_name("verbose")
            .help("Show transpilation progress")
            .short("v")
            .long("verbose"))
        .get_matches();

    let path = Path::new(m.value_of("path").unwrap());
    let verbose = m.is_present("verbose");

    let mut ret_val = ErrorCode::Ok;

    let batch: Vec<PathBuf> = if path.is_file() {
        vec![path.to_path_buf()]
    } else if path.is_dir() {
        match walk_indentex_files(&path) {
            Ok(b) => b,
            Err(e) => {
                ret_val = ErrorCode::WalkError;
                println!("{}", Red.bold().paint(format!("{}", e)));
                Vec::new()
            }
        }
    } else {
        ret_val = ErrorCode::FileTypeError;
        println!("{}",
                 Red.bold().paint(format!("Error: path '{}' is neither a file nor a directory",
                                          path.display())));
        Vec::new()
    };

    let ret_val_transpilation = batch.par_iter()
        .map(|p| match transpile_file(&p) {
            Ok(_) => {
                if verbose {
                    println!("Transpiling file '{}'... {}",
                             p.display(),
                             Green.paint("ok"));
                }
                return ErrorCode::Ok as i32;
            }
            Err(e) => {
                if verbose {
                    println!("Transpiling file '{}'... {}",
                             p.display(),
                             Red.paint("failed"));
                }
                println!("{}",
                         Red.bold().paint(format!("Could not transpile '{}': {}", p.display(), e)));
                return ErrorCode::TranspilationError as i32;
            }
        })
        .max()
        .unwrap();

    process::exit(cmp::max(ret_val as i32, ret_val_transpilation));
}
<|endoftext|>


<commit_before>use std::env;

extern crate zwreec;

fn main() {
    println!("#main started");

    // handling commandline parameters
    let args: Vec<String> = env::args().collect();

    let input_file_name: &str;
    let output_file_name: &str;

    if args.len() == 1 {
        input_file_name = "README.md";
        output_file_name = "file_does_not_exist.zcode";
    } else if args.len() == 3 {
        println!("argument 1 {}", args[1]);
        println!("argument 2 {}", args[2]);

        input_file_name = &args[1];
        output_file_name = &args[2];
    } else {
        help();
        return;
    }

    // call library
    zwreec::compile(input_file_name, output_file_name);

    // only for testing
    println!("(1) {}", zwreec::frontend::temp_hello());
    println!("(2) {}", zwreec::backend::temp_hello());
    println!("(3) {}", zwreec::file::temp_hello());

    println!("#main finished");
}

fn help() {
    println!("usage:\n    zwreec <input_file> <output_file>");
}
<commit_msg>Use match! \O/<commit_after>use std::env;

extern crate zwreec;

fn main() {
    println!("#main started");

    // handling commandline parameters
    let args: Vec<String> = env::args().collect();

    let mut input_file_name: &str;
    let mut output_file_name: &str;

    match args.len() {
        1 => {
            input_file_name = "a.in";
            output_file_name = "a.out";
        },
        3 => {
            input_file_name = &args[1];
            output_file_name = &args[2];
        },
        _ => help()
    }

    // call library
    zwreec::compile(input_file_name, output_file_name);

    // only for testing
    println!("(1) {}", zwreec::frontend::temp_hello());
    println!("(2) {}", zwreec::backend::temp_hello());
    println!("(3) {}", zwreec::file::temp_hello());

    println!("#main finished");
}

fn help() {
    println!("usage:\n    zwreec <input_file> <output_file>");
}
<|endoftext|>


<commit_before>#[macro_use(handler)]
extern crate chatbot;
extern crate getopts;

use std::env;

use chatbot::Chatbot;
use chatbot::adapter::CliAdapter;
use chatbot::adapter::SlackAdapter;
use chatbot::adapter::IrcAdapter;
use chatbot::handler::GithubIssueLinker;

use getopts::Options;
use getopts::ParsingStyle;

#[allow(dead_code)]
fn main() {

    let args = env::args().collect::<Vec<String>>();
    let mut opts = Options::new();
    opts.parsing_style(ParsingStyle::StopAtFirstFree);
    opts.optopt("a", "adapter", "Chat Adapter to use", "slack|cli");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string())
    };

    let adapter_name = matches.opt_str("a").unwrap_or("cli".to_owned());

    let mut bot = Chatbot::new("chatbotbot");

    // Add adapter based on command line argument
    match adapter_name.as_ref() {
        "slack" => bot.add_adapter(SlackAdapter::new()),
        "cli" => bot.add_adapter(CliAdapter::new()),
        "irc" => {
            let config = chatbot::adapter::irc::Config {
                nickname: Some(format!("chatbotbot")),
                alt_nicks: Some(vec![format!("chatbot_rs"), format!("chatbotrs")]),
                server: Some(format!("irc.mozilla.org")),
                channels: Some(vec![format!("#chatbot")]),
                .. Default::default()
            };
            bot.add_adapter(IrcAdapter::new(config))
        },
        _ => panic!("Unexpected adapter name. Use 'cli' or 'slack'.")
    };

    let ping = handler!("PingHandler", r"ping", |_, _| { Some("pong".to_owned()) });

    let robot_name = "Mr. T";
    let trout = handler!("TroutSlap", r"slap (?P<user>.+)", move |matches, _| {
        match matches.name("user") {
            Some(user) => {
                Some(format!("{} slaps {} around a bit with a large trout",
                             robot_name, user))
            },
            None => None
        }
    });

    let echo = handler!("EchoHandler", r"echo (?P<msg>.+)", |matches, _| {
        matches.name("msg").map(|msg| { msg.to_owned() })
    });

    bot.add_handler(ping);
    bot.add_handler(trout);
    bot.add_handler(echo);
    bot.add_handler(GithubIssueLinker::new());

    bot.run();
}
<commit_msg>Clean up name use in main.rs<commit_after>#[macro_use(handler)]
extern crate chatbot;
extern crate getopts;

use std::env;

use chatbot::Chatbot;
use chatbot::adapter::CliAdapter;
use chatbot::adapter::SlackAdapter;
use chatbot::adapter::IrcAdapter;
use chatbot::handler::GithubIssueLinker;

use getopts::Options;
use getopts::ParsingStyle;

#[allow(dead_code)]
fn main() {

    let args = env::args().collect::<Vec<String>>();
    let mut opts = Options::new();
    opts.parsing_style(ParsingStyle::StopAtFirstFree);
    opts.optopt("a", "adapter", "Chat Adapter to use", "slack|cli");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string())
    };

    let adapter_name = matches.opt_str("a").unwrap_or("cli".to_owned());

    let name = "chatbotbot";
    let mut bot = Chatbot::new(name);

    // Add adapter based on command line argument
    match adapter_name.as_ref() {
        "slack" => bot.add_adapter(SlackAdapter::new()),
        "cli" => bot.add_adapter(CliAdapter::new()),
        "irc" => {
            let config = chatbot::adapter::irc::Config {
                nickname: Some(format!("{}", name)),
                alt_nicks: Some(vec![format!("{}_", name), format!("{}__", name)]),
                server: Some(format!("irc.mozilla.org")),
                channels: Some(vec![format!("#chatbot")]),
                .. Default::default()
            };
            bot.add_adapter(IrcAdapter::new(config))
        },
        _ => panic!("Unexpected adapter name. Use 'cli' or 'slack'.")
    };

    let ping = handler!("PingHandler", r"ping", |_, _| { Some("pong".to_owned()) });

    let trout = handler!("TroutSlap", r"slap (?P<user>.+)", move |matches, _| {
        match matches.name("user") {
            Some(user) => {
                Some(format!("{} slaps {} around a bit with a large trout",
                             name, user))
            },
            None => None
        }
    });

    let echo = handler!("EchoHandler", r"echo (?P<msg>.+)", |matches, _| {
        matches.name("msg").map(|msg| { msg.to_owned() })
    });

    bot.add_handler(ping);
    bot.add_handler(trout);
    bot.add_handler(echo);
    bot.add_handler(GithubIssueLinker::new());

    bot.run();
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate walkdir;
extern crate twox_hash;

use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use std::path::Path;

use clap::{App, Arg};
use walkdir::WalkDir;
use twox_hash::XxHash;
use std::hash::Hasher;

fn main() {
    let matches = App::new("dupr")
        .version("0.1.0")
        .author("Mike Sampson")
        .about("Duplicate file finder")
        .arg(
            Arg::with_name("DIR")
                .help("Directory to process")
                .required(true),
        )
        .get_matches();

    let mut files = HashMap::new();

    for entry in WalkDir::new(matches.value_of("DIR").unwrap()) {
        let entry = entry.unwrap();

        if !entry.file_type().is_file() {
            continue;
        }

        let metadata = entry.metadata().unwrap();

        let path = entry.path().to_path_buf();

        files.entry(metadata.len()).or_insert(Vec::new()).push(path);
    }

    for f_paths in files.values() {
        if f_paths.len() == 1 {
            continue;
        }

        let mut hashes = HashMap::new();
        for f_path in f_paths {
            hashes
                .entry(hash_file(&f_path))
                .or_insert(Vec::new())
                .push(f_path);
        }

        for (hash, paths) in hashes {
            if paths.len() < 2 {
                continue;
            }
            for p in paths {
                println!("{:016x} - {}", hash, p.display());
            }
        }
    }
}

fn hash_file<P: AsRef<Path>>(path: P) -> u64 {
    let mut file = File::open(path).unwrap();

    let mut hash = XxHash::with_seed(0);

    let mut buffer = [0; 32 * 1024];

    while let Ok(n) = file.read(&mut buffer) {
        if n == 0 {
            break;
        }
        hash.write(&buffer[0..n]);
    }

    hash.finish()
}
<commit_msg>Add summary option.<commit_after>extern crate clap;
extern crate walkdir;
extern crate twox_hash;

use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use std::path::Path;

use clap::{App, Arg};
use walkdir::WalkDir;
use twox_hash::XxHash;
use std::hash::Hasher;

struct Stats {
    file_count: u64,
    total_size: u64,
    duplicate_count: u64,
}

fn main() {
    let matches = App::new("dupr")
        .version("0.1.0")
        .author("Mike Sampson")
        .about("Duplicate file finder")
        .arg(
            Arg::with_name("DIR")
                .help("Directory to process")
                .required(true),
        )
        .arg(
            Arg::with_name("summary")
                .short("s")
                .long("summary")
                .help("Print out summary information after duplicates")
        )
        .get_matches();

    let mut files = HashMap::new();

    let mut stats = Stats {
        file_count: 0,
        total_size: 0,
        duplicate_count: 0,
    };

    for entry in WalkDir::new(matches.value_of("DIR").unwrap()) {
        let entry = entry.unwrap();

        if !entry.file_type().is_file() {
            continue;
        }

        let metadata = entry.metadata().unwrap();

        stats.file_count += 1;
        stats.total_size += metadata.len();

        let path = entry.path().to_path_buf();

        files.entry(metadata.len()).or_insert(Vec::new()).push(path);
    }

    for f_paths in files.values() {
        if f_paths.len() == 1 {
            continue;
        }

        let mut hashes = HashMap::new();
        for f_path in f_paths {
            hashes
                .entry(hash_file(&f_path))
                .or_insert(Vec::new())
                .push(f_path);
        }

        for (hash, paths) in hashes {
            if paths.len() < 2 {
                continue;
            }
            stats.duplicate_count += paths.len() as u64;
            for p in paths {
                println!("{:016x} - {}", hash, p.display());
            }
        }
    }

    if matches.is_present("summary")  {
        println!(
            "Processed {} files with a total size of {} bytes. {} duplicates found.",
            stats.file_count,
            stats.total_size,
            stats.duplicate_count
        );
    }
}

fn hash_file<P: AsRef<Path>>(path: P) -> u64 {
    let mut file = File::open(path).unwrap();

    let mut hash = XxHash::with_seed(0);

    let mut buffer = [0; 32 * 1024];

    while let Ok(n) = file.read(&mut buffer) {
        if n == 0 {
            break;
        }
        hash.write(&buffer[0..n]);
    }

    hash.finish()
}
<|endoftext|>


<commit_before>extern crate loc;

#[macro_use]
extern crate clap;
extern crate deque;
extern crate num_cpus;
extern crate regex;
extern crate ignore;

use clap::{Arg, App, AppSettings};
use ignore::WalkBuilder;

use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::thread;

use deque::{Stealer, Stolen};
use regex::Regex;

use loc::*;

enum Work {
    File(String),
    Quit,
}

struct Worker {
    chan: Stealer<Work>,
}

#[derive(Clone)]
struct FileCount {
    path: String,
    lang: Lang,
    count: Count,
}

// This concurrency pattern ripped directly from ripgrep
impl Worker {
    fn run(self) -> Vec<FileCount> {
        let mut v: Vec<FileCount> = vec![];
        loop {
            match self.chan.steal() {
                // What causes these?
                Stolen::Empty | Stolen::Abort => continue,
                Stolen::Data(Work::Quit) => break,
                Stolen::Data(Work::File(path)) => {
                    let lang = lang_from_ext(&path);
                    if lang != Lang::Unrecognized {
                        let count = count(&path);
                        v.push(FileCount {
                            lang: lang,
                            path: String::from(path),
                            count: count,
                        });
                    }
                }
            };
        }
        v
    }
}

fn main() {

    let matches = App::new("count")
        .global_settings(&[AppSettings::ColoredHelp])
        .version(crate_version!())
        .author("Curtis Gagliardi <curtis@curtis.io>")
        .about("counts things quickly hopefully")
        // TODO(cgag): actually implement filtering
        .arg(Arg::with_name("exclude")
            .required(false)
            .multiple(true)
            .long("exclude")
            .value_name("REGEX")
            .takes_value(true)
            .help("Rust regex of files to exclude"))
        .arg(Arg::with_name("include")
            .required(false)
            .multiple(true)
            .long("include")
            .value_name("REGEX")
            .takes_value(true)
            .help("Rust regex matching files to include. Anything not matched will be excluded"))
        .arg(Arg::with_name("files")
             .required(false)
             .long("files")
             .takes_value(false)
             .help("Show stats for individual files"))
        .arg(Arg::with_name("sort")
            .required(false)
            .long("sort")
            .value_name("COLUMN")
            .takes_value(true)
            .help("Column to sort by"))
        .arg(Arg::with_name("unrestricted")
             .required(false)
             .multiple(true)
             .long("unrestricted")
             .short("u")
             .takes_value(false)
             .help("A single -u won't respect .gitignore (etc.) files. Two -u flags will additionally count hidden files and directories."))
        .arg(Arg::with_name("target")
            .multiple(true)
            .help("File or directory to count (multiple arguments accepted)"))
        .get_matches();

    let targets = match matches.values_of("target") {
        Some(targets) => targets.collect(),
        None => vec!["."]
    };
    let sort = matches.value_of("sort").unwrap_or("code");
    let by_file = matches.is_present("files");
    let (use_ignore, ignore_hidden) = match matches.occurrences_of("unrestricted") {
        0 => (true,  true),
        1 => (false, true),
        2 => (false, false),
        _ => (false, false),
    };
    let exclude_regex = match matches.values_of("exclude") {
        Some(regex_strs) => {
            let combined_regex = regex_strs.map(|r| format!("({})", r)).collect::<Vec<String>>().join("|");
            match Regex::new(&combined_regex) {
                Ok(r) => Some(r),
                Err(e) => {
                    println!("Error processing exclude regex: {}", e);
                    std::process::exit(1);
                }
            }
        }
        None => None,
    };
    let include_regex = match matches.values_of("include") {
        Some(regex_strs) => {
            let combined_regex = regex_strs.map(|r| format!("({})", r)).collect::<Vec<String>>().join("|");
            match Regex::new(&combined_regex) {
                Ok(r) => Some(r),
                Err(e) => {
                    println!("Error processing include regex: {}", e);
                    std::process::exit(1);
                }
            }
        }
        None => None,
    };

    let threads = num_cpus::get();
    let mut workers = vec![];
    let (workq, stealer) = deque::new();
    for _ in 0..threads {
        let worker = Worker { chan: stealer.clone() };
        workers.push(thread::spawn(|| worker.run()));
    }

    for target in targets {
        // TODO(cgag): use WalkParallel?
        let walker = WalkBuilder::new(target).ignore(use_ignore)
                                             .git_ignore(use_ignore)
                                             .git_exclude(use_ignore)
                                             .hidden(ignore_hidden)
                                             .build();
        let files = walker
            .into_iter()
            .filter_map(|entry| entry.ok())
            .filter(|entry| entry.file_type().expect("no filetype").is_file())
            .map(|entry| String::from(entry.path().to_str().unwrap()))
            .filter(|path| match include_regex {
                None => true,
                Some(ref include) => include.is_match(path),
            })
            .filter(|path| match exclude_regex {
                None => true,
                Some(ref exclude) => !exclude.is_match(path),
            });

        for path in files {
            workq.push(Work::File(path));
        }
    }

    for _ in 0..workers.len() {
        workq.push(Work::Quit);
    }

    let mut filecounts: Vec<FileCount> = Vec::new();
    for worker in workers {
        filecounts.extend(worker.join().unwrap().iter().cloned())
    }

    // TODO(cgag): use insecure hashmaps or something
    let mut by_lang: HashMap<Lang, Vec<FileCount>> = HashMap::new();
    for fc in filecounts {
        match by_lang.entry(fc.lang) {
            Entry::Occupied(mut elem) => elem.get_mut().push(fc),
            Entry::Vacant(elem) => {
                elem.insert(vec![fc]);
            }
        };
    }

    let linesep = str_repeat("-", 80);

    if by_file {
        // print breakdown for each individual file
        println!("{}", linesep);
        println!(" {0: <17} {1: >8} {2: >12} {3: >12} {4: >12} {5: >12}",
                 "Language",
                 "Files",
                 "Lines",
                 "Blank",
                 "Comment",
                 "Code");
        println!("{}", linesep);

        // TODO(cgag): do the summing first, so we can do additional sorting
        // by totals.
        for (lang, mut filecounts) in by_lang {
            let mut total = Count::default();
            for fc in &filecounts {
                total.merge(&fc.count);
            }

            println!("{}", linesep);
            println!(" {0: <17} {1: >8} {2: >12} {3: >12} {4: >12} {5: >12}",
                     lang,
                     filecounts.len(),
                     total.lines,
                     total.blank,
                     total.comment,
                     total.code);

            match sort {
                "code" => filecounts.sort_by(|fc1, fc2| fc2.count.code.cmp(&fc1.count.code)),
                "comment" => {
                    filecounts.sort_by(|fc1, fc2| fc2.count.comment.cmp(&fc1.count.comment))
                }
                "blank" => filecounts.sort_by(|fc1, fc2| fc2.count.blank.cmp(&fc1.count.blank)),
                "lines" => filecounts.sort_by(|fc1, fc2| fc2.count.lines.cmp(&fc1.count.lines)),
                // No sorting by language or files here. Need to do it at a higher level.
                _ => (),
            }

            println!("{}", linesep);
            for fc in filecounts {
                println!("|{0: <25} {1: >12} {2: >12} {3: >12} {4: >12}",
                         last_n_chars(&fc.path, 25),
                         fc.count.lines,
                         fc.count.blank,
                         fc.count.comment,
                         fc.count.code);
            }
        }
    } else {
        // print summary by language
        let mut lang_totals: HashMap<&Lang, LangTotal> = HashMap::new();
        for (lang, filecounts) in &by_lang {
            let mut lang_total = Count::default();
            for fc in filecounts {
                lang_total.merge(&fc.count);
            }
            lang_totals.insert(lang,
                               LangTotal {
                                   files: filecounts.len() as u32,
                                   count: lang_total,
                               });
        }

        let mut totals_by_lang = lang_totals.iter().collect::<Vec<(&&Lang, &LangTotal)>>();
        match sort {
            "language" => totals_by_lang.sort_by(|&(l1, _), &(l2, _)| l1.to_s().cmp(l2.to_s())),
            "files" => totals_by_lang.sort_by(|&(_, c1), &(_, c2)| c2.files.cmp(&c1.files)),
            "code" => {
                totals_by_lang.sort_by(|&(_, c1), &(_, c2)| c2.count.code.cmp(&c1.count.code))
            }
            "comment" => {
                totals_by_lang.sort_by(|&(_, c1), &(_, c2)| c2.count.comment.cmp(&c1.count.comment))
            }
            "blank" => {
                totals_by_lang.sort_by(|&(_, c1), &(_, c2)| c2.count.blank.cmp(&c1.count.blank))
            }
            "lines" => {
                totals_by_lang.sort_by(|&(_, c1), &(_, c2)| c2.count.lines.cmp(&c1.count.lines))
            }
            _ => {
                println!("invalid sort option {}, sorting by code", sort);
                totals_by_lang.sort_by(|&(_, c1), &(_, c2)| c2.count.code.cmp(&c1.count.code))
            }
        }

        print_totals_by_lang(&linesep, &totals_by_lang);
    }

}

fn last_n_chars(s: &str, n: usize) -> String {
    if s.len() <= n {
        return String::from(s);
    }
    s.chars().skip(s.len() - n).collect::<String>()
}

fn str_repeat(s: &str, n: usize) -> String {
    std::iter::repeat(s).take(n).collect::<Vec<_>>().join("")
}

fn print_totals_by_lang(linesep: &str, totals_by_lang: &[(&&Lang, &LangTotal)]) {
    println!("{}", linesep);
    println!(" {0: <17} {1: >8} {2: >12} {3: >12} {4: >12} {5: >12}",
             "Language",
             "Files",
             "Lines",
             "Blank",
             "Comment",
             "Code");
    println!("{}", linesep);

    for &(lang, total) in totals_by_lang {
        println!(" {0: <17} {1: >8} {2: >12} {3: >12} {4: >12} {5: >12}",
                 lang,
                 total.files,
                 total.count.lines,
                 total.count.blank,
                 total.count.comment,
                 total.count.code);
    }

    let mut totals = LangTotal {
        files: 0,
        count: Count::default(),
    };
    for &(_, total) in totals_by_lang {
        totals.files += total.files;
        totals.count.code += total.count.code;
        totals.count.blank += total.count.blank;
        totals.count.comment += total.count.comment;
        totals.count.lines += total.count.lines;
    }

    println!("{}", linesep);
    println!(" {0: <17} {1: >8} {2: >12} {3: >12} {4: >12} {5: >12}",
             "Total",
             totals.files,
             totals.count.lines,
             totals.count.blank,
             totals.count.comment,
             totals.count.code);
    println!("{}", linesep);
}
<commit_msg>Set the name of the tool to `loc`<commit_after>extern crate loc;

#[macro_use]
extern crate clap;
extern crate deque;
extern crate num_cpus;
extern crate regex;
extern crate ignore;

use clap::{Arg, App, AppSettings};
use ignore::WalkBuilder;

use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::thread;

use deque::{Stealer, Stolen};
use regex::Regex;

use loc::*;

enum Work {
    File(String),
    Quit,
}

struct Worker {
    chan: Stealer<Work>,
}

#[derive(Clone)]
struct FileCount {
    path: String,
    lang: Lang,
    count: Count,
}

// This concurrency pattern ripped directly from ripgrep
impl Worker {
    fn run(self) -> Vec<FileCount> {
        let mut v: Vec<FileCount> = vec![];
        loop {
            match self.chan.steal() {
                // What causes these?
                Stolen::Empty | Stolen::Abort => continue,
                Stolen::Data(Work::Quit) => break,
                Stolen::Data(Work::File(path)) => {
                    let lang = lang_from_ext(&path);
                    if lang != Lang::Unrecognized {
                        let count = count(&path);
                        v.push(FileCount {
                            lang: lang,
                            path: String::from(path),
                            count: count,
                        });
                    }
                }
            };
        }
        v
    }
}

fn main() {

    let matches = App::new("loc")
        .global_settings(&[AppSettings::ColoredHelp])
        .version(crate_version!())
        .author("Curtis Gagliardi <curtis@curtis.io>")
        .about("counts things quickly hopefully")
        // TODO(cgag): actually implement filtering
        .arg(Arg::with_name("exclude")
            .required(false)
            .multiple(true)
            .long("exclude")
            .value_name("REGEX")
            .takes_value(true)
            .help("Rust regex of files to exclude"))
        .arg(Arg::with_name("include")
            .required(false)
            .multiple(true)
            .long("include")
            .value_name("REGEX")
            .takes_value(true)
            .help("Rust regex matching files to include. Anything not matched will be excluded"))
        .arg(Arg::with_name("files")
             .required(false)
             .long("files")
             .takes_value(false)
             .help("Show stats for individual files"))
        .arg(Arg::with_name("sort")
            .required(false)
            .long("sort")
            .value_name("COLUMN")
            .takes_value(true)
            .help("Column to sort by"))
        .arg(Arg::with_name("unrestricted")
             .required(false)
             .multiple(true)
             .long("unrestricted")
             .short("u")
             .takes_value(false)
             .help("A single -u won't respect .gitignore (etc.) files. Two -u flags will additionally count hidden files and directories."))
        .arg(Arg::with_name("target")
            .multiple(true)
            .help("File or directory to count (multiple arguments accepted)"))
        .get_matches();

    let targets = match matches.values_of("target") {
        Some(targets) => targets.collect(),
        None => vec!["."]
    };
    let sort = matches.value_of("sort").unwrap_or("code");
    let by_file = matches.is_present("files");
    let (use_ignore, ignore_hidden) = match matches.occurrences_of("unrestricted") {
        0 => (true,  true),
        1 => (false, true),
        2 => (false, false),
        _ => (false, false),
    };
    let exclude_regex = match matches.values_of("exclude") {
        Some(regex_strs) => {
            let combined_regex = regex_strs.map(|r| format!("({})", r)).collect::<Vec<String>>().join("|");
            match Regex::new(&combined_regex) {
                Ok(r) => Some(r),
                Err(e) => {
                    println!("Error processing exclude regex: {}", e);
                    std::process::exit(1);
                }
            }
        }
        None => None,
    };
    let include_regex = match matches.values_of("include") {
        Some(regex_strs) => {
            let combined_regex = regex_strs.map(|r| format!("({})", r)).collect::<Vec<String>>().join("|");
            match Regex::new(&combined_regex) {
                Ok(r) => Some(r),
                Err(e) => {
                    println!("Error processing include regex: {}", e);
                    std::process::exit(1);
                }
            }
        }
        None => None,
    };

    let threads = num_cpus::get();
    let mut workers = vec![];
    let (workq, stealer) = deque::new();
    for _ in 0..threads {
        let worker = Worker { chan: stealer.clone() };
        workers.push(thread::spawn(|| worker.run()));
    }

    for target in targets {
        // TODO(cgag): use WalkParallel?
        let walker = WalkBuilder::new(target).ignore(use_ignore)
                                             .git_ignore(use_ignore)
                                             .git_exclude(use_ignore)
                                             .hidden(ignore_hidden)
                                             .build();
        let files = walker
            .into_iter()
            .filter_map(|entry| entry.ok())
            .filter(|entry| entry.file_type().expect("no filetype").is_file())
            .map(|entry| String::from(entry.path().to_str().unwrap()))
            .filter(|path| match include_regex {
                None => true,
                Some(ref include) => include.is_match(path),
            })
            .filter(|path| match exclude_regex {
                None => true,
                Some(ref exclude) => !exclude.is_match(path),
            });

        for path in files {
            workq.push(Work::File(path));
        }
    }

    for _ in 0..workers.len() {
        workq.push(Work::Quit);
    }

    let mut filecounts: Vec<FileCount> = Vec::new();
    for worker in workers {
        filecounts.extend(worker.join().unwrap().iter().cloned())
    }

    // TODO(cgag): use insecure hashmaps or something
    let mut by_lang: HashMap<Lang, Vec<FileCount>> = HashMap::new();
    for fc in filecounts {
        match by_lang.entry(fc.lang) {
            Entry::Occupied(mut elem) => elem.get_mut().push(fc),
            Entry::Vacant(elem) => {
                elem.insert(vec![fc]);
            }
        };
    }

    let linesep = str_repeat("-", 80);

    if by_file {
        // print breakdown for each individual file
        println!("{}", linesep);
        println!(" {0: <17} {1: >8} {2: >12} {3: >12} {4: >12} {5: >12}",
                 "Language",
                 "Files",
                 "Lines",
                 "Blank",
                 "Comment",
                 "Code");
        println!("{}", linesep);

        // TODO(cgag): do the summing first, so we can do additional sorting
        // by totals.
        for (lang, mut filecounts) in by_lang {
            let mut total = Count::default();
            for fc in &filecounts {
                total.merge(&fc.count);
            }

            println!("{}", linesep);
            println!(" {0: <17} {1: >8} {2: >12} {3: >12} {4: >12} {5: >12}",
                     lang,
                     filecounts.len(),
                     total.lines,
                     total.blank,
                     total.comment,
                     total.code);

            match sort {
                "code" => filecounts.sort_by(|fc1, fc2| fc2.count.code.cmp(&fc1.count.code)),
                "comment" => {
                    filecounts.sort_by(|fc1, fc2| fc2.count.comment.cmp(&fc1.count.comment))
                }
                "blank" => filecounts.sort_by(|fc1, fc2| fc2.count.blank.cmp(&fc1.count.blank)),
                "lines" => filecounts.sort_by(|fc1, fc2| fc2.count.lines.cmp(&fc1.count.lines)),
                // No sorting by language or files here. Need to do it at a higher level.
                _ => (),
            }

            println!("{}", linesep);
            for fc in filecounts {
                println!("|{0: <25} {1: >12} {2: >12} {3: >12} {4: >12}",
                         last_n_chars(&fc.path, 25),
                         fc.count.lines,
                         fc.count.blank,
                         fc.count.comment,
                         fc.count.code);
            }
        }
    } else {
        // print summary by language
        let mut lang_totals: HashMap<&Lang, LangTotal> = HashMap::new();
        for (lang, filecounts) in &by_lang {
            let mut lang_total = Count::default();
            for fc in filecounts {
                lang_total.merge(&fc.count);
            }
            lang_totals.insert(lang,
                               LangTotal {
                                   files: filecounts.len() as u32,
                                   count: lang_total,
                               });
        }

        let mut totals_by_lang = lang_totals.iter().collect::<Vec<(&&Lang, &LangTotal)>>();
        match sort {
            "language" => totals_by_lang.sort_by(|&(l1, _), &(l2, _)| l1.to_s().cmp(l2.to_s())),
            "files" => totals_by_lang.sort_by(|&(_, c1), &(_, c2)| c2.files.cmp(&c1.files)),
            "code" => {
                totals_by_lang.sort_by(|&(_, c1), &(_, c2)| c2.count.code.cmp(&c1.count.code))
            }
            "comment" => {
                totals_by_lang.sort_by(|&(_, c1), &(_, c2)| c2.count.comment.cmp(&c1.count.comment))
            }
            "blank" => {
                totals_by_lang.sort_by(|&(_, c1), &(_, c2)| c2.count.blank.cmp(&c1.count.blank))
            }
            "lines" => {
                totals_by_lang.sort_by(|&(_, c1), &(_, c2)| c2.count.lines.cmp(&c1.count.lines))
            }
            _ => {
                println!("invalid sort option {}, sorting by code", sort);
                totals_by_lang.sort_by(|&(_, c1), &(_, c2)| c2.count.code.cmp(&c1.count.code))
            }
        }

        print_totals_by_lang(&linesep, &totals_by_lang);
    }

}

fn last_n_chars(s: &str, n: usize) -> String {
    if s.len() <= n {
        return String::from(s);
    }
    s.chars().skip(s.len() - n).collect::<String>()
}

fn str_repeat(s: &str, n: usize) -> String {
    std::iter::repeat(s).take(n).collect::<Vec<_>>().join("")
}

fn print_totals_by_lang(linesep: &str, totals_by_lang: &[(&&Lang, &LangTotal)]) {
    println!("{}", linesep);
    println!(" {0: <17} {1: >8} {2: >12} {3: >12} {4: >12} {5: >12}",
             "Language",
             "Files",
             "Lines",
             "Blank",
             "Comment",
             "Code");
    println!("{}", linesep);

    for &(lang, total) in totals_by_lang {
        println!(" {0: <17} {1: >8} {2: >12} {3: >12} {4: >12} {5: >12}",
                 lang,
                 total.files,
                 total.count.lines,
                 total.count.blank,
                 total.count.comment,
                 total.count.code);
    }

    let mut totals = LangTotal {
        files: 0,
        count: Count::default(),
    };
    for &(_, total) in totals_by_lang {
        totals.files += total.files;
        totals.count.code += total.count.code;
        totals.count.blank += total.count.blank;
        totals.count.comment += total.count.comment;
        totals.count.lines += total.count.lines;
    }

    println!("{}", linesep);
    println!(" {0: <17} {1: >8} {2: >12} {3: >12} {4: >12} {5: >12}",
             "Total",
             totals.files,
             totals.count.lines,
             totals.count.blank,
             totals.count.comment,
             totals.count.code);
    println!("{}", linesep);
}
<|endoftext|>


<commit_before>extern crate combine;
extern crate nix; // As in *nix
extern crate rand;
extern crate signal;

mod commander;

// Reference: https://github.com/NixOS/nixpkgs/blob/b07051ce6c3ba3039c89b6755da279002b0c3ace/nixos/modules/virtualisation/nixos-container.pl

use nix::sys::signal::{kill, SIGTERM, SIGINT};
use std::io;
use std::io::Write;
use std::process::Command;

#[derive(Debug)]
enum Error {
    Usage(String),
    NonZero(String, i32),
    ControlError(String),
    IoError(io::Error),
    NixError(nix::Error),
    Interrupted(commander::Signal)
}

use Error::*;

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

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

impl From<commander::NonZero> for Error {
    fn from(e: commander::NonZero) -> Self { match e { commander::NonZero::NonZero(s, o) => NonZero(s, o), commander::NonZero::Interrupted(n) => Interrupted(n) } }
}

type Result<T> = std::result::Result<T, Error>;

// TODO: why can't this be inferred away
fn run<T>(cmd: &mut Command) -> Result<T> where T: commander::CommandOut { commander::run(cmd) }

fn probably_unique_name(length: usize) -> String {
    use rand::distributions::{IndependentSample, Range};
    let between = Range::new(0, 36);
    let mut rng = rand::thread_rng();
    unsafe {
        std::str::from_utf8_unchecked(&(0..length).map(|_| between.ind_sample(&mut rng)).map(|v| if v < 10 { v + 48 } else { v + 87 }).collect::<Vec<u8>>()).to_string()
    }
}

// Result is of the form "10.233.x"
fn unused_ip_block() -> Result<String> {
    // Can bail early anywhere here because we've got the lock
    let mut used = std::collections::HashSet::new();
    for entry in try!(std::fs::read_dir("/etc/containers")) {
        let entry = try!(entry);
        if try!(entry.file_type()).is_file() {
            let mut str = String::new();
            use std::io::Read;
            try!(try!(std::fs::File::open(entry.path())).read_to_string(&mut str));
            for line in str.lines() {
                use combine::{string, many1, choice, digit, token, eof, Parser, ParserExt};
                let ip_component = || { string("10.233.").with(many1::<String, _>(digit())).skip(token('.')).skip(many1::<String, _>(digit())) };
                if let Ok((c, _)) = choice([string("HOST_ADDRESS="), string("LOCAL_ADDRESS=")]).with(ip_component()).skip(eof()).parse(line) {
                    used.insert(c);
                }
            }
        }
    }
    (0..255).map(|p| p.to_string()).find(|s| !used.contains(s)).map(|s| format!("10.233.{}", s)).ok_or(ControlError("Out of IP addresses".to_string()))
}

fn mkpath(mode : u32, path : &str) -> Result<()> {
    use std::os::unix::fs::DirBuilderExt;
    match std::fs::DirBuilder::new().mode(mode).create(path) {
        Err(e) => if let std::io::ErrorKind::AlreadyExists = e.kind() { Ok(()) } else { Err(IoError(e)) },
        Ok(()) => Ok(())
    }
}

fn system_init() -> Result<()> {
    try!(mkpath(0o0755, "/etc/containers"));
    try!(mkpath(0o0700, "/var/lib/containers"));
    Ok(())
}

fn profile_dir(name : &str) -> String {
    format!("/nix/var/nix/profiles/per-container/{}", name)
}

fn container_root(name : &str) -> String {
    format!("/var/lib/containers/{}", name)
}

fn conf_file(name : &str) -> String {
    format!("/etc/containers/{}.conf", name)
}

fn create(config_file : &str) -> Result<String> {
    fn init_container_unsynced(name : &str) -> Result<()> {
        let ip_block = try!(unused_ip_block());
        // Use create_new to avoid clobbering existing container
        let mut f = try!(std::fs::OpenOptions::new().create_new(true).write(true).open(conf_file(name)));
        write!(&mut f, "PRIVATE_NETWORK=1\nHOST_ADDRESS={0}.1\nLOCAL_ADDRESS={0}.2\nAUTO_START=0\n", ip_block).map_err(IoError)
    }

    fn populate_container(name : &str, config_file : &str) -> Result<()> {
        use std::os::unix::fs::DirBuilderExt;
        // The per-container directory is restricted to prevent users on
        // the host from messing with guest users who happen to have the
        // same uid.
        try!(mkpath(0o0700, "/nix/var/nix/profiles/per-container"));
        let profile_dir = profile_dir(name);
        let mut dir_builder = std::fs::DirBuilder::new();
        try!(dir_builder.mode(0o0755).create(&profile_dir));

        let container_root = container_root(name);
        dir_builder.recursive(true);
        try!(dir_builder.mode(0o0755).create(format!("{}/etc/nixos", container_root)));

        // Write configuration.nix into the container fs
        let configuration_nix = format!("{}/etc/nixos/configuration.nix", container_root);
        {
            let mut f = try!(std::fs::File::create(format!("{}/etc/nixos/configuration.nix", container_root)));
            try!(write!(&mut f, "
    {{ config, lib, pkgs, ... }}:
    with lib;

    {{
        boot.isContainer = true;
        networking.hostName = mkDefault \"{}\";
        networking.useDHCP = false;
        imports = [ {} ];
    }}", name, config_file));
            // TODO: guaranteed sync at end of block?
        }

        // Build the container config
        run(Command::new("nix-env")
            .arg("-p").arg(format!("{}/system", profile_dir))
            .arg("-I").arg(format!("nixos-config={}", configuration_nix))
            .arg("-f").arg("<nixpkgs/nixos>")
            .arg("--set")
            .arg("-A").arg("system")
            .arg("--show-trace"))
    }

    let name = format!("fr{}", probably_unique_name(9));

    // Acquire lock
    use std::os::unix::io::IntoRawFd;
    let lock_fd = try!(std::fs::OpenOptions::new().create(true).append(true).open("/run/lock/nixos-container")).into_raw_fd();
    try!(nix::fcntl::flock(lock_fd, nix::fcntl::FlockArg::LockExclusive));

    let res = init_container_unsynced(&name);

    // Drop lock
    unsafe {
        let _ = nix::libc::close(lock_fd);  // If closing fails, ¯\_(ツ)_/¯
    }

    match res.and_then(|()| populate_container(&name, config_file)) {
        Ok(()) => Ok(name),
        // Don't leave a half-baked container behind
        Err(e) => { let _ = destroy(&name); Err(e) }
    }
}

fn run_test(container_name : &str, args : &[String]) -> Result<i32> {
    try!(run(Command::new("systemctl").arg("start").arg(format!("container@{}", container_name))));
    let leader_pid = try!(run(Command::new("machinectl").arg("show").arg(container_name).arg("-p").arg("Leader")).and_then(|o : String| {
        use combine::{string, many1, digit, spaces, eof, Parser, ParserExt};
        string("Leader=").with(many1::<String, _>(digit())).skip(spaces()).skip(eof()).parse(o.as_str())
            .map(|(ds, _)| ds)
            .map_err(|_| ControlError(format!("Bad machinectl output {}", o)))
    }));
    let su_cmd = args.iter().map(|s| format!("'{}'", s.replace("'", "'\\''"))).collect::<Vec<String>>().join(" ");
    match run(Command::new("nsenter")
            .arg("-t").arg(leader_pid)
            .arg("-m").arg("-u").arg("-i").arg("-n").arg("-p")
            .arg("--").arg("/var/setuid-wrappers/su").arg("root").arg("-l").arg("-c").arg(format!("exec {}", su_cmd))) {
        Ok(()) => Ok(0),
        Err(NonZero(_, code)) => Ok(code),
        Err(e) => Err(e)
    }
}

// Remove a directory while recursively unmounting all mounted filesystems within
// that directory and unmounting/removing that directory afterwards as well.
//
// Specified path shouldn't be a mountpoint.
fn safe_remove_tree(path : &str) -> Result<()> {
    if !std::path::Path::new(path).is_dir() {
        return Ok(());
    }
    try!(run(Command::new("find").arg(path)
        .arg("-mindepth").arg("1")
        .arg("-xdev")
        .arg("(")
            .arg("-type").arg("d")
            .arg("-exec").arg("mountpoint").arg("-q").arg("{}").arg(";")
        .arg(")")
        .arg("-exec").arg("umount").arg("-fR").arg("{}").arg("+")));
    try!(run(Command::new("rm").arg("--one-file-system").arg("-rf").arg(path)));

    Ok(())
}

fn stop(container_name : &str) -> Result<()> {
    match run(Command::new("systemctl").arg("stop").arg(format!("container@{}", container_name))) {
        Ok(()) => Ok(()),
        Err(Interrupted(_)) => {
            return run(Command::new("systemctl").arg("kill").arg(format!("container@{}", container_name)))
        },
        Err(e) => Err(e)
    }
}

fn destroy(container_name : &str) -> Result<()> {
    fn log_err(e: Error) {
        println!("Error while destroying container: {:?}", e);
    }

    safe_remove_tree(&profile_dir(container_name)).unwrap_or_else(log_err);
    safe_remove_tree(&format!("/nix/var/nix/gcroots/per-container/{}", container_name)).unwrap_or_else(log_err);
    safe_remove_tree(&container_root(container_name)).unwrap_or_else(log_err);

    std::fs::remove_file(conf_file(container_name)).or_else(|e| match e.kind() {
        std::io::ErrorKind::NotFound => Ok(()),
        _ => Err(IoError(e))
    })
}

fn go() -> Result<i32> {
    // Skip program name
    let mut args = std::env::args().skip(1).peekable();

    let no_destroy = args.peek().map_or(false, |arg| arg == "--no_destroy");
    if no_destroy { args.next(); }

    let config_file = try!(args.next().ok_or(Error::Usage("Missing config file argument".to_string())));

    let test_args : Vec<String> = args.collect();
    if test_args.len() == 0 {
        return Err(Usage("Missing test args".to_string()));
    }

    let config_file = {
        let p = try!(std::fs::canonicalize(config_file).map_err(|e| Usage(format!("For config file: {}", e))));
        try!(p.to_str().ok_or(Usage("Config file path is invalid UTF-8".to_string()))).to_string()
    };

    try!(system_init());
    // We suppress signals in this process and allow the commander module to deal with child processes
    let trap = signal::trap::Trap::trap(&[SIGTERM, SIGINT]);
    let container_name = try!(create(&config_file));
    let res = match trap.wait(std::time::Instant::now()) {
        None => run_test(&container_name, &test_args),
        Some(n) => Err(Interrupted(n))
    };

    if no_destroy {
        try!(stop(&container_name));
    } else {
        try!(destroy(&container_name));
    }

    res
}

fn main() {
    std::process::exit(go().unwrap_or_else(|e| match e {
        NonZero(cmd, code) => { writeln!(io::stderr(), "Command {:?} failed with code {:?}", cmd, code).unwrap(); 2 },
        Usage(msg) => { writeln!(io::stderr(), "Usage error: {}", msg).unwrap(); 1 },
        IoError(err) => { writeln!(io::stderr(), "{}", err).unwrap(); 2 },
        NixError(err) => { writeln!(io::stderr(), "{}", err).unwrap(); 2 },
        ControlError(msg) => { writeln!(io::stderr(), "{}", msg).unwrap(); 2 },
        Interrupted(SIGINT) => { kill(nix::unistd::getpid(), SIGINT).unwrap(); loop {} }
        Interrupted(n) => 128 + n
    }));
}
<commit_msg>Use --no-destroy instead of -no_destroy<commit_after>extern crate combine;
extern crate nix; // As in *nix
extern crate rand;
extern crate signal;

mod commander;

// Reference: https://github.com/NixOS/nixpkgs/blob/b07051ce6c3ba3039c89b6755da279002b0c3ace/nixos/modules/virtualisation/nixos-container.pl

use nix::sys::signal::{kill, SIGTERM, SIGINT};
use std::io;
use std::io::Write;
use std::process::Command;

#[derive(Debug)]
enum Error {
    Usage(String),
    NonZero(String, i32),
    ControlError(String),
    IoError(io::Error),
    NixError(nix::Error),
    Interrupted(commander::Signal)
}

use Error::*;

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

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

impl From<commander::NonZero> for Error {
    fn from(e: commander::NonZero) -> Self { match e { commander::NonZero::NonZero(s, o) => NonZero(s, o), commander::NonZero::Interrupted(n) => Interrupted(n) } }
}

type Result<T> = std::result::Result<T, Error>;

// TODO: why can't this be inferred away
fn run<T>(cmd: &mut Command) -> Result<T> where T: commander::CommandOut { commander::run(cmd) }

fn probably_unique_name(length: usize) -> String {
    use rand::distributions::{IndependentSample, Range};
    let between = Range::new(0, 36);
    let mut rng = rand::thread_rng();
    unsafe {
        std::str::from_utf8_unchecked(&(0..length).map(|_| between.ind_sample(&mut rng)).map(|v| if v < 10 { v + 48 } else { v + 87 }).collect::<Vec<u8>>()).to_string()
    }
}

// Result is of the form "10.233.x"
fn unused_ip_block() -> Result<String> {
    // Can bail early anywhere here because we've got the lock
    let mut used = std::collections::HashSet::new();
    for entry in try!(std::fs::read_dir("/etc/containers")) {
        let entry = try!(entry);
        if try!(entry.file_type()).is_file() {
            let mut str = String::new();
            use std::io::Read;
            try!(try!(std::fs::File::open(entry.path())).read_to_string(&mut str));
            for line in str.lines() {
                use combine::{string, many1, choice, digit, token, eof, Parser, ParserExt};
                let ip_component = || { string("10.233.").with(many1::<String, _>(digit())).skip(token('.')).skip(many1::<String, _>(digit())) };
                if let Ok((c, _)) = choice([string("HOST_ADDRESS="), string("LOCAL_ADDRESS=")]).with(ip_component()).skip(eof()).parse(line) {
                    used.insert(c);
                }
            }
        }
    }
    (0..255).map(|p| p.to_string()).find(|s| !used.contains(s)).map(|s| format!("10.233.{}", s)).ok_or(ControlError("Out of IP addresses".to_string()))
}

fn mkpath(mode : u32, path : &str) -> Result<()> {
    use std::os::unix::fs::DirBuilderExt;
    match std::fs::DirBuilder::new().mode(mode).create(path) {
        Err(e) => if let std::io::ErrorKind::AlreadyExists = e.kind() { Ok(()) } else { Err(IoError(e)) },
        Ok(()) => Ok(())
    }
}

fn system_init() -> Result<()> {
    try!(mkpath(0o0755, "/etc/containers"));
    try!(mkpath(0o0700, "/var/lib/containers"));
    Ok(())
}

fn profile_dir(name : &str) -> String {
    format!("/nix/var/nix/profiles/per-container/{}", name)
}

fn container_root(name : &str) -> String {
    format!("/var/lib/containers/{}", name)
}

fn conf_file(name : &str) -> String {
    format!("/etc/containers/{}.conf", name)
}

fn create(config_file : &str) -> Result<String> {
    fn init_container_unsynced(name : &str) -> Result<()> {
        let ip_block = try!(unused_ip_block());
        // Use create_new to avoid clobbering existing container
        let mut f = try!(std::fs::OpenOptions::new().create_new(true).write(true).open(conf_file(name)));
        write!(&mut f, "PRIVATE_NETWORK=1\nHOST_ADDRESS={0}.1\nLOCAL_ADDRESS={0}.2\nAUTO_START=0\n", ip_block).map_err(IoError)
    }

    fn populate_container(name : &str, config_file : &str) -> Result<()> {
        use std::os::unix::fs::DirBuilderExt;
        // The per-container directory is restricted to prevent users on
        // the host from messing with guest users who happen to have the
        // same uid.
        try!(mkpath(0o0700, "/nix/var/nix/profiles/per-container"));
        let profile_dir = profile_dir(name);
        let mut dir_builder = std::fs::DirBuilder::new();
        try!(dir_builder.mode(0o0755).create(&profile_dir));

        let container_root = container_root(name);
        dir_builder.recursive(true);
        try!(dir_builder.mode(0o0755).create(format!("{}/etc/nixos", container_root)));

        // Write configuration.nix into the container fs
        let configuration_nix = format!("{}/etc/nixos/configuration.nix", container_root);
        {
            let mut f = try!(std::fs::File::create(format!("{}/etc/nixos/configuration.nix", container_root)));
            try!(write!(&mut f, "
    {{ config, lib, pkgs, ... }}:
    with lib;

    {{
        boot.isContainer = true;
        networking.hostName = mkDefault \"{}\";
        networking.useDHCP = false;
        imports = [ {} ];
    }}", name, config_file));
            // TODO: guaranteed sync at end of block?
        }

        // Build the container config
        run(Command::new("nix-env")
            .arg("-p").arg(format!("{}/system", profile_dir))
            .arg("-I").arg(format!("nixos-config={}", configuration_nix))
            .arg("-f").arg("<nixpkgs/nixos>")
            .arg("--set")
            .arg("-A").arg("system")
            .arg("--show-trace"))
    }

    let name = format!("fr{}", probably_unique_name(9));

    // Acquire lock
    use std::os::unix::io::IntoRawFd;
    let lock_fd = try!(std::fs::OpenOptions::new().create(true).append(true).open("/run/lock/nixos-container")).into_raw_fd();
    try!(nix::fcntl::flock(lock_fd, nix::fcntl::FlockArg::LockExclusive));

    let res = init_container_unsynced(&name);

    // Drop lock
    unsafe {
        let _ = nix::libc::close(lock_fd);  // If closing fails, ¯\_(ツ)_/¯
    }

    match res.and_then(|()| populate_container(&name, config_file)) {
        Ok(()) => Ok(name),
        // Don't leave a half-baked container behind
        Err(e) => { let _ = destroy(&name); Err(e) }
    }
}

fn run_test(container_name : &str, args : &[String]) -> Result<i32> {
    try!(run(Command::new("systemctl").arg("start").arg(format!("container@{}", container_name))));
    let leader_pid = try!(run(Command::new("machinectl").arg("show").arg(container_name).arg("-p").arg("Leader")).and_then(|o : String| {
        use combine::{string, many1, digit, spaces, eof, Parser, ParserExt};
        string("Leader=").with(many1::<String, _>(digit())).skip(spaces()).skip(eof()).parse(o.as_str())
            .map(|(ds, _)| ds)
            .map_err(|_| ControlError(format!("Bad machinectl output {}", o)))
    }));
    let su_cmd = args.iter().map(|s| format!("'{}'", s.replace("'", "'\\''"))).collect::<Vec<String>>().join(" ");
    match run(Command::new("nsenter")
            .arg("-t").arg(leader_pid)
            .arg("-m").arg("-u").arg("-i").arg("-n").arg("-p")
            .arg("--").arg("/var/setuid-wrappers/su").arg("root").arg("-l").arg("-c").arg(format!("exec {}", su_cmd))) {
        Ok(()) => Ok(0),
        Err(NonZero(_, code)) => Ok(code),
        Err(e) => Err(e)
    }
}

// Remove a directory while recursively unmounting all mounted filesystems within
// that directory and unmounting/removing that directory afterwards as well.
//
// Specified path shouldn't be a mountpoint.
fn safe_remove_tree(path : &str) -> Result<()> {
    if !std::path::Path::new(path).is_dir() {
        return Ok(());
    }
    try!(run(Command::new("find").arg(path)
        .arg("-mindepth").arg("1")
        .arg("-xdev")
        .arg("(")
            .arg("-type").arg("d")
            .arg("-exec").arg("mountpoint").arg("-q").arg("{}").arg(";")
        .arg(")")
        .arg("-exec").arg("umount").arg("-fR").arg("{}").arg("+")));
    try!(run(Command::new("rm").arg("--one-file-system").arg("-rf").arg(path)));

    Ok(())
}

fn stop(container_name : &str) -> Result<()> {
    match run(Command::new("systemctl").arg("stop").arg(format!("container@{}", container_name))) {
        Ok(()) => Ok(()),
        Err(Interrupted(_)) => {
            return run(Command::new("systemctl").arg("kill").arg(format!("container@{}", container_name)))
        },
        Err(e) => Err(e)
    }
}

fn destroy(container_name : &str) -> Result<()> {
    fn log_err(e: Error) {
        println!("Error while destroying container: {:?}", e);
    }

    safe_remove_tree(&profile_dir(container_name)).unwrap_or_else(log_err);
    safe_remove_tree(&format!("/nix/var/nix/gcroots/per-container/{}", container_name)).unwrap_or_else(log_err);
    safe_remove_tree(&container_root(container_name)).unwrap_or_else(log_err);

    std::fs::remove_file(conf_file(container_name)).or_else(|e| match e.kind() {
        std::io::ErrorKind::NotFound => Ok(()),
        _ => Err(IoError(e))
    })
}

fn go() -> Result<i32> {
    // Skip program name
    let mut args = std::env::args().skip(1).peekable();

    let no_destroy = args.peek().map_or(false, |arg| arg == "--no-destroy");
    if no_destroy { args.next(); }

    let config_file = try!(args.next().ok_or(Error::Usage("Missing config file argument".to_string())));

    let test_args : Vec<String> = args.collect();
    if test_args.len() == 0 {
        return Err(Usage("Missing test args".to_string()));
    }

    let config_file = {
        let p = try!(std::fs::canonicalize(config_file).map_err(|e| Usage(format!("For config file: {}", e))));
        try!(p.to_str().ok_or(Usage("Config file path is invalid UTF-8".to_string()))).to_string()
    };

    try!(system_init());
    // We suppress signals in this process and allow the commander module to deal with child processes
    let trap = signal::trap::Trap::trap(&[SIGTERM, SIGINT]);
    let container_name = try!(create(&config_file));
    let res = match trap.wait(std::time::Instant::now()) {
        None => run_test(&container_name, &test_args),
        Some(n) => Err(Interrupted(n))
    };

    if no_destroy {
        try!(stop(&container_name));
    } else {
        try!(destroy(&container_name));
    }

    res
}

fn main() {
    std::process::exit(go().unwrap_or_else(|e| match e {
        NonZero(cmd, code) => { writeln!(io::stderr(), "Command {:?} failed with code {:?}", cmd, code).unwrap(); 2 },
        Usage(msg) => { writeln!(io::stderr(), "Usage error: {}", msg).unwrap(); 1 },
        IoError(err) => { writeln!(io::stderr(), "{}", err).unwrap(); 2 },
        NixError(err) => { writeln!(io::stderr(), "{}", err).unwrap(); 2 },
        ControlError(msg) => { writeln!(io::stderr(), "{}", msg).unwrap(); 2 },
        Interrupted(SIGINT) => { kill(nix::unistd::getpid(), SIGINT).unwrap(); loop {} }
        Interrupted(n) => 128 + n
    }));
}
<|endoftext|>


<commit_before>extern crate getopts;

mod diff;

use std::io;
use std::io::Write;
use std::env;
use getopts::Options;

static CurrentAddress: Address = Address::Nth(0);

struct Range {
    start: Box<Address>,
    end: Box<Address>,
}

enum Address {
    Current,
    Last,
    Nth(i32),
    Previous,
    NthPrevious(i32),
    Next,
    NthNext(i32),
    All,
    Range(Range),
    CurrentToEnd,
    RENext(String),
    REPrevious(String),
    Mark(String),
}

trait Operation {
    fn apply(Address, &diff::Diff) -> Result<diff::Diff, String>;
}

struct nop;

impl Operation for nop {
    fn apply(addr: Address, diff: &diff::Diff) -> Result<diff::Diff, String> {
        panic!("Command didn't get set. This shouldn't have happened");
        Err("Command didn't get set. This shouldn't have happened".to_string())
    }
}

struct append_after;

impl Operation for append_after {
    fn apply(addr: Address, diff: &diff::Diff) -> Result<diff::Diff, String> {
        let line_num: i32 = match addr {
            Address::Nth(ln) => ln,
            _ => 0,
        };
        input()
            .map_err(|err| err.to_string())
            .map(|lines| {
                diff.add_lines(line_num, lines)
            })
    }
}

struct undo;

impl Operation for undo {
    fn apply(addr: Address, diff: &diff::Diff) -> Result<diff::Diff, String> {
        Ok(diff.undo())
    }
}

struct quit;

impl Operation for quit {
    fn apply(addr: Address, diff: &diff::Diff) -> Result<diff::Diff, String> {
        std::process::exit(0);
        Err("Should never get here".to_string())
    }
}

struct Command {
    address: Address,
    operation: Operation,
}

fn input() -> Result<Vec<String>, io::Error> {
    let mut lines: Vec<String> = Vec::new();
    loop {
        let mut line = String::new();
        match io::stdin().read_line(&mut line) {
            Err(error) => return Err(error),
            Ok(_) => {},
        }
        match line.trim() {
            "." => {
                break
            },
            _ => {
                lines.push(line)
            }
        }
    }
    Ok(lines)
}

fn parse_cmd(line: &String) -> Result<Command, String> {
    let mut command = Command {
        address: Address::Current,
        command: nop,
    };
    command.command = match line.trim() {
        "a" => append_after,
        "q" => quit,
        "u" => undo,
        _ => {
            return Err("Uknown command".to_string())
        },
    };
    Ok(command)
}

fn command(line: &String, line_num: i32, diff: &diff::Diff) -> Result<diff::Diff, String> {
    parse_cmd(line)
        .map_err(|err| err.to_string())
        .and_then(|cmd| {
            cmd.operation.apply(cmd.address, diff)
        })
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let mut opts = Options::new();
    opts.optopt("p", "prompt", "prompt to show while in command mode", "PROMPT");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => panic!(f.to_string())
    };

    let prompt = matches.opt_str("p")
        .unwrap_or_else(|| "".to_string());
    let mut diff = diff::Diff::new();
    let mut line_num: i32 = 1;
    loop {
        print!("{}", prompt);
        match io::stdout().flush() {
            Ok(_) => (),
            Err(err) => {
                println!("Error flushing stdout: {}", err);
                continue
            }
        }
        let mut line = String::new();
        match io::stdin().read_line(&mut line) {
            Ok(_) => (),
            Err(_) => {
                println!("?");
                continue
            },
        }
        match command(&line, line_num, &mut diff) {
            Ok(d) => {
                diff = d;
            },
            Err(err) => println!("Error: {}", err),

        }
    }
}

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

}
<commit_msg>Use fn pointers instead of a trait<commit_after>extern crate getopts;

mod diff;

use std::io;
use std::io::Write;
use std::env;
use getopts::Options;

static CurrentAddress: Address = Address::Nth(0);

struct Range {
    start: Box<Address>,
    end: Box<Address>,
}

enum Address {
    Current,
    Last,
    Nth(i32),
    Previous,
    NthPrevious(i32),
    Next,
    NthNext(i32),
    All,
    Range(Range),
    CurrentToEnd,
    RENext(String),
    REPrevious(String),
    Mark(String),
}

fn nop(addr: Address, diff: &diff::Diff) -> Result<diff::Diff, String> {
    panic!("Command didn't get set. This shouldn't have happened");
    Err("Command didn't get set. This shouldn't have happened".to_string())
}

fn append_after(addr: Address, diff: &diff::Diff) -> Result<diff::Diff, String> {
    let line_num: i32 = match addr {
        Address::Nth(ln) => ln,
        _ => 0,
    };
    input()
        .map_err(|err| err.to_string())
        .map(|lines| {
            diff.add_lines(line_num, lines)
        })
}

fn undo(addr: Address, diff: &diff::Diff) -> Result<diff::Diff, String> {
    Ok(diff.undo())
}

fn quit(addr: Address, diff: &diff::Diff) -> Result<diff::Diff, String> {
    std::process::exit(0);
    Err("Should never get here".to_string())
}

struct Command {
    address: Address,
    operation: fn(Address, &diff::Diff) -> Result<diff::Diff, String>,
}

fn input() -> Result<Vec<String>, io::Error> {
    let mut lines: Vec<String> = Vec::new();
    loop {
        let mut line = String::new();
        match io::stdin().read_line(&mut line) {
            Err(error) => return Err(error),
            Ok(_) => {},
        }
        match line.trim() {
            "." => {
                break
            },
            _ => {
                lines.push(line)
            }
        }
    }
    Ok(lines)
}

fn parse_cmd(line: &String) -> Result<Command, String> {
    let mut command = Command {
        address: Address::Current,
        operation: nop,
    };
    command.operation = match line.trim() {
        "a" => append_after,
        "q" => quit,
        "u" => undo,
        _ => {
            return Err("Uknown command".to_string())
        },
    };
    Ok(command)
}

fn command(line: &String, line_num: i32, diff: &diff::Diff) -> Result<diff::Diff, String> {
    parse_cmd(line)
        .map_err(|err| err.to_string())
        .and_then(|cmd| {
            (cmd.operation)(cmd.address, diff)
        })
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let mut opts = Options::new();
    opts.optopt("p", "prompt", "prompt to show while in command mode", "PROMPT");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => panic!(f.to_string())
    };

    let prompt = matches.opt_str("p")
        .unwrap_or_else(|| "".to_string());
    let mut diff = diff::Diff::new();
    let mut line_num: i32 = 1;
    loop {
        print!("{}", prompt);
        match io::stdout().flush() {
            Ok(_) => (),
            Err(err) => {
                println!("Error flushing stdout: {}", err);
                continue
            }
        }
        let mut line = String::new();
        match io::stdin().read_line(&mut line) {
            Ok(_) => (),
            Err(_) => {
                println!("?");
                continue
            },
        }
        match command(&line, line_num, &mut diff) {
            Ok(d) => {
                diff = d;
            },
            Err(err) => println!("Error: {}", err),

        }
    }
}

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

}
<|endoftext|>