<commit_before>extern crate byteorder;
extern crate regex;
extern crate flate2;
extern crate getopts;
extern crate tempdir;

use std::fs::File;
use std::io;
use std::io::{BufReader, BufWriter};
use std::path::Path;
use std::env;
use std::process::exit;
use std::process::Command;
use std::error::Error;

use getopts::Options;

use tempdir::TempDir;

/// Create an InvalidData io::Error with the description being a
/// formatted string
macro_rules! io_error {
    ($fmtstr:tt) => { io_error!($fmtstr,) };
    ($fmtstr:tt, $( $args:expr ),* ) => {
        Err(::std::io::Error::new(::std::io::ErrorKind::InvalidData,
                                  format!($fmtstr, $( $args ),* )));
    };
}

pub mod data;
mod write;
mod read;
mod string_write;
mod string_read;
#[cfg(test)]
mod tests;

/// Trait used in place of unwrap for printing to stdout/stderr, since if that
/// errors the program should simply exit with no further output
trait SilentUnwrap<T> {
    fn silent_unwrap(self) -> T;
}
impl<T> SilentUnwrap<T> for io::Result<T> {
    fn silent_unwrap(self) -> T {
        match self {
            Ok(t) => t,
            Err(_) => std::process::exit(1),
        }
    }
}

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

    let mut opts = Options::new();
    opts.optflagopt("e", "edit", "edit a NBT file with your $EDITOR.
    If [FILE] is specified, then that file is edited in place, but specifying --input and/or --output will override the input/output.
    If no file is specified, default to read from --input and writing to --output.", "FILE");
    opts.optflagopt("p", "print", "print NBT file to text format. Adding an argument to this is the same as specifying --input", "FILE");
    opts.optflagopt("r", "reverse", "reverse a file in text format to NBT format. Adding an argument to this is the same as specifying --input", "FILE");
    opts.optopt("i",
                "input",
                "specify the input file, defaults to stdin",
                "FILE");
    opts.optopt("o",
                "output",
                "specify the output file, defaults to stdout",
                "FILE");
    opts.optflag("h", "help", "print the help menu");
    opts.optflag("", "version", "print program version");

    let matches = match opts.parse(&args[1..]) {
        Ok(x) => x,
        Err(e) => {
            error(&format!("Error parsing options: {:?}", e.description()))
        },
    };

    if matches.opt_present("h") {
        print_usage(opts);
        return;
    }

    if matches.opt_present("version") {
        println!("{} {} {}",
                 env!("CARGO_PKG_NAME"),
                 env!("CARGO_PKG_VERSION"),
                 /* See build.rs for the git-revision.txt file */
                 include!(concat!(env!("OUT_DIR"), "/git-revision.txt")));
        println!("https://github.com/C4K3/nbted");
        return;
    }

    let is_print: bool = matches.opt_present("print");
    let is_reverse: bool = matches.opt_present("reverse");
    let is_edit: bool = if matches.opt_present("edit") {
        true
    } else {
        /* If edit is not explicitly defined, it is the default action and is
         * selected if no other action is specified */
        (!(is_reverse || is_print))
    };

    /* Hopefully this is a simpler way of ensuring that only one action can be
     * taken than having a long logical expression */
    let mut action_count = 0;
    if is_print {
        action_count += 1;
    }
    if is_reverse {
        action_count += 1;
    }
    if is_edit {
        action_count += 1;
    }
    if action_count > 1 {
        error("You can only specify one action a time.");
    }

    /* Figure out the input file, by trying to read the arguments for all of
     * --input, --edit, --print and --reverse, prioritizing --input over the
     * other arguments, if none of the arguments are specified but there is a
     * free argument, use that, else we finally default to - (stdin) */
    let input = if let Some(x) = matches.opt_str("input") {
        x
    } else if let Some(x) = matches.opt_str("edit") {
        x
    } else if let Some(x) = matches.opt_str("print") {
        x
    } else if let Some(x) = matches.opt_str("reverse") {
        x
    } else if matches.free.len() == 1 {
        matches.free[0].clone()
    } else {
        /* stdin */
        "-".to_string()
    };

    let output = if let Some(x) = matches.opt_str("output") {
        x
    } else if let Some(x) = matches.opt_str("edit") {
        x
    } else if is_edit && matches.free.len() == 1 {
        /* Only want to default to the free argument if we're editing
         * (DO NOT WRITE BACK TO THE READ FILE UNLESS EDITING!) */
        matches.free[0].clone()
    } else {
        /* stdout */
        "-".to_string()
    };

    if matches.free.len() > 1 {
        error("nbted was given multiple arguments, but only supports editing one file at a time.");
    }

    if is_print {
        print(&input, &output);
    } else if is_reverse {
        reverse(&input, &output);
    } else if is_edit {
        edit(&input, &output);
    } else {
        eprintln!("Internal error: No action selected. (Please report this.)");
    }
}

/// When the user wants to edit a specific file in place
fn edit(input: &str, output: &str) {

    /* First we read the NBT data from the input */
    let nbt = if input == "-" {
        // let mut f = BufReader::new(io::stdin());
        let f = io::stdin();
        let mut f = f.lock();
        match read::read_file(&mut f) {
            Ok(x) => x,
            Err(_) => {
                error(&format!("Unable to parse {}, are you sure it's an NBT file?",
                              input))
            },
        }
    } else {
        let path: &Path = Path::new(input);
        let f = match File::open(path) {
            Ok(f) => f,
            Err(_) => error(&format!("Unable to open file {}", input)),
        };
        let mut f = BufReader::new(f);

        match read::read_file(&mut f) {
            Ok(x) => x,
            Err(_) => {
                error(&format!("Unable to parse {}, are you sure it's an NBT file?",
                              input))
            },
        }
    };

    /* Then we create a temporary file and write the NBT data in text format
     * to the temporary file */
    let tmpdir = match TempDir::new("nbted") {
        Ok(x) => x,
        Err(e) => {
            error(&format!("Unable to create temporary directory: {:?}",
                          e.description()))
        },
    };

    let tmp = match Path::new(input).file_name() {
        Some(x) => {
            let mut x = x.to_os_string();
            x.push(".txt");
            x
        },
        None => error(&format!("Error reading file name")),
    };
    let tmp_path = tmpdir.path().join(tmp);

    {
        let mut f = match File::create(&tmp_path) {
            Ok(x) => x,
            Err(e) => {
                error(&format!("Unable to create temporary file: {:?}",
                              e.description()))
            },
        };

        match string_write::write_file(&mut f, &nbt) {
            Ok(()) => (),
            Err(e) => {
                error(&format!("Unable to write temporary file: {:?}",
                              e.description()))
            },
        };

        match f.sync_all() {
            Ok(()) => (),
            Err(e) => {
                error(&format!("Unable to synchronize file: {:?}",
                              e.description()))
            },
        }
    }

    let new_nbt = {
        let mut new_nbt = open_editor(&tmp_path);

        while let Err(e) = new_nbt {
            eprintln!("Unable to parse edited file: {}.\n\
            Do you want to open the file for editing again? (y/N)",
                        e.description());

            let mut line = String::new();
            match io::stdin().read_line(&mut line) {
                Ok(_) => (),
                Err(e) => {
                    error(&format!("Error reading from stdin: {}\n\
                                   Nothing was changed.",
                                  e.description()))
                },
            }

            if line.trim() == "y" {
                new_nbt = open_editor(&tmp_path);
            } else {
                eprintln!("Exiting ... File is unchanged.");
                std::process::exit(1);
            }
        }

        new_nbt.expect("new_nbt was Error")
    };

    if nbt == new_nbt {
        eprintln!("No changes, will do nothing.");
        return;
    }

    /* And finally we write the edited nbt (new_nbt) into the output file */
    if output == "-" {
        let f = io::stdout();
        let mut f = f.lock();
        write::write_file(&mut f, &new_nbt).silent_unwrap();
    } else {
        let path: &Path = Path::new(output);
        let f = match File::create(&path) {
            Ok(x) => x,
            Err(e) => {
                error(&format!("Unable to write to output NBT file {}: {:?}.\n\
            Nothing was changed",
                              output,
                              e.description()))
            },
        };
        let mut f = BufWriter::new(f);

        match write::write_file(&mut f, &new_nbt) {
            Ok(()) => (),
            Err(e) => {
                error(&format!("Error writing NBT file {}: {:?}.\n\
            State of NBT file is unknown, consider restoring it from a backup.",
                              output,
                              e.description()))
            },
        }
    }

    eprintln!("File edited successfully.");
}

/// Open the user's $EDITOR on the temporary file, wait until the editor is
/// closed again, read the temporary file and attempt to parse it into NBT,
/// returning the result.
fn open_editor(tmp_path: &Path) -> io::Result<data::NBTFile> {
    let editor = match env::var("EDITOR") {
        Ok(x) => x,
        Err(_) => {
            match env::var("VISUAL") {
                Ok(x) => x,
                Err(_) => error("Unable to find $EDITOR"),
            }
        },
    };

    let mut cmd = Command::new(editor);
    cmd.arg(&tmp_path.as_os_str());
    let mut cmd = match cmd.spawn() {
        Ok(x) => x,
        Err(e) => {
            error(&format!("Error opening editor: {:?}", e.description()))
        },
    };

    match cmd.wait() {
        Ok(x) if x.success() => (),
        Err(e) => {
            error(&format!("Error executing editor: {:?}", e.description()))
        },
        _ => error("Editor did not exit correctly"),
    }

    /* Then we parse the text format in the temporary file into NBT */
    let mut f = match File::open(&tmp_path) {
        Ok(x) => x,
        Err(e) => {
            error(&format!("Unable to read temporary file: {:?}.\n\
        Nothing was changed",
                          e.description()))
        },
    };

    string_read::read_file(&mut f)
}

/// When the user wants to print an NBT file to text format
fn print(input: &str, output: &str) {
    /* First we read a NBTFile from the input */
    let nbt = if input == "-" {
        let f = io::stdin();
        let mut f = f.lock();
        match read::read_file(&mut f) {
            Ok(x) => x,
            Err(_) => {
                error(&format!("Unable to parse {}, are you sure it's an NBT file?",
                              input))
            },
        }
    } else {
        let path: &Path = Path::new(input);
        let f = match File::open(path) {
            Ok(f) => f,
            Err(_) => error(&format!("Unable to open file {}", input)),
        };
        let mut f = BufReader::new(f);

        match read::read_file(&mut f) {
            Ok(x) => x,
            Err(_) => {
                error(&format!("Unable to parse {}, are you sure it's an NBT file?",
                              input))
            },
        }
    };

    /* Then we write the NBTFile to the output in text format */
    if output == "-" {
        let f = io::stdout();
        let mut f = f.lock();
        string_write::write_file(&mut f, &nbt).silent_unwrap();
    } else {
        let path: &Path = Path::new(output);
        let f = match File::create(&path) {
            Ok(x) => x,
            Err(e) => {
                error(&format!("Unable to write to output NBT file {}: {:?}.\n\
            Nothing was changed",
                              output,
                              e.description()))
            },
        };
        let mut f = BufWriter::new(f);

        match string_write::write_file(&mut f, &nbt) {
            Ok(()) => (),
            Err(e) => {
                error(&format!("Error writing NBT file {}: {:?}.\n\
            State of NBT file is unknown, consider restoring it from a backup.",
                              output,
                              e.description()))
            },
        }
    }
}

/// When the user wants to convert a text format file into an NBT file
fn reverse(input: &str, output: &str) {
    /* First we read the input file in the text format */
    let path: &Path = Path::new(input);
    let mut f = match File::open(&path) {
        Ok(x) => x,
        Err(e) => {
            error(&format!("Unable to read text file {}: {:?}",
                          input,
                          e.description()))
        },
    };

    let nbt = match string_read::read_file(&mut f) {
        Ok(x) => x,
        Err(e) => {
            error(&format!("Unable to parse text file {}: {:?}",
                          input,
                          e.description()))
        },
    };

    /* Then we write the parsed NBT to the output file in NBT format */
    if output == "-" {
        let f = io::stdout();
        let mut f = f.lock();
        write::write_file(&mut f, &nbt).silent_unwrap();
    } else {
        let path: &Path = Path::new(output);
        let f = match File::create(&path) {
            Ok(x) => x,
            Err(e) => {
                error(&format!("Unable to write to output NBT file {}: {:?}.\n\
            Nothing was changed",
                              output,
                              e.description()))
            },
        };
        let mut f = BufWriter::new(f);

        match write::write_file(&mut f, &nbt) {
            Ok(()) => (),
            Err(e) => {
                error(&format!("Error writing NBT file {}: {:?}.\n\
            State of NBT file is unknown, consider restoring it from a backup.",
                              output,
                              e.description()))
            },
        }
    }
}

/// Print the given message and exit with status code 1
fn error(message: &str) -> ! {
    eprintln!("Error: {}", message);
    eprintln!("Run with --help for help, or read the manpage.");
    exit(1);
}

fn print_usage(opts: Options) {
    let brief = "Usage: nbted [options] FILE";
    print!("{}", opts.usage(&brief));
    println!("\nFor detailed usage information, read the manpage.");
}
<commit_msg>Specify the default behavior in --help<commit_after>extern crate byteorder;
extern crate regex;
extern crate flate2;
extern crate getopts;
extern crate tempdir;

use std::fs::File;
use std::io;
use std::io::{BufReader, BufWriter};
use std::path::Path;
use std::env;
use std::process::exit;
use std::process::Command;
use std::error::Error;

use getopts::Options;

use tempdir::TempDir;

/// Create an InvalidData io::Error with the description being a
/// formatted string
macro_rules! io_error {
    ($fmtstr:tt) => { io_error!($fmtstr,) };
    ($fmtstr:tt, $( $args:expr ),* ) => {
        Err(::std::io::Error::new(::std::io::ErrorKind::InvalidData,
                                  format!($fmtstr, $( $args ),* )));
    };
}

pub mod data;
mod write;
mod read;
mod string_write;
mod string_read;
#[cfg(test)]
mod tests;

/// Trait used in place of unwrap for printing to stdout/stderr, since if that
/// errors the program should simply exit with no further output
trait SilentUnwrap<T> {
    fn silent_unwrap(self) -> T;
}
impl<T> SilentUnwrap<T> for io::Result<T> {
    fn silent_unwrap(self) -> T {
        match self {
            Ok(t) => t,
            Err(_) => std::process::exit(1),
        }
    }
}

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

    let mut opts = Options::new();
    opts.optflagopt("e", "edit", "edit a NBT file with your $EDITOR.
    If [FILE] is specified, then that file is edited in place, but specifying --input and/or --output will override the input/output.
    If no file is specified, default to read from --input and writing to --output.", "FILE");
    opts.optflagopt("p", "print", "print NBT file to text format. Adding an argument to this is the same as specifying --input", "FILE");
    opts.optflagopt("r", "reverse", "reverse a file in text format to NBT format. Adding an argument to this is the same as specifying --input", "FILE");
    opts.optopt("i",
                "input",
                "specify the input file, defaults to stdin",
                "FILE");
    opts.optopt("o",
                "output",
                "specify the output file, defaults to stdout",
                "FILE");
    opts.optflag("h", "help", "print the help menu");
    opts.optflag("", "version", "print program version");

    let matches = match opts.parse(&args[1..]) {
        Ok(x) => x,
        Err(e) => {
            error(&format!("Error parsing options: {:?}", e.description()))
        },
    };

    if matches.opt_present("h") {
        print_usage(opts);
        return;
    }

    if matches.opt_present("version") {
        println!("{} {} {}",
                 env!("CARGO_PKG_NAME"),
                 env!("CARGO_PKG_VERSION"),
                 /* See build.rs for the git-revision.txt file */
                 include!(concat!(env!("OUT_DIR"), "/git-revision.txt")));
        println!("https://github.com/C4K3/nbted");
        return;
    }

    let is_print: bool = matches.opt_present("print");
    let is_reverse: bool = matches.opt_present("reverse");
    let is_edit: bool = if matches.opt_present("edit") {
        true
    } else {
        /* If edit is not explicitly defined, it is the default action and is
         * selected if no other action is specified */
        (!(is_reverse || is_print))
    };

    /* Hopefully this is a simpler way of ensuring that only one action can be
     * taken than having a long logical expression */
    let mut action_count = 0;
    if is_print {
        action_count += 1;
    }
    if is_reverse {
        action_count += 1;
    }
    if is_edit {
        action_count += 1;
    }
    if action_count > 1 {
        error("You can only specify one action a time.");
    }

    /* Figure out the input file, by trying to read the arguments for all of
     * --input, --edit, --print and --reverse, prioritizing --input over the
     * other arguments, if none of the arguments are specified but there is a
     * free argument, use that, else we finally default to - (stdin) */
    let input = if let Some(x) = matches.opt_str("input") {
        x
    } else if let Some(x) = matches.opt_str("edit") {
        x
    } else if let Some(x) = matches.opt_str("print") {
        x
    } else if let Some(x) = matches.opt_str("reverse") {
        x
    } else if matches.free.len() == 1 {
        matches.free[0].clone()
    } else {
        /* stdin */
        "-".to_string()
    };

    let output = if let Some(x) = matches.opt_str("output") {
        x
    } else if let Some(x) = matches.opt_str("edit") {
        x
    } else if is_edit && matches.free.len() == 1 {
        /* Only want to default to the free argument if we're editing
         * (DO NOT WRITE BACK TO THE READ FILE UNLESS EDITING!) */
        matches.free[0].clone()
    } else {
        /* stdout */
        "-".to_string()
    };

    if matches.free.len() > 1 {
        error("nbted was given multiple arguments, but only supports editing one file at a time.");
    }

    if is_print {
        print(&input, &output);
    } else if is_reverse {
        reverse(&input, &output);
    } else if is_edit {
        edit(&input, &output);
    } else {
        eprintln!("Internal error: No action selected. (Please report this.)");
    }
}

/// When the user wants to edit a specific file in place
fn edit(input: &str, output: &str) {

    /* First we read the NBT data from the input */
    let nbt = if input == "-" {
        // let mut f = BufReader::new(io::stdin());
        let f = io::stdin();
        let mut f = f.lock();
        match read::read_file(&mut f) {
            Ok(x) => x,
            Err(_) => {
                error(&format!("Unable to parse {}, are you sure it's an NBT file?",
                              input))
            },
        }
    } else {
        let path: &Path = Path::new(input);
        let f = match File::open(path) {
            Ok(f) => f,
            Err(_) => error(&format!("Unable to open file {}", input)),
        };
        let mut f = BufReader::new(f);

        match read::read_file(&mut f) {
            Ok(x) => x,
            Err(_) => {
                error(&format!("Unable to parse {}, are you sure it's an NBT file?",
                              input))
            },
        }
    };

    /* Then we create a temporary file and write the NBT data in text format
     * to the temporary file */
    let tmpdir = match TempDir::new("nbted") {
        Ok(x) => x,
        Err(e) => {
            error(&format!("Unable to create temporary directory: {:?}",
                          e.description()))
        },
    };

    let tmp = match Path::new(input).file_name() {
        Some(x) => {
            let mut x = x.to_os_string();
            x.push(".txt");
            x
        },
        None => error(&format!("Error reading file name")),
    };
    let tmp_path = tmpdir.path().join(tmp);

    {
        let mut f = match File::create(&tmp_path) {
            Ok(x) => x,
            Err(e) => {
                error(&format!("Unable to create temporary file: {:?}",
                              e.description()))
            },
        };

        match string_write::write_file(&mut f, &nbt) {
            Ok(()) => (),
            Err(e) => {
                error(&format!("Unable to write temporary file: {:?}",
                              e.description()))
            },
        };

        match f.sync_all() {
            Ok(()) => (),
            Err(e) => {
                error(&format!("Unable to synchronize file: {:?}",
                              e.description()))
            },
        }
    }

    let new_nbt = {
        let mut new_nbt = open_editor(&tmp_path);

        while let Err(e) = new_nbt {
            eprintln!("Unable to parse edited file: {}.\n\
            Do you want to open the file for editing again? (y/N)",
                        e.description());

            let mut line = String::new();
            match io::stdin().read_line(&mut line) {
                Ok(_) => (),
                Err(e) => {
                    error(&format!("Error reading from stdin: {}\n\
                                   Nothing was changed.",
                                  e.description()))
                },
            }

            if line.trim() == "y" {
                new_nbt = open_editor(&tmp_path);
            } else {
                eprintln!("Exiting ... File is unchanged.");
                std::process::exit(1);
            }
        }

        new_nbt.expect("new_nbt was Error")
    };

    if nbt == new_nbt {
        eprintln!("No changes, will do nothing.");
        return;
    }

    /* And finally we write the edited nbt (new_nbt) into the output file */
    if output == "-" {
        let f = io::stdout();
        let mut f = f.lock();
        write::write_file(&mut f, &new_nbt).silent_unwrap();
    } else {
        let path: &Path = Path::new(output);
        let f = match File::create(&path) {
            Ok(x) => x,
            Err(e) => {
                error(&format!("Unable to write to output NBT file {}: {:?}.\n\
            Nothing was changed",
                              output,
                              e.description()))
            },
        };
        let mut f = BufWriter::new(f);

        match write::write_file(&mut f, &new_nbt) {
            Ok(()) => (),
            Err(e) => {
                error(&format!("Error writing NBT file {}: {:?}.\n\
            State of NBT file is unknown, consider restoring it from a backup.",
                              output,
                              e.description()))
            },
        }
    }

    eprintln!("File edited successfully.");
}

/// Open the user's $EDITOR on the temporary file, wait until the editor is
/// closed again, read the temporary file and attempt to parse it into NBT,
/// returning the result.
fn open_editor(tmp_path: &Path) -> io::Result<data::NBTFile> {
    let editor = match env::var("EDITOR") {
        Ok(x) => x,
        Err(_) => {
            match env::var("VISUAL") {
                Ok(x) => x,
                Err(_) => error("Unable to find $EDITOR"),
            }
        },
    };

    let mut cmd = Command::new(editor);
    cmd.arg(&tmp_path.as_os_str());
    let mut cmd = match cmd.spawn() {
        Ok(x) => x,
        Err(e) => {
            error(&format!("Error opening editor: {:?}", e.description()))
        },
    };

    match cmd.wait() {
        Ok(x) if x.success() => (),
        Err(e) => {
            error(&format!("Error executing editor: {:?}", e.description()))
        },
        _ => error("Editor did not exit correctly"),
    }

    /* Then we parse the text format in the temporary file into NBT */
    let mut f = match File::open(&tmp_path) {
        Ok(x) => x,
        Err(e) => {
            error(&format!("Unable to read temporary file: {:?}.\n\
        Nothing was changed",
                          e.description()))
        },
    };

    string_read::read_file(&mut f)
}

/// When the user wants to print an NBT file to text format
fn print(input: &str, output: &str) {
    /* First we read a NBTFile from the input */
    let nbt = if input == "-" {
        let f = io::stdin();
        let mut f = f.lock();
        match read::read_file(&mut f) {
            Ok(x) => x,
            Err(_) => {
                error(&format!("Unable to parse {}, are you sure it's an NBT file?",
                              input))
            },
        }
    } else {
        let path: &Path = Path::new(input);
        let f = match File::open(path) {
            Ok(f) => f,
            Err(_) => error(&format!("Unable to open file {}", input)),
        };
        let mut f = BufReader::new(f);

        match read::read_file(&mut f) {
            Ok(x) => x,
            Err(_) => {
                error(&format!("Unable to parse {}, are you sure it's an NBT file?",
                              input))
            },
        }
    };

    /* Then we write the NBTFile to the output in text format */
    if output == "-" {
        let f = io::stdout();
        let mut f = f.lock();
        string_write::write_file(&mut f, &nbt).silent_unwrap();
    } else {
        let path: &Path = Path::new(output);
        let f = match File::create(&path) {
            Ok(x) => x,
            Err(e) => {
                error(&format!("Unable to write to output NBT file {}: {:?}.\n\
            Nothing was changed",
                              output,
                              e.description()))
            },
        };
        let mut f = BufWriter::new(f);

        match string_write::write_file(&mut f, &nbt) {
            Ok(()) => (),
            Err(e) => {
                error(&format!("Error writing NBT file {}: {:?}.\n\
            State of NBT file is unknown, consider restoring it from a backup.",
                              output,
                              e.description()))
            },
        }
    }
}

/// When the user wants to convert a text format file into an NBT file
fn reverse(input: &str, output: &str) {
    /* First we read the input file in the text format */
    let path: &Path = Path::new(input);
    let mut f = match File::open(&path) {
        Ok(x) => x,
        Err(e) => {
            error(&format!("Unable to read text file {}: {:?}",
                          input,
                          e.description()))
        },
    };

    let nbt = match string_read::read_file(&mut f) {
        Ok(x) => x,
        Err(e) => {
            error(&format!("Unable to parse text file {}: {:?}",
                          input,
                          e.description()))
        },
    };

    /* Then we write the parsed NBT to the output file in NBT format */
    if output == "-" {
        let f = io::stdout();
        let mut f = f.lock();
        write::write_file(&mut f, &nbt).silent_unwrap();
    } else {
        let path: &Path = Path::new(output);
        let f = match File::create(&path) {
            Ok(x) => x,
            Err(e) => {
                error(&format!("Unable to write to output NBT file {}: {:?}.\n\
            Nothing was changed",
                              output,
                              e.description()))
            },
        };
        let mut f = BufWriter::new(f);

        match write::write_file(&mut f, &nbt) {
            Ok(()) => (),
            Err(e) => {
                error(&format!("Error writing NBT file {}: {:?}.\n\
            State of NBT file is unknown, consider restoring it from a backup.",
                              output,
                              e.description()))
            },
        }
    }
}

/// Print the given message and exit with status code 1
fn error(message: &str) -> ! {
    eprintln!("Error: {}", message);
    eprintln!("Run with --help for help, or read the manpage.");
    exit(1);
}

fn print_usage(opts: Options) {
    let brief = "Usage: nbted [options] FILE";
    print!("{}", opts.usage(&brief));
    println!("\nThe default action, taken if no action is explicitly selected, is to --edit.");
    println!("\nFor detailed usage information, read the manpage.");
}
<|endoftext|>


<commit_before>mod vm;
mod command;
mod parse;
mod program;

use vm::VM;
use std::io::Write;

fn main() {
    let mut vm = VM::new();
    loop {
        print!("{}", "bf> ");
        std::io::stdout().flush().unwrap();

        let mut line = String::new();
        let commands = match std::io::stdin().read_line(&mut line) {
            Ok(0) => break,
            Ok(_) => parse::parse(&line),
            Err(error) => panic!(error),
        };
        vm.apply_many(&commands);
    }
}
<commit_msg>Change runtime of REPL to proxy through program<commit_after>mod vm;
mod command;
mod parse;
mod program;

use vm::VM;
use program::Program;
use std::io::Write;

fn main() {
    let mut vm = VM::new();
    let mut program = Program::new();
    loop {
        print!("{}", "bf> ");
        std::io::stdout().flush().unwrap();

        let mut line = String::new();
        let commands = match std::io::stdin().read_line(&mut line) {
            Ok(0) => break,
            Ok(_) => parse::parse(&line),
            Err(error) => panic!(error),
        };
        program.append(&commands);
        program.execute(&mut vm);
    }
}
<|endoftext|>


<commit_before>extern crate irc;
extern crate slack;
#[macro_use] extern crate lazy_static;
extern crate regex;
extern crate rustc_serialize;

use std::collections::HashMap;
use std::thread;
use std::sync::mpsc::{Sender, channel};
use std::default::Default;
use irc::client::prelude::*;
use regex::Regex;
use rustc_serialize::json;

#[derive(Debug)]
enum SlackToIrc {
    Message { to: String, msg: String },
    MeMessage { to: String, msg: String },
    Raw(String),
    Away(bool),
}

#[derive(Debug)]
enum IrcToSlack {
    Message { to: String, from: Option<String>, msg: String },
    Topic { chan: String, topic: Option<String> },
    Kick { by: Option<String>, chans: String, nicks: String, reason: Option<String> },
    Join { nick: String, chans: String },
    Part { nick: String, chans: String, reason: Option<String> },
    Quit { nick: String, reason: Option<String> },
    Nick { old_nick: String, new_nick: String },
    Mode { by: Option<String>, name: String, modes: String, params: Option<String> },
    Error(String),
}

struct SlackHandler<'a, T: 'a> {
    chan: &'a Sender<T>,
}

fn get_channel_with_id(id: &str, channels: Vec<slack::Channel>) -> Option<slack::Channel> {
    channels.into_iter().find(|channel| {
        channel.id == id
    })
}

fn get_user_with_id(id: &str, users: Vec<slack::User>) -> Option<slack::User> {
    users.into_iter().find(|user| {
        user.id == id
    })
}

fn parse_slack_text(text: &str, cli: &slack::RtmClient) -> String {
    lazy_static!{
        static ref REPLACEMENTS: Vec<(Regex, &'static str)> = vec![
            (Regex::new(r"\n|\r\n|\r").unwrap(), " "),
            (Regex::new(r"&").unwrap(), "&"),
            (Regex::new(r"<").unwrap(), "<"),
            (Regex::new(r">").unwrap(), ">"),
            (Regex::new(r"<!channel>").unwrap(), "@channel"),
            (Regex::new(r"<!group>").unwrap(), "@group"),
            (Regex::new(r"<!everyone>").unwrap(), "@everyone"),
        ];
        static ref EMOJIS: HashMap<&'static str, &'static str> = {
            let mut m = HashMap::new();
            m.insert("smile", ":)");
            m.insert("simple_smile", ":)");
            m.insert("smiley", ":-)");
            m.insert("grin", ":D");
            m.insert("wink", ";)");
            m.insert("smirk", ";)");
            m.insert("blush", ":$");
            m.insert("stuck_out_tongue", ":P");
            m.insert("stuck_out_tongue_winking_eye", ";P");
            m.insert("stuck_out_tongue_closed_eyes", "xP");
            m.insert("disappointed", ":(");
            m.insert("astonished", ":O");
            m.insert("open_mouth", ":O");
            m.insert("heart", "<3");
            m.insert("broken_heart", "</3");
            m.insert("confused", ":S");
            m.insert("angry", ">:(");
            m.insert("cry", ":,(");
            m.insert("frowning", ":(");
            m.insert("imp", "]:(");
            m.insert("innocent", "o:)");
            m.insert("joy", ":,)");
            m.insert("kissing", ":*");
            m.insert("laughing", "x)");
            m.insert("neutral_face", ":|");
            m.insert("no_mouth", ":-");
            m.insert("rage", ":@");
            m.insert("smiling_imp", "]:)");
            m.insert("sob", ":,'(");
            m.insert("sunglasses", "8)");
            m.insert("sweat", ",:(");
            m.insert("sweat_smile", ",:)");
            m.insert("unamused", ":$");
            m
        };
        static ref FUNCTIONS: Vec<(Regex, Box<Fn(&regex::Captures, &slack::RtmClient) -> String + Sync>)> = vec![
            (Regex::new(r"<#(C\w+)\|?(\w+)?>").unwrap(), Box::new(|captures: &regex::Captures, cli: &slack::RtmClient| {
                if let Some(id) = captures.at(1) {
                    if let Some(channel) = get_channel_with_id(id, cli.get_channels()) {
                        return format!("#{}", channel.name).to_owned()
                    }
                }
                captures.at(2).unwrap_or("").to_owned()
            })),
            (Regex::new(r"<@(U\w+)\|?(\w+)?>").unwrap(), Box::new(|captures: &regex::Captures, cli: &slack::RtmClient| {
                if let Some(id) = captures.at(1) {
                    if let Some(user) = get_user_with_id(id, cli.get_users()) {
                        return format!("@{}", user.name).to_owned()
                    }
                }
                captures.at(2).unwrap_or("").to_owned()
            })),
            (Regex::new(r"<(?!!)(\S+)>").unwrap(), Box::new(|captures: &regex::Captures, _| {
                captures.at(1).unwrap_or("").to_owned()
            })),
            (Regex::new(r"<!(\w+)\|?(\w+)?>").unwrap(), Box::new(|captures: &regex::Captures, _| {
                format!("<{}>", captures.at(2).or(captures.at(1)).unwrap_or("")).to_owned()
            })),
            (Regex::new(r"\:(\w+)\:").unwrap(), Box::new(|captures: &regex::Captures, _| {
                if let Some(&emoji) = EMOJIS.get(captures.at(1).unwrap()) {
                    emoji.to_owned()
                } else {
                    captures.at(0).unwrap().to_owned()
                }
            })),
        ];
    }

    let text = REPLACEMENTS.iter().fold(text.to_owned(), |text, &(ref re, replacement)| {
        re.replace_all(&text, replacement)
    });

    FUNCTIONS.iter().fold(text, |text, &(ref re, ref replacement)| {
        re.replace_all(&text, |captures: &regex::Captures| {
            replacement(captures, cli)
        })
    })
}

impl<'a> slack::EventHandler for SlackHandler<'a, SlackToIrc> {
    fn on_event(&mut self, cli: &mut slack::RtmClient, event: Result<&slack::Event, slack::Error>, _raw_json: &str) {
        let slack_user = &"TEMP_REPLACE_ME";
        let bot_channel = &"TEMP_REPLACE_ME";
        match event {
            Ok(&slack::Event::Message(ref message)) => match message {
                &slack::Message::Standard { channel: Some(ref channel), user: Some(ref user), text: Some(ref text), .. } if user == slack_user => {
                    lazy_static!{
                        static ref PM_RE: Regex = Regex::new(r"^(\S+):\s+(.+)").unwrap();
                    }

                    let text = parse_slack_text(&text, cli);

                    if text.starts_with("%") {
                        self.chan.send(SlackToIrc::Raw(text[1..].to_owned())).unwrap();
                        cli.send_message(channel, "_sent raw command_").unwrap();
                    } else if channel == bot_channel {
                        if let Some(captures) = PM_RE.captures(&text) {
                            self.chan.send(SlackToIrc::Message { to: captures.at(1).unwrap().to_owned(), msg: captures.at(2).unwrap().to_owned() }).unwrap();
                        } else {
                            cli.send_message(channel, "_no message sent, are you missing a `user: ` prefix?_").unwrap();
                        }
                    } else {
                        self.chan.send(SlackToIrc::Message { to: channel.clone(), msg: text.clone() }).unwrap();
                    }
                },
                &slack::Message::MeMessage { ref channel, ref user, ref text, .. } if user == slack_user => {
                    self.chan.send(SlackToIrc::MeMessage { to: channel.clone(), msg: text.clone() }).unwrap();
                },
                _ => (),
            },
            Ok(&slack::Event::PresenceChange { ref user, ref presence }) if user == slack_user => {
                self.chan.send(SlackToIrc::Away(presence == "active")).unwrap();
            },
            _ => (),
        }
    }

    fn on_ping(&mut self, _cli: &mut slack::RtmClient) {}

    fn on_close(&mut self, _cli: &mut slack::RtmClient) {
        println!("Disconnected from Slack");
    }

    fn on_connect(&mut self, _cli: &mut slack::RtmClient) {
        println!("Connected to Slack");
    }
}

fn make_slack_action(username: Option<String>, text: String) -> String {
    #[derive(RustcEncodable)]
    struct SlackAction {
        text: String,
        username: Option<String>,
        parse: String,
        icon_url: Option<String>,
    }

    json::encode(&SlackAction {
        text: text,
        username: username.clone(),
        parse: "full".to_owned(),
        icon_url: username.map(|username| { format!("http://api.adorable.io/avatars/48/{}.png", username) }),
    }).unwrap()
}

fn main() {
    let (slack_tx, slack_rx) = channel();
    let (irc_tx, irc_rx) = channel();

    let slack_thread = thread::Builder::new().name("slack".to_owned()).spawn(move || {
        let mut handler = SlackHandler { chan: &irc_tx };
        let mut cli = slack::RtmClient::new(&"SLACK_TOKEN_REPLACE_ME");

        cli.login_and_run(&mut handler).unwrap();

        let all_channels = ["ALL_CHANNELS_REPLACE_ME"];
        let dm_channel = "SLACK_DM_CHANNEL_TODO_REPLACE_ME";

        for msg in slack_rx {
            match msg {
                IrcToSlack::Message { to, from, msg } => {
                    let to: &str = if to.starts_with("#") {
                        dm_channel
                    } else {
                        &to
                    };

                    let msg = make_slack_action(from, msg);

                    cli.post_message(to, &msg, None).unwrap();
                },
                IrcToSlack::Topic { chan, topic } => {
                    cli.set_topic(&chan, &topic.unwrap_or("".to_owned())).unwrap();
                },
                IrcToSlack::Kick { by, chans, nicks, reason } => {
                    let by = &by.unwrap_or("server".to_owned());
                    let reason = &reason.unwrap_or("".to_owned());

                    for chan in chans.split(",") {
                        for nick in nicks.split(",") {
                            cli.send_message(&chan, &format!("*{}* has kicked *{}* (_{}_)", by, nick, reason)).unwrap();
                        }
                    }
                },
                IrcToSlack::Join { nick, chans } => {
                    for chan in chans.split(",") {
                        cli.send_message(&chan, &format!("*{}* has joined", nick)).unwrap();
                    }
                },
                IrcToSlack::Part { nick, chans, reason } => {
                    let reason = &reason.unwrap_or("".to_owned());

                    for chan in chans.split(",") {
                        cli.send_message(&chan, &format!("*{}* has left (_{}_)", nick, reason)).unwrap();
                    }
                },
                IrcToSlack::Quit { nick, reason } => {
                    let reason = &reason.unwrap_or("".to_owned());

                    for chan in &all_channels {
                        cli.send_message(&chan, &format!("*{}* has quit (_{}_)", nick, reason)).unwrap();
                    }
                },
                IrcToSlack::Nick { old_nick, new_nick } => {
                    for chan in &all_channels {
                        cli.send_message(&chan, &format!("*{}* is now known as *{}*", old_nick, new_nick)).unwrap();
                    }
                },
                IrcToSlack::Mode { by, name, modes, params } => {
                    if name.starts_with("#") {
                        cli.send_message(&name, &format!("*{}* sets mode *{}* on _{}_", by.unwrap_or("server".to_owned()), modes, params.unwrap_or(name.clone()))).unwrap();
                    }
                },
                IrcToSlack::Error(msg) => {
                    cli.send_message(dm_channel, &msg).unwrap();
                },
            }
        }
    }).unwrap();

    let irc_thread = thread::Builder::new().name("irc".to_owned()).spawn(move || {
        let config = Config {
            nickname: Some("NICK_REPLACE_ME".to_owned()),
            server: Some("SERVER_REPLACE_ME".to_owned()),
            channels: Some(vec!["CHANNELS_REPLACE_ME".to_owned()]),
            port: Some(6697),
            use_ssl: Some(true),
            ..Default::default()
        };

        let server = IrcServer::from_config(config).unwrap();

        server.identify().unwrap();

        let server_thread = {
            let server = server.clone();
            thread::Builder::new().name("irc_server".to_owned()).spawn(move || {
                for message in server.iter() {
                    let message = message.unwrap();
                    match message.command {
                        Command::ERROR(msg) => {
                            slack_tx.send(IrcToSlack::Error(msg)).unwrap();
                        },
                        Command::PRIVMSG(to, text) | Command::NOTICE(to, text) => {
                            lazy_static!{
                                static ref ACTION_RE: Regex = Regex::new("^\x01ACTION (.+)\x01$").unwrap();
                            }

                            if ACTION_RE.is_match(&text) {
                                slack_tx.send(IrcToSlack::Message {
                                    to: to,
                                    from: message.prefix,
                                    msg: format!("_{}_", ACTION_RE.captures(&text).unwrap().at(1).unwrap())
                                }).unwrap();
                            } else {
                                slack_tx.send(IrcToSlack::Message { to: to, from: message.prefix, msg: text }).unwrap();
                            }
                        },
                        Command::TOPIC(channel, topic) => {
                            slack_tx.send(IrcToSlack::Topic { chan: channel, topic: topic }).unwrap();
                        },
                        Command::KICK(channels, users, reason) => {
                            slack_tx.send(IrcToSlack::Kick { by: message.prefix, chans: channels, nicks: users, reason: reason }).unwrap();
                        },
                        Command::JOIN(channels, _, _) => {
                            slack_tx.send(IrcToSlack::Join { nick: message.prefix.unwrap(), chans: channels }).unwrap();
                        },
                        Command::PART(channels, reason) => {
                            slack_tx.send(IrcToSlack::Part { nick: message.prefix.unwrap(), chans: channels, reason: reason }).unwrap();
                        },
                        Command::QUIT(reason) => {
                            slack_tx.send(IrcToSlack::Quit { nick: message.prefix.unwrap(), reason: reason }).unwrap();
                        },
                        Command::NICK(nick) => {
                            slack_tx.send(IrcToSlack::Nick { old_nick: message.prefix.unwrap(), new_nick: nick }).unwrap();
                        },
                        Command::MODE(name, modes, params) => {
                            slack_tx.send(IrcToSlack::Mode { by: message.prefix, name: name, modes: modes, params: params }).unwrap();
                        }
                        _ => (),
                    }
                }
            }).unwrap()
        };

        for msg in irc_rx {
            match msg {
                SlackToIrc::Message { to, msg } => {
                    server.send_privmsg(&to, &msg).unwrap();
                },
                SlackToIrc::MeMessage { to, msg } => {
                    server.send_privmsg(&to, &format!("\x01ACTION {}\x01", msg)).unwrap();
                },
                SlackToIrc::Raw(msg) => {
                    server.send(Command::Raw(msg, vec![], None)).unwrap();
                },
                SlackToIrc::Away(is_away) => {
                    server.send(Command::AWAY(Some(if is_away { " ".to_owned() } else { "".to_owned() }))).unwrap();
                },
            }
        }

        server_thread.join().unwrap();
    }).unwrap();

    slack_thread.join().unwrap();
    irc_thread.join().unwrap();
}
<commit_msg>split Slack recv onto a separate thread, connect with another slack::RtmClient instance to send messages<commit_after>extern crate irc;
extern crate slack;
#[macro_use] extern crate lazy_static;
extern crate regex;
extern crate rustc_serialize;

use std::collections::HashMap;
use std::thread;
use std::sync::mpsc::{Sender, channel};
use std::default::Default;
use irc::client::prelude::*;
use regex::Regex;
use rustc_serialize::json;

#[derive(Debug)]
enum SlackToIrc {
    Message { to: String, msg: String },
    MeMessage { to: String, msg: String },
    Raw(String),
    Away(bool),
}

#[derive(Debug)]
enum IrcToSlack {
    Message { to: String, from: Option<String>, msg: String },
    Topic { chan: String, topic: Option<String> },
    Kick { by: Option<String>, chans: String, nicks: String, reason: Option<String> },
    Join { nick: String, chans: String },
    Part { nick: String, chans: String, reason: Option<String> },
    Quit { nick: String, reason: Option<String> },
    Nick { old_nick: String, new_nick: String },
    Mode { by: Option<String>, name: String, modes: String, params: Option<String> },
    Error(String),
}

struct SlackHandler<'a, T: 'a> {
    chan: &'a Sender<T>,
}

fn get_channel_with_id(id: &str, channels: Vec<slack::Channel>) -> Option<slack::Channel> {
    channels.into_iter().find(|channel| {
        channel.id == id
    })
}

fn get_user_with_id(id: &str, users: Vec<slack::User>) -> Option<slack::User> {
    users.into_iter().find(|user| {
        user.id == id
    })
}

fn parse_slack_text(text: &str, cli: &slack::RtmClient) -> String {
    lazy_static!{
        static ref REPLACEMENTS: Vec<(Regex, &'static str)> = vec![
            (Regex::new(r"\n|\r\n|\r").unwrap(), " "),
            (Regex::new(r"&").unwrap(), "&"),
            (Regex::new(r"<").unwrap(), "<"),
            (Regex::new(r">").unwrap(), ">"),
            (Regex::new(r"<!channel>").unwrap(), "@channel"),
            (Regex::new(r"<!group>").unwrap(), "@group"),
            (Regex::new(r"<!everyone>").unwrap(), "@everyone"),
        ];
        static ref EMOJIS: HashMap<&'static str, &'static str> = {
            let mut m = HashMap::new();
            m.insert("smile", ":)");
            m.insert("simple_smile", ":)");
            m.insert("smiley", ":-)");
            m.insert("grin", ":D");
            m.insert("wink", ";)");
            m.insert("smirk", ";)");
            m.insert("blush", ":$");
            m.insert("stuck_out_tongue", ":P");
            m.insert("stuck_out_tongue_winking_eye", ";P");
            m.insert("stuck_out_tongue_closed_eyes", "xP");
            m.insert("disappointed", ":(");
            m.insert("astonished", ":O");
            m.insert("open_mouth", ":O");
            m.insert("heart", "<3");
            m.insert("broken_heart", "</3");
            m.insert("confused", ":S");
            m.insert("angry", ">:(");
            m.insert("cry", ":,(");
            m.insert("frowning", ":(");
            m.insert("imp", "]:(");
            m.insert("innocent", "o:)");
            m.insert("joy", ":,)");
            m.insert("kissing", ":*");
            m.insert("laughing", "x)");
            m.insert("neutral_face", ":|");
            m.insert("no_mouth", ":-");
            m.insert("rage", ":@");
            m.insert("smiling_imp", "]:)");
            m.insert("sob", ":,'(");
            m.insert("sunglasses", "8)");
            m.insert("sweat", ",:(");
            m.insert("sweat_smile", ",:)");
            m.insert("unamused", ":$");
            m
        };
        static ref FUNCTIONS: Vec<(Regex, Box<Fn(&regex::Captures, &slack::RtmClient) -> String + Sync>)> = vec![
            (Regex::new(r"<#(C\w+)\|?(\w+)?>").unwrap(), Box::new(|captures: &regex::Captures, cli: &slack::RtmClient| {
                if let Some(id) = captures.at(1) {
                    if let Some(channel) = get_channel_with_id(id, cli.get_channels()) {
                        return format!("#{}", channel.name).to_owned()
                    }
                }
                captures.at(2).unwrap_or("").to_owned()
            })),
            (Regex::new(r"<@(U\w+)\|?(\w+)?>").unwrap(), Box::new(|captures: &regex::Captures, cli: &slack::RtmClient| {
                if let Some(id) = captures.at(1) {
                    if let Some(user) = get_user_with_id(id, cli.get_users()) {
                        return format!("@{}", user.name).to_owned()
                    }
                }
                captures.at(2).unwrap_or("").to_owned()
            })),
            (Regex::new(r"<(?!!)(\S+)>").unwrap(), Box::new(|captures: &regex::Captures, _| {
                captures.at(1).unwrap_or("").to_owned()
            })),
            (Regex::new(r"<!(\w+)\|?(\w+)?>").unwrap(), Box::new(|captures: &regex::Captures, _| {
                format!("<{}>", captures.at(2).or(captures.at(1)).unwrap_or("")).to_owned()
            })),
            (Regex::new(r"\:(\w+)\:").unwrap(), Box::new(|captures: &regex::Captures, _| {
                if let Some(&emoji) = EMOJIS.get(captures.at(1).unwrap()) {
                    emoji.to_owned()
                } else {
                    captures.at(0).unwrap().to_owned()
                }
            })),
        ];
    }

    let text = REPLACEMENTS.iter().fold(text.to_owned(), |text, &(ref re, replacement)| {
        re.replace_all(&text, replacement)
    });

    FUNCTIONS.iter().fold(text, |text, &(ref re, ref replacement)| {
        re.replace_all(&text, |captures: &regex::Captures| {
            replacement(captures, cli)
        })
    })
}

impl<'a> slack::EventHandler for SlackHandler<'a, SlackToIrc> {
    fn on_event(&mut self, cli: &mut slack::RtmClient, event: Result<&slack::Event, slack::Error>, _raw_json: &str) {
        let slack_user = &"TEMP_REPLACE_ME";
        let bot_channel = &"TEMP_REPLACE_ME";
        match event {
            Ok(&slack::Event::Message(ref message)) => match message {
                &slack::Message::Standard { channel: Some(ref channel), user: Some(ref user), text: Some(ref text), .. } if user == slack_user => {
                    lazy_static!{
                        static ref PM_RE: Regex = Regex::new(r"^(\S+):\s+(.+)").unwrap();
                    }

                    let text = parse_slack_text(&text, cli);

                    if text.starts_with("%") {
                        self.chan.send(SlackToIrc::Raw(text[1..].to_owned())).unwrap();
                        cli.send_message(channel, "_sent raw command_").unwrap();
                    } else if channel == bot_channel {
                        if let Some(captures) = PM_RE.captures(&text) {
                            self.chan.send(SlackToIrc::Message { to: captures.at(1).unwrap().to_owned(), msg: captures.at(2).unwrap().to_owned() }).unwrap();
                        } else {
                            cli.send_message(channel, "_no message sent, are you missing a `user: ` prefix?_").unwrap();
                        }
                    } else {
                        self.chan.send(SlackToIrc::Message { to: channel.clone(), msg: text.clone() }).unwrap();
                    }
                },
                &slack::Message::MeMessage { ref channel, ref user, ref text, .. } if user == slack_user => {
                    self.chan.send(SlackToIrc::MeMessage { to: channel.clone(), msg: text.clone() }).unwrap();
                },
                _ => (),
            },
            Ok(&slack::Event::PresenceChange { ref user, ref presence }) if user == slack_user => {
                self.chan.send(SlackToIrc::Away(presence == "active")).unwrap();
            },
            _ => (),
        }
    }

    fn on_ping(&mut self, _cli: &mut slack::RtmClient) {}

    fn on_close(&mut self, _cli: &mut slack::RtmClient) {
        println!("Disconnected from Slack");
    }

    fn on_connect(&mut self, _cli: &mut slack::RtmClient) {
        println!("Connected to Slack");
    }
}

fn make_slack_action(username: Option<String>, text: String) -> String {
    #[derive(RustcEncodable)]
    struct SlackAction {
        text: String,
        username: Option<String>,
        parse: String,
        icon_url: Option<String>,
    }

    json::encode(&SlackAction {
        text: text,
        username: username.clone(),
        parse: "full".to_owned(),
        icon_url: username.map(|username| { format!("http://api.adorable.io/avatars/48/{}.png", username) }),
    }).unwrap()
}

fn main() {
    let (slack_tx, slack_rx) = channel();
    let (irc_tx, irc_rx) = channel();

    let slack_thread = thread::Builder::new().name("slack".to_owned()).spawn(move || {
        let slack_token = "SLACK_TOKEN_REPLACE_ME";

        let recv_thread = thread::Builder::new().name("slack_recv".to_owned()).spawn(move || {
            let mut handler = SlackHandler { chan: &irc_tx };
            let mut cli = slack::RtmClient::new(&slack_token);
            cli.login_and_run(&mut handler).unwrap();
        }).unwrap();

        let mut cli = slack::RtmClient::new(&slack_token);
        cli.login().unwrap();

        let all_channels = ["ALL_CHANNELS_REPLACE_ME"];
        let dm_channel = "SLACK_DM_CHANNEL_TODO_REPLACE_ME";

        for msg in slack_rx {
            match msg {
                IrcToSlack::Message { to, from, msg } => {
                    let to: &str = if to.starts_with("#") {
                        dm_channel
                    } else {
                        &to
                    };

                    let msg = make_slack_action(from, msg);

                    cli.post_message(to, &msg, None).unwrap();
                },
                IrcToSlack::Topic { chan, topic } => {
                    cli.set_topic(&chan, &topic.unwrap_or("".to_owned())).unwrap();
                },
                IrcToSlack::Kick { by, chans, nicks, reason } => {
                    let by = &by.unwrap_or("server".to_owned());
                    let reason = &reason.unwrap_or("".to_owned());

                    for chan in chans.split(",") {
                        for nick in nicks.split(",") {
                            cli.send_message(&chan, &format!("*{}* has kicked *{}* (_{}_)", by, nick, reason)).unwrap();
                        }
                    }
                },
                IrcToSlack::Join { nick, chans } => {
                    for chan in chans.split(",") {
                        cli.send_message(&chan, &format!("*{}* has joined", nick)).unwrap();
                    }
                },
                IrcToSlack::Part { nick, chans, reason } => {
                    let reason = &reason.unwrap_or("".to_owned());

                    for chan in chans.split(",") {
                        cli.send_message(&chan, &format!("*{}* has left (_{}_)", nick, reason)).unwrap();
                    }
                },
                IrcToSlack::Quit { nick, reason } => {
                    let reason = &reason.unwrap_or("".to_owned());

                    for chan in &all_channels {
                        cli.send_message(&chan, &format!("*{}* has quit (_{}_)", nick, reason)).unwrap();
                    }
                },
                IrcToSlack::Nick { old_nick, new_nick } => {
                    for chan in &all_channels {
                        cli.send_message(&chan, &format!("*{}* is now known as *{}*", old_nick, new_nick)).unwrap();
                    }
                },
                IrcToSlack::Mode { by, name, modes, params } => {
                    if name.starts_with("#") {
                        cli.send_message(&name, &format!("*{}* sets mode *{}* on _{}_", by.unwrap_or("server".to_owned()), modes, params.unwrap_or(name.clone()))).unwrap();
                    }
                },
                IrcToSlack::Error(msg) => {
                    cli.send_message(dm_channel, &msg).unwrap();
                },
            }
        }

        recv_thread.join().unwrap();
    }).unwrap();

    let irc_thread = thread::Builder::new().name("irc".to_owned()).spawn(move || {
        let config = Config {
            nickname: Some("NICK_REPLACE_ME".to_owned()),
            server: Some("SERVER_REPLACE_ME".to_owned()),
            channels: Some(vec!["CHANNELS_REPLACE_ME".to_owned()]),
            port: Some(6697),
            use_ssl: Some(true),
            ..Default::default()
        };

        let server = IrcServer::from_config(config).unwrap();

        server.identify().unwrap();

        let server_thread = {
            let server = server.clone();
            thread::Builder::new().name("irc_server".to_owned()).spawn(move || {
                for message in server.iter() {
                    let message = message.unwrap();
                    match message.command {
                        Command::ERROR(msg) => {
                            slack_tx.send(IrcToSlack::Error(msg)).unwrap();
                        },
                        Command::PRIVMSG(to, text) | Command::NOTICE(to, text) => {
                            lazy_static!{
                                static ref ACTION_RE: Regex = Regex::new("^\x01ACTION (.+)\x01$").unwrap();
                            }

                            if ACTION_RE.is_match(&text) {
                                slack_tx.send(IrcToSlack::Message {
                                    to: to,
                                    from: message.prefix,
                                    msg: format!("_{}_", ACTION_RE.captures(&text).unwrap().at(1).unwrap())
                                }).unwrap();
                            } else {
                                slack_tx.send(IrcToSlack::Message { to: to, from: message.prefix, msg: text }).unwrap();
                            }
                        },
                        Command::TOPIC(channel, topic) => {
                            slack_tx.send(IrcToSlack::Topic { chan: channel, topic: topic }).unwrap();
                        },
                        Command::KICK(channels, users, reason) => {
                            slack_tx.send(IrcToSlack::Kick { by: message.prefix, chans: channels, nicks: users, reason: reason }).unwrap();
                        },
                        Command::JOIN(channels, _, _) => {
                            slack_tx.send(IrcToSlack::Join { nick: message.prefix.unwrap(), chans: channels }).unwrap();
                        },
                        Command::PART(channels, reason) => {
                            slack_tx.send(IrcToSlack::Part { nick: message.prefix.unwrap(), chans: channels, reason: reason }).unwrap();
                        },
                        Command::QUIT(reason) => {
                            slack_tx.send(IrcToSlack::Quit { nick: message.prefix.unwrap(), reason: reason }).unwrap();
                        },
                        Command::NICK(nick) => {
                            slack_tx.send(IrcToSlack::Nick { old_nick: message.prefix.unwrap(), new_nick: nick }).unwrap();
                        },
                        Command::MODE(name, modes, params) => {
                            slack_tx.send(IrcToSlack::Mode { by: message.prefix, name: name, modes: modes, params: params }).unwrap();
                        }
                        _ => (),
                    }
                }
            }).unwrap()
        };

        for msg in irc_rx {
            match msg {
                SlackToIrc::Message { to, msg } => {
                    server.send_privmsg(&to, &msg).unwrap();
                },
                SlackToIrc::MeMessage { to, msg } => {
                    server.send_privmsg(&to, &format!("\x01ACTION {}\x01", msg)).unwrap();
                },
                SlackToIrc::Raw(msg) => {
                    server.send(Command::Raw(msg, vec![], None)).unwrap();
                },
                SlackToIrc::Away(is_away) => {
                    server.send(Command::AWAY(Some(if is_away { " ".to_owned() } else { "".to_owned() }))).unwrap();
                },
            }
        }

        server_thread.join().unwrap();
    }).unwrap();

    slack_thread.join().unwrap();
    irc_thread.join().unwrap();
}
<|endoftext|>


<commit_before>use chrono::NaiveDateTime;
use crossbeam_channel::bounded;
use std::collections::{HashMap, HashSet};
use std::env;
use std::fs::File;
use std::panic;
use std::rc::Rc;
use std::thread;

mod event;
mod logcat;
mod reader;
mod ui;

use crate::event::Event;
use crate::logcat::{LogcatEntry, LoggerEntry};
use crate::reader::reader_thread;
use crate::ui::{input_thread, UserInterface};

struct Context {
    tags: HashSet<Rc<String>>,
    process_names: HashMap<i32, Option<Rc<String>>>,
    logcat_entries: Vec<LogcatEntry>,
}

impl Context {
    fn new() -> Context {
        Context {
            tags: HashSet::new(),
            process_names: HashMap::new(),
            logcat_entries: Vec::new(),
        }
    }

    fn on_logger_entry(&mut self, raw: LoggerEntry) -> &LogcatEntry {
        if !self.tags.contains(&raw.tag) {
            self.tags.insert(Rc::new(raw.tag.clone()));
        }
        self.process_names.entry(raw.pid).or_insert(None); // TODO: call 'adb shell' instead
        let e = LogcatEntry {
            pid: raw.pid,
            tid: raw.tid,
            process_name: self
                .process_names
                .get(&raw.pid)
                .unwrap()
                .as_ref()
                .map(|v| Rc::clone(v)),
            timestamp: NaiveDateTime::from_timestamp(i64::from(raw.sec), raw.nsec),
            level: raw.level,
            tag: Rc::clone(self.tags.get(&raw.tag).unwrap()),
            text: raw.text,
        };
        self.logcat_entries.push(e);
        self.logcat_entries.last().unwrap()
    }
}

fn main() {
    let path = env::args_os().nth(1).expect("usage: lokatt <path-to-file>");

    panic::set_hook(Box::new(|panic_info| {
        UserInterface::shutdown();

        match panic_info.location() {
            Some(l) => eprintln!("{}:{}: panic:", l.file(), l.line()),
            None => eprintln!("panic:"),
        }
        eprintln!("{:#?}", panic_info);
    }));

    let (sender, receiver) = bounded(32);

    let s = sender.clone();
    let ui = UserInterface::init(s);

    let s = sender.clone();
    thread::Builder::new()
        .name("input".to_string())
        .spawn(move || {
            input_thread(s);
        })
        .unwrap();

    let mut src = File::open(path).unwrap();
    let s = sender.clone();
    thread::Builder::new()
        .name("reader".to_string())
        .spawn(move || {
            reader_thread(s, &mut src);
        })
        .unwrap();

    drop(sender);

    let mut ctx = Context::new();
    loop {
        match receiver.recv().unwrap() {
            Event::Command(cmd) => {
                if cmd == "quit" {
                    break;
                }
            }
            Event::KeyCode(ch) => {
                ui.on_key(ch);
            }
            Event::LoggerEntry(le) => {
                let logcat = ctx.on_logger_entry(le);
                ui.on_logcat(logcat);
            }
            Event::Refresh => {
                ui.on_refresh();
            }
        }
    }

    UserInterface::shutdown();
}
<commit_msg>Include backtrace in panic output<commit_after>use backtrace::Backtrace;
use chrono::NaiveDateTime;
use crossbeam_channel::bounded;
use std::collections::{HashMap, HashSet};
use std::env;
use std::fs::File;
use std::panic;
use std::rc::Rc;
use std::thread;

mod event;
mod logcat;
mod reader;
mod ui;

use crate::event::Event;
use crate::logcat::{LogcatEntry, LoggerEntry};
use crate::reader::reader_thread;
use crate::ui::{input_thread, UserInterface};

struct Context {
    tags: HashSet<Rc<String>>,
    process_names: HashMap<i32, Option<Rc<String>>>,
    logcat_entries: Vec<LogcatEntry>,
}

impl Context {
    fn new() -> Context {
        Context {
            tags: HashSet::new(),
            process_names: HashMap::new(),
            logcat_entries: Vec::new(),
        }
    }

    fn on_logger_entry(&mut self, raw: LoggerEntry) -> &LogcatEntry {
        if !self.tags.contains(&raw.tag) {
            self.tags.insert(Rc::new(raw.tag.clone()));
        }
        self.process_names.entry(raw.pid).or_insert(None); // TODO: call 'adb shell' instead
        let e = LogcatEntry {
            pid: raw.pid,
            tid: raw.tid,
            process_name: self
                .process_names
                .get(&raw.pid)
                .unwrap()
                .as_ref()
                .map(|v| Rc::clone(v)),
            timestamp: NaiveDateTime::from_timestamp(i64::from(raw.sec), raw.nsec),
            level: raw.level,
            tag: Rc::clone(self.tags.get(&raw.tag).unwrap()),
            text: raw.text,
        };
        self.logcat_entries.push(e);
        self.logcat_entries.last().unwrap()
    }
}

fn main() {
    let path = env::args_os().nth(1).expect("usage: lokatt <path-to-file>");

    panic::set_hook(Box::new(|panic_info| {
        UserInterface::shutdown();

        let bt = Backtrace::new();
        match panic_info.location() {
            Some(l) => eprintln!("{}:{}: panic:", l.file(), l.line()),
            None => eprintln!("panic:"),
        }
        eprintln!("{:#?}\n{:#?}", panic_info, bt);
    }));

    let (sender, receiver) = bounded(32);

    let s = sender.clone();
    let ui = UserInterface::init(s);

    let s = sender.clone();
    thread::Builder::new()
        .name("input".to_string())
        .spawn(move || {
            input_thread(s);
        })
        .unwrap();

    let mut src = File::open(path).unwrap();
    let s = sender.clone();
    thread::Builder::new()
        .name("reader".to_string())
        .spawn(move || {
            reader_thread(s, &mut src);
        })
        .unwrap();

    drop(sender);

    let mut ctx = Context::new();
    loop {
        match receiver.recv().unwrap() {
            Event::Command(cmd) => {
                if cmd == "quit" {
                    break;
                }
            }
            Event::KeyCode(ch) => {
                ui.on_key(ch);
            }
            Event::LoggerEntry(le) => {
                let logcat = ctx.on_logger_entry(le);
                ui.on_logcat(logcat);
            }
            Event::Refresh => {
                ui.on_refresh();
            }
        }
    }

    UserInterface::shutdown();
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::str::FromStr;
use std::path::{PathBuf, Path, Component};

use id_map::{IdMap, mapping_validator};
use range::Range;
use quire::validate::{Sequence, Mapping, Scalar, Numeric};
use quire::validate::{Structure};
use serde::de::{Deserializer, Deserialize, Error};

#[derive(Deserialize)]
pub struct SandboxConfig {
    pub config_file: Option<PathBuf>,
    pub image_dir: PathBuf,
    pub image_dir_levels: u32,
    pub used_images_list: Option<PathBuf>,
    pub log_file: Option<PathBuf>,
    pub log_level: Option<String>,
    pub readonly_paths: BTreeMap<PathBuf, PathBuf>,
    pub writable_paths: BTreeMap<PathBuf, PathBuf>,
    pub allow_users: Vec<Range>,
    pub allow_groups: Vec<Range>,
    pub allow_tcp_ports: Vec<Range>,
    pub additional_hosts: BTreeMap<String, String>,
    pub uid_map: Vec<IdMap>,
    pub gid_map: Vec<IdMap>,
    pub auto_clean: bool,
}

impl SandboxConfig {
    pub fn check_path<P: AsRef<Path>>(&self, path: P) -> bool {
        let mut num = 0;
        for component in path.as_ref().components() {
            match component {
                Component::Normal(x) if x.to_str().is_some() => num += 1,
                _ => return false,
            }
        }
        return num == self.image_dir_levels;
    }
    pub fn validator<'x>() -> Structure<'x> {
        Structure::new()
        .member("config_file", Scalar::new().optional())
        .member("image_dir", Scalar::new().optional()
            .default("/var/lib/lithos/containers"))
        .member("image_dir_levels",
            Numeric::new().min(1).max(16).default(1))
        .member("used_images_list", Scalar::new().optional())
        .member("log_file", Scalar::new().optional())
        .member("log_level", Scalar::new().optional())
        .member("readonly_paths", Mapping::new(
            Scalar::new(),
            Scalar::new()))
        .member("writable_paths", Mapping::new(
            Scalar::new(),
            Scalar::new()))
        .member("allow_users", Sequence::new(Scalar::new()))
        .member("allow_groups", Sequence::new(Scalar::new()))
        .member("allow_tcp_ports", Sequence::new(Scalar::new()))
        .member("uid_map", mapping_validator())
        .member("gid_map", mapping_validator())
        .member("additional_hosts", Mapping::new(
            Scalar::new(),
            Scalar::new()))
        .member("auto_clean", Scalar::new().default("true").optional())
    }
}
<commit_msg>Fix unused import warnings<commit_after>use std::collections::BTreeMap;
use std::path::{PathBuf, Path, Component};

use id_map::{IdMap, mapping_validator};
use range::Range;
use quire::validate::{Sequence, Mapping, Scalar, Numeric};
use quire::validate::{Structure};

#[derive(Deserialize)]
pub struct SandboxConfig {
    pub config_file: Option<PathBuf>,
    pub image_dir: PathBuf,
    pub image_dir_levels: u32,
    pub used_images_list: Option<PathBuf>,
    pub log_file: Option<PathBuf>,
    pub log_level: Option<String>,
    pub readonly_paths: BTreeMap<PathBuf, PathBuf>,
    pub writable_paths: BTreeMap<PathBuf, PathBuf>,
    pub allow_users: Vec<Range>,
    pub allow_groups: Vec<Range>,
    pub allow_tcp_ports: Vec<Range>,
    pub additional_hosts: BTreeMap<String, String>,
    pub uid_map: Vec<IdMap>,
    pub gid_map: Vec<IdMap>,
    pub auto_clean: bool,
}

impl SandboxConfig {
    pub fn check_path<P: AsRef<Path>>(&self, path: P) -> bool {
        let mut num = 0;
        for component in path.as_ref().components() {
            match component {
                Component::Normal(x) if x.to_str().is_some() => num += 1,
                _ => return false,
            }
        }
        return num == self.image_dir_levels;
    }
    pub fn validator<'x>() -> Structure<'x> {
        Structure::new()
        .member("config_file", Scalar::new().optional())
        .member("image_dir", Scalar::new().optional()
            .default("/var/lib/lithos/containers"))
        .member("image_dir_levels",
            Numeric::new().min(1).max(16).default(1))
        .member("used_images_list", Scalar::new().optional())
        .member("log_file", Scalar::new().optional())
        .member("log_level", Scalar::new().optional())
        .member("readonly_paths", Mapping::new(
            Scalar::new(),
            Scalar::new()))
        .member("writable_paths", Mapping::new(
            Scalar::new(),
            Scalar::new()))
        .member("allow_users", Sequence::new(Scalar::new()))
        .member("allow_groups", Sequence::new(Scalar::new()))
        .member("allow_tcp_ports", Sequence::new(Scalar::new()))
        .member("uid_map", mapping_validator())
        .member("gid_map", mapping_validator())
        .member("additional_hosts", Mapping::new(
            Scalar::new(),
            Scalar::new()))
        .member("auto_clean", Scalar::new().default("true").optional())
    }
}
<|endoftext|>


<commit_before>use std::{
    io::{self, ErrorKind, Write},
    net::SocketAddr,
    sync::Arc,
    task::{Context, Poll, Waker},
    time::{Duration, Instant},
};

use futures::future;
use kcp::{Error as KcpError, Kcp, KcpResult};
use log::{error, trace};
use tokio::{net::UdpSocket, sync::mpsc};

use crate::{utils::now_millis, KcpConfig};

/// Writer for sending packets to the underlying UdpSocket
struct UdpOutput {
    socket: Arc<UdpSocket>,
    target_addr: SocketAddr,
    delay_tx: mpsc::UnboundedSender<Vec<u8>>,
}

impl UdpOutput {
    /// Create a new Writer for writing packets to UdpSocket
    pub fn new(socket: Arc<UdpSocket>, target_addr: SocketAddr) -> UdpOutput {
        let (delay_tx, mut delay_rx) = mpsc::unbounded_channel::<Vec<u8>>();

        {
            let socket = socket.clone();
            tokio::spawn(async move {
                while let Some(buf) = delay_rx.recv().await {
                    if let Err(err) = socket.send_to(&buf, target_addr).await {
                        error!("[SEND] UDP delayed send failed, error: {}", err);
                    }
                }
            });
        }

        UdpOutput {
            socket,
            target_addr,
            delay_tx,
        }
    }
}

impl Write for UdpOutput {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self.socket.try_send_to(buf, self.target_addr) {
            Ok(n) => Ok(n),
            Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
                // send return EAGAIN
                // ignored as packet was lost in transmission
                trace!("[SEND] UDP send EAGAIN, packet.size: {} bytes, delayed send", buf.len());

                self.delay_tx.send(buf.to_owned()).expect("channel closed unexpectly");

                Ok(buf.len())
            }
            Err(err) => Err(err),
        }
    }

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

pub struct KcpSocket {
    kcp: Kcp<UdpOutput>,
    last_update: Instant,
    socket: Arc<UdpSocket>,
    flush_write: bool,
    flush_ack_input: bool,
    sent_first: bool,
    pending_sender: Option<Waker>,
    pending_receiver: Option<Waker>,
    closed: bool,
}

impl KcpSocket {
    pub fn new(
        c: &KcpConfig,
        conv: u32,
        socket: Arc<UdpSocket>,
        target_addr: SocketAddr,
        stream: bool,
    ) -> KcpResult<KcpSocket> {
        let output = UdpOutput::new(socket.clone(), target_addr);
        let mut kcp = if stream {
            Kcp::new_stream(conv, output)
        } else {
            Kcp::new(conv, output)
        };
        c.apply_config(&mut kcp);

        // Ask server to allocate one
        if conv == 0 {
            kcp.input_conv();
        }

        kcp.update(now_millis())?;

        Ok(KcpSocket {
            kcp,
            last_update: Instant::now(),
            socket,
            flush_write: c.flush_write,
            flush_ack_input: c.flush_acks_input,
            sent_first: false,
            pending_sender: None,
            pending_receiver: None,
            closed: false,
        })
    }

    /// Call every time you got data from transmission
    pub fn input(&mut self, buf: &[u8]) -> KcpResult<bool> {
        match self.kcp.input(buf) {
            Ok(..) => {}
            Err(KcpError::ConvInconsistent(expected, actual)) => {
                trace!("[INPUT] Conv expected={} actual={} ignored", expected, actual);
                return Ok(false);
            }
            Err(err) => return Err(err),
        }
        self.last_update = Instant::now();

        if self.flush_ack_input {
            self.kcp.flush_ack()?;
        }

        Ok(self.try_wake_pending_waker())
    }

    /// Call if you want to send some data
    pub fn poll_send(&mut self, cx: &mut Context<'_>, mut buf: &[u8]) -> Poll<KcpResult<usize>> {
        if self.closed {
            return Ok(0).into();
        }

        // If:
        //     1. Have sent the first packet (asking for conv)
        //     2. Too many pending packets
        if self.sent_first
            && (self.kcp.wait_snd() >= self.kcp.snd_wnd() as usize
                || self.kcp.wait_snd() >= self.kcp.rmt_wnd() as usize
                || self.kcp.waiting_conv())
        {
            trace!(
                "[SEND] waitsnd={} sndwnd={} excceeded or waiting conv={}",
                self.kcp.wait_snd(),
                self.kcp.snd_wnd(),
                self.kcp.waiting_conv()
            );

            if let Some(waker) = self.pending_sender.replace(cx.waker().clone()) {
                if !cx.waker().will_wake(&waker) {
                    waker.wake();
                }
            }
            return Poll::Pending;
        }

        if !self.sent_first && self.kcp.waiting_conv() && buf.len() > self.kcp.mss() as usize {
            buf = &buf[..self.kcp.mss() as usize];
        }

        let n = self.kcp.send(buf)?;
        self.sent_first = true;

        if self.kcp.wait_snd() >= self.kcp.snd_wnd() as usize || self.kcp.wait_snd() >= self.kcp.rmt_wnd() as usize {
            self.kcp.flush()?;
        }

        self.last_update = Instant::now();

        if self.flush_write {
            self.kcp.flush()?;
        }

        Ok(n).into()
    }

    /// Call if you want to send some data
    #[allow(dead_code)]
    pub async fn send(&mut self, buf: &[u8]) -> KcpResult<usize> {
        future::poll_fn(|cx| self.poll_send(cx, buf)).await
    }

    #[allow(dead_code)]
    pub fn try_recv(&mut self, buf: &mut [u8]) -> KcpResult<usize> {
        if self.closed {
            return Ok(0);
        }
        self.kcp.recv(buf)
    }

    pub fn poll_recv(&mut self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<KcpResult<usize>> {
        if self.closed {
            return Ok(0).into();
        }

        match self.kcp.recv(buf) {
            Ok(0) | Err(KcpError::RecvQueueEmpty) | Err(KcpError::ExpectingFragment) => {
                if let Some(waker) = self.pending_receiver.replace(cx.waker().clone()) {
                    if !cx.waker().will_wake(&waker) {
                        waker.wake();
                    }
                }
                Poll::Pending
            }
            Err(err) => Err(err).into(),
            Ok(n) => Ok(n).into(),
        }
    }

    #[allow(dead_code)]
    pub async fn recv(&mut self, buf: &mut [u8]) -> KcpResult<usize> {
        future::poll_fn(|cx| self.poll_recv(cx, buf)).await
    }

    pub fn flush(&mut self) -> KcpResult<()> {
        self.kcp.flush()?;
        self.last_update = Instant::now();
        Ok(())
    }

    fn try_wake_pending_waker(&mut self) -> bool {
        let mut waked = false;

        if self.pending_sender.is_some()
            && self.kcp.wait_snd() < self.kcp.snd_wnd() as usize
            && self.kcp.wait_snd() < self.kcp.rmt_wnd() as usize
            && !self.kcp.waiting_conv()
        {
            let waker = self.pending_sender.take().unwrap();
            waker.wake();

            waked = true;
        }

        if self.pending_receiver.is_some() {
            if let Ok(peek) = self.kcp.peeksize() {
                if peek > 0 {
                    let waker = self.pending_receiver.take().unwrap();
                    waker.wake();

                    waked = true;
                }
            }
        }

        waked
    }

    pub fn update(&mut self) -> KcpResult<Instant> {
        let now = now_millis();
        self.kcp.update(now)?;
        let next = self.kcp.check(now);

        self.try_wake_pending_waker();

        Ok(Instant::now() + Duration::from_millis(next as u64))
    }

    pub fn close(&mut self) {
        self.closed = true;
        if let Some(w) = self.pending_sender.take() {
            w.wake();
        }
        if let Some(w) = self.pending_receiver.take() {
            w.wake();
        }
    }

    pub fn udp_socket(&self) -> &Arc<UdpSocket> {
        &self.socket
    }

    pub fn can_close(&self) -> bool {
        self.kcp.wait_snd() == 0
    }

    pub fn conv(&self) -> u32 {
        self.kcp.conv()
    }

    pub fn set_conv(&mut self, conv: u32) {
        self.kcp.set_conv(conv);
    }

    pub fn waiting_conv(&self) -> bool {
        self.kcp.waiting_conv()
    }

    pub fn peek_size(&self) -> KcpResult<usize> {
        self.kcp.peeksize()
    }

    pub fn last_update_time(&self) -> Instant {
        self.last_update
    }
}

#[cfg(test)]
mod test {

    use kcp::Error as KcpError;
    use log::trace;
    use std::sync::Arc;
    use tokio::{
        net::UdpSocket,
        sync::Mutex,
        time::{self, Instant},
    };

    use super::KcpSocket;
    use crate::config::KcpConfig;

    #[tokio::test]
    async fn kcp_echo() {
        let _ = env_logger::try_init();

        static CONV: u32 = 0xdeadbeef;

        // s1 connects s2
        let s1 = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let s2 = UdpSocket::bind("127.0.0.1:0").await.unwrap();

        let s1_addr = s1.local_addr().unwrap();
        let s2_addr = s2.local_addr().unwrap();

        let s1 = Arc::new(s1);
        let s2 = Arc::new(s2);

        let config = KcpConfig::default();
        let kcp1 = KcpSocket::new(&config, 0, s1.clone(), s2_addr, true).unwrap();
        let kcp2 = KcpSocket::new(&config, CONV, s2.clone(), s1_addr, true).unwrap();

        let kcp1 = Arc::new(Mutex::new(kcp1));
        let kcp2 = Arc::new(Mutex::new(kcp2));

        let kcp1_task = {
            let kcp1 = kcp1.clone();
            tokio::spawn(async move {
                loop {
                    let mut kcp = kcp1.lock().await;
                    let next = kcp.update().expect("update");
                    trace!("kcp1 next tick {:?}", next);
                    time::sleep_until(Instant::from_std(next)).await;
                }
            })
        };

        let kcp2_task = {
            let kcp2 = kcp2.clone();
            tokio::spawn(async move {
                loop {
                    let mut kcp = kcp2.lock().await;
                    let next = kcp.update().expect("update");
                    trace!("kcp2 next tick {:?}", next);
                    time::sleep_until(Instant::from_std(next)).await;
                }
            })
        };

        const SEND_BUFFER: &[u8] = b"HELLO WORLD";

        {
            let n = kcp1.lock().await.send(SEND_BUFFER).await.unwrap();
            assert_eq!(n, SEND_BUFFER.len());
        }

        let echo_task = tokio::spawn(async move {
            let mut buf = [0u8; 1024];

            loop {
                let n = s2.recv(&mut buf).await.unwrap();

                let packet = &mut buf[..n];

                let conv = kcp::get_conv(packet);
                if conv == 0 {
                    kcp::set_conv(packet, CONV);
                }

                let mut kcp2 = kcp2.lock().await;
                kcp2.input(packet).unwrap();

                match kcp2.try_recv(&mut buf) {
                    Ok(n) => {
                        let received = &buf[..n];
                        kcp2.send(received).await.unwrap();
                    }
                    Err(KcpError::RecvQueueEmpty) => {
                        continue;
                    }
                    Err(err) => {
                        panic!("kcp.recv error: {:?}", err);
                    }
                }
            }
        });

        {
            let mut buf = [0u8; 1024];

            loop {
                let n = s1.recv(&mut buf).await.unwrap();

                let packet = &buf[..n];

                let mut kcp1 = kcp1.lock().await;
                kcp1.input(packet).unwrap();

                match kcp1.try_recv(&mut buf) {
                    Ok(n) => {
                        let received = &buf[..n];
                        assert_eq!(received, SEND_BUFFER);
                        break;
                    }
                    Err(KcpError::RecvQueueEmpty) => {
                        continue;
                    }
                    Err(err) => {
                        panic!("kcp.recv error: {:?}", err);
                    }
                }
            }
        }

        echo_task.abort();
        kcp1_task.abort();
        kcp2_task.abort();
    }
}
<commit_msg>write to closed socket, return BrokenPipe<commit_after>use std::{
    io::{self, ErrorKind, Write},
    net::SocketAddr,
    sync::Arc,
    task::{Context, Poll, Waker},
    time::{Duration, Instant},
};

use futures::future;
use kcp::{Error as KcpError, Kcp, KcpResult};
use log::{error, trace};
use tokio::{net::UdpSocket, sync::mpsc};

use crate::{utils::now_millis, KcpConfig};

/// Writer for sending packets to the underlying UdpSocket
struct UdpOutput {
    socket: Arc<UdpSocket>,
    target_addr: SocketAddr,
    delay_tx: mpsc::UnboundedSender<Vec<u8>>,
}

impl UdpOutput {
    /// Create a new Writer for writing packets to UdpSocket
    pub fn new(socket: Arc<UdpSocket>, target_addr: SocketAddr) -> UdpOutput {
        let (delay_tx, mut delay_rx) = mpsc::unbounded_channel::<Vec<u8>>();

        {
            let socket = socket.clone();
            tokio::spawn(async move {
                while let Some(buf) = delay_rx.recv().await {
                    if let Err(err) = socket.send_to(&buf, target_addr).await {
                        error!("[SEND] UDP delayed send failed, error: {}", err);
                    }
                }
            });
        }

        UdpOutput {
            socket,
            target_addr,
            delay_tx,
        }
    }
}

impl Write for UdpOutput {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self.socket.try_send_to(buf, self.target_addr) {
            Ok(n) => Ok(n),
            Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
                // send return EAGAIN
                // ignored as packet was lost in transmission
                trace!("[SEND] UDP send EAGAIN, packet.size: {} bytes, delayed send", buf.len());

                self.delay_tx.send(buf.to_owned()).expect("channel closed unexpectly");

                Ok(buf.len())
            }
            Err(err) => Err(err),
        }
    }

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

pub struct KcpSocket {
    kcp: Kcp<UdpOutput>,
    last_update: Instant,
    socket: Arc<UdpSocket>,
    flush_write: bool,
    flush_ack_input: bool,
    sent_first: bool,
    pending_sender: Option<Waker>,
    pending_receiver: Option<Waker>,
    closed: bool,
}

impl KcpSocket {
    pub fn new(
        c: &KcpConfig,
        conv: u32,
        socket: Arc<UdpSocket>,
        target_addr: SocketAddr,
        stream: bool,
    ) -> KcpResult<KcpSocket> {
        let output = UdpOutput::new(socket.clone(), target_addr);
        let mut kcp = if stream {
            Kcp::new_stream(conv, output)
        } else {
            Kcp::new(conv, output)
        };
        c.apply_config(&mut kcp);

        // Ask server to allocate one
        if conv == 0 {
            kcp.input_conv();
        }

        kcp.update(now_millis())?;

        Ok(KcpSocket {
            kcp,
            last_update: Instant::now(),
            socket,
            flush_write: c.flush_write,
            flush_ack_input: c.flush_acks_input,
            sent_first: false,
            pending_sender: None,
            pending_receiver: None,
            closed: false,
        })
    }

    /// Call every time you got data from transmission
    pub fn input(&mut self, buf: &[u8]) -> KcpResult<bool> {
        match self.kcp.input(buf) {
            Ok(..) => {}
            Err(KcpError::ConvInconsistent(expected, actual)) => {
                trace!("[INPUT] Conv expected={} actual={} ignored", expected, actual);
                return Ok(false);
            }
            Err(err) => return Err(err),
        }
        self.last_update = Instant::now();

        if self.flush_ack_input {
            self.kcp.flush_ack()?;
        }

        Ok(self.try_wake_pending_waker())
    }

    /// Call if you want to send some data
    pub fn poll_send(&mut self, cx: &mut Context<'_>, mut buf: &[u8]) -> Poll<KcpResult<usize>> {
        if self.closed {
            return Err(io::Error::from(ErrorKind::BrokenPipe).into()).into();
        }

        // If:
        //     1. Have sent the first packet (asking for conv)
        //     2. Too many pending packets
        if self.sent_first
            && (self.kcp.wait_snd() >= self.kcp.snd_wnd() as usize
                || self.kcp.wait_snd() >= self.kcp.rmt_wnd() as usize
                || self.kcp.waiting_conv())
        {
            trace!(
                "[SEND] waitsnd={} sndwnd={} rmtwnd={} excceeded or waiting conv={}",
                self.kcp.wait_snd(),
                self.kcp.snd_wnd(),
                self.kcp.rmt_wnd(),
                self.kcp.waiting_conv()
            );

            if let Some(waker) = self.pending_sender.replace(cx.waker().clone()) {
                if !cx.waker().will_wake(&waker) {
                    waker.wake();
                }
            }
            return Poll::Pending;
        }

        if !self.sent_first && self.kcp.waiting_conv() && buf.len() > self.kcp.mss() as usize {
            buf = &buf[..self.kcp.mss() as usize];
        }

        let n = self.kcp.send(buf)?;
        self.sent_first = true;

        if self.kcp.wait_snd() >= self.kcp.snd_wnd() as usize || self.kcp.wait_snd() >= self.kcp.rmt_wnd() as usize {
            self.kcp.flush()?;
        }

        self.last_update = Instant::now();

        if self.flush_write {
            self.kcp.flush()?;
        }

        Ok(n).into()
    }

    /// Call if you want to send some data
    #[allow(dead_code)]
    pub async fn send(&mut self, buf: &[u8]) -> KcpResult<usize> {
        future::poll_fn(|cx| self.poll_send(cx, buf)).await
    }

    #[allow(dead_code)]
    pub fn try_recv(&mut self, buf: &mut [u8]) -> KcpResult<usize> {
        if self.closed {
            return Ok(0);
        }
        self.kcp.recv(buf)
    }

    pub fn poll_recv(&mut self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<KcpResult<usize>> {
        if self.closed {
            return Ok(0).into();
        }

        match self.kcp.recv(buf) {
            Ok(0) | Err(KcpError::RecvQueueEmpty) | Err(KcpError::ExpectingFragment) => {
                if let Some(waker) = self.pending_receiver.replace(cx.waker().clone()) {
                    if !cx.waker().will_wake(&waker) {
                        waker.wake();
                    }
                }
                Poll::Pending
            }
            Err(err) => Err(err).into(),
            Ok(n) => Ok(n).into(),
        }
    }

    #[allow(dead_code)]
    pub async fn recv(&mut self, buf: &mut [u8]) -> KcpResult<usize> {
        future::poll_fn(|cx| self.poll_recv(cx, buf)).await
    }

    pub fn flush(&mut self) -> KcpResult<()> {
        self.kcp.flush()?;
        self.last_update = Instant::now();
        Ok(())
    }

    fn try_wake_pending_waker(&mut self) -> bool {
        let mut waked = false;

        if self.pending_sender.is_some()
            && self.kcp.wait_snd() < self.kcp.snd_wnd() as usize
            && self.kcp.wait_snd() < self.kcp.rmt_wnd() as usize
            && !self.kcp.waiting_conv()
        {
            let waker = self.pending_sender.take().unwrap();
            waker.wake();

            waked = true;
        }

        if self.pending_receiver.is_some() {
            if let Ok(peek) = self.kcp.peeksize() {
                if peek > 0 {
                    let waker = self.pending_receiver.take().unwrap();
                    waker.wake();

                    waked = true;
                }
            }
        }

        waked
    }

    pub fn update(&mut self) -> KcpResult<Instant> {
        let now = now_millis();
        self.kcp.update(now)?;
        let next = self.kcp.check(now);

        self.try_wake_pending_waker();

        Ok(Instant::now() + Duration::from_millis(next as u64))
    }

    pub fn close(&mut self) {
        self.closed = true;
        if let Some(w) = self.pending_sender.take() {
            w.wake();
        }
        if let Some(w) = self.pending_receiver.take() {
            w.wake();
        }
    }

    pub fn udp_socket(&self) -> &Arc<UdpSocket> {
        &self.socket
    }

    pub fn can_close(&self) -> bool {
        self.kcp.wait_snd() == 0
    }

    pub fn conv(&self) -> u32 {
        self.kcp.conv()
    }

    pub fn set_conv(&mut self, conv: u32) {
        self.kcp.set_conv(conv);
    }

    pub fn waiting_conv(&self) -> bool {
        self.kcp.waiting_conv()
    }

    pub fn peek_size(&self) -> KcpResult<usize> {
        self.kcp.peeksize()
    }

    pub fn last_update_time(&self) -> Instant {
        self.last_update
    }
}

#[cfg(test)]
mod test {

    use kcp::Error as KcpError;
    use log::trace;
    use std::sync::Arc;
    use tokio::{
        net::UdpSocket,
        sync::Mutex,
        time::{self, Instant},
    };

    use super::KcpSocket;
    use crate::config::KcpConfig;

    #[tokio::test]
    async fn kcp_echo() {
        let _ = env_logger::try_init();

        static CONV: u32 = 0xdeadbeef;

        // s1 connects s2
        let s1 = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let s2 = UdpSocket::bind("127.0.0.1:0").await.unwrap();

        let s1_addr = s1.local_addr().unwrap();
        let s2_addr = s2.local_addr().unwrap();

        let s1 = Arc::new(s1);
        let s2 = Arc::new(s2);

        let config = KcpConfig::default();
        let kcp1 = KcpSocket::new(&config, 0, s1.clone(), s2_addr, true).unwrap();
        let kcp2 = KcpSocket::new(&config, CONV, s2.clone(), s1_addr, true).unwrap();

        let kcp1 = Arc::new(Mutex::new(kcp1));
        let kcp2 = Arc::new(Mutex::new(kcp2));

        let kcp1_task = {
            let kcp1 = kcp1.clone();
            tokio::spawn(async move {
                loop {
                    let mut kcp = kcp1.lock().await;
                    let next = kcp.update().expect("update");
                    trace!("kcp1 next tick {:?}", next);
                    time::sleep_until(Instant::from_std(next)).await;
                }
            })
        };

        let kcp2_task = {
            let kcp2 = kcp2.clone();
            tokio::spawn(async move {
                loop {
                    let mut kcp = kcp2.lock().await;
                    let next = kcp.update().expect("update");
                    trace!("kcp2 next tick {:?}", next);
                    time::sleep_until(Instant::from_std(next)).await;
                }
            })
        };

        const SEND_BUFFER: &[u8] = b"HELLO WORLD";

        {
            let n = kcp1.lock().await.send(SEND_BUFFER).await.unwrap();
            assert_eq!(n, SEND_BUFFER.len());
        }

        let echo_task = tokio::spawn(async move {
            let mut buf = [0u8; 1024];

            loop {
                let n = s2.recv(&mut buf).await.unwrap();

                let packet = &mut buf[..n];

                let conv = kcp::get_conv(packet);
                if conv == 0 {
                    kcp::set_conv(packet, CONV);
                }

                let mut kcp2 = kcp2.lock().await;
                kcp2.input(packet).unwrap();

                match kcp2.try_recv(&mut buf) {
                    Ok(n) => {
                        let received = &buf[..n];
                        kcp2.send(received).await.unwrap();
                    }
                    Err(KcpError::RecvQueueEmpty) => {
                        continue;
                    }
                    Err(err) => {
                        panic!("kcp.recv error: {:?}", err);
                    }
                }
            }
        });

        {
            let mut buf = [0u8; 1024];

            loop {
                let n = s1.recv(&mut buf).await.unwrap();

                let packet = &buf[..n];

                let mut kcp1 = kcp1.lock().await;
                kcp1.input(packet).unwrap();

                match kcp1.try_recv(&mut buf) {
                    Ok(n) => {
                        let received = &buf[..n];
                        assert_eq!(received, SEND_BUFFER);
                        break;
                    }
                    Err(KcpError::RecvQueueEmpty) => {
                        continue;
                    }
                    Err(err) => {
                        panic!("kcp.recv error: {:?}", err);
                    }
                }
            }
        }

        echo_task.abort();
        kcp1_task.abort();
        kcp2_task.abort();
    }
}
<|endoftext|>