<commit_before>use std::convert::From;
use std::error::Error;
use std::result;
use std::fmt;
use snap;
use lz4_compress as lz4;

type Result<T> = result::Result<T, CompressionError>;

pub const LZ4: &'static str = "lz4";
pub const SNAPPY: &'static str = "snappy";

/// It's an error which may occure during encoding or deconding
/// frame body. As there are only two types of compressors it
/// contains two related enum options.
#[derive(Debug)]
pub enum CompressionError {
    /// Snappy error.
    Snappy(Box<Error>),
    /// Lz4 error.
    Lz4(String),
}

impl fmt::Display for CompressionError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &CompressionError::Snappy(ref err) => write!(f, "Snappy Error: {:?}", err),
            &CompressionError::Lz4(ref s) => write!(f, "Lz4 Error: {:?}", s),
        }
    }
}

impl Error for CompressionError {
    fn description(&self) -> &str {
        match self {
            &CompressionError::Snappy(ref err) => err.description(),
            &CompressionError::Lz4(ref s) => s.as_str(),
        }
    }
}

/// Compressor trait that defines functionality
/// which should be provided by typical compressor.
pub trait Compressor {
    /// Encodes given bytes and returns `Result` that contains either
    /// encoded data or an error which occures during the transformation.
    fn encode(&self, bytes: Vec<u8>) -> Result<Vec<u8>>;
    /// Encodes given encoded data and returns `Result` that contains either
    /// encoded bytes or an error which occures during the transformation.
    fn decode(&self, bytes: Vec<u8>) -> Result<Vec<u8>>;
    /// Returns a string which is a name of a compressor. This name should be
    /// exactly the same as one which returns a server in a response to
    /// `Options` request.
    fn into_string(&self) -> Option<String>;
}

/// Enum which represents a type of compression. Only non-startup frame's body can be compressen.
#[derive(Debug, PartialEq, Clone, Copy, Eq, Ord, PartialOrd)]
pub enum Compression {
    /// [lz4](https://code.google.com/p/lz4/) compression
    Lz4,
    /// [snappy](https://code.google.com/p/snappy/) compression
    Snappy,
    /// Non compression
    None,
}

impl Compression {
    /// It encodes `bytes` basing on type of `Compression`..
    ///
    /// # Examples
    ///
    /// ```
    ///    use cdrs::compression::Compression;
    ///
    ///   let snappy_compression = Compression::Snappy;
    ///   let bytes = String::from("Hello World").into_bytes().to_vec();
    ///   let encoded = snappy_compression.encode(bytes.clone()).unwrap();
    ///   assert_eq!(snappy_compression.decode(encoded).unwrap(), bytes);
    ///
    /// ```
    pub fn encode(&self, bytes: Vec<u8>) -> Result<Vec<u8>> {
        match self {
            &Compression::Lz4 => Compression::encode_lz4(bytes),
            &Compression::Snappy => Compression::encode_snappy(bytes),
            &Compression::None => Ok(bytes),
        }
    }

    /// It decodes `bytes` basing on type of compression.
    ///
    /// # Examples
    ///
    /// ```
    ///    use cdrs::compression::Compression;
    ///     let lz4_compression = Compression::Lz4;
    ///     let bytes = String::from("Hello World").into_bytes().to_vec();
    ///     let encoded = lz4_compression.encode(bytes.clone()).unwrap();
    ///     let len = encoded.len() as u8;
    ///     let mut input = vec![0, 0, 0, len];
    ///     input.extend_from_slice(encoded.as_slice());
    ///     assert_eq!(lz4_compression.decode(input).unwrap(), bytes);
    /// ```
    pub fn decode(&self, bytes: Vec<u8>) -> Result<Vec<u8>> {
        match self {
            &Compression::Lz4 => Compression::decode_lz4(bytes),
            &Compression::Snappy => Compression::decode_snappy(bytes),
            &Compression::None => Ok(bytes),
        }
    }

    /// It transforms compression method into a `&str`.
    pub fn as_str(&self) -> Option<&'static str> {
        match self {
            &Compression::Lz4 => Some(LZ4),
            &Compression::Snappy => Some(SNAPPY),
            &Compression::None => None,
        }
    }

    fn encode_snappy(bytes: Vec<u8>) -> Result<Vec<u8>> {
        let mut encoder = snap::Encoder::new();
        encoder.compress_vec(bytes.as_slice())
            .map_err(|err| CompressionError::Snappy(Box::new(err)))
    }

    fn decode_snappy(bytes: Vec<u8>) -> Result<Vec<u8>> {
        let mut decoder = snap::Decoder::new();
        decoder.decompress_vec(bytes.as_slice())
            .map_err(|err| CompressionError::Snappy(Box::new(err)))
    }

    fn encode_lz4(bytes: Vec<u8>) -> Result<Vec<u8>> {
        return Ok(lz4::compress(bytes.as_slice()));
    }

    fn decode_lz4(bytes: Vec<u8>) -> Result<Vec<u8>> {
        // skip first 4 bytes in accordance to
        // https://github.com/apache/cassandra/blob/trunk/doc/native_protocol_v4.spec#L805
        lz4::decompress(&bytes[4..])
            .map_err(|err| CompressionError::Lz4(err.description().to_string()))
    }
}

impl From<String> for Compression {
    /// It converts `String` into `Compression`. If string is neither `lz4` nor `snappy` then
    /// `Compression::None` will be returned
    fn from(compression_string: String) -> Compression {
        return Compression::from(compression_string.as_str());
    }
}

impl<'a> From<&'a str> for Compression {
    /// It converts `str` into `Compression`. If string is neither `lz4` nor `snappy` then
    /// `Compression::None` will be returned
    fn from(compression_str: &'a str) -> Compression {
        match compression_str {
            LZ4 => Compression::Lz4,
            SNAPPY => Compression::Snappy,
            _ => Compression::None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_compression_from_str() {
        let lz4 = "lz4";
        assert_eq!(Compression::from(lz4), Compression::Lz4);
        let snappy = "snappy";
        assert_eq!(Compression::from(snappy), Compression::Snappy);
        let none = "x";
        assert_eq!(Compression::from(none), Compression::None);
    }

    #[test]
    fn test_compression_from_string() {
        let lz4 = "lz4".to_string();
        assert_eq!(Compression::from(lz4), Compression::Lz4);
        let snappy = "snappy".to_string();
        assert_eq!(Compression::from(snappy), Compression::Snappy);
        let none = "x".to_string();
        assert_eq!(Compression::from(none), Compression::None);
    }

    #[test]
    fn test_compression_encode_snappy() {
        let snappy_compression = Compression::Snappy;
        let bytes = String::from("Hello World").into_bytes().to_vec();
        snappy_compression.encode(bytes.clone()).expect("Should work without exceptions");
    }

    #[test]
    fn test_compression_decode_snappy() {
        let snappy_compression = Compression::Snappy;
        let bytes = String::from("Hello World").into_bytes().to_vec();
        let encoded = snappy_compression.encode(bytes.clone()).unwrap();
        assert_eq!(snappy_compression.decode(encoded).unwrap(), bytes);
    }

    #[test]
    fn test_compression_encode_lz4() {
        let snappy_compression = Compression::Lz4;
        let bytes = String::from("Hello World").into_bytes().to_vec();
        snappy_compression.encode(bytes.clone()).expect("Should work without exceptions");
    }

    #[test]
    fn test_compression_decode_lz4() {
        let lz4_compression = Compression::Lz4;
        let bytes = String::from("Hello World").into_bytes().to_vec();
        let encoded = lz4_compression.encode(bytes.clone()).unwrap();
        let len = encoded.len() as u8;
        let mut input = vec![0, 0, 0, len];
        input.extend_from_slice(encoded.as_slice());
        assert_eq!(lz4_compression.decode(input).unwrap(), bytes);
    }

    #[test]
    fn test_compression_encode_none() {
        let none_compression = Compression::None;
        let bytes = String::from("Hello World").into_bytes().to_vec();
        none_compression.encode(bytes.clone()).expect("Should work without exceptions");
    }

    #[test]
    fn test_compression_decode_none() {
        let none_compression = Compression::None;
        let bytes = String::from("Hello World").into_bytes().to_vec();
        let encoded = none_compression.encode(bytes.clone()).unwrap();
        assert_eq!(none_compression.decode(encoded).unwrap(), bytes);
    }

    #[test]
    fn test_compression_encode_lz4_with_invalid_input() {
        let lz4_compression = Compression::Lz4;
        let bytes: Vec<u8> = vec![0x7f, 0x7f, 0x7f, 0x7f, 0x7f];
        let encoded = lz4_compression.encode(bytes.clone()).unwrap();
        let decode = lz4_compression.decode(encoded);
        assert_eq!(decode.is_err(), true);
    }

    #[test]
    fn test_compression_encode_snappy_with_non_utf8() {
        let snappy_compression = Compression::Snappy;
        let v = vec![0xff, 0xff];
        let encoded = snappy_compression.encode(v.clone()).expect("Should work without exceptions");
        assert_eq!(snappy_compression.decode(encoded).unwrap(), v);
    }

}
<commit_msg>Use upstream error types in CompressionError<commit_after>use std::convert::From;
use std::error::Error;
use std::result;
use std::fmt;
use std::io;
use snap;
use lz4_compress as lz4;

type Result<T> = result::Result<T, CompressionError>;

pub const LZ4: &'static str = "lz4";
pub const SNAPPY: &'static str = "snappy";

/// It's an error which may occure during encoding or deconding
/// frame body. As there are only two types of compressors it
/// contains two related enum options.
#[derive(Debug)]
pub enum CompressionError {
    /// Snappy error.
    Snappy(snap::Error),
    /// Lz4 error.
    Lz4(io::Error),
}

impl fmt::Display for CompressionError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &CompressionError::Snappy(ref err) => write!(f, "Snappy Error: {:?}", err),
            &CompressionError::Lz4(ref err) => write!(f, "Lz4 Error: {:?}", err),
        }
    }
}

impl Error for CompressionError {
    fn description(&self) -> &str {
        match self {
            &CompressionError::Snappy(ref err) => err.description(),
            &CompressionError::Lz4(ref err) => err.description(),
        }
    }
}

/// Compressor trait that defines functionality
/// which should be provided by typical compressor.
pub trait Compressor {
    /// Encodes given bytes and returns `Result` that contains either
    /// encoded data or an error which occures during the transformation.
    fn encode(&self, bytes: Vec<u8>) -> Result<Vec<u8>>;
    /// Encodes given encoded data and returns `Result` that contains either
    /// encoded bytes or an error which occures during the transformation.
    fn decode(&self, bytes: Vec<u8>) -> Result<Vec<u8>>;
    /// Returns a string which is a name of a compressor. This name should be
    /// exactly the same as one which returns a server in a response to
    /// `Options` request.
    fn into_string(&self) -> Option<String>;
}

/// Enum which represents a type of compression. Only non-startup frame's body can be compressen.
#[derive(Debug, PartialEq, Clone, Copy, Eq, Ord, PartialOrd)]
pub enum Compression {
    /// [lz4](https://code.google.com/p/lz4/) compression
    Lz4,
    /// [snappy](https://code.google.com/p/snappy/) compression
    Snappy,
    /// Non compression
    None,
}

impl Compression {
    /// It encodes `bytes` basing on type of `Compression`..
    ///
    /// # Examples
    ///
    /// ```
    ///    use cdrs::compression::Compression;
    ///
    ///   let snappy_compression = Compression::Snappy;
    ///   let bytes = String::from("Hello World").into_bytes().to_vec();
    ///   let encoded = snappy_compression.encode(bytes.clone()).unwrap();
    ///   assert_eq!(snappy_compression.decode(encoded).unwrap(), bytes);
    ///
    /// ```
    pub fn encode(&self, bytes: Vec<u8>) -> Result<Vec<u8>> {
        match self {
            &Compression::Lz4 => Compression::encode_lz4(bytes),
            &Compression::Snappy => Compression::encode_snappy(bytes),
            &Compression::None => Ok(bytes),
        }
    }

    /// It decodes `bytes` basing on type of compression.
    ///
    /// # Examples
    ///
    /// ```
    ///    use cdrs::compression::Compression;
    ///     let lz4_compression = Compression::Lz4;
    ///     let bytes = String::from("Hello World").into_bytes().to_vec();
    ///     let encoded = lz4_compression.encode(bytes.clone()).unwrap();
    ///     let len = encoded.len() as u8;
    ///     let mut input = vec![0, 0, 0, len];
    ///     input.extend_from_slice(encoded.as_slice());
    ///     assert_eq!(lz4_compression.decode(input).unwrap(), bytes);
    /// ```
    pub fn decode(&self, bytes: Vec<u8>) -> Result<Vec<u8>> {
        match self {
            &Compression::Lz4 => Compression::decode_lz4(bytes),
            &Compression::Snappy => Compression::decode_snappy(bytes),
            &Compression::None => Ok(bytes),
        }
    }

    /// It transforms compression method into a `&str`.
    pub fn as_str(&self) -> Option<&'static str> {
        match self {
            &Compression::Lz4 => Some(LZ4),
            &Compression::Snappy => Some(SNAPPY),
            &Compression::None => None,
        }
    }

    fn encode_snappy(bytes: Vec<u8>) -> Result<Vec<u8>> {
        let mut encoder = snap::Encoder::new();
        encoder.compress_vec(bytes.as_slice())
            .map_err(CompressionError::Snappy)
    }

    fn decode_snappy(bytes: Vec<u8>) -> Result<Vec<u8>> {
        let mut decoder = snap::Decoder::new();
        decoder.decompress_vec(bytes.as_slice())
            .map_err(CompressionError::Snappy)
    }

    fn encode_lz4(bytes: Vec<u8>) -> Result<Vec<u8>> {
        return Ok(lz4::compress(bytes.as_slice()));
    }

    fn decode_lz4(bytes: Vec<u8>) -> Result<Vec<u8>> {
        // skip first 4 bytes in accordance to
        // https://github.com/apache/cassandra/blob/trunk/doc/native_protocol_v4.spec#L805
        lz4::decompress(&bytes[4..]).map_err(CompressionError::Lz4)
    }
}

impl From<String> for Compression {
    /// It converts `String` into `Compression`. If string is neither `lz4` nor `snappy` then
    /// `Compression::None` will be returned
    fn from(compression_string: String) -> Compression {
        return Compression::from(compression_string.as_str());
    }
}

impl<'a> From<&'a str> for Compression {
    /// It converts `str` into `Compression`. If string is neither `lz4` nor `snappy` then
    /// `Compression::None` will be returned
    fn from(compression_str: &'a str) -> Compression {
        match compression_str {
            LZ4 => Compression::Lz4,
            SNAPPY => Compression::Snappy,
            _ => Compression::None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_compression_from_str() {
        let lz4 = "lz4";
        assert_eq!(Compression::from(lz4), Compression::Lz4);
        let snappy = "snappy";
        assert_eq!(Compression::from(snappy), Compression::Snappy);
        let none = "x";
        assert_eq!(Compression::from(none), Compression::None);
    }

    #[test]
    fn test_compression_from_string() {
        let lz4 = "lz4".to_string();
        assert_eq!(Compression::from(lz4), Compression::Lz4);
        let snappy = "snappy".to_string();
        assert_eq!(Compression::from(snappy), Compression::Snappy);
        let none = "x".to_string();
        assert_eq!(Compression::from(none), Compression::None);
    }

    #[test]
    fn test_compression_encode_snappy() {
        let snappy_compression = Compression::Snappy;
        let bytes = String::from("Hello World").into_bytes().to_vec();
        snappy_compression.encode(bytes.clone()).expect("Should work without exceptions");
    }

    #[test]
    fn test_compression_decode_snappy() {
        let snappy_compression = Compression::Snappy;
        let bytes = String::from("Hello World").into_bytes().to_vec();
        let encoded = snappy_compression.encode(bytes.clone()).unwrap();
        assert_eq!(snappy_compression.decode(encoded).unwrap(), bytes);
    }

    #[test]
    fn test_compression_encode_lz4() {
        let snappy_compression = Compression::Lz4;
        let bytes = String::from("Hello World").into_bytes().to_vec();
        snappy_compression.encode(bytes.clone()).expect("Should work without exceptions");
    }

    #[test]
    fn test_compression_decode_lz4() {
        let lz4_compression = Compression::Lz4;
        let bytes = String::from("Hello World").into_bytes().to_vec();
        let encoded = lz4_compression.encode(bytes.clone()).unwrap();
        let len = encoded.len() as u8;
        let mut input = vec![0, 0, 0, len];
        input.extend_from_slice(encoded.as_slice());
        assert_eq!(lz4_compression.decode(input).unwrap(), bytes);
    }

    #[test]
    fn test_compression_encode_none() {
        let none_compression = Compression::None;
        let bytes = String::from("Hello World").into_bytes().to_vec();
        none_compression.encode(bytes.clone()).expect("Should work without exceptions");
    }

    #[test]
    fn test_compression_decode_none() {
        let none_compression = Compression::None;
        let bytes = String::from("Hello World").into_bytes().to_vec();
        let encoded = none_compression.encode(bytes.clone()).unwrap();
        assert_eq!(none_compression.decode(encoded).unwrap(), bytes);
    }

    #[test]
    fn test_compression_encode_lz4_with_invalid_input() {
        let lz4_compression = Compression::Lz4;
        let bytes: Vec<u8> = vec![0x7f, 0x7f, 0x7f, 0x7f, 0x7f];
        let encoded = lz4_compression.encode(bytes.clone()).unwrap();
        let decode = lz4_compression.decode(encoded);
        assert_eq!(decode.is_err(), true);
    }

    #[test]
    fn test_compression_encode_snappy_with_non_utf8() {
        let snappy_compression = Compression::Snappy;
        let v = vec![0xff, 0xff];
        let encoded = snappy_compression.encode(v.clone()).expect("Should work without exceptions");
        assert_eq!(snappy_compression.decode(encoded).unwrap(), v);
    }

}
<|endoftext|>


<commit_before>use clap::{App, Arg};
use env_logger::{Env, DEFAULT_FILTER_ENV};
use godcoin::{blockchain::ReindexOpts, prelude::*};
use log::info;
use serde::Deserialize;
use std::{
    env, fs,
    path::{Path, PathBuf},
};
use tokio::{prelude::*, runtime::Runtime};

#[derive(Debug, Deserialize)]
struct Config {
    minter_key: String,
    enable_stale_production: bool,
    bind_address: Option<String>,
}

fn main() {
    env_logger::init_from_env(Env::new().filter_or(DEFAULT_FILTER_ENV, "godcoin=info"));
    godcoin::init().unwrap();

    let home = {
        match env::var("GODCOIN_HOME") {
            Ok(s) => PathBuf::from(s),
            Err(_) => Path::join(&dirs::data_local_dir().unwrap(), "godcoin"),
        }
    };

    let home = home.to_string_lossy();
    let args = App::new("godcoin-server")
        .about("GODcoin core server daemon")
        .version(env!("CARGO_PKG_VERSION"))
        .arg(
            Arg::with_name("home")
                .long("home")
                .default_value(&home)
                .empty_values(false)
                .help("Home directory which defaults to env var GODCOIN_HOME"),
        )
        .arg(
            Arg::with_name("reindex")
                .long("reindex")
                .help("Reindexes the block log"),
        )
        .arg(
            Arg::with_name("auto_trim")
                .long("reindex-trim-corrupt")
                .help("Trims any corruption detected in the block log during reindexing"),
        )
        .get_matches();

    let home = PathBuf::from(args.value_of("home").expect("Failed to obtain home path"));
    let (blocklog_loc, index_loc) = {
        if !Path::is_dir(&home) {
            let res = std::fs::create_dir(&home);
            res.unwrap_or_else(|_| panic!("Failed to create dir at {:?}", &home));
            info!("Created GODcoin home at {:?}", &home);
        } else {
            info!("Found GODcoin home at {:?}", &home);
        }
        let blocklog_loc = Path::join(&home, "blklog");
        let index_loc = Path::join(&home, "index");
        (blocklog_loc, index_loc)
    };

    let config_file = Path::join(&home, "config.toml");
    info!("Opening configuration file at {:?}", config_file);
    let config_file = fs::read(config_file).expect("Failed to open config");
    let config: Config = toml::from_str(&String::from_utf8(config_file).unwrap()).unwrap();

    let minter_key =
        PrivateKey::from_wif(&config.minter_key).expect("Provided minter key is invalid");
    let bind_addr = config
        .bind_address
        .unwrap_or_else(|| "127.0.0.1:7777".to_owned());

    let reindex = if args.is_present("reindex") {
        info!("User requested reindexing");
        if Path::exists(&index_loc) {
            info!("Deleting current index");
            fs::remove_dir_all(&index_loc)
                .expect("Failed to delete the blockchain index directory");
        } else {
            info!("Current index does not exist");
        }
        let auto_trim = args.is_present("auto_trim");
        Some(ReindexOpts { auto_trim })
    } else {
        None
    };

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

    let enable_stale_production = config.enable_stale_production;
    rt.spawn(future::lazy(move || {
        godcoin_server::start(godcoin_server::ServerOpts {
            blocklog_loc,
            index_loc,
            minter_key,
            bind_addr,
            reindex,
            enable_stale_production,
        });
        Ok(())
    }));

    let future = tokio_signal::ctrl_c()
        .flatten_stream()
        .into_future()
        .and_then(|_| {
            info!("Received ctrl-c, shutting down...");
            Ok(())
        });
    rt.block_on(future).map_err(|(e, _)| e).unwrap();
}
<commit_msg>Abort process when a thread panics<commit_after>use clap::{App, Arg};
use env_logger::{Env, DEFAULT_FILTER_ENV};
use godcoin::{blockchain::ReindexOpts, prelude::*};
use log::info;
use serde::Deserialize;
use std::{
    env, fs,
    path::{Path, PathBuf},
};
use tokio::{prelude::*, runtime::Builder};

#[derive(Debug, Deserialize)]
struct Config {
    minter_key: String,
    enable_stale_production: bool,
    bind_address: Option<String>,
}

fn main() {
    env_logger::init_from_env(Env::new().filter_or(DEFAULT_FILTER_ENV, "godcoin=info"));
    godcoin::init().unwrap();

    let home = {
        match env::var("GODCOIN_HOME") {
            Ok(s) => PathBuf::from(s),
            Err(_) => Path::join(&dirs::data_local_dir().unwrap(), "godcoin"),
        }
    };

    let home = home.to_string_lossy();
    let args = App::new("godcoin-server")
        .about("GODcoin core server daemon")
        .version(env!("CARGO_PKG_VERSION"))
        .arg(
            Arg::with_name("home")
                .long("home")
                .default_value(&home)
                .empty_values(false)
                .help("Home directory which defaults to env var GODCOIN_HOME"),
        )
        .arg(
            Arg::with_name("reindex")
                .long("reindex")
                .help("Reindexes the block log"),
        )
        .arg(
            Arg::with_name("auto_trim")
                .long("reindex-trim-corrupt")
                .help("Trims any corruption detected in the block log during reindexing"),
        )
        .get_matches();

    let home = PathBuf::from(args.value_of("home").expect("Failed to obtain home path"));
    let (blocklog_loc, index_loc) = {
        if !Path::is_dir(&home) {
            let res = std::fs::create_dir(&home);
            res.unwrap_or_else(|_| panic!("Failed to create dir at {:?}", &home));
            info!("Created GODcoin home at {:?}", &home);
        } else {
            info!("Found GODcoin home at {:?}", &home);
        }
        let blocklog_loc = Path::join(&home, "blklog");
        let index_loc = Path::join(&home, "index");
        (blocklog_loc, index_loc)
    };

    let config_file = Path::join(&home, "config.toml");
    info!("Opening configuration file at {:?}", config_file);
    let config_file = fs::read(config_file).expect("Failed to open config");
    let config: Config = toml::from_str(&String::from_utf8(config_file).unwrap()).unwrap();

    let minter_key =
        PrivateKey::from_wif(&config.minter_key).expect("Provided minter key is invalid");
    let bind_addr = config
        .bind_address
        .unwrap_or_else(|| "127.0.0.1:7777".to_owned());

    let reindex = if args.is_present("reindex") {
        info!("User requested reindexing");
        if Path::exists(&index_loc) {
            info!("Deleting current index");
            fs::remove_dir_all(&index_loc)
                .expect("Failed to delete the blockchain index directory");
        } else {
            info!("Current index does not exist");
        }
        let auto_trim = args.is_present("auto_trim");
        Some(ReindexOpts { auto_trim })
    } else {
        None
    };

    let mut rt = Builder::new()
        .panic_handler(|_| std::process::abort())
        .build()
        .unwrap();

    let enable_stale_production = config.enable_stale_production;
    rt.spawn(future::lazy(move || {
        godcoin_server::start(godcoin_server::ServerOpts {
            blocklog_loc,
            index_loc,
            minter_key,
            bind_addr,
            reindex,
            enable_stale_production,
        });
        Ok(())
    }));

    let future = tokio_signal::ctrl_c()
        .flatten_stream()
        .into_future()
        .and_then(|_| {
            info!("Received ctrl-c, shutting down...");
            Ok(())
        });
    rt.block_on(future).map_err(|(e, _)| e).unwrap();
}
<|endoftext|>


<commit_before>use std::{fs, io, time};
use std::io::Read;
use std::sync::atomic;
use std::collections::HashMap;

use slog::Logger;
use chrono::Utc;

use {rpc, tracker, disk, listener, CONFIG, SHUTDOWN, PEER_ID};
use util::{io_err, io_err_val, id_to_hash, hash_to_id};
use torrent::{self, peer, Torrent};
use throttle::Throttler;

pub mod cio;
pub mod acio;
mod job;

/// Tracker update job interval
const TRK_JOB_SECS: u64 = 60;
/// Unchoke rotation job interval
const UNCHK_JOB_SECS: u64 = 15;
/// Session serialization job interval
const SES_JOB_SECS: u64 = 10;
/// Interval to update RPC of transfer stats
const TX_JOB_MS: u64 = 500;

/// Interval to requery all jobs and execute if needed
const JOB_INT_MS: usize = 500;

pub struct Control<T: cio::CIO> {
    throttler: Throttler,
    cio: T,
    tid_cnt: usize,
    job_timer: usize,
    tx_rates: Option<(u64, u64)>,
    jobs: job::JobManager<T>,
    torrents: HashMap<usize, Torrent<T>>,
    peers: HashMap<usize, usize>,
    hash_idx: HashMap<[u8; 20], usize>,
    l: Logger,
}

impl<T: cio::CIO> Control<T> {
    pub fn new(mut cio: T, throttler: Throttler, l: Logger) -> io::Result<Control<T>> {
        let torrents = HashMap::new();
        let peers = HashMap::new();
        let hash_idx = HashMap::new();
        // Every minute check to update trackers;
        let mut jobs = job::JobManager::new();
        jobs.add_job(job::TrackerUpdate, time::Duration::from_secs(TRK_JOB_SECS));
        jobs.add_job(
            job::UnchokeUpdate,
            time::Duration::from_secs(UNCHK_JOB_SECS),
        );
        jobs.add_job(job::SessionUpdate, time::Duration::from_secs(SES_JOB_SECS));
        jobs.add_job(
            job::TorrentTxUpdate::new(),
            time::Duration::from_millis(TX_JOB_MS),
        );
        let job_timer = cio.set_timer(JOB_INT_MS).map_err(
            |_| io_err_val("timer failure!"),
        )?;
        // 5 MiB max bucket
        Ok(Control {
            throttler,
            cio,
            tid_cnt: 0,
            job_timer,
            jobs,
            torrents,
            peers,
            hash_idx,
            tx_rates: None,
            l,
        })
    }

    pub fn run(&mut self) {
        if self.deserialize().is_err() {
            warn!(self.l, "Session deserialization failed!");
        }
        debug!(self.l, "Initialized!");
        self.send_rpc_info();
        let mut events = Vec::with_capacity(20);
        loop {
            self.cio.poll(&mut events);
            for event in events.drain(..) {
                if self.handle_event(event) {
                    self.serialize();
                    return;
                }
            }
            if SHUTDOWN.load(atomic::Ordering::SeqCst) == true {
                break;
            }
        }
    }

    fn serialize(&mut self) {
        debug!(self.l, "Serializing torrents!");
        for (_, torrent) in self.torrents.iter_mut() {
            torrent.serialize();
        }
    }

    fn deserialize(&mut self) -> io::Result<()> {
        debug!(self.l, "Deserializing torrents!");
        let sd = &CONFIG.disk.session;
        for entry in fs::read_dir(sd)? {
            if let Err(e) = self.deserialize_torrent(entry) {
                warn!(self.l, "Failed to deserialize torrent file: {:?}!", e);
            }
        }
        Ok(())
    }

    fn deserialize_torrent(&mut self, entry: io::Result<fs::DirEntry>) -> io::Result<()> {
        let dir = entry?;
        // TODO: We probably should improve this heuristic with and not rely
        // on directory entries, but this is good enough for now.
        if dir.file_name().len() != 40 {
            return Ok(());
        }
        trace!(self.l, "Attempting to deserialize file {:?}", dir);
        let mut f = fs::File::open(dir.path())?;
        let mut data = Vec::new();
        f.read_to_end(&mut data)?;
        trace!(self.l, "Succesfully read file");

        let tid = self.tid_cnt;
        let throttle = self.throttler.get_throttle(tid);
        let log = self.l.new(o!("torrent" => tid));
        if let Ok(t) = Torrent::deserialize(tid, &data, throttle, self.cio.new_handle(), log) {
            trace!(self.l, "Succesfully parsed torrent file {:?}", dir.path());
            self.hash_idx.insert(t.info().hash, tid);
            self.tid_cnt += 1;
            self.torrents.insert(tid, t);
        } else {
            return io_err("Torrent data invalid!");
        }
        Ok(())
    }

    fn handle_event(&mut self, event: cio::Event) -> bool {
        match event {
            cio::Event::Tracker(Ok(e)) => {
                self.handle_trk_ev(e);
            }
            cio::Event::Tracker(Err(e)) => {
                error!(self.l, "tracker error: {:?}", e);
                trace!(self.l, "tracker error bt: {:?}", e.backtrace());
            }
            cio::Event::Disk(Ok(e)) => {
                self.handle_disk_ev(e);
            }
            cio::Event::Disk(Err(e)) => {
                error!(self.l, "disk error: {:?}", e);
                trace!(self.l, "disk error: {:?}", e.backtrace());
            }
            cio::Event::RPC(Ok(e)) => {
                return self.handle_rpc_ev(e);
            }
            cio::Event::RPC(Err(e)) => {
                error!(self.l, "rpc error: {:?}", e);
                trace!(self.l, "rpc error: {:?}", e.backtrace());
            }
            cio::Event::Listener(Ok(e)) => {
                self.handle_lst_ev(e);
            }
            cio::Event::Listener(Err(e)) => {
                error!(self.l, "listener error: {:?}", e);
                trace!(self.l, "listener error: {:?}", e.backtrace());
            }
            cio::Event::Timer(t) => {
                if t == self.throttler.id() {
                    if let Some(p) = self.throttler.update() {
                        self.tx_rates = Some(p);
                    }
                } else if t == self.throttler.fid() {
                    self.flush_blocked_peers();
                } else if t == self.job_timer {
                    self.update_jobs();
                    self.update_rpc_tx();
                } else {
                    error!(self.l, "unknown timer id {} reported", t);
                }
            }
            cio::Event::Peer { peer, event } => {
                self.handle_peer_ev(peer, event);
            }
        }
        false
    }

    fn handle_trk_ev(&mut self, tr: tracker::Response) {
        debug!(self.l, "Handling tracker response");
        let id = tr.0;
        let resp = tr.1;
        {
            if let Some(torrent) = self.torrents.get_mut(&id) {
                torrent.set_tracker_response(&resp);
            } else {
                return;
            }
        }
        trace!(self.l, "Adding peers!");
        if let Ok(r) = resp {
            for ip in r.peers.iter() {
                trace!(self.l, "Adding peer({:?})!", ip);
                if let Ok(peer) = peer::PeerConn::new_outgoing(ip) {
                    trace!(self.l, "Added peer({:?})!", ip);
                    self.add_peer(id, peer);
                }
            }
            if let Some(torrent) = self.torrents.get_mut(&id) {
                torrent.update_rpc_peers();
            }
        }
    }

    fn update_jobs(&mut self) {
        trace!(self.l, "Handling job timer");
        self.jobs.update(&mut self.torrents);
    }

    fn handle_disk_ev(&mut self, resp: disk::Response) {
        trace!(self.l, "Got disk response {:?}!", resp);
        if let Some(torrent) = self.torrents.get_mut(&resp.tid()) {
            torrent.handle_disk_resp(resp);
        }
    }

    fn handle_lst_ev(&mut self, msg: listener::Message) {
        debug!(self.l, "Adding peer for torrent with hash {:?}!", msg.hash);
        if let Some(tid) = self.hash_idx.get(&msg.hash).cloned() {
            self.add_inc_peer(tid, msg.peer, msg.id, msg.rsv);
        } else {
            warn!(
                self.l,
                "Couldn't add peer, torrent with hash {:?} doesn't exist",
                msg.hash
            );
        }
    }

    fn handle_peer_ev(&mut self, peer: cio::PID, ev: cio::Result<torrent::Message>) {
        let ref mut p = self.peers;
        let ref mut t = self.torrents;

        p.get(&peer).cloned().and_then(|id| t.get_mut(&id)).map(
            |torrent| if torrent.peer_ev(peer, ev).is_err() {
                p.remove(&peer);
                torrent.update_rpc_peers();
            },
        );
    }

    fn flush_blocked_peers(&mut self) {
        trace!(self.l, "Flushing blocked peers!");
        self.cio.flush_peers(self.throttler.flush_dl());
        self.cio.flush_peers(self.throttler.flush_ul());
    }

    fn add_torrent(&mut self, info: torrent::Info, path: Option<String>) {
        debug!(self.l, "Adding {:?}!", info);
        if self.hash_idx.contains_key(&info.hash) {
            warn!(self.l, "Torrent already exists!");
            return;
        }
        let tid = self.tid_cnt;
        let throttle = self.throttler.get_throttle(tid);
        let log = self.l.new(o!("torrent" => tid));
        let t = Torrent::new(tid, path, info, throttle, self.cio.new_handle(), log);
        self.hash_idx.insert(t.info().hash, tid);
        self.tid_cnt += 1;
        self.torrents.insert(tid, t);
    }

    fn handle_rpc_ev(&mut self, req: rpc::Message) -> bool {
        debug!(self.l, "Handling rpc reqest!");
        match req {
            rpc::Message::UpdateTorrent(u) => {
                let hash_idx = &self.hash_idx;
                let torrents = &mut self.torrents;
                let res = id_to_hash(&u.id)
                    .and_then(|d| hash_idx.get(d.as_ref()))
                    .and_then(|i| torrents.get_mut(i));
                if let Some(t) = res {
                    t.rpc_update(u);
                }
            }
            rpc::Message::Torrent{ info, path } => self.add_torrent(info, path),
            rpc::Message::UpdateFile {
                id,
                torrent_id,
                priority,
            } => {
                let hash_idx = &self.hash_idx;
                let torrents = &mut self.torrents;
                let res = id_to_hash(&torrent_id)
                    .and_then(|d| hash_idx.get(d.as_ref()))
                    .and_then(|i| torrents.get_mut(i));
                if let Some(t) = res {
                    t.rpc_update_file(id, priority);
                }
            }
            rpc::Message::UpdateServer {
                id,
                throttle_up,
                throttle_down,
            } => {
                let tu = throttle_up.unwrap_or(self.throttler.ul_rate() as u32);
                let td = throttle_down.unwrap_or(self.throttler.dl_rate() as u32);
                self.throttler.set_ul_rate(tu as usize);
                self.throttler.set_dl_rate(td as usize);
                self.cio.msg_rpc(rpc::CtlMessage::Update(vec![
                    rpc::resource::SResourceUpdate::Throttle {
                        id,
                        throttle_up: tu,
                        throttle_down: td,
                    },
                ]));
            }
            rpc::Message::RemoveTorrent(id) => {
                let hash_idx = &mut self.hash_idx;
                let torrents = &mut self.torrents;
                id_to_hash(&id)
                    .and_then(|d| hash_idx.remove(d.as_ref()))
                    .and_then(|i| torrents.remove(&i))
                    .map(|mut t| t.delete());
            }
            rpc::Message::RemovePeer { id, torrent_id } => {
                let hash_idx = &self.hash_idx;
                let torrents = &mut self.torrents;
                id_to_hash(&torrent_id)
                    .and_then(|d| hash_idx.get(d.as_ref()))
                    .and_then(|i| torrents.get_mut(i))
                    .map(|t| t.remove_peer(&id));
            }
            rpc::Message::RemoveTracker { id, torrent_id } => {
                let hash_idx = &self.hash_idx;
                let torrents = &mut self.torrents;
                id_to_hash(&torrent_id)
                    .and_then(|d| hash_idx.get(d.as_ref()))
                    .and_then(|i| torrents.get_mut(i))
                    .map(|t| t.remove_tracker(&id));
            }
        }
        false
    }

    fn add_peer(&mut self, id: usize, peer: peer::PeerConn) {
        trace!(self.l, "Adding peer to torrent {:?}!", id);
        if let Some(torrent) = self.torrents.get_mut(&id) {
            if let Some(pid) = torrent.add_peer(peer) {
                self.peers.insert(pid, id);
            }
        }
    }

    fn add_inc_peer(&mut self, id: usize, peer: peer::PeerConn, cid: [u8; 20], rsv: [u8; 8]) {
        trace!(self.l, "Adding peer to torrent {:?}!", id);
        if let Some(torrent) = self.torrents.get_mut(&id) {
            if let Some(pid) = torrent.add_inc_peer(peer, cid, rsv) {
                self.peers.insert(pid, id);
            }
        }
    }

    fn update_rpc_tx(&mut self) {
        if let Some((rate_up, rate_down)) = self.tx_rates {
            self.cio.msg_rpc(rpc::CtlMessage::Update(vec![
                rpc::resource::SResourceUpdate::Rate {
                    id: hash_to_id(&PEER_ID[..]),
                    rate_up,
                    rate_down,
                },
            ]));
            self.tx_rates = None;
        }
    }

    fn send_rpc_info(&mut self) {
        let res = rpc::resource::Resource::Server(rpc::resource::Server {
            id: hash_to_id(&PEER_ID[..]),
            rate_up: 0,
            rate_down: 0,
            throttle_up: 0,
            throttle_down: 0,
            started: Utc::now(),
        });
        self.cio.msg_rpc(rpc::CtlMessage::Extant(vec![res]));
    }
}

impl<T: cio::CIO> Drop for Control<T> {
    fn drop(&mut self) {
        debug!(self.l, "Triggering thread shutdown sequence!");
        self.cio.msg_disk(disk::Request::shutdown());
        self.cio.msg_rpc(rpc::CtlMessage::Shutdown);
        self.cio.msg_trk(tracker::Request::Shutdown);
        self.cio.msg_listener(listener::Request::Shutdown);
    }
}
<commit_msg>Increase session serialization interval<commit_after>use std::{fs, io, time};
use std::io::Read;
use std::sync::atomic;
use std::collections::HashMap;

use slog::Logger;
use chrono::Utc;

use {rpc, tracker, disk, listener, CONFIG, SHUTDOWN, PEER_ID};
use util::{io_err, io_err_val, id_to_hash, hash_to_id};
use torrent::{self, peer, Torrent};
use throttle::Throttler;

pub mod cio;
pub mod acio;
mod job;

/// Tracker update job interval
const TRK_JOB_SECS: u64 = 60;
/// Unchoke rotation job interval
const UNCHK_JOB_SECS: u64 = 15;
/// Session serialization job interval
const SES_JOB_SECS: u64 = 60;
/// Interval to update RPC of transfer stats
const TX_JOB_MS: u64 = 500;

/// Interval to requery all jobs and execute if needed
const JOB_INT_MS: usize = 500;

pub struct Control<T: cio::CIO> {
    throttler: Throttler,
    cio: T,
    tid_cnt: usize,
    job_timer: usize,
    tx_rates: Option<(u64, u64)>,
    jobs: job::JobManager<T>,
    torrents: HashMap<usize, Torrent<T>>,
    peers: HashMap<usize, usize>,
    hash_idx: HashMap<[u8; 20], usize>,
    l: Logger,
}

impl<T: cio::CIO> Control<T> {
    pub fn new(mut cio: T, throttler: Throttler, l: Logger) -> io::Result<Control<T>> {
        let torrents = HashMap::new();
        let peers = HashMap::new();
        let hash_idx = HashMap::new();
        // Every minute check to update trackers;
        let mut jobs = job::JobManager::new();
        jobs.add_job(job::TrackerUpdate, time::Duration::from_secs(TRK_JOB_SECS));
        jobs.add_job(
            job::UnchokeUpdate,
            time::Duration::from_secs(UNCHK_JOB_SECS),
        );
        jobs.add_job(job::SessionUpdate, time::Duration::from_secs(SES_JOB_SECS));
        jobs.add_job(
            job::TorrentTxUpdate::new(),
            time::Duration::from_millis(TX_JOB_MS),
        );
        let job_timer = cio.set_timer(JOB_INT_MS).map_err(
            |_| io_err_val("timer failure!"),
        )?;
        // 5 MiB max bucket
        Ok(Control {
            throttler,
            cio,
            tid_cnt: 0,
            job_timer,
            jobs,
            torrents,
            peers,
            hash_idx,
            tx_rates: None,
            l,
        })
    }

    pub fn run(&mut self) {
        if self.deserialize().is_err() {
            warn!(self.l, "Session deserialization failed!");
        }
        debug!(self.l, "Initialized!");
        self.send_rpc_info();
        let mut events = Vec::with_capacity(20);
        loop {
            self.cio.poll(&mut events);
            for event in events.drain(..) {
                if self.handle_event(event) {
                    self.serialize();
                    return;
                }
            }
            if SHUTDOWN.load(atomic::Ordering::SeqCst) == true {
                break;
            }
        }
    }

    fn serialize(&mut self) {
        debug!(self.l, "Serializing torrents!");
        for (_, torrent) in self.torrents.iter_mut() {
            torrent.serialize();
        }
    }

    fn deserialize(&mut self) -> io::Result<()> {
        debug!(self.l, "Deserializing torrents!");
        let sd = &CONFIG.disk.session;
        for entry in fs::read_dir(sd)? {
            if let Err(e) = self.deserialize_torrent(entry) {
                warn!(self.l, "Failed to deserialize torrent file: {:?}!", e);
            }
        }
        Ok(())
    }

    fn deserialize_torrent(&mut self, entry: io::Result<fs::DirEntry>) -> io::Result<()> {
        let dir = entry?;
        // TODO: We probably should improve this heuristic with and not rely
        // on directory entries, but this is good enough for now.
        if dir.file_name().len() != 40 {
            return Ok(());
        }
        trace!(self.l, "Attempting to deserialize file {:?}", dir);
        let mut f = fs::File::open(dir.path())?;
        let mut data = Vec::new();
        f.read_to_end(&mut data)?;
        trace!(self.l, "Succesfully read file");

        let tid = self.tid_cnt;
        let throttle = self.throttler.get_throttle(tid);
        let log = self.l.new(o!("torrent" => tid));
        if let Ok(t) = Torrent::deserialize(tid, &data, throttle, self.cio.new_handle(), log) {
            trace!(self.l, "Succesfully parsed torrent file {:?}", dir.path());
            self.hash_idx.insert(t.info().hash, tid);
            self.tid_cnt += 1;
            self.torrents.insert(tid, t);
        } else {
            return io_err("Torrent data invalid!");
        }
        Ok(())
    }

    fn handle_event(&mut self, event: cio::Event) -> bool {
        match event {
            cio::Event::Tracker(Ok(e)) => {
                self.handle_trk_ev(e);
            }
            cio::Event::Tracker(Err(e)) => {
                error!(self.l, "tracker error: {:?}", e);
                trace!(self.l, "tracker error bt: {:?}", e.backtrace());
            }
            cio::Event::Disk(Ok(e)) => {
                self.handle_disk_ev(e);
            }
            cio::Event::Disk(Err(e)) => {
                error!(self.l, "disk error: {:?}", e);
                trace!(self.l, "disk error: {:?}", e.backtrace());
            }
            cio::Event::RPC(Ok(e)) => {
                return self.handle_rpc_ev(e);
            }
            cio::Event::RPC(Err(e)) => {
                error!(self.l, "rpc error: {:?}", e);
                trace!(self.l, "rpc error: {:?}", e.backtrace());
            }
            cio::Event::Listener(Ok(e)) => {
                self.handle_lst_ev(e);
            }
            cio::Event::Listener(Err(e)) => {
                error!(self.l, "listener error: {:?}", e);
                trace!(self.l, "listener error: {:?}", e.backtrace());
            }
            cio::Event::Timer(t) => {
                if t == self.throttler.id() {
                    if let Some(p) = self.throttler.update() {
                        self.tx_rates = Some(p);
                    }
                } else if t == self.throttler.fid() {
                    self.flush_blocked_peers();
                } else if t == self.job_timer {
                    self.update_jobs();
                    self.update_rpc_tx();
                } else {
                    error!(self.l, "unknown timer id {} reported", t);
                }
            }
            cio::Event::Peer { peer, event } => {
                self.handle_peer_ev(peer, event);
            }
        }
        false
    }

    fn handle_trk_ev(&mut self, tr: tracker::Response) {
        debug!(self.l, "Handling tracker response");
        let id = tr.0;
        let resp = tr.1;
        {
            if let Some(torrent) = self.torrents.get_mut(&id) {
                torrent.set_tracker_response(&resp);
            } else {
                return;
            }
        }
        trace!(self.l, "Adding peers!");
        if let Ok(r) = resp {
            for ip in r.peers.iter() {
                trace!(self.l, "Adding peer({:?})!", ip);
                if let Ok(peer) = peer::PeerConn::new_outgoing(ip) {
                    trace!(self.l, "Added peer({:?})!", ip);
                    self.add_peer(id, peer);
                }
            }
            if let Some(torrent) = self.torrents.get_mut(&id) {
                torrent.update_rpc_peers();
            }
        }
    }

    fn update_jobs(&mut self) {
        trace!(self.l, "Handling job timer");
        self.jobs.update(&mut self.torrents);
    }

    fn handle_disk_ev(&mut self, resp: disk::Response) {
        trace!(self.l, "Got disk response {:?}!", resp);
        if let Some(torrent) = self.torrents.get_mut(&resp.tid()) {
            torrent.handle_disk_resp(resp);
        }
    }

    fn handle_lst_ev(&mut self, msg: listener::Message) {
        debug!(self.l, "Adding peer for torrent with hash {:?}!", msg.hash);
        if let Some(tid) = self.hash_idx.get(&msg.hash).cloned() {
            self.add_inc_peer(tid, msg.peer, msg.id, msg.rsv);
        } else {
            warn!(
                self.l,
                "Couldn't add peer, torrent with hash {:?} doesn't exist",
                msg.hash
            );
        }
    }

    fn handle_peer_ev(&mut self, peer: cio::PID, ev: cio::Result<torrent::Message>) {
        let ref mut p = self.peers;
        let ref mut t = self.torrents;

        p.get(&peer).cloned().and_then(|id| t.get_mut(&id)).map(
            |torrent| if torrent.peer_ev(peer, ev).is_err() {
                p.remove(&peer);
                torrent.update_rpc_peers();
            },
        );
    }

    fn flush_blocked_peers(&mut self) {
        trace!(self.l, "Flushing blocked peers!");
        self.cio.flush_peers(self.throttler.flush_dl());
        self.cio.flush_peers(self.throttler.flush_ul());
    }

    fn add_torrent(&mut self, info: torrent::Info, path: Option<String>) {
        debug!(self.l, "Adding {:?}!", info);
        if self.hash_idx.contains_key(&info.hash) {
            warn!(self.l, "Torrent already exists!");
            return;
        }
        let tid = self.tid_cnt;
        let throttle = self.throttler.get_throttle(tid);
        let log = self.l.new(o!("torrent" => tid));
        let t = Torrent::new(tid, path, info, throttle, self.cio.new_handle(), log);
        self.hash_idx.insert(t.info().hash, tid);
        self.tid_cnt += 1;
        self.torrents.insert(tid, t);
    }

    fn handle_rpc_ev(&mut self, req: rpc::Message) -> bool {
        debug!(self.l, "Handling rpc reqest!");
        match req {
            rpc::Message::UpdateTorrent(u) => {
                let hash_idx = &self.hash_idx;
                let torrents = &mut self.torrents;
                let res = id_to_hash(&u.id)
                    .and_then(|d| hash_idx.get(d.as_ref()))
                    .and_then(|i| torrents.get_mut(i));
                if let Some(t) = res {
                    t.rpc_update(u);
                }
            }
            rpc::Message::Torrent{ info, path } => self.add_torrent(info, path),
            rpc::Message::UpdateFile {
                id,
                torrent_id,
                priority,
            } => {
                let hash_idx = &self.hash_idx;
                let torrents = &mut self.torrents;
                let res = id_to_hash(&torrent_id)
                    .and_then(|d| hash_idx.get(d.as_ref()))
                    .and_then(|i| torrents.get_mut(i));
                if let Some(t) = res {
                    t.rpc_update_file(id, priority);
                }
            }
            rpc::Message::UpdateServer {
                id,
                throttle_up,
                throttle_down,
            } => {
                let tu = throttle_up.unwrap_or(self.throttler.ul_rate() as u32);
                let td = throttle_down.unwrap_or(self.throttler.dl_rate() as u32);
                self.throttler.set_ul_rate(tu as usize);
                self.throttler.set_dl_rate(td as usize);
                self.cio.msg_rpc(rpc::CtlMessage::Update(vec![
                    rpc::resource::SResourceUpdate::Throttle {
                        id,
                        throttle_up: tu,
                        throttle_down: td,
                    },
                ]));
            }
            rpc::Message::RemoveTorrent(id) => {
                let hash_idx = &mut self.hash_idx;
                let torrents = &mut self.torrents;
                id_to_hash(&id)
                    .and_then(|d| hash_idx.remove(d.as_ref()))
                    .and_then(|i| torrents.remove(&i))
                    .map(|mut t| t.delete());
            }
            rpc::Message::RemovePeer { id, torrent_id } => {
                let hash_idx = &self.hash_idx;
                let torrents = &mut self.torrents;
                id_to_hash(&torrent_id)
                    .and_then(|d| hash_idx.get(d.as_ref()))
                    .and_then(|i| torrents.get_mut(i))
                    .map(|t| t.remove_peer(&id));
            }
            rpc::Message::RemoveTracker { id, torrent_id } => {
                let hash_idx = &self.hash_idx;
                let torrents = &mut self.torrents;
                id_to_hash(&torrent_id)
                    .and_then(|d| hash_idx.get(d.as_ref()))
                    .and_then(|i| torrents.get_mut(i))
                    .map(|t| t.remove_tracker(&id));
            }
        }
        false
    }

    fn add_peer(&mut self, id: usize, peer: peer::PeerConn) {
        trace!(self.l, "Adding peer to torrent {:?}!", id);
        if let Some(torrent) = self.torrents.get_mut(&id) {
            if let Some(pid) = torrent.add_peer(peer) {
                self.peers.insert(pid, id);
            }
        }
    }

    fn add_inc_peer(&mut self, id: usize, peer: peer::PeerConn, cid: [u8; 20], rsv: [u8; 8]) {
        trace!(self.l, "Adding peer to torrent {:?}!", id);
        if let Some(torrent) = self.torrents.get_mut(&id) {
            if let Some(pid) = torrent.add_inc_peer(peer, cid, rsv) {
                self.peers.insert(pid, id);
            }
        }
    }

    fn update_rpc_tx(&mut self) {
        if let Some((rate_up, rate_down)) = self.tx_rates {
            self.cio.msg_rpc(rpc::CtlMessage::Update(vec![
                rpc::resource::SResourceUpdate::Rate {
                    id: hash_to_id(&PEER_ID[..]),
                    rate_up,
                    rate_down,
                },
            ]));
            self.tx_rates = None;
        }
    }

    fn send_rpc_info(&mut self) {
        let res = rpc::resource::Resource::Server(rpc::resource::Server {
            id: hash_to_id(&PEER_ID[..]),
            rate_up: 0,
            rate_down: 0,
            throttle_up: 0,
            throttle_down: 0,
            started: Utc::now(),
        });
        self.cio.msg_rpc(rpc::CtlMessage::Extant(vec![res]));
    }
}

impl<T: cio::CIO> Drop for Control<T> {
    fn drop(&mut self) {
        debug!(self.l, "Triggering thread shutdown sequence!");
        self.cio.msg_disk(disk::Request::shutdown());
        self.cio.msg_rpc(rpc::CtlMessage::Shutdown);
        self.cio.msg_trk(tracker::Request::Shutdown);
        self.cio.msg_listener(listener::Request::Shutdown);
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::fs;
use std::fs::File;
use std::io::Write;
use std::path::Path;

fn main() {
    let out_dir = env::var("OUT_DIR").unwrap();
    let dest_path = Path::new(&out_dir).join("fonts.rs");
    let mut f = File::create(&dest_path).unwrap();

    writeln!(f,
             "pub fn fonts() -> HashMap<&'static str, &'static [u8]> {{")
            .unwrap();
    writeln!(f,
             "let mut fonts = HashMap::<&'static str, &'static [u8]>::new();")
            .unwrap();
    for l1 in fs::read_dir("./src/static/fonts/").unwrap() {
        let l1fn = l1.unwrap().path();
        if l1fn.is_dir() {
            for pbf in fs::read_dir(l1fn).unwrap() {
                let pbfpath = pbf.unwrap().path();
                let mut pbfcomp = pbfpath.components();
                pbfcomp.next();
                pbfcomp.next();
                let inclpath = pbfcomp.as_path();
                pbfcomp.next();
                let keypath = pbfcomp.as_path();
                writeln!(f,
                         "fonts.insert(\"{}\", include_bytes!(\"{}\"));",
                         keypath.display(),
                         inclpath.display())
                        .unwrap();
            }
        }
    }
    writeln!(f, "fonts").unwrap();
    writeln!(f, "}}").unwrap();
}
<commit_msg>Force Unix path separators in pbf lookup map<commit_after>use std::env;
use std::fs;
use std::fs::File;
use std::io::Write;
use std::path::Path;

fn main() {
    let out_dir = env::var("OUT_DIR").unwrap();
    let dest_path = Path::new(&out_dir).join("fonts.rs");
    let mut f = File::create(&dest_path).unwrap();

    writeln!(f,
             "pub fn fonts() -> HashMap<&'static str, &'static [u8]> {{")
            .unwrap();
    writeln!(f,
             "let mut fonts = HashMap::<&'static str, &'static [u8]>::new();")
            .unwrap();
    for l1 in fs::read_dir("./src/static/fonts/").unwrap() {
        let l1fn = l1.unwrap().path();
        if l1fn.is_dir() {
            for pbf in fs::read_dir(l1fn).unwrap() {
                let pbfpath = pbf.unwrap().path();
                let mut pbfcomp = pbfpath.components();
                pbfcomp.next();
                pbfcomp.next();
                let inclpath = pbfcomp.as_path();
                pbfcomp.next();
                let keypath = pbfcomp.as_path();
                writeln!(f,
                         "fonts.insert(\"{}\", include_bytes!(\"{}\"));",
                         keypath.display().to_string().replace(std::path::MAIN_SEPARATOR, "/"),
                         inclpath.display().to_string().replace(std::path::MAIN_SEPARATOR, "/"))
                        .unwrap();
            }
        }
    }
    writeln!(f, "fonts").unwrap();
    writeln!(f, "}}").unwrap();
}
<|endoftext|>


<commit_before>#![feature(proc_macro, custom_attribute)]

#[macro_use]
extern crate smelter;

#[derive(PartialEq, Debug, Builder, Default)]
struct Point {
    x: u32,
    #[smelter(field_name="y_axis")]
    y: u32,
}

#[derive(PartialEq, Debug, Builder, Default)]
struct Container<T>
    where T: PartialEq  + Default {
    item: T,
}

#[test]
fn can_generate_builder_methods() {
    let point = Point::default().x(1).y_axis(2);
    let expected = Point {x: 1, y: 2};
    assert_eq!(point, expected);
}

#[test]
fn can_generate_generic_builder_methods() {
    let container: Container<u32> = Container::default().item(1u32);
    let expected = Container { item: 1u32 };
    assert_eq!(container, expected);
}

#[test]
fn can_generate_mutable_methods() {
    let mut point = Point::default();
    point.x_mut(1).y_axis_mut(2);
    let expected = Point { x: 1, y: 2};
    assert_eq!(point, expected);
}

#[derive(PartialEq, Builder, Default, Debug)]
#[smelter(prefix="with_")]
struct ContainerWith<T>
    where T: PartialEq + Default {
        item: T,
        #[smelter(field_name = "id")]
        item_id: u64,
}

#[test]
fn can_generate_container_with_prefix() {
    let container: ContainerWith<u32> = ContainerWith::default()
                                            .with_item(1u32)
                                            .with_id(5u64);
    let expected = ContainerWith { item: 1u32 , item_id: 5u64 };
    assert_eq!(container, expected);
}

#[derive(Builder, PartialEq, Debug, Default)]
struct LotsOfFields<T: PartialEq> {
    pub this: String,
    structure: u32,
    has: i32,
    pub a: String,
    lot: T,
    pub of: &'static str,
    fields: String,
}

#[test]
fn public_fields_work() {
    let this = "this".to_string();
    let structure = 1u32;
    let has = -10i32;
    let a = "a".to_string();
    let lot = 7u32;
    let of = &"of";
    let fields = "fields".to_string();
    let expected: LotsOfFields<u32> = LotsOfFields {
        this: this.clone(),
        structure: structure.clone(),
        has: has.clone(),
        a: a.clone(),
        lot: lot.clone(),
        of: of,
        fields: fields.clone(),
    };
    let lof: LotsOfFields<u32> = LotsOfFields::default()
        .this(this)
        .structure(structure)
        .has(has)
        .a(a)
        .lot(lot)
        .of(of)
        .fields(fields);

    assert_eq!(lof, expected);
}

#[derive(Builder, PartialEq, Debug)]
struct WithLifeTime<'a> {
    l: &'a String,
}

#[test]
fn with_lifetime() {
    let s1 = "hello".to_string();
    let s2 = "hello".to_string();
    let s3 = "world".to_string();
    let with_lifetime = WithLifeTime {
        l: &s1,
    };

    let expected = WithLifeTime {
        l: &s3,
    }.l(&s2);

    assert_eq!(with_lifetime, expected);
}

mod pub_definition;
use pub_definition::*;

#[test]
fn exported_public_method() {
    let mut t: PubTest = Default::default();
    // We can access r mutable here
    t.r_mut(1u32);
    let not_expected: PubTest = Default::default();
    assert!(&not_expected != &t);
}
<commit_msg>Derived a struct with a collection.<commit_after>#![feature(proc_macro, custom_attribute)]

#[macro_use]
extern crate smelter;

#[derive(PartialEq, Debug, Builder, Default)]
struct Point {
    x: u32,
    #[smelter(field_name="y_axis")]
    y: u32,
}

#[derive(PartialEq, Debug, Builder, Default)]
struct Container<T>
    where T: PartialEq  + Default {
    item: T,
}

#[test]
fn can_generate_builder_methods() {
    let point = Point::default().x(1).y_axis(2);
    let expected = Point {x: 1, y: 2};
    assert_eq!(point, expected);
}

#[test]
fn can_generate_generic_builder_methods() {
    let container: Container<u32> = Container::default().item(1u32);
    let expected = Container { item: 1u32 };
    assert_eq!(container, expected);
}

#[test]
fn can_generate_mutable_methods() {
    let mut point = Point::default();
    point.x_mut(1).y_axis_mut(2);
    let expected = Point { x: 1, y: 2};
    assert_eq!(point, expected);
}

#[derive(PartialEq, Builder, Default, Debug)]
#[smelter(prefix="with_")]
struct ContainerWith<T>
    where T: PartialEq + Default {
        item: T,
        #[smelter(field_name = "id")]
        item_id: u64,
}

#[test]
fn can_generate_container_with_prefix() {
    let container: ContainerWith<u32> = ContainerWith::default()
                                            .with_item(1u32)
                                            .with_id(5u64);
    let expected = ContainerWith { item: 1u32 , item_id: 5u64 };
    assert_eq!(container, expected);
}

#[derive(Builder, PartialEq, Debug, Default)]
struct LotsOfFields<T: PartialEq> {
    pub this: String,
    structure: u32,
    has: i32,
    pub a: String,
    lot: T,
    pub of: &'static str,
    fields: String,
}

#[test]
fn public_fields_work() {
    let this = "this".to_string();
    let structure = 1u32;
    let has = -10i32;
    let a = "a".to_string();
    let lot = 7u32;
    let of = &"of";
    let fields = "fields".to_string();
    let expected: LotsOfFields<u32> = LotsOfFields {
        this: this.clone(),
        structure: structure.clone(),
        has: has.clone(),
        a: a.clone(),
        lot: lot.clone(),
        of: of,
        fields: fields.clone(),
    };
    let lof: LotsOfFields<u32> = LotsOfFields::default()
        .this(this)
        .structure(structure)
        .has(has)
        .a(a)
        .lot(lot)
        .of(of)
        .fields(fields);

    assert_eq!(lof, expected);
}

#[derive(Builder, PartialEq, Debug)]
struct WithLifeTime<'a> {
    l: &'a String,
}

#[test]
fn with_lifetime() {
    let s1 = "hello".to_string();
    let s2 = "hello".to_string();
    let s3 = "world".to_string();
    let with_lifetime = WithLifeTime {
        l: &s1,
    };

    let expected = WithLifeTime {
        l: &s3,
    }.l(&s2);

    assert_eq!(with_lifetime, expected);
}

mod pub_definition;
use pub_definition::*;

#[test]
fn exported_public_method() {
    let mut t: PubTest = Default::default();
    // We can access r mutable here
    t.r_mut(1u32);
    let not_expected: PubTest = Default::default();
    assert!(&not_expected != &t);
}

#[derive(PartialEq, Builder, Default, Debug, Clone)]
#[smelter(prefix="with_")]
pub struct User {
    pub uid: u64,
    pub email: String,
    pub alias: String,
    pub friends: Vec<User>,
}

#[test]
fn can_derive_collection() {
    let mut u1 = User::default();
    let u2 = User::default()
                .with_email("email@example.com".to_string())
                .with_alias("Ed".to_string())
                .with_uid(10u64);
    u1.with_email_mut("email@example.com".to_string())
      .with_alias_mut("Ed".to_string())
      .with_uid_mut(10u64);

    assert_eq!(u1, u2);
    let u3 = User::default().with_friends(vec![u1.clone(), u2.clone()]);
    assert_eq!(vec![u1, u2], u3.friends);

}
<|endoftext|>


<commit_before>// This module defines constants and functions for working with zakocmp
// configuration files.

extern crate yaml_rust;

use std::io::{Error, ErrorKind};
use std::result::Result;
use std::string::String;
use std::vec::Vec;
use yaml_rust::{Yaml, YamlLoader};

const DEFAULT_POLICY_KEY: &'static str = "default-policy";
const POLICIES_KEY: &'static str = "policies";

// Enumerates the string representations of known policies.
const POLICY_REPR_IGNORE: &'static str = "ignore";
const POLICY_REPR_NOADD: &'static str = "noadd";
const POLICY_REPR_NODELETE: &'static str = "nodelete";
const POLICY_REPR_NOMODIFY: &'static str = "nomodify";
const POLICY_REPR_IMMUTABLE: &'static str = "immutable";

// Represents known policies as integers.
pub const POLICY_IGNORE: i32 = 0;
pub const POLICY_NOADD: i32 = 1 << 0;
pub const POLICY_NODELETE: i32 = 1 << 1;
pub const POLICY_NOMODIFY: i32 = 1 << 2;
pub const POLICY_IMMUTABLE: i32 = POLICY_NOADD | POLICY_NODELETE | POLICY_NOMODIFY;

// Represents a sorted vector of zakopane config rules, each mapping a
// path (prefix) to a policy. This type alias is provided for ease of
// coding.
type Policies = Vec<(String, i32)>;

// Represents a zakopane config. Please consult the documentation.
pub struct Config {
    default_policy: i32,
    policies: Policies,
}

// Borrows the string representation of one policy `token` and returns
// the equivalent integral representation.
fn policy_token_as_int(token: &str) -> Result<i32, Error> {
    match token {
        POLICY_REPR_IGNORE => Ok(POLICY_IGNORE),
        POLICY_REPR_NOADD => Ok(POLICY_NOADD),
        POLICY_REPR_NODELETE => Ok(POLICY_NODELETE),
        POLICY_REPR_NOMODIFY => Ok(POLICY_NOMODIFY),
        POLICY_REPR_IMMUTABLE => Ok(POLICY_IMMUTABLE),
        _ => Err(Error::new(
            ErrorKind::InvalidInput,
            format!("bad token: ``{}''", token),
        )),
    }
}

// Borrows the string representation of a combined `policy` and returns
// the equivalent integral representation. This function expects
// `policy` to comprise one or more policy tokens separated by commas.
fn policy_tokens_as_int(policy: &str) -> Result<i32, Error> {
    let policy_ints: Vec<i32> = policy
        .split(",")
        .map(|tok| policy_token_as_int(tok))
        .collect::<Result<Vec<i32>, Error>>()?;
    return Ok(policy_ints
        .iter()
        .fold(POLICY_IGNORE, |accum, elem| accum | elem));
}

// Borrows yaml representations of one line of zakopane policy and
// returns the corresponding valid tuple suitable for use in building a
// Policies object.
fn extract_policy(ypath: &Yaml, policy_tokens: &Yaml) -> Result<(String, i32), Error> {
    let path: String = match ypath.as_str() {
        Some(string) => string.to_owned(),
        None => return Err(Error::new(ErrorKind::InvalidData, "malformed path?")),
    };
    let policy: i32 = match policy_tokens.as_str() {
        Some(string) => policy_tokens_as_int(string)?,
        None => return Err(Error::new(ErrorKind::InvalidData, "malformed policy?")),
    };
    Ok((path, policy))
}

// Borrows the YAML representation of a zakopane config and returns the
// corresponding Policies. The return value can be benignly
// empty (e.g. if the present config elects not to specify any rules).
fn policies_from_yaml(doc: &Yaml) -> Result<Policies, Error> {
    let policies_map_yaml = &doc[POLICIES_KEY];
    if policies_map_yaml.is_badvalue() {
        // Assumes the config may be benignly devoid of specific
        // policies, returning Ok.
        return Ok(vec![]);
    }
    // Otherwise, iterates over the policies map. Each entry in the
    // policies map correlates a path prefix to a comma-separated list
    // of policies.
    let policies_map: &yaml_rust::yaml::Hash = match policies_map_yaml.as_hash() {
        Some(map) => map,
        None => return Err(Error::new(ErrorKind::InvalidData, "malformed policies")),
    };
    let mut policies: Policies = policies_map
        .into_iter()
        .map(|pair| extract_policy(&pair.0, &pair.1))
        .collect::<Result<Policies, Error>>()?;
    policies.sort_unstable_by_key(|pair| pair.0.to_owned());
    Ok(policies)
}

// Borrows the YAML representation of a zakopane config and returns the
// integral default-policy defined within.
fn default_policy_from_yaml(doc: &Yaml) -> Result<i32, Error> {
    let default_policy_yaml = &doc[DEFAULT_POLICY_KEY];
    if default_policy_yaml.is_badvalue() {
        return Err(Error::new(ErrorKind::InvalidData, DEFAULT_POLICY_KEY));
    }
    let default_policy: i32 = match default_policy_yaml.as_str() {
        None => return Err(Error::new(ErrorKind::InvalidData, DEFAULT_POLICY_KEY)),
        Some(token) => policy_tokens_as_int(&token),
    }?;
    Ok(default_policy)
}

impl Config {
    // Borrows the string representation of a zakopane config and
    // returns a corresponding Config.
    pub fn new(config: &str) -> Result<Config, Error> {
        let docs: Vec<Yaml> = match YamlLoader::load_from_str(config) {
            Ok(val) => val,
            Err(scan_error) => return Err(Error::new(ErrorKind::InvalidData, scan_error)),
        };
        if docs.len() == 0 {
            return Err(Error::new(ErrorKind::InvalidData, "empty zakopane config"));
        }
        let doc = &docs[0];

        let default_policy = default_policy_from_yaml(&doc)?;
        let policies: Policies = policies_from_yaml(&doc)?;

        Ok(Config {
            default_policy: default_policy,
            policies: policies,
        })
    }

    // Returns how many rules this config contains.
    // This is never less than 1 as a default-policy is always required.
    pub fn rules(&self) -> usize {
        1 + self.policies.len()
    }

    // Borrows a `path` and returns the best-matched policy that
    // applies. This function returns an owned tuple of the
    // (closest-matched path expression, integral policy).
    //
    // This function represents the default-policy fallback by
    // returning the tuple consisting of an empty &str and the
    // default policy.
    pub fn match_policy(&self, path: &str) -> (&str, i32) {
        let mut best_match_path: &str = "";
        let mut best_match_policy: i32 = 0;
        for (prefix, policy) in self.policies.iter() {
            if path.starts_with(prefix) && prefix.len() > best_match_path.len() {
                best_match_path = prefix;
                best_match_policy = *policy;
            }
        }
        if best_match_path.len() == 0 {
            return ("", self.default_policy);
        }
        return (best_match_path, best_match_policy);
    }
}

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

    #[test]
    fn policy_token_bare() {
        let policy: i32 = policy_tokens_as_int(&"noadd").unwrap();
        assert!(policy == POLICY_NOADD);

        let policy: i32 = policy_tokens_as_int(&"nodelete").unwrap();
        assert!(policy == POLICY_NODELETE);

        let policy: i32 = policy_tokens_as_int(&"nomodify").unwrap();
        assert!(policy == POLICY_NOMODIFY);
    }

    #[test]
    fn policy_tokens_can_combo() {
        let policy: i32 = policy_tokens_as_int(&"noadd,nodelete").unwrap();
        assert!(policy == POLICY_NOADD | POLICY_NODELETE);
    }

    #[test]
    fn policy_tokens_can_repeat() {
        let policy: i32 =
            policy_tokens_as_int(&"noadd,noadd,noadd,noadd,nodelete,nodelete,nodelete,noadd")
                .unwrap();
        assert!(policy == POLICY_NOADD | POLICY_NODELETE);
    }

    #[test]
    fn config_must_not_be_empty() {
        let config = "";
        assert!(!Config::new(&config).is_ok());
    }

    #[test]
    fn config_must_not_be_obviously_malformed() {
        let config = r#"
This is not a zakopane config -
rather, it's two lines of text.
        "#;
        assert!(!Config::new(&config).is_ok());
    }

    #[test]
    fn config_requires_default_policy() {
        let config = r#"
policies:
    hello-there: nomodify
        "#;
        assert!(!Config::new(&config).is_ok());

        let mut config_with_default_policy: String = r#"
default-policy: immutable
        "#
        .to_string();
        config_with_default_policy.push_str(&config);
        assert!(Config::new(&config_with_default_policy).is_ok())
    }

    #[test]
    fn config_might_not_have_specific_policies() {
        let config = r#"
default-policy: nodelete
one-irrelevant-key: it doesn't matter what we put here
another-irrelevant-key: this doesn't invalidate the YAML
third-irrelevant-key: so long as it contains a default-policy
        "#;
        assert!(Config::new(&config).is_ok());
    }

    #[test]
    fn config_policies_must_be_a_map() {
        let config = r#"
default-policy: noadd
policies:
    -   eh?
    -   this ain't a map
        "#;
        assert!(!Config::new(&config).is_ok());
    }

    #[test]
    fn config_can_have_several_policies() {
        let config = r#"
default-policy: immutable
policies:
    hello-there: noadd
    general-kenobi: nodelete
        "#;
        assert!(Config::new(&config).is_ok());
    }

    #[test]
    fn match_default_policy() {
        let config_yaml = r#"
default-policy: noadd
        "#;
        let config = Config::new(&config_yaml).unwrap();

        // With only a default policy, this config has just 1 rule.
        assert!(config.rules() == 1);

        // Any path prefix we throw at match_policy() shall come up
        // as the default policy.
        let (_path, policy) = config.match_policy("./Documents/hello/there.txt");
        assert!(policy == POLICY_NOADD);
        let (_path, policy) = config.match_policy("./Music/general/kenobi.txt");
        assert!(policy == POLICY_NOADD);
    }

    #[test]
    fn match_nondefault_policies() {
        let config_yaml = r#"
default-policy: immutable
policies:
    ./Pictures/: noadd
    ./Pictures/2019/third-party/: nodelete
    ./Pictures/2020/: nomodify
    ./Pictures/2020/food/: nodelete,nomodify
        "#;
        let config = Config::new(&config_yaml).unwrap();

        assert!(config.rules() == 5);

        // Falls back on the default-policy absent any specific policy
        // defined for this file.
        let (_path, policy) = config.match_policy("./Documents/catch-me-senpai.txt");
        assert!(policy == POLICY_IMMUTABLE);
        // Matches only ``./Pictures.''
        let (_path, policy) = config.match_policy("./Pictures/2016/yano.jpg");
        assert!(policy == POLICY_NOADD);
        // As above and does _not_ match ``./Pictures/2019/third-party/.''
        let (_path, policy) = config.match_policy("./Pictures/2019/first-party.jpg");
        assert!(policy == POLICY_NOADD);
        // Does match ``./Pictures/2019/third-party/.''
        let (_path, policy) = config.match_policy("./Pictures/2019/third-party/yano.jpg");
        assert!(policy == POLICY_NODELETE);

        // Path prefix matching is done strictly and exactly;
        // ``food.md'' doesn't match ``food/,'' so there's no risk of
        // zakopane confusing cohabiting entities with similar basenames.
        let (path, policy) = config.match_policy("./Pictures/2020/food.md");
        assert!(policy == POLICY_NOMODIFY);
        assert!(path == "./Pictures/2020/");
        let (path, policy) = config.match_policy("./Pictures/2020/food/tacos.jpg");
        assert!(policy == POLICY_NODELETE | POLICY_NOMODIFY);
        assert!(path == "./Pictures/2020/food/");
    }
}
<commit_msg>zakocmp: use is_err() in unit tests<commit_after>// This module defines constants and functions for working with zakocmp
// configuration files.

extern crate yaml_rust;

use std::io::{Error, ErrorKind};
use std::result::Result;
use std::string::String;
use std::vec::Vec;
use yaml_rust::{Yaml, YamlLoader};

const DEFAULT_POLICY_KEY: &'static str = "default-policy";
const POLICIES_KEY: &'static str = "policies";

// Enumerates the string representations of known policies.
const POLICY_REPR_IGNORE: &'static str = "ignore";
const POLICY_REPR_NOADD: &'static str = "noadd";
const POLICY_REPR_NODELETE: &'static str = "nodelete";
const POLICY_REPR_NOMODIFY: &'static str = "nomodify";
const POLICY_REPR_IMMUTABLE: &'static str = "immutable";

// Represents known policies as integers.
pub const POLICY_IGNORE: i32 = 0;
pub const POLICY_NOADD: i32 = 1 << 0;
pub const POLICY_NODELETE: i32 = 1 << 1;
pub const POLICY_NOMODIFY: i32 = 1 << 2;
pub const POLICY_IMMUTABLE: i32 = POLICY_NOADD | POLICY_NODELETE | POLICY_NOMODIFY;

// Represents a sorted vector of zakopane config rules, each mapping a
// path (prefix) to a policy. This type alias is provided for ease of
// coding.
type Policies = Vec<(String, i32)>;

// Represents a zakopane config. Please consult the documentation.
pub struct Config {
    default_policy: i32,
    policies: Policies,
}

// Borrows the string representation of one policy `token` and returns
// the equivalent integral representation.
fn policy_token_as_int(token: &str) -> Result<i32, Error> {
    match token {
        POLICY_REPR_IGNORE => Ok(POLICY_IGNORE),
        POLICY_REPR_NOADD => Ok(POLICY_NOADD),
        POLICY_REPR_NODELETE => Ok(POLICY_NODELETE),
        POLICY_REPR_NOMODIFY => Ok(POLICY_NOMODIFY),
        POLICY_REPR_IMMUTABLE => Ok(POLICY_IMMUTABLE),
        _ => Err(Error::new(
            ErrorKind::InvalidInput,
            format!("bad token: ``{}''", token),
        )),
    }
}

// Borrows the string representation of a combined `policy` and returns
// the equivalent integral representation. This function expects
// `policy` to comprise one or more policy tokens separated by commas.
fn policy_tokens_as_int(policy: &str) -> Result<i32, Error> {
    let policy_ints: Vec<i32> = policy
        .split(",")
        .map(|tok| policy_token_as_int(tok))
        .collect::<Result<Vec<i32>, Error>>()?;
    return Ok(policy_ints
        .iter()
        .fold(POLICY_IGNORE, |accum, elem| accum | elem));
}

// Borrows yaml representations of one line of zakopane policy and
// returns the corresponding valid tuple suitable for use in building a
// Policies object.
fn extract_policy(ypath: &Yaml, policy_tokens: &Yaml) -> Result<(String, i32), Error> {
    let path: String = match ypath.as_str() {
        Some(string) => string.to_owned(),
        None => return Err(Error::new(ErrorKind::InvalidData, "malformed path?")),
    };
    let policy: i32 = match policy_tokens.as_str() {
        Some(string) => policy_tokens_as_int(string)?,
        None => return Err(Error::new(ErrorKind::InvalidData, "malformed policy?")),
    };
    Ok((path, policy))
}

// Borrows the YAML representation of a zakopane config and returns the
// corresponding Policies. The return value can be benignly
// empty (e.g. if the present config elects not to specify any rules).
fn policies_from_yaml(doc: &Yaml) -> Result<Policies, Error> {
    let policies_map_yaml = &doc[POLICIES_KEY];
    if policies_map_yaml.is_badvalue() {
        // Assumes the config may be benignly devoid of specific
        // policies, returning Ok.
        return Ok(vec![]);
    }
    // Otherwise, iterates over the policies map. Each entry in the
    // policies map correlates a path prefix to a comma-separated list
    // of policies.
    let policies_map: &yaml_rust::yaml::Hash = match policies_map_yaml.as_hash() {
        Some(map) => map,
        None => return Err(Error::new(ErrorKind::InvalidData, "malformed policies")),
    };
    let mut policies: Policies = policies_map
        .into_iter()
        .map(|pair| extract_policy(&pair.0, &pair.1))
        .collect::<Result<Policies, Error>>()?;
    policies.sort_unstable_by_key(|pair| pair.0.to_owned());
    Ok(policies)
}

// Borrows the YAML representation of a zakopane config and returns the
// integral default-policy defined within.
fn default_policy_from_yaml(doc: &Yaml) -> Result<i32, Error> {
    let default_policy_yaml = &doc[DEFAULT_POLICY_KEY];
    if default_policy_yaml.is_badvalue() {
        return Err(Error::new(ErrorKind::InvalidData, DEFAULT_POLICY_KEY));
    }
    let default_policy: i32 = match default_policy_yaml.as_str() {
        None => return Err(Error::new(ErrorKind::InvalidData, DEFAULT_POLICY_KEY)),
        Some(token) => policy_tokens_as_int(&token),
    }?;
    Ok(default_policy)
}

impl Config {
    // Borrows the string representation of a zakopane config and
    // returns a corresponding Config.
    pub fn new(config: &str) -> Result<Config, Error> {
        let docs: Vec<Yaml> = match YamlLoader::load_from_str(config) {
            Ok(val) => val,
            Err(scan_error) => return Err(Error::new(ErrorKind::InvalidData, scan_error)),
        };
        if docs.len() == 0 {
            return Err(Error::new(ErrorKind::InvalidData, "empty zakopane config"));
        }
        let doc = &docs[0];

        let default_policy = default_policy_from_yaml(&doc)?;
        let policies: Policies = policies_from_yaml(&doc)?;

        Ok(Config {
            default_policy: default_policy,
            policies: policies,
        })
    }

    // Returns how many rules this config contains.
    // This is never less than 1 as a default-policy is always required.
    pub fn rules(&self) -> usize {
        1 + self.policies.len()
    }

    // Borrows a `path` and returns the best-matched policy that
    // applies. This function returns an owned tuple of the
    // (closest-matched path expression, integral policy).
    //
    // This function represents the default-policy fallback by
    // returning the tuple consisting of an empty &str and the
    // default policy.
    pub fn match_policy(&self, path: &str) -> (&str, i32) {
        let mut best_match_path: &str = "";
        let mut best_match_policy: i32 = 0;
        for (prefix, policy) in self.policies.iter() {
            if path.starts_with(prefix) && prefix.len() > best_match_path.len() {
                best_match_path = prefix;
                best_match_policy = *policy;
            }
        }
        if best_match_path.len() == 0 {
            return ("", self.default_policy);
        }
        return (best_match_path, best_match_policy);
    }
}

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

    #[test]
    fn policy_token_bare() {
        let policy: i32 = policy_tokens_as_int(&"noadd").unwrap();
        assert!(policy == POLICY_NOADD);

        let policy: i32 = policy_tokens_as_int(&"nodelete").unwrap();
        assert!(policy == POLICY_NODELETE);

        let policy: i32 = policy_tokens_as_int(&"nomodify").unwrap();
        assert!(policy == POLICY_NOMODIFY);
    }

    #[test]
    fn policy_tokens_can_combo() {
        let policy: i32 = policy_tokens_as_int(&"noadd,nodelete").unwrap();
        assert!(policy == POLICY_NOADD | POLICY_NODELETE);
    }

    #[test]
    fn policy_tokens_can_repeat() {
        let policy: i32 =
            policy_tokens_as_int(&"noadd,noadd,noadd,noadd,nodelete,nodelete,nodelete,noadd")
                .unwrap();
        assert!(policy == POLICY_NOADD | POLICY_NODELETE);
    }

    #[test]
    fn config_must_not_be_empty() {
        let config = "";
        assert!(Config::new(&config).is_err());
    }

    #[test]
    fn config_must_not_be_obviously_malformed() {
        let config = r#"
This is not a zakopane config -
rather, it's two lines of text.
        "#;
        assert!(Config::new(&config).is_err());
    }

    #[test]
    fn config_requires_default_policy() {
        let config = r#"
policies:
    hello-there: nomodify
        "#;
        assert!(Config::new(&config).is_err());

        let mut config_with_default_policy: String = r#"
default-policy: immutable
        "#
        .to_string();
        config_with_default_policy.push_str(&config);
        assert!(Config::new(&config_with_default_policy).is_ok())
    }

    #[test]
    fn config_might_not_have_specific_policies() {
        let config = r#"
default-policy: nodelete
one-irrelevant-key: it doesn't matter what we put here
another-irrelevant-key: this doesn't invalidate the YAML
third-irrelevant-key: so long as it contains a default-policy
        "#;
        assert!(Config::new(&config).is_ok());
    }

    #[test]
    fn config_policies_must_be_a_map() {
        let config = r#"
default-policy: noadd
policies:
    -   eh?
    -   this ain't a map
        "#;
        assert!(Config::new(&config).is_err());
    }

    #[test]
    fn config_can_have_several_policies() {
        let config = r#"
default-policy: immutable
policies:
    hello-there: noadd
    general-kenobi: nodelete
        "#;
        assert!(Config::new(&config).is_ok());
    }

    #[test]
    fn match_default_policy() {
        let config_yaml = r#"
default-policy: noadd
        "#;
        let config = Config::new(&config_yaml).unwrap();

        // With only a default policy, this config has just 1 rule.
        assert!(config.rules() == 1);

        // Any path prefix we throw at match_policy() shall come up
        // as the default policy.
        let (_path, policy) = config.match_policy("./Documents/hello/there.txt");
        assert!(policy == POLICY_NOADD);
        let (_path, policy) = config.match_policy("./Music/general/kenobi.txt");
        assert!(policy == POLICY_NOADD);
    }

    #[test]
    fn match_nondefault_policies() {
        let config_yaml = r#"
default-policy: immutable
policies:
    ./Pictures/: noadd
    ./Pictures/2019/third-party/: nodelete
    ./Pictures/2020/: nomodify
    ./Pictures/2020/food/: nodelete,nomodify
        "#;
        let config = Config::new(&config_yaml).unwrap();

        assert!(config.rules() == 5);

        // Falls back on the default-policy absent any specific policy
        // defined for this file.
        let (_path, policy) = config.match_policy("./Documents/catch-me-senpai.txt");
        assert!(policy == POLICY_IMMUTABLE);
        // Matches only ``./Pictures.''
        let (_path, policy) = config.match_policy("./Pictures/2016/yano.jpg");
        assert!(policy == POLICY_NOADD);
        // As above and does _not_ match ``./Pictures/2019/third-party/.''
        let (_path, policy) = config.match_policy("./Pictures/2019/first-party.jpg");
        assert!(policy == POLICY_NOADD);
        // Does match ``./Pictures/2019/third-party/.''
        let (_path, policy) = config.match_policy("./Pictures/2019/third-party/yano.jpg");
        assert!(policy == POLICY_NODELETE);

        // Path prefix matching is done strictly and exactly;
        // ``food.md'' doesn't match ``food/,'' so there's no risk of
        // zakopane confusing cohabiting entities with similar basenames.
        let (path, policy) = config.match_policy("./Pictures/2020/food.md");
        assert!(policy == POLICY_NOMODIFY);
        assert!(path == "./Pictures/2020/");
        let (path, policy) = config.match_policy("./Pictures/2020/food/tacos.jpg");
        assert!(policy == POLICY_NODELETE | POLICY_NOMODIFY);
        assert!(path == "./Pictures/2020/food/");
    }
}
<|endoftext|>