<commit_before>use std::collections::HashMap;
use std::fmt;
use std::error;
use rustc_serialize;

use super::{Data, StrVal, Bool, VecVal, Map, OptVal};
use self::Error::*;

#[derive(Default)]
pub struct Encoder {
     data: Vec<Data>,
}

impl Encoder {
    pub fn new() -> Encoder {
        Encoder::default()
    }

    pub fn stack(&self) -> &[Data] {
        &self.data
    }
}

/// Error type to represent encoding failure.
///
/// This type is not intended to be matched exhaustively as new variants
/// may be added in future without a version bump.
#[derive(Debug)]
pub enum Error {
    NestedOptions,
    UnsupportedType,
    MissingElements,
    KeyIsNotString,
    NoDataToEncode,
    MultipleRootsFound,

    #[doc(hidden)]
    __Nonexhaustive,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use std::error::Error;
        self.description().fmt(f)
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            NestedOptions => "nested Option types are not supported",
            UnsupportedType => "unsupported type",
            MissingElements => "no elements in value",
            KeyIsNotString => "key is not a string",
            NoDataToEncode => "the encodable type created no data",
            MultipleRootsFound => {
                "the encodable type emitted data that was not tree-like in structure"
            }
            Error::__Nonexhaustive => unreachable!(),
        }
    }
}

pub type EncoderResult = Result<(), Error>;

impl Encoder {
    fn push(&mut self, data: Data) {
        self.data.push(data);
    }

    fn push_ok(&mut self, data: Data) -> EncoderResult {
        self.push(data);
        Ok(())
    }

    fn push_str_ok<T: ToString>(&mut self, data: T) -> EncoderResult {
        self.push(StrVal(data.to_string()));
        Ok(())
    }
}

impl rustc_serialize::Encoder for Encoder {
    type Error = Error;

    fn emit_nil(&mut self) -> EncoderResult {
        Err(UnsupportedType)
    }

    fn emit_isize(&mut self, v: isize) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_usize(&mut self, v: usize) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_u64(&mut self, v: u64) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_u32(&mut self, v: u32) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_u16(&mut self, v: u16) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_u8(&mut self, v: u8) -> EncoderResult {
        self.push_str_ok(v)
    }

    fn emit_i64(&mut self, v: i64) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_i32(&mut self, v: i32) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_i16(&mut self, v: i16) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_i8(&mut self, v: i8) -> EncoderResult {
        self.push_str_ok(v)
    }

    fn emit_bool(&mut self, v: bool) -> EncoderResult {
        self.push_ok(Bool(v))
    }

    fn emit_f64(&mut self, v: f64) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_f32(&mut self, v: f32) -> EncoderResult {
        self.push_str_ok(v)
    }

    fn emit_char(&mut self, v: char) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_str(&mut self, v: &str) -> EncoderResult {
        self.push_str_ok(v)
    }

    fn emit_enum<F>(&mut self, _name: &str, _f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        Err(UnsupportedType)
    }

    fn emit_enum_variant<F>(&mut self, _name: &str, _id: usize, _len: usize, _f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        Err(UnsupportedType)
    }

    fn emit_enum_variant_arg<F>(&mut self, _a_idx: usize, _f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        Err(UnsupportedType)
    }

    fn emit_enum_struct_variant<F>(&mut self,
                                   _v_name: &str,
                                   _v_id: usize,
                                   _len: usize,
                                   _f: F)
                                   -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        Err(UnsupportedType)
    }

    fn emit_enum_struct_variant_field<F>(&mut self,
                                         _f_name: &str,
                                         _f_idx: usize,
                                         _f: F)
                                         -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        Err(UnsupportedType)
    }

    fn emit_struct<F>(&mut self, _name: &str, _len: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.push(Map(HashMap::new()));
        f(self)
    }

    fn emit_struct_field<F>(&mut self, name: &str, _idx: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        let mut m = match self.data.pop() {
            Some(Map(m)) => m,
            _ => {
                return Err(UnsupportedType);
            }
        };
        try!(f(self));
        let data = match self.data.pop() {
            Some(d) => d,
            _ => {
                return Err(UnsupportedType);
            }
        };
        m.insert(name.to_string(), data);
        self.push_ok(Map(m))
    }

    fn emit_tuple<F>(&mut self, len: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.emit_seq(len, f)
    }

    fn emit_tuple_arg<F>(&mut self, idx: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.emit_seq_elt(idx, f)
    }

    fn emit_tuple_struct<F>(&mut self, _name: &str, len: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.emit_seq(len, f)
    }

    fn emit_tuple_struct_arg<F>(&mut self, idx: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.emit_seq_elt(idx, f)
    }

    // Specialized types:
    fn emit_option<F>(&mut self, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        f(self)
    }

    fn emit_option_none(&mut self) -> EncoderResult {
        self.push_ok(OptVal(None))
    }

    fn emit_option_some<F>(&mut self, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        try!(f(self));
        let val = match self.data.pop() {
            Some(OptVal(_)) => return Err(NestedOptions),
            Some(d) => d,
            _ => {
                return Err(UnsupportedType);
            }
        };
        self.push_ok(OptVal(Some(Box::new(val))))
    }

    fn emit_seq<F>(&mut self, _len: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.push(VecVal(Vec::new()));
        f(self)
    }

    fn emit_seq_elt<F>(&mut self, _idx: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        let mut v = match self.data.pop() {
            Some(VecVal(v)) => v,
            _ => {
                return Err(UnsupportedType);
            }
        };
        try!(f(self));
        let data = match self.data.pop() {
            Some(d) => d,
            _ => {
                return Err(UnsupportedType);
            }
        };
        v.push(data);
        self.push_ok(VecVal(v))
    }

    fn emit_map<F>(&mut self, _len: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.push(Map(HashMap::new()));
        f(self)
    }

    fn emit_map_elt_key<F>(&mut self, _idx: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        try!(f(self));
        let last = match self.data.last() {
            Some(d) => d,
            None => {
                return Err(MissingElements);
            }
        };
        match *last {
            StrVal(_) => Ok(()),
            _ => Err(KeyIsNotString),
        }
    }

    fn emit_map_elt_val<F>(&mut self, _idx: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        let k = match self.data.pop() {
            Some(StrVal(s)) => s,
            _ => {
                return Err(KeyIsNotString);
            }
        };
        let mut m = match self.data.pop() {
            Some(Map(m)) => m,
            _ => bug!("Expected a map"),
        };
        try!(f(self));
        let popped = match self.data.pop() {
            Some(p) => p,
            None => bug!("Nothing to pop!"),
        };
        m.insert(k, popped);
        self.push(Map(m));
        Ok(())
    }
}

pub fn encode<T: rustc_serialize::Encodable>(data: &T) -> Result<Data, Error> {
    let mut encoder = Encoder::new();
    try!(data.encode(&mut encoder));
    match encoder.data.len() {
        1 => Ok(encoder.data.pop().unwrap()),
        // These two errors are really down to either the type implementing Encodable wrong
        // or using some data which is incompatible with the way mustache *currently* works.
        // It does feel like `bug!` more than Error but maybe it covers some usecases so
        // lets not panic here.
        0 => Err(NoDataToEncode),
        _ => Err(MultipleRootsFound),
    }
}

#[cfg(test)]
mod tests {
    use super::{encode, Error};
    use rustc_serialize::{Encodable, Encoder};

    struct NoData;

    impl Encodable for NoData {
        fn encode<S: Encoder>(&self, _: &mut S) -> Result<(), S::Error> {
            Ok(())
        }
    }

    struct Multiroot;

    impl Encodable for Multiroot {
        fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
            try!(s.emit_u32(4));
            s.emit_u32(1)
        }
    }

    #[test]
    fn encodable_with_no_data() {
        assert_let!(Err(Error::NoDataToEncode) = encode(&NoData));
    }

    #[test]
    fn encodable_with_multiple_roots() {
        assert_let!(Err(Error::MultipleRootsFound) = encode(&Multiroot));
    }
}
<commit_msg>Don't import the Error variants into the namespace<commit_after>use std::collections::HashMap;
use std::fmt;
use std::error;
use rustc_serialize;

use super::{Data, StrVal, Bool, VecVal, Map, OptVal};

#[derive(Default)]
pub struct Encoder {
     data: Vec<Data>,
}

impl Encoder {
    pub fn new() -> Encoder {
        Encoder::default()
    }

    pub fn stack(&self) -> &[Data] {
        &self.data
    }
}

/// Error type to represent encoding failure.
///
/// This type is not intended to be matched exhaustively as new variants
/// may be added in future without a version bump.
#[derive(Debug)]
pub enum Error {
    NestedOptions,
    UnsupportedType,
    MissingElements,
    KeyIsNotString,
    NoDataToEncode,
    MultipleRootsFound,

    #[doc(hidden)]
    __Nonexhaustive,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use std::error::Error;
        self.description().fmt(f)
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::NestedOptions => "nested Option types are not supported",
            Error::UnsupportedType => "unsupported type",
            Error::MissingElements => "no elements in value",
            Error::KeyIsNotString => "key is not a string",
            Error::NoDataToEncode => "the encodable type created no data",
            Error::MultipleRootsFound => {
                "the encodable type emitted data that was not tree-like in structure"
            }
            Error::__Nonexhaustive => unreachable!(),
        }
    }
}

pub type EncoderResult = Result<(), Error>;

impl Encoder {
    fn push(&mut self, data: Data) {
        self.data.push(data);
    }

    fn push_ok(&mut self, data: Data) -> EncoderResult {
        self.push(data);
        Ok(())
    }

    fn push_str_ok<T: ToString>(&mut self, data: T) -> EncoderResult {
        self.push(StrVal(data.to_string()));
        Ok(())
    }
}

impl rustc_serialize::Encoder for Encoder {
    type Error = Error;

    fn emit_nil(&mut self) -> EncoderResult {
        Err(Error::UnsupportedType)
    }

    fn emit_isize(&mut self, v: isize) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_usize(&mut self, v: usize) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_u64(&mut self, v: u64) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_u32(&mut self, v: u32) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_u16(&mut self, v: u16) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_u8(&mut self, v: u8) -> EncoderResult {
        self.push_str_ok(v)
    }

    fn emit_i64(&mut self, v: i64) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_i32(&mut self, v: i32) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_i16(&mut self, v: i16) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_i8(&mut self, v: i8) -> EncoderResult {
        self.push_str_ok(v)
    }

    fn emit_bool(&mut self, v: bool) -> EncoderResult {
        self.push_ok(Bool(v))
    }

    fn emit_f64(&mut self, v: f64) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_f32(&mut self, v: f32) -> EncoderResult {
        self.push_str_ok(v)
    }

    fn emit_char(&mut self, v: char) -> EncoderResult {
        self.push_str_ok(v)
    }
    fn emit_str(&mut self, v: &str) -> EncoderResult {
        self.push_str_ok(v)
    }

    fn emit_enum<F>(&mut self, _name: &str, _f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        Err(Error::UnsupportedType)
    }

    fn emit_enum_variant<F>(&mut self, _name: &str, _id: usize, _len: usize, _f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        Err(Error::UnsupportedType)
    }

    fn emit_enum_variant_arg<F>(&mut self, _a_idx: usize, _f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        Err(Error::UnsupportedType)
    }

    fn emit_enum_struct_variant<F>(&mut self,
                                   _v_name: &str,
                                   _v_id: usize,
                                   _len: usize,
                                   _f: F)
                                   -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        Err(Error::UnsupportedType)
    }

    fn emit_enum_struct_variant_field<F>(&mut self,
                                         _f_name: &str,
                                         _f_idx: usize,
                                         _f: F)
                                         -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        Err(Error::UnsupportedType)
    }

    fn emit_struct<F>(&mut self, _name: &str, _len: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.push(Map(HashMap::new()));
        f(self)
    }

    fn emit_struct_field<F>(&mut self, name: &str, _idx: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        let mut m = match self.data.pop() {
            Some(Map(m)) => m,
            _ => {
                return Err(Error::UnsupportedType);
            }
        };
        try!(f(self));
        let data = match self.data.pop() {
            Some(d) => d,
            _ => {
                return Err(Error::UnsupportedType);
            }
        };
        m.insert(name.to_string(), data);
        self.push_ok(Map(m))
    }

    fn emit_tuple<F>(&mut self, len: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.emit_seq(len, f)
    }

    fn emit_tuple_arg<F>(&mut self, idx: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.emit_seq_elt(idx, f)
    }

    fn emit_tuple_struct<F>(&mut self, _name: &str, len: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.emit_seq(len, f)
    }

    fn emit_tuple_struct_arg<F>(&mut self, idx: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.emit_seq_elt(idx, f)
    }

    // Specialized types:
    fn emit_option<F>(&mut self, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        f(self)
    }

    fn emit_option_none(&mut self) -> EncoderResult {
        self.push_ok(OptVal(None))
    }

    fn emit_option_some<F>(&mut self, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        try!(f(self));
        let val = match self.data.pop() {
            Some(OptVal(_)) => return Err(Error::NestedOptions),
            Some(d) => d,
            _ => {
                return Err(Error::UnsupportedType);
            }
        };
        self.push_ok(OptVal(Some(Box::new(val))))
    }

    fn emit_seq<F>(&mut self, _len: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.push(VecVal(Vec::new()));
        f(self)
    }

    fn emit_seq_elt<F>(&mut self, _idx: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        let mut v = match self.data.pop() {
            Some(VecVal(v)) => v,
            _ => {
                return Err(Error::UnsupportedType);
            }
        };
        try!(f(self));
        let data = match self.data.pop() {
            Some(d) => d,
            _ => {
                return Err(Error::UnsupportedType);
            }
        };
        v.push(data);
        self.push_ok(VecVal(v))
    }

    fn emit_map<F>(&mut self, _len: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        self.push(Map(HashMap::new()));
        f(self)
    }

    fn emit_map_elt_key<F>(&mut self, _idx: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        try!(f(self));
        let last = match self.data.last() {
            Some(d) => d,
            None => {
                return Err(Error::MissingElements);
            }
        };
        match *last {
            StrVal(_) => Ok(()),
            _ => Err(Error::KeyIsNotString),
        }
    }

    fn emit_map_elt_val<F>(&mut self, _idx: usize, f: F) -> EncoderResult
    where F: FnOnce(&mut Encoder) -> EncoderResult
    {
        let k = match self.data.pop() {
            Some(StrVal(s)) => s,
            _ => {
                return Err(Error::KeyIsNotString);
            }
        };
        let mut m = match self.data.pop() {
            Some(Map(m)) => m,
            _ => bug!("Expected a map"),
        };
        try!(f(self));
        let popped = match self.data.pop() {
            Some(p) => p,
            None => bug!("Nothing to pop!"),
        };
        m.insert(k, popped);
        self.push(Map(m));
        Ok(())
    }
}

pub fn encode<T: rustc_serialize::Encodable>(data: &T) -> Result<Data, Error> {
    let mut encoder = Encoder::new();
    try!(data.encode(&mut encoder));
    match encoder.data.len() {
        1 => Ok(encoder.data.pop().unwrap()),
        // These two errors are really down to either the type implementing Encodable wrong
        // or using some data which is incompatible with the way mustache *currently* works.
        // It does feel like `bug!` more than Error but maybe it covers some usecases so
        // lets not panic here.
        0 => Err(Error::NoDataToEncode),
        _ => Err(Error::MultipleRootsFound),
    }
}

#[cfg(test)]
mod tests {
    use super::{encode, Error};
    use rustc_serialize::{Encodable, Encoder};

    struct NoData;

    impl Encodable for NoData {
        fn encode<S: Encoder>(&self, _: &mut S) -> Result<(), S::Error> {
            Ok(())
        }
    }

    struct Multiroot;

    impl Encodable for Multiroot {
        fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
            try!(s.emit_u32(4));
            s.emit_u32(1)
        }
    }

    #[test]
    fn encodable_with_no_data() {
        assert_let!(Err(Error::NoDataToEncode) = encode(&NoData));
    }

    #[test]
    fn encodable_with_multiple_roots() {
        assert_let!(Err(Error::MultipleRootsFound) = encode(&Multiroot));
    }
}
<|endoftext|>


<commit_before>extern crate crypto;
extern crate base64;
extern crate byteorder;
extern crate docopt;
extern crate rustc_serialize;
extern crate rpassword;
extern crate ring;
extern crate untrusted;
#[macro_use]
extern crate error_chain;

use std::process;
use std::mem;
use std::io::prelude::*;
use std::io::BufReader;
use std::fs::{OpenOptions, File};
use std::path::Path;

use ring::rand::SystemRandom;
use ring::signature::Ed25519KeyPair;
use ring::digest;
use crypto::bcrypt_pbkdf::bcrypt_pbkdf;

use docopt::Docopt;

mod structs;
mod errors;

use structs::*;
use errors::*;

const USAGE: &'static str = "
signify-rs

Usage:
  signify -h
  signify -G [-n] [-c <comment>] -p <pubkey> -s <seckey>
  signify -S [-x <sigfile>] -s <seckey> -m <message>
  signify -V [-x <sigfile>] -p <pubkey> -m <message>

Options:
  -h --help     Show this screen.
  -c <comment>  Specify the comment to be added during key generation.
  -m <message>  When signing, the file containing the message to sign.  When verifying, the file containing the
                message to verify.  When verifying with -e, the file to create.
  -n            Do not ask for a passphrase during key generation. Otherwise, signify will prompt the user for a
                passphrase to protect the secret key.
  -p <pubkey>   Public key produced by -G, and used by -V to check a signature.
  -s <seckey>   Secret (private) key produced by -G, and used by -S to sign a message.
  -x <sigfile>  The signature file to create or verify.  The default is <message>.sig.
";

#[allow(non_snake_case)]
#[derive(Debug, RustcDecodable)]
struct Args {
    flag_G: bool,
    flag_S: bool,
    flag_V: bool,

    flag_x: Option<String>,
    flag_c: Option<String>,

    flag_p: String,
    flag_s: String,
    flag_m: String,
    flag_n: bool,
}

enum FileContent {
    PublicKey(PublicKey),
    PrivateKey(PrivateKey),
    Signature(Signature),
}

fn write_base64_file<P: AsRef<Path>>(file: P, comment: &str, buf: &[u8]) -> Result<()> {
    let mut f = try!(OpenOptions::new().write(true).create_new(true).open(file));

    try!(write!(f, "{}", COMMENTHDR));
    try!(write!(f, "{}\n", comment));
    let out = base64::encode(buf);
    try!(write!(f, "{}\n", out));

    Ok(())
}

fn read_base64_file<P: AsRef<Path>>(file: P) -> Result<FileContent> {
    let file_display = format!("{}", file.as_ref().display());
    let f = try!(File::open(file));
    let mut reader = BufReader::new(f);

    let mut comment_line = String::new();
    let len = try!(reader.read_line(&mut comment_line));

    if len == 0 || len < COMMENTHDRLEN || !comment_line.starts_with(COMMENTHDR) {
        return Err(format!("invalid comment in {}; must start with '{}'", file_display, COMMENTHDR).into());
    }

    if &comment_line[len-1..len] != "\n" {
        return Err(format!("missing new line after comment in {}", file_display).into());
    }

    if len > COMMENTHDRLEN + COMMENTMAXLEN {
        return Err("comment too long".into());
    }

    let mut base64_line = String::new();
    let len = try!(reader.read_line(&mut base64_line));

    if len == 0 {
        return Err(format!("missing line in {}", file_display).into());
    }

    if &base64_line[len-1..len] != "\n" {
        return Err(format!("missing new line after comment in {}", file_display).into());
    }

    let base64_line = &base64_line[0..len-1];

    let data = try!(base64::decode(base64_line)
                    .chain_err(|| format!("invalid base64 encoding in {}", file_display)));

    if &data[0..2] != PKGALG {
        return Err(format!("unsupported file {}", file_display).into());
    }

    match data.len() {
        x if x == mem::size_of::<PublicKey>() => {
            PublicKey::from_buf(&data)
                .map(FileContent::PublicKey)
        }
        x if x == mem::size_of::<PrivateKey>() => {
            PrivateKey::from_buf(&data)
                .map(FileContent::PrivateKey)
        }
        x if x == mem::size_of::<Signature>() => {
            Signature::from_buf(&data)
                .map(FileContent::Signature)
        },
        _ => {
            Err(format!("unsupported file {}", file_display).into())
        }
    }
}

fn verify(pubkey_path: String, msg_path: String, signature_path: Option<String>) -> Result<()> {
    let pkey = match read_base64_file(&pubkey_path) {
        Ok(FileContent::PublicKey(pkey)) => pkey,
        _ => return Err("an error occured.".into()),
    };

    let signature_path = match signature_path {
        Some(path) => path,
        None => format!("{}.sig", msg_path)
    };

    let signature = match read_base64_file(&signature_path) {
        Ok(FileContent::Signature(sig)) => sig,
        _ => return Err(format!("Can't read signature from {}", signature_path).into()),
    };

    let mut msgfile = try!(File::open(&msg_path).chain_err(|| read_error(&msg_path)));
    let mut msg = vec![];
    try!(msgfile.read_to_end(&mut msg).chain_err(|| read_error(&msg_path)));

    if signature.keynum != pkey.keynum {
        return Err("signature verification failed: checked against wrong key".into());
    }

    if signature.verify(&msg, &pkey) {
        println!("Signature Verified");
        Ok(())
    } else {
        Err("signature verification failed".into())
    }
}

fn sign(seckey_path: String, msg_path: String, signature_path: Option<String>) -> Result<()> {
    let mut skey = match read_base64_file(&seckey_path) {
        Ok(FileContent::PrivateKey(skey)) => skey,
        _ => return Err("an error occured.".into()),
    };

    let rounds = skey.kdfrounds;
    let xorkey = try!(kdf(&skey.salt, rounds, false, SECRETBYTES));

    for (prv, xor) in skey.seckey.iter_mut().zip(xorkey.iter()) {
        *prv = *prv ^ xor;
    }
    let skey = skey;

    let mut msgfile = try!(File::open(&msg_path)
                           .chain_err(|| read_error(&msg_path)).into());
    let mut msg = vec![];
    try!(msgfile.read_to_end(&mut msg).chain_err(|| read_error(&msg_path)));

    let signature_path = match signature_path {
        Some(path) => path,
        None => format!("{}.sig", msg_path)
    };

    let sig = try!(skey.sign(&msg).chain_err(|| "Failed to sign message"));

    let mut out = vec![];
    try!(sig.write(&mut out).chain_err(|| "Can't write to internal buffer"));

    let sig_comment = "signature from signify secret key";

    write_base64_file(&signature_path, sig_comment, &out)
        .chain_err(|| "Failed to write signature file")
}

fn read_password(prompt: &str) -> Result<String> {
    let mut stdout = std::io::stdout();
    try!(stdout.write_all(prompt.as_bytes()).chain_err(|| "Write to stdout failed"));
    try!(stdout.flush().chain_err(|| "Flushing stdout failed"));

    rpassword::read_password()
        .chain_err(|| "unable to read passphrase")
}

fn kdf(salt: &[u8], rounds: u32, confirm: bool, keylen: usize) -> Result<Vec<u8>> {
    let mut result = vec![0; keylen];
    if rounds == 0 {
        return Ok(result);
    }

    let passphrase = try!(read_password("passphrase: "));

    if confirm {
        let confirm_passphrase = try!(read_password("confirm passphrase: "));

        if passphrase != confirm_passphrase {
            return Err("passwords don't match".into());
        }
    }

    bcrypt_pbkdf(passphrase.as_bytes(), salt, rounds, &mut result);
    Ok(result)
}

fn generate(pubkey_path: String, privkey_path: String, comment: Option<String>, kdfrounds: u32) -> Result<()> {
    let comment = match comment {
        Some(s) => s,
        None    => "signify".into()
    };

    let mut keynum = [0; KEYNUMLEN];
    try!(SystemRandom.fill(&mut keynum).chain_err(|| "Can't fill keynum randomly"));

    let (_, keypair_bytes) = try!(Ed25519KeyPair::generate_serializable(&SystemRandom));
    let mut skey = keypair_bytes.private_key;
    let pkey = keypair_bytes.public_key;

    let mut salt = [0; 16];
    try!(SystemRandom.fill(&mut salt).chain_err(|| "Can't fill salt randomly"));

    let xorkey = try!(kdf(&salt, kdfrounds, true, SECRETBYTES));

    for (prv, xor) in skey.iter_mut().zip(xorkey.iter()) {
        *prv = *prv ^ xor;
    }

    // signify stores the extended key as the private key,
    // that is the 32 byte of the secret key, followed by the 32 byte of the public key,
    // summing up to 64 byte.
    //
    //  *ring* separates them, so we need to stick them together again.
    let mut complete_key = [0; 64];
    complete_key[0..32].copy_from_slice(&skey[0..32]);
    complete_key[32..].copy_from_slice(&pkey);

    // Store private key
    let digest = digest::digest(&digest::SHA512, &complete_key);
    let mut checksum = [0; 8];
    checksum.copy_from_slice(&digest.as_ref()[0..8]);

    let private_key = PrivateKey {
        pkgalg: PKGALG,
        kdfalg: KDFALG,
        kdfrounds: kdfrounds,
        salt: salt,
        checksum: checksum,
        keynum: keynum,
        seckey: complete_key,
    };

    let mut out = vec![];
    try!(private_key.write(&mut out).chain_err(|| "Can't write to internal buffer"));

    let priv_comment = format!("{} secret key", comment);
    try!(write_base64_file(&privkey_path, &priv_comment, &out)
        .chain_err(|| write_error(&privkey_path))
    );

    // Store public key
    let public_key = PublicKey::with_key_and_keynum(pkey, keynum);

    let mut out = vec![];
    try!(public_key.write(&mut out).chain_err(|| "Can't write to internal buffer"));

    let pub_comment = format!("{} public key", comment);
    write_base64_file(&pubkey_path, &pub_comment, &out)
        .chain_err(|| write_error(&pubkey_path))
}

fn human(res: Result<()>) {
    match res {
        Err(e) => {
            let mut it = e.iter();
            let e = it.next().unwrap();
            let next_err = it.next();
            if let Some(next_err) = next_err {
                println!("{}: {}", e, next_err);
            } else {
                println!("{}", e);
            }

            process::exit(1);
        },
        Ok(()) => {}
    }
}

fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    if args.flag_V {
        human(verify(args.flag_p, args.flag_m, args.flag_x));
    } else if args.flag_G {
        let rounds = if args.flag_n { 0 } else { 42 };
        human(generate(args.flag_p, args.flag_s, args.flag_c, rounds));
    } else if args.flag_S {
        human(sign(args.flag_s, args.flag_m, args.flag_x));
    }
}
<commit_msg>refactor: Note why unwrap is safe<commit_after>extern crate crypto;
extern crate base64;
extern crate byteorder;
extern crate docopt;
extern crate rustc_serialize;
extern crate rpassword;
extern crate ring;
extern crate untrusted;
#[macro_use]
extern crate error_chain;

use std::process;
use std::mem;
use std::io::prelude::*;
use std::io::BufReader;
use std::fs::{OpenOptions, File};
use std::path::Path;

use ring::rand::SystemRandom;
use ring::signature::Ed25519KeyPair;
use ring::digest;
use crypto::bcrypt_pbkdf::bcrypt_pbkdf;

use docopt::Docopt;

mod structs;
mod errors;

use structs::*;
use errors::*;

const USAGE: &'static str = "
signify-rs

Usage:
  signify -h
  signify -G [-n] [-c <comment>] -p <pubkey> -s <seckey>
  signify -S [-x <sigfile>] -s <seckey> -m <message>
  signify -V [-x <sigfile>] -p <pubkey> -m <message>

Options:
  -h --help     Show this screen.
  -c <comment>  Specify the comment to be added during key generation.
  -m <message>  When signing, the file containing the message to sign.  When verifying, the file containing the
                message to verify.  When verifying with -e, the file to create.
  -n            Do not ask for a passphrase during key generation. Otherwise, signify will prompt the user for a
                passphrase to protect the secret key.
  -p <pubkey>   Public key produced by -G, and used by -V to check a signature.
  -s <seckey>   Secret (private) key produced by -G, and used by -S to sign a message.
  -x <sigfile>  The signature file to create or verify.  The default is <message>.sig.
";

#[allow(non_snake_case)]
#[derive(Debug, RustcDecodable)]
struct Args {
    flag_G: bool,
    flag_S: bool,
    flag_V: bool,

    flag_x: Option<String>,
    flag_c: Option<String>,

    flag_p: String,
    flag_s: String,
    flag_m: String,
    flag_n: bool,
}

enum FileContent {
    PublicKey(PublicKey),
    PrivateKey(PrivateKey),
    Signature(Signature),
}

fn write_base64_file<P: AsRef<Path>>(file: P, comment: &str, buf: &[u8]) -> Result<()> {
    let mut f = try!(OpenOptions::new().write(true).create_new(true).open(file));

    try!(write!(f, "{}", COMMENTHDR));
    try!(write!(f, "{}\n", comment));
    let out = base64::encode(buf);
    try!(write!(f, "{}\n", out));

    Ok(())
}

fn read_base64_file<P: AsRef<Path>>(file: P) -> Result<FileContent> {
    let file_display = format!("{}", file.as_ref().display());
    let f = try!(File::open(file));
    let mut reader = BufReader::new(f);

    let mut comment_line = String::new();
    let len = try!(reader.read_line(&mut comment_line));

    if len == 0 || len < COMMENTHDRLEN || !comment_line.starts_with(COMMENTHDR) {
        return Err(format!("invalid comment in {}; must start with '{}'", file_display, COMMENTHDR).into());
    }

    if &comment_line[len-1..len] != "\n" {
        return Err(format!("missing new line after comment in {}", file_display).into());
    }

    if len > COMMENTHDRLEN + COMMENTMAXLEN {
        return Err("comment too long".into());
    }

    let mut base64_line = String::new();
    let len = try!(reader.read_line(&mut base64_line));

    if len == 0 {
        return Err(format!("missing line in {}", file_display).into());
    }

    if &base64_line[len-1..len] != "\n" {
        return Err(format!("missing new line after comment in {}", file_display).into());
    }

    let base64_line = &base64_line[0..len-1];

    let data = try!(base64::decode(base64_line)
                    .chain_err(|| format!("invalid base64 encoding in {}", file_display)));

    if &data[0..2] != PKGALG {
        return Err(format!("unsupported file {}", file_display).into());
    }

    match data.len() {
        x if x == mem::size_of::<PublicKey>() => {
            PublicKey::from_buf(&data)
                .map(FileContent::PublicKey)
        }
        x if x == mem::size_of::<PrivateKey>() => {
            PrivateKey::from_buf(&data)
                .map(FileContent::PrivateKey)
        }
        x if x == mem::size_of::<Signature>() => {
            Signature::from_buf(&data)
                .map(FileContent::Signature)
        },
        _ => {
            Err(format!("unsupported file {}", file_display).into())
        }
    }
}

fn verify(pubkey_path: String, msg_path: String, signature_path: Option<String>) -> Result<()> {
    let pkey = match read_base64_file(&pubkey_path) {
        Ok(FileContent::PublicKey(pkey)) => pkey,
        _ => return Err("an error occured.".into()),
    };

    let signature_path = match signature_path {
        Some(path) => path,
        None => format!("{}.sig", msg_path)
    };

    let signature = match read_base64_file(&signature_path) {
        Ok(FileContent::Signature(sig)) => sig,
        _ => return Err(format!("Can't read signature from {}", signature_path).into()),
    };

    let mut msgfile = try!(File::open(&msg_path).chain_err(|| read_error(&msg_path)));
    let mut msg = vec![];
    try!(msgfile.read_to_end(&mut msg).chain_err(|| read_error(&msg_path)));

    if signature.keynum != pkey.keynum {
        return Err("signature verification failed: checked against wrong key".into());
    }

    if signature.verify(&msg, &pkey) {
        println!("Signature Verified");
        Ok(())
    } else {
        Err("signature verification failed".into())
    }
}

fn sign(seckey_path: String, msg_path: String, signature_path: Option<String>) -> Result<()> {
    let mut skey = match read_base64_file(&seckey_path) {
        Ok(FileContent::PrivateKey(skey)) => skey,
        _ => return Err("an error occured.".into()),
    };

    let rounds = skey.kdfrounds;
    let xorkey = try!(kdf(&skey.salt, rounds, false, SECRETBYTES));

    for (prv, xor) in skey.seckey.iter_mut().zip(xorkey.iter()) {
        *prv = *prv ^ xor;
    }
    let skey = skey;

    let mut msgfile = try!(File::open(&msg_path)
                           .chain_err(|| read_error(&msg_path)).into());
    let mut msg = vec![];
    try!(msgfile.read_to_end(&mut msg).chain_err(|| read_error(&msg_path)));

    let signature_path = match signature_path {
        Some(path) => path,
        None => format!("{}.sig", msg_path)
    };

    let sig = try!(skey.sign(&msg).chain_err(|| "Failed to sign message"));

    let mut out = vec![];
    try!(sig.write(&mut out).chain_err(|| "Can't write to internal buffer"));

    let sig_comment = "signature from signify secret key";

    write_base64_file(&signature_path, sig_comment, &out)
        .chain_err(|| "Failed to write signature file")
}

fn read_password(prompt: &str) -> Result<String> {
    let mut stdout = std::io::stdout();
    try!(stdout.write_all(prompt.as_bytes()).chain_err(|| "Write to stdout failed"));
    try!(stdout.flush().chain_err(|| "Flushing stdout failed"));

    rpassword::read_password()
        .chain_err(|| "unable to read passphrase")
}

fn kdf(salt: &[u8], rounds: u32, confirm: bool, keylen: usize) -> Result<Vec<u8>> {
    let mut result = vec![0; keylen];
    if rounds == 0 {
        return Ok(result);
    }

    let passphrase = try!(read_password("passphrase: "));

    if confirm {
        let confirm_passphrase = try!(read_password("confirm passphrase: "));

        if passphrase != confirm_passphrase {
            return Err("passwords don't match".into());
        }
    }

    bcrypt_pbkdf(passphrase.as_bytes(), salt, rounds, &mut result);
    Ok(result)
}

fn generate(pubkey_path: String, privkey_path: String, comment: Option<String>, kdfrounds: u32) -> Result<()> {
    let comment = match comment {
        Some(s) => s,
        None    => "signify".into()
    };

    let mut keynum = [0; KEYNUMLEN];
    try!(SystemRandom.fill(&mut keynum).chain_err(|| "Can't fill keynum randomly"));

    let (_, keypair_bytes) = try!(Ed25519KeyPair::generate_serializable(&SystemRandom));
    let mut skey = keypair_bytes.private_key;
    let pkey = keypair_bytes.public_key;

    let mut salt = [0; 16];
    try!(SystemRandom.fill(&mut salt).chain_err(|| "Can't fill salt randomly"));

    let xorkey = try!(kdf(&salt, kdfrounds, true, SECRETBYTES));

    for (prv, xor) in skey.iter_mut().zip(xorkey.iter()) {
        *prv = *prv ^ xor;
    }

    // signify stores the extended key as the private key,
    // that is the 32 byte of the secret key, followed by the 32 byte of the public key,
    // summing up to 64 byte.
    //
    //  *ring* separates them, so we need to stick them together again.
    let mut complete_key = [0; 64];
    complete_key[0..32].copy_from_slice(&skey[0..32]);
    complete_key[32..].copy_from_slice(&pkey);

    // Store private key
    let digest = digest::digest(&digest::SHA512, &complete_key);
    let mut checksum = [0; 8];
    checksum.copy_from_slice(&digest.as_ref()[0..8]);

    let private_key = PrivateKey {
        pkgalg: PKGALG,
        kdfalg: KDFALG,
        kdfrounds: kdfrounds,
        salt: salt,
        checksum: checksum,
        keynum: keynum,
        seckey: complete_key,
    };

    let mut out = vec![];
    try!(private_key.write(&mut out).chain_err(|| "Can't write to internal buffer"));

    let priv_comment = format!("{} secret key", comment);
    try!(write_base64_file(&privkey_path, &priv_comment, &out)
        .chain_err(|| write_error(&privkey_path))
    );

    // Store public key
    let public_key = PublicKey::with_key_and_keynum(pkey, keynum);

    let mut out = vec![];
    try!(public_key.write(&mut out).chain_err(|| "Can't write to internal buffer"));

    let pub_comment = format!("{} public key", comment);
    write_base64_file(&pubkey_path, &pub_comment, &out)
        .chain_err(|| write_error(&pubkey_path))
}

fn human(res: Result<()>) {
    match res {
        Err(e) => {
            let mut it = e.iter();
            let e = it.next().unwrap(); // We have an error, so we can definitely unwrap it.
            let next_err = it.next();
            if let Some(next_err) = next_err {
                println!("{}: {}", e, next_err);
            } else {
                println!("{}", e);
            }

            process::exit(1);
        },
        Ok(()) => {}
    }
}

fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    if args.flag_V {
        human(verify(args.flag_p, args.flag_m, args.flag_x));
    } else if args.flag_G {
        let rounds = if args.flag_n { 0 } else { 42 };
        human(generate(args.flag_p, args.flag_s, args.flag_c, rounds));
    } else if args.flag_S {
        human(sign(args.flag_s, args.flag_m, args.flag_x));
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]

#[macro_use]
extern crate error_chain;
extern crate fs2;
extern crate libc;
extern crate tempdir;
extern crate walkdir;

use std::hash::{Hash, Hasher};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::{Command, ExitStatus};
use std::{env, mem, process};

mod cargo;
mod errors;
mod flock;
mod fs;
mod io;
mod parse;
mod rustc;
mod sysroot;
mod toml;
mod xargo;

use parse::Args;
use errors::*;

fn main() {
    match run() {
        Err(e) => {
            writeln!(io::stderr(), "{:?}", e).ok();
            process::exit(1);
        }
        Ok(status) => {
            if !status.success() {
                process::exit(status.code().unwrap_or(1))
            }
        }
    }
}

fn run() -> Result<ExitStatus> {
    let Args { mut cmd, target, subcommand, verbose } = try!(parse::args());

    let target = if let Some(target) = target.as_ref() {
        Some(try!(Target::from(target)))
    } else if let Some(target) = try!(parse::target_in_cargo_config()) {
        Some(try!(Target::from(&target)))
    } else {
        None
    };

    let needs_sysroot = match subcommand.as_ref().map(|s| &s[..]) {
        // we need rebuild the sysroot for these subcommands
        None | Some("clean") | Some("init") | Some("new") |
        Some("update") | Some("search") => None,
        _ => {
            if let Some(target) = target.as_ref() {
                if target.needs_sysroot() {
                    try!(sysroot::update(target, verbose));
                    Some(target)
                } else {
                    None
                }
            } else {
                None
            }
        }
    };

    let locks = if let Some(target) = needs_sysroot {
        // TODO(Filesystem.display()) this could be better ...
        let sysroot = format!("{}", try!(xargo::home()).display());

        if subcommand.as_ref().map(|s| &s[..]) == Some("doc") {
            let mut flags = try!(rustc::flags(target, "rustdocflags"));
            flags.push("--sysroot".to_owned());

            flags.push(sysroot.clone());

            cmd.env("RUSTDOCFLAGS", flags.join(" "));
        }

        let mut flags = try!(rustc::flags(target, "rustflags"));
        flags.push("--sysroot".to_owned());

        // TODO(Filesystem.display()) this could be better ...
        flags.push(sysroot);

        cmd.env("RUSTFLAGS", flags.join(" "));

        // Make sure the sysroot is not blown up while the Cargo command is running
        Some((xargo::lock_ro(&try!(rustc::meta()).host),
              xargo::lock_ro(target.triple())))
    } else {
        None
    };

    let status = try!(cmd.status()
        .chain_err(|| "failed to execute `cargo`. Is it not installed?"));

    mem::drop(locks);

    Ok(status)
}

#[derive(Debug)]
pub enum Target {
    BuiltIn { triple: String },
    Custom { triple: String, json: PathBuf },
    Path { triple: String, json: PathBuf },
}

impl Target {
    fn from(target: &str) -> Result<Target> {
        let target_list = try!(rustc::target_list());

        let target = target.to_owned();
        if target_list.iter().any(|t| t == &target) {
            Ok(Target::BuiltIn { triple: target })
        } else if target.ends_with(".json") {
            if let Some(triple) = Path::new(&target)
                .file_stem()
                .and_then(|f| f.to_str()) {
                Ok(Target::Path {
                    json: PathBuf::from(&target),
                    triple: triple.to_owned(),
                })
            } else {
                try!(Err(format!("error extracting triple from {}",
                                 target)))
            }
        } else {
            let json = Path::new(&target).with_extension("json");

            if json.exists() {
                Ok(Target::Custom {
                    json: json,
                    triple: target,
                })
            } else {
                if let Some(target_dir) = env::var_os("RUST_TARGET_PATH") {
                    let json = PathBuf::from(target_dir)
                        .join(&target)
                        .with_extension("json");

                    if json.exists() {
                        return Ok(Target::Custom {
                            json: json,
                            triple: target,
                        });
                    }
                }

                try!(Err(format!("no target specification file found \
                                  for {}",
                                 target)))
            }
        }
    }
}

impl Target {
    fn has_atomics(&self) -> Result<bool> {
        let mut cmd = rustc::rustc();
        cmd.arg("--target");

        match *self {
            Target::BuiltIn { ref triple } |
            Target::Custom { ref triple, .. } => {
                cmd.arg(triple);
            }
            Target::Path { ref json, .. } => {
                cmd.arg(json);
            }
        }

        cmd.args(&["--print", "cfg"]);

        Ok(try!(cmd.run_and_get_stdout()).contains("target_has_atomic"))
    }

    fn hash<H>(&self, hasher: &mut H) -> Result<()>
        where H: Hasher
    {
        match *self {
            Target::BuiltIn { .. } => {}
            Target::Custom { ref json, .. } |
            Target::Path { ref json, .. } => try!(io::read(json)).hash(hasher),
        }

        Ok(())
    }

    fn needs_sysroot(&self) -> bool {
        match *self {
            Target::BuiltIn { ref triple } => {
                match &triple[..] {
                    "thumbv6m-none-eabi" |
                    "thumbv7m-none-eabi" |
                    "thumbv7em-none-eabi" |
                    "thumbv7em-none-eabihf" => true,
                    _ => false,
                }
            }
            _ => true,
        }
    }

    fn triple(&self) -> &str {
        match *self {
            Target::BuiltIn { ref triple } |
            Target::Custom { ref triple, .. } |
            Target::Path { ref triple, .. } => triple,
        }
    }
}

trait CommandExt {
    fn run_and_get_stdout(&mut self) -> Result<String>;
    fn run_or_error(&mut self) -> Result<()>;
}

impl CommandExt for Command {
    fn run_and_get_stdout(&mut self) -> Result<String> {
        let cmd = &format!("`{:?}`", self);

        let output = try!(self.output()
            .chain_err(|| format!("failed to execute {}", cmd)));

        if !output.status.success() {
            try!(Err(format!("{} failed with exit status: {:?}",
                             cmd,
                             output.status.code())))
        }

        Ok(try!(String::from_utf8(output.stdout)
            .chain_err(|| format!("{} output was not UTF-8 encoded", cmd))))
    }

    fn run_or_error(&mut self) -> Result<()> {
        let cmd = &format!("`{:?}`", self);

        let exit = try!(self.status()
            .chain_err(|| format!("failed to execute {}", cmd)));

        if !exit.success() {
            try!(Err(format!("{} failed with exit status {:?}",
                             cmd,
                             exit.code())))
        }

        Ok(())
    }
}
<commit_msg>Auto merge of #75 - japaric:nicer-errors, r=japaric<commit_after>#![allow(dead_code)]

#[macro_use]
extern crate error_chain;
extern crate fs2;
extern crate libc;
extern crate tempdir;
extern crate walkdir;

use std::hash::{Hash, Hasher};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::{Command, ExitStatus};
use std::{env, mem, process};

mod cargo;
mod errors;
mod flock;
mod fs;
mod io;
mod parse;
mod rustc;
mod sysroot;
mod toml;
mod xargo;

use parse::Args;
use errors::*;

fn main() {
    fn show_backtrace() -> bool {
        env::var("RUST_BACKTRACE").as_ref().map(|s| &s[..]) == Ok("1")
    }

    match run() {
        Err(e) => {
            let stderr = io::stderr();
            let mut stderr = stderr.lock();

            writeln!(stderr, "{}", e).ok();

            for e in e.iter().skip(1) {
                writeln!(stderr, "caused by: {}", e).ok();
            }

            if show_backtrace() {
                if let Some(backtrace) = e.backtrace() {
                    writeln!(stderr, "{:?}", backtrace).ok();
                }
            }

            process::exit(1)
        },
        Ok(status) => {
            if !status.success() {
                process::exit(status.code().unwrap_or(1))
            }
        }
    }
}

fn run() -> Result<ExitStatus> {
    let Args { mut cmd, target, subcommand, verbose } = try!(parse::args());

    let target = if let Some(target) = target.as_ref() {
        Some(try!(Target::from(target)))
    } else if let Some(target) = try!(parse::target_in_cargo_config()) {
        Some(try!(Target::from(&target)))
    } else {
        None
    };

    let needs_sysroot = match subcommand.as_ref().map(|s| &s[..]) {
        // we don't need to rebuild the sysroot for these subcommands
        None | Some("clean") | Some("init") | Some("new") |
        Some("update") | Some("search") => None,
        _ => {
            if let Some(target) = target.as_ref() {
                if target.needs_sysroot() {
                    try!(sysroot::update(target, verbose));
                    Some(target)
                } else {
                    None
                }
            } else {
                None
            }
        }
    };

    let locks = if let Some(target) = needs_sysroot {
        // TODO(Filesystem.display()) this could be better ...
        let sysroot = format!("{}", try!(xargo::home()).display());

        if subcommand.as_ref().map(|s| &s[..]) == Some("doc") {
            let mut flags = try!(rustc::flags(target, "rustdocflags"));
            flags.push("--sysroot".to_owned());

            flags.push(sysroot.clone());

            cmd.env("RUSTDOCFLAGS", flags.join(" "));
        }

        let mut flags = try!(rustc::flags(target, "rustflags"));
        flags.push("--sysroot".to_owned());

        // TODO(Filesystem.display()) this could be better ...
        flags.push(sysroot);

        cmd.env("RUSTFLAGS", flags.join(" "));

        // Make sure the sysroot is not blown up while the Cargo command is running
        Some((xargo::lock_ro(&try!(rustc::meta()).host),
              xargo::lock_ro(target.triple())))
    } else {
        None
    };

    let status = try!(cmd.status()
        .chain_err(|| "failed to execute `cargo`. Is it not installed?"));

    mem::drop(locks);

    Ok(status)
}

#[derive(Debug)]
pub enum Target {
    BuiltIn { triple: String },
    Custom { triple: String, json: PathBuf },
    Path { triple: String, json: PathBuf },
}

impl Target {
    fn from(target: &str) -> Result<Target> {
        let target_list = try!(rustc::target_list());

        let target = target.to_owned();
        if target_list.iter().any(|t| t == &target) {
            Ok(Target::BuiltIn { triple: target })
        } else if target.ends_with(".json") {
            if let Some(triple) = Path::new(&target)
                .file_stem()
                .and_then(|f| f.to_str()) {
                Ok(Target::Path {
                    json: PathBuf::from(&target),
                    triple: triple.to_owned(),
                })
            } else {
                try!(Err(format!("error extracting triple from {}",
                                 target)))
            }
        } else {
            let json = Path::new(&target).with_extension("json");

            if json.exists() {
                Ok(Target::Custom {
                    json: json,
                    triple: target,
                })
            } else {
                if let Some(target_dir) = env::var_os("RUST_TARGET_PATH") {
                    let json = PathBuf::from(target_dir)
                        .join(&target)
                        .with_extension("json");

                    if json.exists() {
                        return Ok(Target::Custom {
                            json: json,
                            triple: target,
                        });
                    }
                }

                try!(Err(format!("no target specification file found \
                                  for {}",
                                 target)))
            }
        }
    }
}

impl Target {
    fn has_atomics(&self) -> Result<bool> {
        let mut cmd = rustc::rustc();
        cmd.arg("--target");

        match *self {
            Target::BuiltIn { ref triple } |
            Target::Custom { ref triple, .. } => {
                cmd.arg(triple);
            }
            Target::Path { ref json, .. } => {
                cmd.arg(json);
            }
        }

        cmd.args(&["--print", "cfg"]);

        Ok(try!(cmd.run_and_get_stdout()).contains("target_has_atomic"))
    }

    fn hash<H>(&self, hasher: &mut H) -> Result<()>
        where H: Hasher
    {
        match *self {
            Target::BuiltIn { .. } => {}
            Target::Custom { ref json, .. } |
            Target::Path { ref json, .. } => try!(io::read(json)).hash(hasher),
        }

        Ok(())
    }

    fn needs_sysroot(&self) -> bool {
        match *self {
            Target::BuiltIn { ref triple } => {
                match &triple[..] {
                    "thumbv6m-none-eabi" |
                    "thumbv7m-none-eabi" |
                    "thumbv7em-none-eabi" |
                    "thumbv7em-none-eabihf" => true,
                    _ => false,
                }
            }
            _ => true,
        }
    }

    fn triple(&self) -> &str {
        match *self {
            Target::BuiltIn { ref triple } |
            Target::Custom { ref triple, .. } |
            Target::Path { ref triple, .. } => triple,
        }
    }
}

trait CommandExt {
    fn run_and_get_stdout(&mut self) -> Result<String>;
    fn run_or_error(&mut self) -> Result<()>;
}

impl CommandExt for Command {
    fn run_and_get_stdout(&mut self) -> Result<String> {
        let cmd = &format!("`{:?}`", self);

        let output = try!(self.output()
            .chain_err(|| format!("failed to execute {}", cmd)));

        if !output.status.success() {
            try!(Err(format!("{} failed with exit status: {:?}",
                             cmd,
                             output.status.code())))
        }

        Ok(try!(String::from_utf8(output.stdout)
            .chain_err(|| format!("{} output was not UTF-8 encoded", cmd))))
    }

    fn run_or_error(&mut self) -> Result<()> {
        let cmd = &format!("`{:?}`", self);

        let exit = try!(self.status()
            .chain_err(|| format!("failed to execute {}", cmd)));

        if !exit.success() {
            try!(Err(format!("{} failed with exit status {:?}",
                             cmd,
                             exit.code())))
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>#[macro_use] extern crate log;

use std::num::{Int, FromPrimitive};
use std::os;
use std::io::File;
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::iter::{range_step, count};

// Index source
struct InputSource {
    haystack: Vec<char>,
    indexed_haystack: HashMap<char, Vec<usize>>,
}

fn index_chars(input: &Vec<char>) -> HashMap<char, Vec<usize>> {
    let mut indexed_haystack = HashMap::new();
    for (idx, &c) in input.iter().enumerate() {
        match indexed_haystack.entry(c) {
            Vacant(entry) => { entry.insert(vec![idx]); },
            Occupied(mut entry) => { (*entry.get_mut()).push(idx); }
        }
    }

    indexed_haystack
}

impl InputSource {
    fn new(haystack: Vec<char>) -> InputSource {
        let indexed_haystack = index_chars(&haystack);

        InputSource {
            haystack: haystack,
            indexed_haystack: indexed_haystack,
        }
    }

    fn search_for(&self, needle: Vec<char>) -> Option<EDS> {
        EDS::new(self, needle)
    }
}

// This is a struct for doing the searching
struct EDS<'a> {
    source: &'a InputSource,

    // String to search for
    needle: Vec<char>,

    first_poses: Vec<usize>,
    second_poses: Vec<usize>,

    // keep track of state between runs
    first_poses_idx: usize,
    second_poses_idx: usize,
}

impl<'a> Iterator for EDS<'a> {
    // (start, offset)
    type Item = (usize, i64);

    // This actually does the searching
    fn next(&mut self) -> Option<(usize, i64)> {
        debug!("Start of next call");
        let len_first_poses = self.first_poses.len();
        let len_second_poses = self.second_poses.len();

        if self.first_poses_idx == len_first_poses {
            debug!("Reached end, finishing");
            // reached the end
            return None;
        }

        let len_needle = self.needle.len();
        let len_haystack = self.source.haystack.len();
        let mut have_done_one_loop = false;
        let first_run = (self.first_poses_idx == 0 && self.second_poses_idx == 0);

        loop {
            debug!("Start of loop");
            // Starts off at 0/0, and we don't want to increment it the first time
            if ! (self.first_poses_idx == 0 && self.second_poses_idx == 0 && !have_done_one_loop) {
                debug!("Incrementing counters");
                if self.second_poses_idx == len_second_poses - 1 {
                    self.second_poses_idx = 0;
                    self.first_poses_idx += 1;
                } else{
                    self.second_poses_idx += 1;
                }
            }
            have_done_one_loop = true;
            if self.first_poses_idx == len_first_poses {
                debug!("Reached end, finishing");
                // reached the end
                return None;
            }

            debug!("first_poses_idx {} second_poses_idx {}", self.first_poses_idx, self.second_poses_idx);
            let first_pos = self.first_poses[self.first_poses_idx];
            let second_pos = self.second_poses[self.second_poses_idx];
            debug!("first_pos {} second_pos {}", first_pos, second_pos);

            let step: i64 = (second_pos as i64) - (first_pos as i64);
            if step == 1 {
                continue;
            }
            let positions: Vec<i64> = std::iter::count(first_pos as i64, step as i64).take(len_needle).collect();
            //let positions: Vec<i64> = range_step(first_pos as i64, step*len_needle, step).collect();
            debug!("Start {} step {}", first_pos, step);
            if positions.iter().any(|&x| { x<0 }) {
                // TODO there's probably a better maths way to do this
                continue;
            }
            let positions: Vec<usize> = positions.iter().map(|&x| { FromPrimitive::from_i64(x).unwrap() }).collect();
            debug!("positions {:?}", positions);
            if positions.len() == 0 {
                // WTF
                continue;
            }

            if positions[positions.len()-1] > len_haystack - 1 {
                debug!("Past end");
                continue;
            }
            if positions.iter().map(|&c| { &self.source.haystack[c] }).zip(self.needle.iter()).all(|(&a, &b)| { a == b }) {
                //println!("Found starting at {} with step of {}", first_pos, step);
                return Some((first_pos, step));
            }
        }
        None
    }
}

impl<'a> EDS<'a> {
    fn new(source: &InputSource, needle: Vec<char>) -> Option<EDS> {
        let first_poses;
        match source.indexed_haystack.get(&needle[0]) {
            None => { return None ; }
            Some(ref entries) => {
                first_poses = (**entries).clone();
            }
        }

        let second_poses;
        match source.indexed_haystack.get(&needle[1]) {
            None => { return None ; }
            Some(ref entries) => {
                second_poses = (**entries).clone();
            }
        }

        let result = EDS {
            needle: needle, source: source,
            first_poses: first_poses, second_poses: second_poses,
            first_poses_idx: 0, second_poses_idx: 0,
        };

        return Some(result);

    }
}

fn print_results(source: &InputSource, start: usize, step: i64, len: usize) {
    let total_rows = len + 2;
    let width: usize = 31; // 15 + 1 + 15
    let start_i = start as i64;

    for row_num in range(0, len as i64) {
        println!("{}", count(row_num*step + start_i - 15, 1).take(width).map(|i| { source.haystack[i as usize] }).collect::<String>());
    }

}

fn only_alphanumeric(input: String) -> Vec<char> {
    input.chars().filter(|c| { c.is_alphabetic() }).map(|c| { c.to_lowercase() }).collect()
}

fn main() {
    let file_to_search = &os::args()[1];
    println!("Reading in {}", file_to_search);
    let haystack: Vec<char> = only_alphanumeric(File::open(&Path::new(file_to_search)).read_to_string().unwrap());
    let source = InputSource::new(haystack);

    let needle = (&os::args()[2]).to_string();
    println!("Looking for {}", needle);
    let needle = only_alphanumeric(needle);
    let len_needle = needle.len();

    match source.search_for(needle) {
        None => {
            println!("Cannot look for this string");
        }
        Some(eds_searcher) => {
            for (start, step) in (eds_searcher).take(5) {
                println!("Found starting at {} with step of {}", start, step);
                print_results(&source, start, step, len_needle);
            }
        }
    }

}
<commit_msg>references and such<commit_after>#[macro_use] extern crate log;

use std::num::{Int, FromPrimitive};
use std::os;
use std::io::File;
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::iter::{range_step, count};

// Index source
struct InputSource {
    haystack: Vec<char>,
    indexed_haystack: HashMap<char, Vec<usize>>,
}

fn index_chars(input: &Vec<char>) -> HashMap<char, Vec<usize>> {
    let mut indexed_haystack = HashMap::new();
    for (idx, &c) in input.iter().enumerate() {
        match indexed_haystack.entry(c) {
            Vacant(entry) => { entry.insert(vec![idx]); },
            Occupied(mut entry) => { (*entry.get_mut()).push(idx); }
        }
    }

    indexed_haystack
}

impl InputSource {
    fn new(haystack: Vec<char>) -> InputSource {
        let indexed_haystack = index_chars(&haystack);

        InputSource {
            haystack: haystack,
            indexed_haystack: indexed_haystack,
        }
    }

    fn search_for<'a>(&'a self, needle: &'a Vec<char>) -> Option<EDS> {
        EDS::new(self, needle)
    }
}

// This is a struct for doing the searching
struct EDS<'a, 'b> {
    source: &'a InputSource,

    // String to search for
    needle: &'b Vec<char>,

    first_poses: Vec<usize>,
    second_poses: Vec<usize>,

    // keep track of state between runs
    first_poses_idx: usize,
    second_poses_idx: usize,
}

impl<'a, 'b> Iterator for EDS<'a, 'b> {
    // (start, offset)
    type Item = (usize, i64);

    // This actually does the searching
    fn next(&mut self) -> Option<(usize, i64)> {
        debug!("Start of next call");
        let len_first_poses = self.first_poses.len();
        let len_second_poses = self.second_poses.len();

        if self.first_poses_idx == len_first_poses {
            debug!("Reached end, finishing");
            // reached the end
            return None;
        }

        let len_needle = self.needle.len();
        let len_haystack = self.source.haystack.len();
        let mut have_done_one_loop = false;
        let first_run = (self.first_poses_idx == 0 && self.second_poses_idx == 0);

        loop {
            debug!("Start of loop");
            // Starts off at 0/0, and we don't want to increment it the first time
            if ! (self.first_poses_idx == 0 && self.second_poses_idx == 0 && !have_done_one_loop) {
                debug!("Incrementing counters");
                if self.second_poses_idx == len_second_poses - 1 {
                    self.second_poses_idx = 0;
                    self.first_poses_idx += 1;
                } else{
                    self.second_poses_idx += 1;
                }
            }
            have_done_one_loop = true;
            if self.first_poses_idx == len_first_poses {
                debug!("Reached end, finishing");
                // reached the end
                return None;
            }

            debug!("first_poses_idx {} second_poses_idx {}", self.first_poses_idx, self.second_poses_idx);
            let first_pos = self.first_poses[self.first_poses_idx];
            let second_pos = self.second_poses[self.second_poses_idx];
            debug!("first_pos {} second_pos {}", first_pos, second_pos);

            let step: i64 = (second_pos as i64) - (first_pos as i64);
            if step == 1 {
                continue;
            }
            let positions: Vec<i64> = std::iter::count(first_pos as i64, step as i64).take(len_needle).collect();
            //let positions: Vec<i64> = range_step(first_pos as i64, step*len_needle, step).collect();
            debug!("Start {} step {}", first_pos, step);
            if positions.iter().any(|&x| { x<0 }) {
                // TODO there's probably a better maths way to do this
                continue;
            }
            let positions: Vec<usize> = positions.iter().map(|&x| { FromPrimitive::from_i64(x).unwrap() }).collect();
            debug!("positions {:?}", positions);
            if positions.len() == 0 {
                // WTF
                continue;
            }

            if positions[positions.len()-1] > len_haystack - 1 {
                debug!("Past end");
                continue;
            }
            if positions.iter().map(|&c| { &self.source.haystack[c] }).zip(self.needle.iter()).all(|(&a, &b)| { a == b }) {
                //println!("Found starting at {} with step of {}", first_pos, step);
                return Some((first_pos, step));
            }
        }
        None
    }
}

impl<'a, 'b> EDS<'a, 'b> {
    fn new(source: &'a InputSource, needle: &'b Vec<char>) -> Option<EDS<'a, 'b>> {
        let first_poses;
        match source.indexed_haystack.get(&needle[0]) {
            None => { return None ; }
            Some(ref entries) => {
                first_poses = (**entries).clone();
            }
        }

        let second_poses;
        match source.indexed_haystack.get(&needle[1]) {
            None => { return None ; }
            Some(ref entries) => {
                second_poses = (**entries).clone();
            }
        }

        let result = EDS {
            needle: needle, source: source,
            first_poses: first_poses, second_poses: second_poses,
            first_poses_idx: 0, second_poses_idx: 0,
        };

        return Some(result);

    }
}

fn print_results(source: &InputSource, start: usize, step: i64, len: usize) {
    let total_rows = (len + 2) as i64;
    let width: usize = 31; // 15 + 1 + 15
    let start_i = start as i64;

    for row_num in range(0, total_rows) {
        println!("{}", count((row_num-1)*step + start_i - 15, 1).take(width).map(|i| { if i < 0 { ' ' } else { source.haystack[i as usize] } }).collect::<String>());
    }

}

fn only_alphanumeric(input: String) -> Vec<char> {
    input.chars().filter(|c| { c.is_alphabetic() }).map(|c| { c.to_lowercase() }).collect()
}

fn main() {
    let file_to_search = &os::args()[1];
    println!("Reading in {}", file_to_search);
    let haystack: Vec<char> = only_alphanumeric(File::open(&Path::new(file_to_search)).read_to_string().unwrap());
    let source = InputSource::new(haystack);

    let needle = (&os::args()[2]).to_string();
    println!("Looking for {}", needle);
    let needle = only_alphanumeric(needle);
    let len_needle = needle.len();

    let needle2 = (&os::args()[3]).to_string();
    println!("Looking also for {}", needle2);
    let needle2 = only_alphanumeric(needle2);
    let len_needle2 = needle2.len();

    match source.search_for(&needle) {
        None => {
            println!("Cannot look for this string");
        }
        Some(eds_searcher) => {
            for (start, step) in eds_searcher.take(10) {
                println!("Found starting at {} with step of {}", start, step);
                print_results(&source, start, step, len_needle);
                // now try to find the second needle2
                match source.search_for(&needle2) {
                    None => {
                        println!("Cannot look for {:?}", needle2);
                    }
                    Some(needle2_searcher) => {
                        for (start2, step2) in needle2_searcher.take(1) {
                            println!("Found {:?} starting at {} with step of {}", needle2, start2, step2);

                        }
                    }
                }
            }
        }
    }

}
<|endoftext|>


<commit_before>pub mod api;
pub mod db;
pub mod indexer;
pub mod model;
pub mod opening;
pub mod util;

use crate::{
    api::{
        Error, GameRow, GameRowWithUci, PersonalMoveRow, PersonalQuery, PersonalQueryFilter,
        PersonalResponse,
    },
    db::Database,
    indexer::{IndexerOpt, IndexerStub},
    model::{PersonalKeyBuilder, PersonalKeyPrefix},
    opening::{Opening, Openings},
    util::NdJson,
};
use axum::{
    extract::{Extension, Path, Query},
    handler::get,
    http::StatusCode,
    AddExtensionLayer, Router,
};
use clap::Clap;
use futures_util::stream::Stream;
use shakmaty::{fen::Fen, variant::VariantPosition, zobrist::Zobrist, CastlingMode};
use std::{mem, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};
use tokio::sync::oneshot;

#[derive(Clap)]
struct Opt {
    #[clap(long = "bind", default_value = "127.0.0.1:9000")]
    bind: SocketAddr,
    #[clap(long = "db", default_value = "_db")]
    db: PathBuf,
    #[clap(flatten)]
    indexer: IndexerOpt,
}

#[tokio::main]
async fn main() {
    env_logger::init_from_env(
        env_logger::Env::new()
            .filter("EXPLORER_LOG")
            .write_style("EXPLORER_LOG_STYLE"),
    );

    let opt = Opt::parse();

    let openings: &'static Openings = Box::leak(Box::new(Openings::build_table()));
    let db = Arc::new(Database::open(opt.db).expect("db"));
    let (indexer, join_handles) = IndexerStub::spawn(db.clone(), opt.indexer);

    let app = Router::new()
        .route("/admin/:prop", get(db_property))
        .route("/admin/game/:prop", get(game_property))
        .route("/admin/personal/:prop", get(personal_property))
        .route("/personal", get(personal))
        .route("/test", get(test))
        .layer(AddExtensionLayer::new(openings))
        .layer(AddExtensionLayer::new(db))
        .layer(AddExtensionLayer::new(indexer));

    axum::Server::bind(&opt.bind)
        .serve(app.into_make_service())
        .with_graceful_shutdown(async {
            tokio::signal::ctrl_c().await.expect("wait for ctrl-c");
        })
        .await
        .expect("bind");

    for join_handle in join_handles {
        join_handle.await.expect("indexer");
    }
}

async fn test() -> NdJson<impl futures_util::stream::Stream<Item = u32>> {
    NdJson::new(futures_util::stream::unfold(0, |state| async move {
        tokio::time::sleep(std::time::Duration::from_secs(2)).await;
        Some((state, state + 1))
    }))
}

async fn db_property(
    Extension(db): Extension<Arc<Database>>,
    Path(prop): Path<String>,
) -> Result<String, StatusCode> {
    db.queryable()
        .db_property(&prop)
        .expect("get property")
        .ok_or(StatusCode::NOT_FOUND)
}

async fn game_property(
    Extension(db): Extension<Arc<Database>>,
    Path(prop): Path<String>,
) -> Result<String, StatusCode> {
    db.queryable()
        .game_property(&prop)
        .expect("get property")
        .ok_or(StatusCode::NOT_FOUND)
}

async fn personal_property(
    Extension(db): Extension<Arc<Database>>,
    Path(prop): Path<String>,
) -> Result<String, StatusCode> {
    db.queryable()
        .personal_property(&prop)
        .expect("get property")
        .ok_or(StatusCode::NOT_FOUND)
}

struct PersonalStreamState {
    indexing: Option<oneshot::Receiver<()>>,
    key: PersonalKeyPrefix,
    db: Arc<Database>,
    filter: PersonalQueryFilter,
    pos: VariantPosition,
    opening: Option<&'static Opening>,
    first: bool,
    done: bool,
}

async fn personal(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Extension(indexer): Extension<IndexerStub>,
    Query(query): Query<PersonalQuery>,
) -> Result<NdJson<impl Stream<Item = PersonalResponse>>, Error> {
    let indexing = if query.update {
        Some(indexer.index_player(query.player.clone()).await)
    } else {
        None
    };

    let variant = query.variant.into();

    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => VariantPosition::from_setup(variant, &Fen::from(fen), CastlingMode::Chess960)?,
        None => VariantPosition::new(variant),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;

    let key = PersonalKeyBuilder::with_user_pov(&query.player.into(), query.color)
        .with_zobrist(variant, pos.zobrist_hash());

    let state = PersonalStreamState {
        filter: query.filter,
        db,
        indexing,
        opening,
        key,
        pos: pos.into_inner(),
        first: true,
        done: false,
    };

    Ok(NdJson::new(futures_util::stream::unfold(
        state,
        |mut state| async move {
            if state.done {
                return None;
            }

            let first = mem::replace(&mut state.first, false);
            state.done = match state.indexing {
                Some(ref mut indexing) => {
                    tokio::select! {
                        _ = indexing => true,
                        _ = tokio::time::sleep(Duration::from_millis(if first { 200 } else { 5000 })) => false,
                    }
                }
                None => true,
            };
            let queryable = state.db.queryable();
            let filtered = queryable
                .get_personal(&state.key, state.filter.since, state.filter.until)
                .expect("get personal")
                .prepare(&state.pos, &state.filter);

            Some((
                PersonalResponse {
                    total: filtered.total,
                    moves: filtered
                        .moves
                        .into_iter()
                        .map(|row| PersonalMoveRow {
                            uci: row.uci,
                            san: row.san,
                            stats: row.stats,
                            game: row.game.and_then(|id| {
                                queryable
                                    .get_game_info(id)
                                    .expect("get game")
                                    .map(|info| GameRow { id, info })
                            }),
                        })
                        .collect(),
                    recent_games: filtered
                        .recent_games
                        .into_iter()
                        .flat_map(|(uci, id)| {
                            queryable.get_game_info(id).expect("get game").map(|info| {
                                GameRowWithUci {
                                    uci,
                                    row: GameRow { id, info },
                                }
                            })
                        })
                        .collect(),
                    opening: state.opening,
                },
                state,
            ))
        },
    )))
}
<commit_msg>no longer need ndjson demo<commit_after>pub mod api;
pub mod db;
pub mod indexer;
pub mod model;
pub mod opening;
pub mod util;

use crate::{
    api::{
        Error, GameRow, GameRowWithUci, PersonalMoveRow, PersonalQuery, PersonalQueryFilter,
        PersonalResponse,
    },
    db::Database,
    indexer::{IndexerOpt, IndexerStub},
    model::{PersonalKeyBuilder, PersonalKeyPrefix},
    opening::{Opening, Openings},
    util::NdJson,
};
use axum::{
    extract::{Extension, Path, Query},
    handler::get,
    http::StatusCode,
    AddExtensionLayer, Router,
};
use clap::Clap;
use futures_util::stream::Stream;
use shakmaty::{fen::Fen, variant::VariantPosition, zobrist::Zobrist, CastlingMode};
use std::{mem, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};
use tokio::sync::oneshot;

#[derive(Clap)]
struct Opt {
    #[clap(long = "bind", default_value = "127.0.0.1:9000")]
    bind: SocketAddr,
    #[clap(long = "db", default_value = "_db")]
    db: PathBuf,
    #[clap(flatten)]
    indexer: IndexerOpt,
}

#[tokio::main]
async fn main() {
    env_logger::init_from_env(
        env_logger::Env::new()
            .filter("EXPLORER_LOG")
            .write_style("EXPLORER_LOG_STYLE"),
    );

    let opt = Opt::parse();

    let openings: &'static Openings = Box::leak(Box::new(Openings::build_table()));
    let db = Arc::new(Database::open(opt.db).expect("db"));
    let (indexer, join_handles) = IndexerStub::spawn(db.clone(), opt.indexer);

    let app = Router::new()
        .route("/admin/:prop", get(db_property))
        .route("/admin/game/:prop", get(game_property))
        .route("/admin/personal/:prop", get(personal_property))
        .route("/personal", get(personal))
        .layer(AddExtensionLayer::new(openings))
        .layer(AddExtensionLayer::new(db))
        .layer(AddExtensionLayer::new(indexer));

    axum::Server::bind(&opt.bind)
        .serve(app.into_make_service())
        .with_graceful_shutdown(async {
            tokio::signal::ctrl_c().await.expect("wait for ctrl-c");
        })
        .await
        .expect("bind");

    for join_handle in join_handles {
        join_handle.await.expect("indexer");
    }
}

async fn db_property(
    Extension(db): Extension<Arc<Database>>,
    Path(prop): Path<String>,
) -> Result<String, StatusCode> {
    db.queryable()
        .db_property(&prop)
        .expect("get property")
        .ok_or(StatusCode::NOT_FOUND)
}

async fn game_property(
    Extension(db): Extension<Arc<Database>>,
    Path(prop): Path<String>,
) -> Result<String, StatusCode> {
    db.queryable()
        .game_property(&prop)
        .expect("get property")
        .ok_or(StatusCode::NOT_FOUND)
}

async fn personal_property(
    Extension(db): Extension<Arc<Database>>,
    Path(prop): Path<String>,
) -> Result<String, StatusCode> {
    db.queryable()
        .personal_property(&prop)
        .expect("get property")
        .ok_or(StatusCode::NOT_FOUND)
}

struct PersonalStreamState {
    indexing: Option<oneshot::Receiver<()>>,
    key: PersonalKeyPrefix,
    db: Arc<Database>,
    filter: PersonalQueryFilter,
    pos: VariantPosition,
    opening: Option<&'static Opening>,
    first: bool,
    done: bool,
}

async fn personal(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Extension(indexer): Extension<IndexerStub>,
    Query(query): Query<PersonalQuery>,
) -> Result<NdJson<impl Stream<Item = PersonalResponse>>, Error> {
    let indexing = if query.update {
        Some(indexer.index_player(query.player.clone()).await)
    } else {
        None
    };

    let variant = query.variant.into();

    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => VariantPosition::from_setup(variant, &Fen::from(fen), CastlingMode::Chess960)?,
        None => VariantPosition::new(variant),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;

    let key = PersonalKeyBuilder::with_user_pov(&query.player.into(), query.color)
        .with_zobrist(variant, pos.zobrist_hash());

    let state = PersonalStreamState {
        filter: query.filter,
        db,
        indexing,
        opening,
        key,
        pos: pos.into_inner(),
        first: true,
        done: false,
    };

    Ok(NdJson::new(futures_util::stream::unfold(
        state,
        |mut state| async move {
            if state.done {
                return None;
            }

            let first = mem::replace(&mut state.first, false);
            state.done = match state.indexing {
                Some(ref mut indexing) => {
                    tokio::select! {
                        _ = indexing => true,
                        _ = tokio::time::sleep(Duration::from_millis(if first { 200 } else { 5000 })) => false,
                    }
                }
                None => true,
            };
            let queryable = state.db.queryable();
            let filtered = queryable
                .get_personal(&state.key, state.filter.since, state.filter.until)
                .expect("get personal")
                .prepare(&state.pos, &state.filter);

            Some((
                PersonalResponse {
                    total: filtered.total,
                    moves: filtered
                        .moves
                        .into_iter()
                        .map(|row| PersonalMoveRow {
                            uci: row.uci,
                            san: row.san,
                            stats: row.stats,
                            game: row.game.and_then(|id| {
                                queryable
                                    .get_game_info(id)
                                    .expect("get game")
                                    .map(|info| GameRow { id, info })
                            }),
                        })
                        .collect(),
                    recent_games: filtered
                        .recent_games
                        .into_iter()
                        .flat_map(|(uci, id)| {
                            queryable.get_game_info(id).expect("get game").map(|info| {
                                GameRowWithUci {
                                    uci,
                                    row: GameRow { id, info },
                                }
                            })
                        })
                        .collect(),
                    opening: state.opening,
                },
                state,
            ))
        },
    )))
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate chrono;
extern crate curl;
extern crate env_logger;
extern crate flate2;
extern crate rustc_version;
extern crate tar;
extern crate tempdir;

#[macro_use]
extern crate log;

use std::env;
use std::fs::{self, File};
use std::hash::{Hash, Hasher, SipHasher};
use std::io::Read;
use std::path::{Path, PathBuf};
use std::process::{self, Command};

// TODO proper error handling/reporting
macro_rules! try {
    ($e:expr) => {
        $e.unwrap_or_else(|e| panic!("{} with {}", stringify!($e), e))
    }
}

mod sysroot;

fn main() {
    init_logger();

    let (mut cargo, target) = parse_args();

    if let Some(target) = target {
        let sysroot = sysroot::create(&target);

        if let Some(mut rustflags) = env::var("RUSTFLAGS").ok() {
            rustflags.push_str(&format!(" --sysroot {}", sysroot.display()));
            cargo.env("RUSTFLAGS", rustflags);
        } else {
            cargo.env("RUSTFLAGS", format!("--sysroot {}", sysroot.display()));
        }
    }

    if let Some(code) = try!(try!(cargo.spawn()).wait()).code() {
        process::exit(code);
    }
}

fn init_logger() {
    try!(env_logger::init());
}

/// Custom target with specification file
#[derive(Debug)]
pub struct Target {
    hash: u64,
    path: PathBuf,
    triple: String,
}

impl Target {
    fn from(s: &str) -> Option<Target> {
        let path = &PathBuf::from(format!("{}.json", s));
        if path.is_file() {
            return Some(Target::from_path(path));
        }

        let target_path = &env::var_os("RUST_TARGET_PATH").unwrap_or_default();

        for dir in env::split_paths(target_path) {
            let path = &dir.join(path);

            if path.is_file() {
                return Some(Target::from_path(path));
            }
        }

        None
    }

    fn from_path(path: &Path) -> Target {
        fn hash(path: &Path) -> u64 {
            let h = &mut SipHasher::new();
            let contents = &mut String::new();
            try!(try!(File::open(path)).read_to_string(contents));
            contents.hash(h);
            h.finish()
        }

        let triple = path.file_stem().unwrap().to_string_lossy().into_owned();
        info!("target: {}", triple);
        Target {
            hash: hash(path),
            path: try!(fs::canonicalize(path)),
            triple: triple,
        }
    }
}

fn parse_args() -> (Command, Option<Target>) {
    let mut cmd = Command::new("cargo");
    let mut target = None;

    let mut next_is_target = false;
    for arg_os in env::args_os().skip(1) {
        if target.is_none() {
            let arg = &*arg_os.to_string_lossy();

            if next_is_target {
                target = Target::from(arg);
            } else {
                if arg == "--target" {
                    next_is_target = true;
                } else if arg.starts_with("--target=") {
                    target = arg.split('=').skip(1).next().and_then(|s| Target::from(s))
                }
            }
        }

        cmd.arg(arg_os);
    }

    (cmd, target)
}
<commit_msg>OsString doesn't implement Default on stable/beta<commit_after>#![deny(warnings)]

extern crate chrono;
extern crate curl;
extern crate env_logger;
extern crate flate2;
extern crate rustc_version;
extern crate tar;
extern crate tempdir;

#[macro_use]
extern crate log;

use std::env;
use std::ffi::OsString;
use std::fs::{self, File};
use std::hash::{Hash, Hasher, SipHasher};
use std::io::Read;
use std::path::{Path, PathBuf};
use std::process::{self, Command};

// TODO proper error handling/reporting
macro_rules! try {
    ($e:expr) => {
        $e.unwrap_or_else(|e| panic!("{} with {}", stringify!($e), e))
    }
}

mod sysroot;

fn main() {
    init_logger();

    let (mut cargo, target) = parse_args();

    if let Some(target) = target {
        let sysroot = sysroot::create(&target);

        if let Some(mut rustflags) = env::var("RUSTFLAGS").ok() {
            rustflags.push_str(&format!(" --sysroot {}", sysroot.display()));
            cargo.env("RUSTFLAGS", rustflags);
        } else {
            cargo.env("RUSTFLAGS", format!("--sysroot {}", sysroot.display()));
        }
    }

    if let Some(code) = try!(try!(cargo.spawn()).wait()).code() {
        process::exit(code);
    }
}

fn init_logger() {
    try!(env_logger::init());
}

/// Custom target with specification file
#[derive(Debug)]
pub struct Target {
    hash: u64,
    path: PathBuf,
    triple: String,
}

impl Target {
    fn from(s: &str) -> Option<Target> {
        let path = &PathBuf::from(format!("{}.json", s));
        if path.is_file() {
            return Some(Target::from_path(path));
        }

        let target_path = &env::var_os("RUST_TARGET_PATH").unwrap_or(OsString::new());

        for dir in env::split_paths(target_path) {
            let path = &dir.join(path);

            if path.is_file() {
                return Some(Target::from_path(path));
            }
        }

        None
    }

    fn from_path(path: &Path) -> Target {
        fn hash(path: &Path) -> u64 {
            let h = &mut SipHasher::new();
            let contents = &mut String::new();
            try!(try!(File::open(path)).read_to_string(contents));
            contents.hash(h);
            h.finish()
        }

        let triple = path.file_stem().unwrap().to_string_lossy().into_owned();
        info!("target: {}", triple);
        Target {
            hash: hash(path),
            path: try!(fs::canonicalize(path)),
            triple: triple,
        }
    }
}

fn parse_args() -> (Command, Option<Target>) {
    let mut cmd = Command::new("cargo");
    let mut target = None;

    let mut next_is_target = false;
    for arg_os in env::args_os().skip(1) {
        if target.is_none() {
            let arg = &*arg_os.to_string_lossy();

            if next_is_target {
                target = Target::from(arg);
            } else {
                if arg == "--target" {
                    next_is_target = true;
                } else if arg.starts_with("--target=") {
                    target = arg.split('=').skip(1).next().and_then(|s| Target::from(s))
                }
            }
        }

        cmd.arg(arg_os);
    }

    (cmd, target)
}
<|endoftext|>


<commit_before>use std::env;
use std::fs::File;
use std::error::Error;
use std::io::prelude::*;
use std::io::BufWriter;
extern crate time;
extern crate toml;
extern crate rawloader;
use rawloader::decoders;
use rawloader::imageops;

fn usage() {
  println!("rawloader <file> [outfile]");
  std::process::exit(1);
}

fn error(err: &str) {
  println!("ERROR: {}", err);
  std::process::exit(2);
}

fn main() {
  let args: Vec<_> = env::args().collect();
  if args.len() < 2 {
    usage();
  }
  let file = &args[1];
  let fileppm = format!("{}.ppm",file);
  let outfile = if args.len() > 2 {
    &args[2]
  } else {
    &fileppm
  };
  println!("Loading file \"{}\" and saving it as \"{}\"", file, outfile);

  let rawloader = decoders::RawLoader::new();
  let from_time = time::precise_time_ns();
  let image = match rawloader.decode_safe(file) {
    Ok(val) => val,
    Err(e) => {error(&e);unreachable!()},
  };
  let to_time = time::precise_time_ns();
  println!("Decoded in {} ms", (to_time - from_time)/1000000);

  println!("Found camera \"{}\" model \"{}\"", image.make, image.model);
  println!("Found canonical named camera \"{}\" model \"{}\"", image.canonical_make, image.canonical_model);
  println!("Image size is {}x{}", image.width, image.height);
  println!("WB coeffs are {:?}", image.wb_coeffs);
  println!("black levels are {:?}", image.blacklevels);
  println!("white levels are {:?}", image.whitelevels);
  println!("color matrix is {:?}", image.color_matrix);
  println!("dcraw filters is {:#x}", image.dcraw_filters);
  println!("crops are {:?}", image.crops);

  let mut sum: u64 = 0;
  for i in 0..(image.width*image.height) {
    sum += image.data[i as usize] as u64;
  }
  println!("Image sum: {}", sum);
  let count: u64 = (image.width as u64) * (image.height as u64);
  println!("Image avg: {}", sum/count);

  let decoded = imageops::simple_decode(&image);

  let uf = match File::create(outfile) {
    Ok(val) => val,
    Err(e) => {error(e.description());unreachable!()},
  };
  let mut f = BufWriter::new(uf);
  let preamble = format!("P6 {} {} {}\n", image.width, image.height, 255).into_bytes();
  if let Err(err) = f.write_all(&preamble) {
    error(err.description());
  }
  for pix in decoded {
    let pixel = ((pix.max(0.0)*255.0).min(255.0)) as u8;
    if let Err(err) = f.write_all(&[pixel]) {
      error(err.description());
    }
  }
}
<commit_msg>Print out the max and min values of the image<commit_after>use std::env;
use std::fs::File;
use std::error::Error;
use std::io::prelude::*;
use std::io::BufWriter;
extern crate time;
extern crate toml;
extern crate rawloader;
use rawloader::decoders;
use rawloader::imageops;

fn usage() {
  println!("rawloader <file> [outfile]");
  std::process::exit(1);
}

fn error(err: &str) {
  println!("ERROR: {}", err);
  std::process::exit(2);
}

fn main() {
  let args: Vec<_> = env::args().collect();
  if args.len() < 2 {
    usage();
  }
  let file = &args[1];
  let fileppm = format!("{}.ppm",file);
  let outfile = if args.len() > 2 {
    &args[2]
  } else {
    &fileppm
  };
  println!("Loading file \"{}\" and saving it as \"{}\"", file, outfile);

  let rawloader = decoders::RawLoader::new();
  let from_time = time::precise_time_ns();
  let image = match rawloader.decode_safe(file) {
    Ok(val) => val,
    Err(e) => {error(&e);unreachable!()},
  };
  let to_time = time::precise_time_ns();
  println!("Decoded in {} ms", (to_time - from_time)/1000000);

  println!("Found camera \"{}\" model \"{}\"", image.make, image.model);
  println!("Found canonical named camera \"{}\" model \"{}\"", image.canonical_make, image.canonical_model);
  println!("Image size is {}x{}", image.width, image.height);
  println!("WB coeffs are {:?}", image.wb_coeffs);
  println!("black levels are {:?}", image.blacklevels);
  println!("white levels are {:?}", image.whitelevels);
  println!("color matrix is {:?}", image.color_matrix);
  println!("dcraw filters is {:#x}", image.dcraw_filters);
  println!("crops are {:?}", image.crops);

  let mut sum: u64 = 0;
  let mut min: u16 = image.data[0];
  let mut max: u16 = 0;
  for i in 0..(image.width*image.height) {
    let val = image.data[i as usize];
    sum += val as u64;
    if val < min { min = val; }
    if val > max { max = val; }
  }
  println!("Image sum: {}", sum);
  let count: u64 = (image.width as u64) * (image.height as u64);
  println!("Image avg: {}", sum/count);
  println!("Image min: {}", min);
  println!("Image max: {}", max);

  let decoded = imageops::simple_decode(&image);

  let uf = match File::create(outfile) {
    Ok(val) => val,
    Err(e) => {error(e.description());unreachable!()},
  };
  let mut f = BufWriter::new(uf);
  let preamble = format!("P6 {} {} {}\n", image.width, image.height, 255).into_bytes();
  if let Err(err) = f.write_all(&preamble) {
    error(err.description());
  }
  for pix in decoded {
    let pixel = ((pix.max(0.0)*255.0).min(255.0)) as u8;
    if let Err(err) = f.write_all(&[pixel]) {
      error(err.description());
    }
  }
}
<|endoftext|>