<commit_before>//! The dump module contains code related for outputing/dumping data.
use std::fmt::Display;
use std::convert::From;
use super::errors::*;

/// Enum which provides all possible output value formats supported by the dump module.
#[derive(Debug, PartialEq, Clone, Copy)]
pub enum OutputFormat {
    Hex,
    Decimal,
    Octal,
    Binary,
}

impl From<String> for OutputFormat {
    fn from(format_string: String) -> Self {
        match format_string.to_lowercase().as_ref() {
            "hex" => OutputFormat::Hex,
            "dec" => OutputFormat::Decimal,
            "oct" => OutputFormat::Octal,
            "bin" => OutputFormat::Binary,
            _ => panic!("Invalid output format"),
        }
    }
}

#[derive(Debug)]
pub struct OutputSettings {
    start_address: u32,
    show_address: bool,
    group_size: usize,
    columns: usize,
    show_interpretation: bool,
    output_fmt: OutputFormat,
}

impl OutputSettings {
    fn new() -> OutputSettings {
        OutputSettings {
            start_address: 0,
            show_address: true,
            group_size: 1,
            columns: 8,
            show_interpretation: true,
            output_fmt: OutputFormat::Hex,
        }
    }

    pub fn start_address(mut self, address: u32) -> Self {
        self.start_address = address;
        self
    }

    pub fn show_address(mut self, show: bool) -> Self {
        self.show_address = show;
        self
    }

    pub fn group_size(mut self, size: usize) -> Self {
        self.group_size = size;
        self
    }

    pub fn columns(mut self, columns: usize) -> Self {
        self.columns = columns;
        self
    }

    pub fn show_interpretation(mut self, show: bool) -> Self {
        self.show_interpretation = show;
        self
    }

    pub fn format(mut self, fmt: OutputFormat) -> Self {
        self.output_fmt = fmt;
        self
    }
}

/// The `OutputLine` struct contains all  information needed to dump/output a single line of data.
#[derive(Debug)]
pub struct OutputLine<'a> {
    output_settings: OutputSettings,
    data: &'a [u8], // TODO NiCo: Add member for data format (output format of data)
}

impl<'a> OutputLine<'a> {
    fn new(data: &[u8]) -> OutputLine {
        OutputLine {
            output_settings: OutputSettings::new(),
            data: data,
        }
    }

    pub fn format(self, settings: OutputSettings) -> Self {
        OutputLine {
            output_settings: settings,
            data: self.data,
        }
    }

    fn write_address(&self, f: &mut ::fmt::Formatter) -> Result<()> {
        write!(f, "{:08.X}: ", self.output_settings.start_address).map_err(|e| e.into())
    }

    fn write_bytes(&self, f: &mut ::fmt::Formatter) -> Result<()> {
        let mut byte_count = 0;
        for b in self.data.iter() {
            byte_count += 1;
            let is_seperator_necessary = byte_count % self.output_settings.group_size == 0;
            if is_seperator_necessary {
                self.write_formated_byte(f, b)?;
                write!(f, " ")?
            } else {
                self.write_formated_byte(f, b)?
            }
        }
        Ok(())
    }

    fn write_formated_byte(&self, f: &mut ::fmt::Formatter, byte: &u8) -> Result<()> {
        match self.output_settings.output_fmt {
            OutputFormat::Hex => write!(f, "{:02.X}", byte).map_err(|e| e.into()),
            OutputFormat::Octal => write!(f, "{:02.o}", byte).map_err(|e| e.into()),
            OutputFormat::Decimal => write!(f, "{:02}", byte).map_err(|e| e.into()),
            OutputFormat::Binary => write!(f, "{:02b}", byte).map_err(|e| e.into()),
        }
    }

    fn write_interpretation(&self, f: &mut ::fmt::Formatter) -> Result<()> {
        write!(f, "   ");
        for b in self.data.iter() {
            match *b {
                character @ 20u8...126u8 => write!(f, "{}", character as char)?,
                _ => write!(f, "{}", ".")?,
            }
        }
        Ok(())
    }
}

impl<'a> ::fmt::Display for OutputLine<'a> {
    fn fmt(&self, f: &mut ::fmt::Formatter) -> ::fmt::Result {
        if self.output_settings.show_address {
            self.write_address(f);
        }
        self.write_bytes(f);
        if self.output_settings.show_interpretation {
            self.write_interpretation(f);
        }
        Ok(())
    }
}

mod test {
    use super::*;
    use std::fmt::Write;

    #[test]
    fn output_settings_can_be_constructed() {
        let output_settings = OutputSettings::new();
        assert!(true);
    }

    #[test]
    fn output_settings_builder() {
        let format = OutputFormat::Binary;
        let start_address = 0xFF00AABB;
        let group_size = 2;
        let show_address = false;
        let show_interpretation = false;

        let mut output_settings = OutputSettings::new()
            .format(format)
            .start_address(start_address)
            .group_size(group_size)
            .show_address(show_address)
            .show_interpretation(show_interpretation);

        assert_eq!(output_settings.output_fmt, format);
        assert_eq!(output_settings.start_address, start_address);
        assert_eq!(output_settings.group_size, group_size);
        assert_eq!(output_settings.show_address, show_address);
        assert_eq!(output_settings.show_interpretation, show_interpretation);
    }

    #[test]
    fn output_settings_from_string() {
        assert_eq!(OutputFormat::Binary, OutputFormat::from("bin".to_string()));
        assert_eq!(OutputFormat::Hex, OutputFormat::from("hex".to_string()));
        assert_eq!(OutputFormat::Octal, OutputFormat::from("oct".to_string()));
        assert_eq!(OutputFormat::Decimal, OutputFormat::from("dec".to_string()));

        assert_eq!(OutputFormat::Binary, OutputFormat::from("Bin".to_string()));
        assert_eq!(OutputFormat::Hex, OutputFormat::from("hEx".to_string()));
        assert_eq!(OutputFormat::Octal, OutputFormat::from("ocT".to_string()));
        assert_eq!(OutputFormat::Decimal, OutputFormat::from("DEC".to_string()));
    }

    #[test]
    fn outputline_can_be_constructed() {
        let data = [1, 2, 3];
        let output_line = OutputLine::new(&data);
        assert!(true);
    }

    #[test]
    fn default_output_format_for_a_single_line() {
        let data = [1, 2, 3, 4, 5, 6, 7, 8];
        let expected_output = "00000000: 01 02 03 04 05 06 07 08    ........";
        let output_line = OutputLine::new(&data);
        let mut buffer = String::new();
        let result = write!(&mut buffer, "{}", output_line);
        assert_eq!(Ok(()), result);
        assert_eq!(expected_output, buffer);
    }

    #[test]
    fn octal_output_format_on_a_single_line() {
        let data = [8, 9, 10, 11, 12, 13, 14, 15];
        let expected_output = "00000000: 10 11 12 13 14 15 16 17    ........";
        let output_settings = OutputSettings::new().format(OutputFormat::Octal);
        let output_line = OutputLine::new(&data).format(output_settings);
        let mut buffer = String::new();
        let result = write!(&mut buffer, "{}", output_line);
        assert_eq!(Ok(()), result);
        assert_eq!(expected_output, buffer);
    }

    #[test]
    fn binary_output_format_on_a_single_line() {
        let data = [65, 66, 67, 68, 126, 124, 60, 46];
        let expected_output = "00000000: 1000001 1000010 1000011 1000100 1111110 1111100 111100 101110    ABCD~|<.";
        let output_settings = OutputSettings::new().format(OutputFormat::Binary);
        let output_line = OutputLine::new(&data).format(output_settings);
        let mut buffer = String::new();
        let result = write!(&mut buffer, "{}", output_line);
        assert_eq!(Ok(()), result);
        assert_eq!(expected_output, buffer);
    }

    #[test]
    fn interpretation_for_default_settings() {
        let data = [65, 66, 67, 68, 126, 124, 60, 46];
        let expected_output = "00000000: 41 42 43 44 7E 7C 3C 2E    ABCD~|<.";
        let output_settings = OutputSettings::new();
        let output_line = OutputLine::new(&data).format(output_settings);
        let mut buffer = String::new();
        let result = write!(&mut buffer, "{}", output_line);
        assert_eq!(Ok(()), result);
        assert_eq!(expected_output, buffer);
    }
}
<commit_msg>Change visibility of OutputSettings constructor<commit_after>//! The dump module contains code related for outputing/dumping data.
use std::fmt::Display;
use std::convert::From;
use super::errors::*;

/// Enum which provides all possible output value formats supported by the dump module.
#[derive(Debug, PartialEq, Clone, Copy)]
pub enum OutputFormat {
    Hex,
    Decimal,
    Octal,
    Binary,
}

impl From<String> for OutputFormat {
    fn from(format_string: String) -> Self {
        match format_string.to_lowercase().as_ref() {
            "hex" => OutputFormat::Hex,
            "dec" => OutputFormat::Decimal,
            "oct" => OutputFormat::Octal,
            "bin" => OutputFormat::Binary,
            _ => panic!("Invalid output format"),
        }
    }
}

#[derive(Debug)]
pub struct OutputSettings {
    start_address: u32,
    show_address: bool,
    group_size: usize,
    columns: usize,
    show_interpretation: bool,
    output_fmt: OutputFormat,
}

impl OutputSettings {
    pub fn new() -> OutputSettings {
        OutputSettings {
            start_address: 0,
            show_address: true,
            group_size: 1,
            columns: 8,
            show_interpretation: true,
            output_fmt: OutputFormat::Hex,
        }
    }

    pub fn start_address(mut self, address: u32) -> Self {
        self.start_address = address;
        self
    }

    pub fn show_address(mut self, show: bool) -> Self {
        self.show_address = show;
        self
    }

    pub fn group_size(mut self, size: usize) -> Self {
        self.group_size = size;
        self
    }

    pub fn columns(mut self, columns: usize) -> Self {
        self.columns = columns;
        self
    }

    pub fn show_interpretation(mut self, show: bool) -> Self {
        self.show_interpretation = show;
        self
    }

    pub fn format(mut self, fmt: OutputFormat) -> Self {
        self.output_fmt = fmt;
        self
    }
}

/// The `OutputLine` struct contains all  information needed to dump/output a single line of data.
#[derive(Debug)]
pub struct OutputLine<'a> {
    output_settings: OutputSettings,
    data: &'a [u8], // TODO NiCo: Add member for data format (output format of data)
}

impl<'a> OutputLine<'a> {
    fn new(data: &[u8]) -> OutputLine {
        OutputLine {
            output_settings: OutputSettings::new(),
            data: data,
        }
    }

    pub fn format(self, settings: OutputSettings) -> Self {
        OutputLine {
            output_settings: settings,
            data: self.data,
        }
    }

    fn write_address(&self, f: &mut ::fmt::Formatter) -> Result<()> {
        write!(f, "{:08.X}: ", self.output_settings.start_address).map_err(|e| e.into())
    }

    fn write_bytes(&self, f: &mut ::fmt::Formatter) -> Result<()> {
        let mut byte_count = 0;
        for b in self.data.iter() {
            byte_count += 1;
            let is_seperator_necessary = byte_count % self.output_settings.group_size == 0;
            if is_seperator_necessary {
                self.write_formated_byte(f, b)?;
                write!(f, " ")?
            } else {
                self.write_formated_byte(f, b)?
            }
        }
        Ok(())
    }

    fn write_formated_byte(&self, f: &mut ::fmt::Formatter, byte: &u8) -> Result<()> {
        match self.output_settings.output_fmt {
            OutputFormat::Hex => write!(f, "{:02.X}", byte).map_err(|e| e.into()),
            OutputFormat::Octal => write!(f, "{:02.o}", byte).map_err(|e| e.into()),
            OutputFormat::Decimal => write!(f, "{:02}", byte).map_err(|e| e.into()),
            OutputFormat::Binary => write!(f, "{:02b}", byte).map_err(|e| e.into()),
        }
    }

    fn write_interpretation(&self, f: &mut ::fmt::Formatter) -> Result<()> {
        write!(f, "   ");
        for b in self.data.iter() {
            match *b {
                character @ 20u8...126u8 => write!(f, "{}", character as char)?,
                _ => write!(f, "{}", ".")?,
            }
        }
        Ok(())
    }
}

impl<'a> ::fmt::Display for OutputLine<'a> {
    fn fmt(&self, f: &mut ::fmt::Formatter) -> ::fmt::Result {
        if self.output_settings.show_address {
            self.write_address(f);
        }
        self.write_bytes(f);
        if self.output_settings.show_interpretation {
            self.write_interpretation(f);
        }
        Ok(())
    }
}

mod test {
    use super::*;
    use std::fmt::Write;

    #[test]
    fn output_settings_can_be_constructed() {
        let output_settings = OutputSettings::new();
        assert!(true);
    }

    #[test]
    fn output_settings_builder() {
        let format = OutputFormat::Binary;
        let start_address = 0xFF00AABB;
        let group_size = 2;
        let show_address = false;
        let show_interpretation = false;

        let mut output_settings = OutputSettings::new()
            .format(format)
            .start_address(start_address)
            .group_size(group_size)
            .show_address(show_address)
            .show_interpretation(show_interpretation);

        assert_eq!(output_settings.output_fmt, format);
        assert_eq!(output_settings.start_address, start_address);
        assert_eq!(output_settings.group_size, group_size);
        assert_eq!(output_settings.show_address, show_address);
        assert_eq!(output_settings.show_interpretation, show_interpretation);
    }

    #[test]
    fn output_settings_from_string() {
        assert_eq!(OutputFormat::Binary, OutputFormat::from("bin".to_string()));
        assert_eq!(OutputFormat::Hex, OutputFormat::from("hex".to_string()));
        assert_eq!(OutputFormat::Octal, OutputFormat::from("oct".to_string()));
        assert_eq!(OutputFormat::Decimal, OutputFormat::from("dec".to_string()));

        assert_eq!(OutputFormat::Binary, OutputFormat::from("Bin".to_string()));
        assert_eq!(OutputFormat::Hex, OutputFormat::from("hEx".to_string()));
        assert_eq!(OutputFormat::Octal, OutputFormat::from("ocT".to_string()));
        assert_eq!(OutputFormat::Decimal, OutputFormat::from("DEC".to_string()));
    }

    #[test]
    fn outputline_can_be_constructed() {
        let data = [1, 2, 3];
        let output_line = OutputLine::new(&data);
        assert!(true);
    }

    #[test]
    fn default_output_format_for_a_single_line() {
        let data = [1, 2, 3, 4, 5, 6, 7, 8];
        let expected_output = "00000000: 01 02 03 04 05 06 07 08    ........";
        let output_line = OutputLine::new(&data);
        let mut buffer = String::new();
        let result = write!(&mut buffer, "{}", output_line);
        assert_eq!(Ok(()), result);
        assert_eq!(expected_output, buffer);
    }

    #[test]
    fn octal_output_format_on_a_single_line() {
        let data = [8, 9, 10, 11, 12, 13, 14, 15];
        let expected_output = "00000000: 10 11 12 13 14 15 16 17    ........";
        let output_settings = OutputSettings::new().format(OutputFormat::Octal);
        let output_line = OutputLine::new(&data).format(output_settings);
        let mut buffer = String::new();
        let result = write!(&mut buffer, "{}", output_line);
        assert_eq!(Ok(()), result);
        assert_eq!(expected_output, buffer);
    }

    #[test]
    fn binary_output_format_on_a_single_line() {
        let data = [65, 66, 67, 68, 126, 124, 60, 46];
        let expected_output = "00000000: 1000001 1000010 1000011 1000100 1111110 1111100 111100 101110    ABCD~|<.";
        let output_settings = OutputSettings::new().format(OutputFormat::Binary);
        let output_line = OutputLine::new(&data).format(output_settings);
        let mut buffer = String::new();
        let result = write!(&mut buffer, "{}", output_line);
        assert_eq!(Ok(()), result);
        assert_eq!(expected_output, buffer);
    }

    #[test]
    fn interpretation_for_default_settings() {
        let data = [65, 66, 67, 68, 126, 124, 60, 46];
        let expected_output = "00000000: 41 42 43 44 7E 7C 3C 2E    ABCD~|<.";
        let output_settings = OutputSettings::new();
        let output_line = OutputLine::new(&data).format(output_settings);
        let mut buffer = String::new();
        let result = write!(&mut buffer, "{}", output_line);
        assert_eq!(Ok(()), result);
        assert_eq!(expected_output, buffer);
    }
}
<|endoftext|>


<commit_before>//! Various trait aliases to simplify generics in other modules. Every trait in this module is
//! implemented for all possible valid types.
use num_traits;
use std::fmt::Debug;

pub trait Float: num_traits::float::Float + num_traits::cast::FromPrimitive + Debug {}
impl<F> Float for F where F: num_traits::float::Float + num_traits::cast::FromPrimitive + Debug {}

pub trait Point<F: Float>: AsRef<[F]> + AsMut<[F]> + Clone {}
impl<F, P> Point<F> for P
        where F: Float,
              P: AsRef<[F]> + AsMut<[F]> + Clone {}

pub trait Vector<F: Float, P: Point<F>>: AsRef<[F]> + AsMut<[F]> + From<P> {}
impl<F, P, V> Vector<F, P> for V
        where F: Float,
              P: Point<F>,
              V: AsRef<[F]> + AsMut<[F]> + From<P> {}<commit_msg>Made `Point` and `Vector` traits more robust<commit_after>//! Various trait aliases to simplify generics in other modules. Every trait in this module is
//! implemented for all possible valid types.
use num_traits;
use std::fmt::Debug;
use std::ops::{Add, Sub, Mul, Div, Neg};

pub trait Float: num_traits::float::Float + num_traits::cast::FromPrimitive + Debug {}
impl<F> Float for F where F: num_traits::float::Float + num_traits::cast::FromPrimitive + Debug {}

pub trait PVOps<F, Other>:
		Add<Other, Output = Self> +
		Sub<Other, Output = Self> +

		Add<Self, Output = Self> +
		Sub<Self, Output = Self> +
		Mul<F, Output = Self> +
		Div<F, Output = Self> +
		Neg<Output = Self>

		where Self: Sized,
			  F: Float,
			  Other: PVOps<F, Self> {}

pub trait Point<F, V>:
		AsRef<[F]> +
		AsMut<[F]> +
		From<V> +
		Clone +
		PVOps<F, V>

		where Self: Sized,
			  F: Float,
			  V: Vector<F, Self> {}

impl<F, P, V> Point<F, V> for P
        where F: Float,
        	  V: Vector<F, P>,
              P: AsRef<[F]> + AsMut<[F]> + From<V> + Clone + PVOps<F, V> {}

pub trait Vector<F: Float, P: Point<F, Self>>:
		AsRef<[F]> +
		AsMut<[F]> +
		From<P> +
		Clone +
		PVOps<F, P> {}

impl<F, P, V> Vector<F, P> for V
        where F: Float,
              P: Point<F, V>,
              V: AsRef<[F]> + AsMut<[F]> + From<P> + Clone + PVOps<F, P> {}<|endoftext|>


<commit_before>#[phase(plugin)]
extern crate regex_macros;
extern crate regex;
extern crate serialize;
extern crate git2;
extern crate docopt;
#[phase(plugin)] extern crate docopt_macros;
#[phase(plugin, link)] extern crate log;
extern crate term;

pub use errors;

use git2::Repository;
use std::os;
use std::io::process::Command;
use utils::say::sayln;
use errors::{DeliveryError};

pub fn get_repository() -> Result<git2::Repository, DeliveryError> {
    let repo = try!(git2::Repository::discover(&os::getcwd()));
    Ok(repo)
}

pub fn get_head(repo: git2::Repository) -> Result<String, DeliveryError> {
    let head = try!(repo.head());
    let shorthand = head.shorthand();
    let result = match shorthand {
        Some(result) => Ok(String::from_str(result)),
        None => Err(DeliveryError{ kind: errors::NotOnABranch, detail: None })
    };
    result
}

pub fn git_push(branch: &str, target: &str) -> Result<String, DeliveryError> {
    let mut command = Command::new("git");
    command.arg("push");
    command.arg("--porcelain");
    command.arg("--progress");
    command.arg("--verbose");
    command.arg("origin");
    command.arg(format!("{}:_for/{}/{}", branch, target, branch));
    debug!("Running: {}", command);
    let output = match command.output() {
        Ok(o) => o,
        Err(e) => { return Err(DeliveryError{ kind: errors::FailedToExecute, detail: Some(format!("failed to execute git: {}", e.desc))}) },
    };
    if !output.status.success() {
        return Err(DeliveryError{ kind: errors::PushFailed, detail: Some(format!("STDOUT: {}\nSTDERR: {}\n", String::from_utf8_lossy(output.output.as_slice()), String::from_utf8_lossy(output.error.as_slice())))});
    }
    let stdout = String::from_utf8_lossy(output.output.as_slice()).into_string();
    let stderr = String::from_utf8_lossy(output.error.as_slice()).into_string();
    debug!("Git stdout: {}", stdout);
    debug!("Git stderr: {}", stderr);
    debug!("Git exited: {}", output.status);
    let output = try!(parse_git_push_output(stdout.as_slice(), stderr.as_slice()));
    for result in output.iter() {
        match result.flag {
            SuccessfulFastForward => sayln("green", format!("Updated change: {}", result.reason).as_slice()),
            SuccessfulForcedUpdate => sayln("green", format!("Force updated change: {}", result.reason).as_slice()),
            SuccessfulDeletedRef => sayln("red", format!("Deleted change: {}", result.reason).as_slice()),
            SuccessfulPushedNewRef => sayln("green", format!("Created change: {}", result.reason).as_slice()),
            Rejected => sayln("red", format!("Rejected change: {}", result.reason).as_slice()),
            UpToDate => sayln("yellow", format!("Nothing added to the existing change").as_slice()),
        }
    }
    Ok(stdout.into_string())
}

pub enum PushResultFlags {
    SuccessfulFastForward,
    SuccessfulForcedUpdate,
    SuccessfulDeletedRef,
    SuccessfulPushedNewRef,
    Rejected,
    UpToDate,
}

pub struct PushResult {
    flag: PushResultFlags,
    from: String,
    to: String,
    reason: String
}

pub fn parse_git_push_output(push_output: &str, push_error: &str) -> Result<Vec<PushResult>, DeliveryError> {
    let mut push_results: Vec<PushResult> = Vec::new();
    for line in push_error.lines_any() {
        debug!("error: {}", line);
        if line.starts_with("remote") {
            let r = regex!(r"remote: (.+)");
            let caps_result = r.captures(line);
            match caps_result {
                Some(caps) => sayln("white", format!("{}", caps.at(1)).as_slice()),
                None => {}
            }
        }
    }
    for line in push_output.lines_any() {
        debug!("output: {}", line);
        if line.starts_with("To") {
            continue;
        } else if line.starts_with("Done") {
            continue;
        }
        let r = regex!(r"(.)\t(.+):(.+)\t\[(.+)\]");
        let caps_result = r.captures(line);
        let caps = match caps_result {
            Some(caps) => caps,
            None => { return Err(DeliveryError{ kind: errors::BadGitOutputMatch, detail: Some(format!("Failed to match: {}", line)) }) }
        };
        let result_flag = match caps.at(1) {
            " " => SuccessfulFastForward,
            "+" => SuccessfulForcedUpdate,
            "-" => SuccessfulDeletedRef,
            "*" => SuccessfulPushedNewRef,
            "!" => Rejected,
            "=" => UpToDate,
            _ => { return Err(DeliveryError{ kind: errors::BadGitOutputMatch, detail: Some(format!("Unknown result flag")) }) }
        };
        push_results.push(
            PushResult{
                flag: result_flag,
                from: String::from_str(caps.at(2)),
                to: String::from_str(caps.at(3)),
                reason: String::from_str(caps.at(4))
            }
        )
    }
    Ok(push_results)
}

#[test]
fn test_parse_git_push_output() {
    let input = "To ssh://adam@127.0.0.1/Users/adam/src/opscode/delivery/opscode/delivery-cli2
=	refs/heads/foo:refs/heads/_for/master/foo	[up to date]
Done";
    let results: Vec<String> = parse_git_push_output(input);
    let mut valid_result: Vec<String> = Vec::new();
    valid_result.push(String::from_str("Review branch _for/master/foo is up to date"));
    assert_eq!(valid_result, results);
}

<commit_msg>Adding a test for the git output parsing.<commit_after>#[phase(plugin)]
extern crate regex_macros;
extern crate regex;
extern crate serialize;
extern crate git2;
extern crate docopt;
#[phase(plugin)] extern crate docopt_macros;
#[phase(plugin, link)] extern crate log;
extern crate term;

pub use errors;

use git2::Repository;
use std::os;
use std::io::process::Command;
use utils::say::sayln;
use errors::{DeliveryError};

pub fn get_repository() -> Result<git2::Repository, DeliveryError> {
    let repo = try!(git2::Repository::discover(&os::getcwd()));
    Ok(repo)
}

pub fn get_head(repo: git2::Repository) -> Result<String, DeliveryError> {
    let head = try!(repo.head());
    let shorthand = head.shorthand();
    let result = match shorthand {
        Some(result) => Ok(String::from_str(result)),
        None => Err(DeliveryError{ kind: errors::NotOnABranch, detail: None })
    };
    result
}

pub fn git_push(branch: &str, target: &str) -> Result<String, DeliveryError> {
    let mut command = Command::new("git");
    command.arg("push");
    command.arg("--porcelain");
    command.arg("--progress");
    command.arg("--verbose");
    command.arg("origin");
    command.arg(format!("{}:_for/{}/{}", branch, target, branch));
    debug!("Running: {}", command);
    let output = match command.output() {
        Ok(o) => o,
        Err(e) => { return Err(DeliveryError{ kind: errors::FailedToExecute, detail: Some(format!("failed to execute git: {}", e.desc))}) },
    };
    if !output.status.success() {
        return Err(DeliveryError{ kind: errors::PushFailed, detail: Some(format!("STDOUT: {}\nSTDERR: {}\n", String::from_utf8_lossy(output.output.as_slice()), String::from_utf8_lossy(output.error.as_slice())))});
    }
    let stdout = String::from_utf8_lossy(output.output.as_slice()).into_string();
    let stderr = String::from_utf8_lossy(output.error.as_slice()).into_string();
    debug!("Git exited: {}", output.status);
    let output = try!(parse_git_push_output(stdout.as_slice(), stderr.as_slice()));
    for result in output.iter() {
        match result.flag {
            SuccessfulFastForward => sayln("green", format!("Updated change: {}", result.reason).as_slice()),
            SuccessfulForcedUpdate => sayln("green", format!("Force updated change: {}", result.reason).as_slice()),
            SuccessfulDeletedRef => sayln("red", format!("Deleted change: {}", result.reason).as_slice()),
            SuccessfulPushedNewRef => sayln("green", format!("Created change: {}", result.reason).as_slice()),
            Rejected => sayln("red", format!("Rejected change: {}", result.reason).as_slice()),
            UpToDate => sayln("yellow", format!("Nothing added to the existing change").as_slice()),
        }
    }
    Ok(stdout.into_string())
}

pub enum PushResultFlags {
    SuccessfulFastForward,
    SuccessfulForcedUpdate,
    SuccessfulDeletedRef,
    SuccessfulPushedNewRef,
    Rejected,
    UpToDate,
}

pub struct PushResult {
    flag: PushResultFlags,
    from: String,
    to: String,
    reason: String
}

pub fn parse_git_push_output(push_output: &str, push_error: &str) -> Result<Vec<PushResult>, DeliveryError> {
    let mut push_results: Vec<PushResult> = Vec::new();
    for line in push_error.lines_any() {
        debug!("error: {}", line);
        if line.starts_with("remote") {
            let r = regex!(r"remote: (.+)");
            let caps_result = r.captures(line);
            match caps_result {
                Some(caps) => sayln("white", format!("{}", caps.at(1)).as_slice()),
                None => {}
            }
        }
    }
    for line in push_output.lines_any() {
        debug!("output: {}", line);
        if line.starts_with("To") {
            continue;
        } else if line.starts_with("Done") {
            continue;
        }
        let r = regex!(r"(.)\t(.+):(.+)\t\[(.+)\]");
        let caps_result = r.captures(line);
        let caps = match caps_result {
            Some(caps) => caps,
            None => { return Err(DeliveryError{ kind: errors::BadGitOutputMatch, detail: Some(format!("Failed to match: {}", line)) }) }
        };
        let result_flag = match caps.at(1) {
            " " => SuccessfulFastForward,
            "+" => SuccessfulForcedUpdate,
            "-" => SuccessfulDeletedRef,
            "*" => SuccessfulPushedNewRef,
            "!" => Rejected,
            "=" => UpToDate,
            _ => { return Err(DeliveryError{ kind: errors::BadGitOutputMatch, detail: Some(format!("Unknown result flag")) }) }
        };
        push_results.push(
            PushResult{
                flag: result_flag,
                from: String::from_str(caps.at(2)),
                to: String::from_str(caps.at(3)),
                reason: String::from_str(caps.at(4))
            }
        )
    }
    Ok(push_results)
}

#[test]
fn test_parse_git_push_output_success() {
    let stdout = "To ssh://adam@127.0.0.1/Users/adam/src/opscode/delivery/opscode/delivery-cli2
=	refs/heads/foo:refs/heads/_for/master/foo	[up to date]
Done";
    let stderr = "Pushing to ssh://adam@Chef@172.31.6.130:8989/Chef/adam_universe/delivery-cli
Total 0 (delta 0), reused 0 (delta 0)
remote: Patchset already up to date, nothing to do
remote: https://172.31.6.130/e/Chef/#/organizations/adam_universe/projects/delivery-cli/changes/146a9573-1bd0-4a27-a106-528347761811
updating local tracking ref 'refs/remotes/origin/_for/master/adam/test6'";
    let result = parse_git_push_output(stdout, stderr);
    match result {
        Ok(pr_vec) => {
            // assert!(r_vec[0].flag, UpToDate);
            assert_eq!(pr_vec[0].from.as_slice(), "refs/heads/foo");
            assert_eq!(pr_vec[0].to.as_slice(), "refs/heads/_for/master/foo");
            assert_eq!(pr_vec[0].reason.as_slice(), "up to date");
        },
        Err(e) => panic!("No result")
    };
}

<|endoftext|>


<commit_before>//! Transports for communicating with the docker daemon

extern crate hyper;

use self::super::{Error, Result};
use hyper::Body;
use hyper::client::{Client, HttpConnector};
use hyper::{Method, Request, Response, StatusCode};
use hyper::header;
use hyper_openssl::HttpsConnector;
use hyper::rt::{Future, Stream};
use hyperlocal::Uri as DomainUri;
use hyperlocal::UnixConnector;
use mime::Mime;
use rustc_serialize::json;
use std::fmt;
use std::cell::{RefMut, RefCell};
use std::io::{BufReader, Cursor};
use std::io::{Read, Write};
use tokio::runtime::Runtime;

trait InteractiveStream : Read + Write { }

pub fn tar() -> Mime {
    "application/tar".parse().unwrap()
}

/// Transports are types which define the means of communication
/// with the docker daemon
pub enum Transport {
    /// A network tcp interface
    Tcp {
        client: Client<HttpConnector>,
        runtime: RefCell<Runtime>,
        host: String,
    },
    /// TCP/TLS.w
    EncryptedTcp {
        client: Client<HttpsConnector<HttpConnector>>,
        runtime: RefCell<Runtime>,
        host: String,
    },
    /// A Unix domain socket
    Unix {
        client: Client<UnixConnector>,
        runtime: RefCell<Runtime>,
        path: String,
    },
}

impl fmt::Debug for Transport {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Transport::Tcp { ref host, .. } => write!(f, "Tcp({})", host),
            Transport::EncryptedTcp { ref host, .. } => write!(f, "EncryptedTcp({})", host),
            Transport::Unix { ref path, .. } => write!(f, "Unix({})", path),
        }
    }
}

impl Transport {
    pub fn request<B>(
        &self,
        method: Method,
        endpoint: &str,
        body: Option<(B, Mime)>,
    ) -> Result<String>
    where
        B: Into<Body>,
    {
        let mut res = self.stream(method, endpoint, body)?;
        let mut body = String::new();
        res.read_to_string(&mut body)?;
        debug!("{} raw response: {}", endpoint, body);
        Ok(body)
    }

    /// Builds an HTTP request.
    fn build_request<B>(
        &self,
        method: Method,
        endpoint: &str,
        body: Option<(B, Mime)>,
    ) -> Result<Request<Body>>
    where
        B: Into<Body>,
    {
        let mut builder = Request::builder();
        let req = match *self {
            Transport::Tcp { ref host, .. } => {
                builder.method(method).uri(&format!("{}{}", host, endpoint))
            }
            Transport::EncryptedTcp { ref host, .. } => {
                builder.method(method).uri(&format!("{}{}", host, endpoint))
            }
            Transport::Unix { ref path, .. } => {
                let uri: hyper::Uri = DomainUri::new(&path, endpoint).into();
                builder.method(method).uri(&uri.to_string())
            }
        };
        let req = req.header(header::HOST, "");

        match body {
            Some((b, c)) => Ok(req
                .header(header::CONTENT_TYPE, &c.to_string()[..])
                .body(b.into())?),
            _ => Ok(req.body(Body::empty())?),
        }
    }

    pub fn stream<B>(
        &self,
        method: Method,
        endpoint: &str,
        body: Option<(B, Mime)>,
    ) -> Result<Box<Read>>
    where
        B: Into<Body>,
    {
        let req = self.build_request(method, endpoint, body)?;
        let res = self.send_request(req)?;

        match res.status() {
            StatusCode::OK |
            StatusCode::CREATED |
            StatusCode::SWITCHING_PROTOCOLS => {
                let chunk = self.runtime().block_on(res.into_body().concat2())?;
                Ok(Box::new(Cursor::new(chunk.into_iter().collect::<Vec<u8>>())))
            },
            StatusCode::NO_CONTENT => Ok(
                Box::new(BufReader::new("".as_bytes())),
            ),
            // todo: constantize these
            StatusCode::BAD_REQUEST => {
                Err(Error::Fault {
                    code: res.status(),
                    message: get_error_message(res).unwrap_or(
                        "bad parameter"
                            .to_owned(),
                    ),
                })
            }
            StatusCode::NOT_FOUND => {
                Err(Error::Fault {
                    code: res.status(),
                    message: get_error_message(res).unwrap_or(
                        "not found".to_owned(),
                    ),
                })
            }
            StatusCode::NOT_MODIFIED => {
                Err(Error::Fault {
                    code: res.status(),
                    message: get_error_message(res).unwrap_or(
                        "not modified"
                            .to_owned(),
                    ),
                })
            },
            StatusCode::NOT_ACCEPTABLE => {
                Err(Error::Fault {
                    code: res.status(),
                    message: get_error_message(res).unwrap_or(
                        "not acceptable"
                            .to_owned(),
                    ),
                })
            }
            StatusCode::CONFLICT => {
                Err(Error::Fault {
                    code: res.status(),
                    message: get_error_message(res).unwrap_or(
                        "conflict found"
                            .to_owned(),
                    ),
                })
            }
            StatusCode::INTERNAL_SERVER_ERROR => {
                Err(Error::Fault {
                    code: res.status(),
                    message: get_error_message(res).unwrap_or(
                        "internal server error"
                            .to_owned(),
                    ),
                })
            }
            _ => unreachable!(),
        }
    }

    fn send_request(&self, req: Request<hyper::Body>) -> Result<hyper::Response<Body>> {
        let req = match self {
            Transport::Tcp { ref client, .. } => client.request(req),
            Transport::EncryptedTcp { ref client, .. } => client.request(req),
            Transport::Unix { ref client, .. } => client.request(req),
        };

        self.runtime().block_on(req).map_err(Error::Hyper)
    }

    fn runtime(&self) -> RefMut<Runtime> {
        match self {
            Transport::Tcp { ref runtime, .. } => runtime.borrow_mut(),
            Transport::EncryptedTcp { ref runtime, .. } => runtime.borrow_mut(),
            Transport::Unix { ref runtime, .. } => runtime.borrow_mut(),
        }
    }
}

/// Extract the error message content from an HTTP response that
/// contains a Docker JSON error structure.
fn get_error_message(res: Response<Body>) -> Option<String> {
    let chunk = match res.into_body().concat2().wait() {
        Ok(c) => c,
        Err(..) => return None,
    };

    match String::from_utf8(chunk.into_iter().collect()) {
        Ok(output) => {
            let json_response = json::Json::from_str(output.as_str()).ok();
            let message = json_response
                .as_ref()
                .and_then(|x| x.as_object())
                .and_then(|x| x.get("message"))
                .and_then(|x| x.as_string())
                .map(|x| x.to_owned());

            message
        }
        Err(..) => {
            None
        },
    }
}
<commit_msg>Remove unused InteractiveStream<commit_after>//! Transports for communicating with the docker daemon

extern crate hyper;

use self::super::{Error, Result};
use hyper::Body;
use hyper::client::{Client, HttpConnector};
use hyper::{Method, Request, Response, StatusCode};
use hyper::header;
use hyper_openssl::HttpsConnector;
use hyper::rt::{Future, Stream};
use hyperlocal::Uri as DomainUri;
use hyperlocal::UnixConnector;
use mime::Mime;
use rustc_serialize::json;
use std::fmt;
use std::cell::{RefMut, RefCell};
use std::io::{BufReader, Cursor};
use std::io::Read;
use tokio::runtime::Runtime;

pub fn tar() -> Mime {
    "application/tar".parse().unwrap()
}

/// Transports are types which define the means of communication
/// with the docker daemon
pub enum Transport {
    /// A network tcp interface
    Tcp {
        client: Client<HttpConnector>,
        runtime: RefCell<Runtime>,
        host: String,
    },
    /// TCP/TLS.w
    EncryptedTcp {
        client: Client<HttpsConnector<HttpConnector>>,
        runtime: RefCell<Runtime>,
        host: String,
    },
    /// A Unix domain socket
    Unix {
        client: Client<UnixConnector>,
        runtime: RefCell<Runtime>,
        path: String,
    },
}

impl fmt::Debug for Transport {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Transport::Tcp { ref host, .. } => write!(f, "Tcp({})", host),
            Transport::EncryptedTcp { ref host, .. } => write!(f, "EncryptedTcp({})", host),
            Transport::Unix { ref path, .. } => write!(f, "Unix({})", path),
        }
    }
}

impl Transport {
    pub fn request<B>(
        &self,
        method: Method,
        endpoint: &str,
        body: Option<(B, Mime)>,
    ) -> Result<String>
    where
        B: Into<Body>,
    {
        let mut res = self.stream(method, endpoint, body)?;
        let mut body = String::new();
        res.read_to_string(&mut body)?;
        debug!("{} raw response: {}", endpoint, body);
        Ok(body)
    }

    /// Builds an HTTP request.
    fn build_request<B>(
        &self,
        method: Method,
        endpoint: &str,
        body: Option<(B, Mime)>,
    ) -> Result<Request<Body>>
    where
        B: Into<Body>,
    {
        let mut builder = Request::builder();
        let req = match *self {
            Transport::Tcp { ref host, .. } => {
                builder.method(method).uri(&format!("{}{}", host, endpoint))
            }
            Transport::EncryptedTcp { ref host, .. } => {
                builder.method(method).uri(&format!("{}{}", host, endpoint))
            }
            Transport::Unix { ref path, .. } => {
                let uri: hyper::Uri = DomainUri::new(&path, endpoint).into();
                builder.method(method).uri(&uri.to_string())
            }
        };
        let req = req.header(header::HOST, "");

        match body {
            Some((b, c)) => Ok(req
                .header(header::CONTENT_TYPE, &c.to_string()[..])
                .body(b.into())?),
            _ => Ok(req.body(Body::empty())?),
        }
    }

    pub fn stream<B>(
        &self,
        method: Method,
        endpoint: &str,
        body: Option<(B, Mime)>,
    ) -> Result<Box<Read>>
    where
        B: Into<Body>,
    {
        let req = self.build_request(method, endpoint, body)?;
        let res = self.send_request(req)?;

        match res.status() {
            StatusCode::OK |
            StatusCode::CREATED |
            StatusCode::SWITCHING_PROTOCOLS => {
                let chunk = self.runtime().block_on(res.into_body().concat2())?;
                Ok(Box::new(Cursor::new(chunk.into_iter().collect::<Vec<u8>>())))
            },
            StatusCode::NO_CONTENT => Ok(
                Box::new(BufReader::new("".as_bytes())),
            ),
            // todo: constantize these
            StatusCode::BAD_REQUEST => {
                Err(Error::Fault {
                    code: res.status(),
                    message: get_error_message(res).unwrap_or(
                        "bad parameter"
                            .to_owned(),
                    ),
                })
            }
            StatusCode::NOT_FOUND => {
                Err(Error::Fault {
                    code: res.status(),
                    message: get_error_message(res).unwrap_or(
                        "not found".to_owned(),
                    ),
                })
            }
            StatusCode::NOT_MODIFIED => {
                Err(Error::Fault {
                    code: res.status(),
                    message: get_error_message(res).unwrap_or(
                        "not modified"
                            .to_owned(),
                    ),
                })
            },
            StatusCode::NOT_ACCEPTABLE => {
                Err(Error::Fault {
                    code: res.status(),
                    message: get_error_message(res).unwrap_or(
                        "not acceptable"
                            .to_owned(),
                    ),
                })
            }
            StatusCode::CONFLICT => {
                Err(Error::Fault {
                    code: res.status(),
                    message: get_error_message(res).unwrap_or(
                        "conflict found"
                            .to_owned(),
                    ),
                })
            }
            StatusCode::INTERNAL_SERVER_ERROR => {
                Err(Error::Fault {
                    code: res.status(),
                    message: get_error_message(res).unwrap_or(
                        "internal server error"
                            .to_owned(),
                    ),
                })
            }
            _ => unreachable!(),
        }
    }

    fn send_request(&self, req: Request<hyper::Body>) -> Result<hyper::Response<Body>> {
        let req = match self {
            Transport::Tcp { ref client, .. } => client.request(req),
            Transport::EncryptedTcp { ref client, .. } => client.request(req),
            Transport::Unix { ref client, .. } => client.request(req),
        };

        self.runtime().block_on(req).map_err(Error::Hyper)
    }

    fn runtime(&self) -> RefMut<Runtime> {
        match self {
            Transport::Tcp { ref runtime, .. } => runtime.borrow_mut(),
            Transport::EncryptedTcp { ref runtime, .. } => runtime.borrow_mut(),
            Transport::Unix { ref runtime, .. } => runtime.borrow_mut(),
        }
    }
}

/// Extract the error message content from an HTTP response that
/// contains a Docker JSON error structure.
fn get_error_message(res: Response<Body>) -> Option<String> {
    let chunk = match res.into_body().concat2().wait() {
        Ok(c) => c,
        Err(..) => return None,
    };

    match String::from_utf8(chunk.into_iter().collect()) {
        Ok(output) => {
            let json_response = json::Json::from_str(output.as_str()).ok();
            let message = json_response
                .as_ref()
                .and_then(|x| x.as_object())
                .and_then(|x| x.get("message"))
                .and_then(|x| x.as_string())
                .map(|x| x.to_owned());

            message
        }
        Err(..) => {
            None
        },
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utils for forming gzchunked streams.

/// Encodes the given iterator over protobuf messages into the gzchunked format.
///
/// This is a streaming encoder that lazily encodes the data and can be used to
/// effectively process megabytes of data. The function returns an iterator that
/// yields parts of the gzchunked file.
///
/// This variant uses default encoding settings. To customize them, one can use
/// the [`encode_with_opts`] function.
///
/// [`encode_with_opts`]: fn.encode_with_opts.html
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
/// use std::io::Write as _;
///
/// let items = vec! {
///     String::from("foo"),
///     String::from("bar"),
///     String::from("baz"),
/// };
///
/// let chunks = rrg::gzchunked::encode(items.into_iter());
/// for (idx, chunk) in chunks.enumerate() {
///     let mut file = File::open(format!("output.gzc.{}", idx)).unwrap();
///     file.write_all(chunk.unwrap().as_slice());
/// }
/// ```
pub fn encode<I>(iter: I) -> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    encode_with_opts(iter, EncodeOpts::default())
}

/// Encodes the given iterator over protobuf messages into the gzchunked format.
///
/// This is a variant of the [`encode`] function that allows customization of
/// encoding parameters. Refer to its documentation for more details.
///
/// [`encode`]: fn.encode.html
pub fn encode_with_opts<I>(iter: I, opts: EncodeOpts) -> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    Encode::with_opts(iter, opts)
}

/// Decodes an iterator over gzchunked file parts into a stream of messages.
///
/// This is a streaming decoder that performs the decoding in a lazy way and can
/// be used to effectively process megabytes of data.
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let paths = ["foo.gzc.1", "foo.gzc.2", "foo.gzc.3"];
/// let files = paths.iter().map(|path| File::open(path).unwrap());
///
/// for (idx, msg) in rrg::gzchunked::decode(files).enumerate() {
///     let msg: String = msg.unwrap();
///     println!("item #{}: {:?}", idx, msg);
/// }
/// ```
pub fn decode<I, M>(iter: I) -> impl Iterator<Item=std::io::Result<M>>
where
    I: Iterator,
    I::Item: std::io::Read,
    M: prost::Message + Default,
{
    let parts = iter.map(flate2::read::GzDecoder::new);
    crate::chunked::decode(crate::io::IterReader::new(parts))
}

/// A wrapper type for gzip compression level.
#[derive(Clone, Copy, Debug)]
pub struct Compression(flate2::Compression);

/// A type for defining gzip compression level for gzchunked.
impl Compression {

    /// Creates a new compression descriptor at the specified level.
    ///
    /// The lower the number the worse the compression (with 0 meaning no
    /// compression at all).
    pub fn new(level: u32) -> Compression {
        Compression(flate2::Compression::new(level))
    }

    /// Creates a new compression descriptor with disabled compression.
    pub fn none() -> Compression {
        Compression(flate2::Compression::none())
    }

    /// Creates a new compression descriptor with highest compression.
    pub fn best() -> Compression {
        Compression(flate2::Compression::best())
    }
}

impl Default for Compression {

    fn default() -> Compression {
        Compression(flate2::Compression::new(5))
    }
}

/// Options and flags that configure encoding into the gzchuned format.
#[derive(Clone, Copy, Debug)]
pub struct EncodeOpts {
    /// Compression level used for the gzip encoding.
    pub compression: Compression,
    /// A rough file size limit for parts of the output file.
    pub part_size: u64,
}

impl Default for EncodeOpts {

    fn default() -> EncodeOpts {
        EncodeOpts {
            compression: Compression::default(),
            part_size: 1 * 1024 * 1024, // 1 MiB.
        }
    }
}

/// Streaming encoder for the gzchunked format.
///
/// It implements the `Iterator` trait, lazily polling the underlying iterator
/// over Protocol Buffers messages as more parts are needed.
///
/// Instances of this type can be constructed using the [`encode_with_opts`] or
/// [`encode`] function.
///
/// [`encode`]: fn.encode.html
/// [`encode_with_opts`]: fn.encode_with_opts.html
pub struct Encode<I> {
    chunked: crate::chunked::Encode<I>,
    opts: EncodeOpts,
}

impl<I> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    /// Creates a new encoder instance with the specified options.
    fn with_opts(iter: I, opts: EncodeOpts) -> Encode<I> {
        Encode {
            chunked: crate::chunked::encode(iter),
            opts: opts,
        }
    }

    /// Obtains the next part of the output file (if available).
    fn next_part(&mut self) -> std::io::Result<Option<Vec<u8>>> {
        use crate::io::copy_until;

        let compression = self.opts.compression.0;
        let part_size = self.opts.part_size;

        let mut encoder = flate2::write::GzEncoder::new(vec!(), compression);
        let len = copy_until(&mut self.chunked, &mut encoder, |_, encoder| {
            encoder.get_ref().len() as u64 >= part_size
        })?;

        if len == 0 {
            Ok(None)
        } else {
            Ok(Some(encoder.finish()?))
        }
    }
}

impl<I> Iterator for Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    type Item = std::io::Result<Vec<u8>>;

    fn next(&mut self) -> Option<std::io::Result<Vec<u8>>> {
        self.next_part().transpose()
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_encode_with_empty_iter() {
        let mut iter = encode(std::iter::empty::<()>())
            .map(Result::unwrap);

        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_decode_with_empty_iter() {
        let mut iter = decode::<_, ()>(std::iter::empty::<&[u8]>())
            .map(Result::unwrap);

        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_single_item_iter() {
        let chunks = encode(std::iter::once(String::from("foo")))
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(String::from("foo")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_multiple_items_iter() {
        let data = vec! {
            String::from("foo"),
            String::from("bar"),
            String::from("baz"),
        };

        let chunks = encode(data.into_iter())
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(String::from("foo")));
        assert_eq!(iter.next(), Some(String::from("bar")));
        assert_eq!(iter.next(), Some(String::from("baz")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_empty_items_iter() {
        let data = vec!((), (), ());

        let chunks = encode(data.into_iter())
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_many_items_iter() {
        let sample = rand::random::<[u8; 32]>().to_vec();
        let items = std::iter::repeat(sample.clone()).take(32 * 1024);

        let opts = EncodeOpts {
            compression: Compression::default(),
            part_size: 4 * 1024,
        };

        let chunks = encode_with_opts(items, opts)
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode::<_, Vec<u8>>(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert!(iter.all(|item| item == sample));
    }

    #[test]
    fn test_encode_and_decode_with_no_compression() {
        let sample = rand::random::<[u8; 32]>().to_vec();
        let items = std::iter::repeat(sample.clone()).take(32 * 1024);

        let opts = EncodeOpts {
            compression: Compression::none(),
            part_size: 4 * 1024,
        };

        let chunks = encode_with_opts(items, opts)
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode::<_, Vec<u8>>(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert!(iter.all(|item| item == sample));
    }

    #[test]
    fn test_encode_and_decode_with_best_compression() {
        let sample = rand::random::<[u8; 32]>().to_vec();
        let items = std::iter::repeat(sample.clone()).take(32 * 1024);

        let opts = EncodeOpts {
            compression: Compression::best(),
            part_size: 4 * 1024,
        };

        let chunks = encode_with_opts(items, opts)
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode::<_, Vec<u8>>(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert!(iter.all(|item| item == sample));
    }
}
<commit_msg>Improve documentation for the gzchunked module.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utilities for working with streams in the gzchunked format.
//!
//! gzchunked is a simple file format used for storing large sequences of
//! Protocol Buffers messages. A gzchunked file consists of multiple parts where
//! each part is a gzipped fragment of a stream encoded in the [chunked] format.
//!
//! A high-level pseudocode for encoding and decoding procedures of the
//! gzchunked format can be described using the following formulae:
//!
//!   * _encode(protos) = map(gzip, partition(chunk(protos)))_
//!   * _decode(parts) = unchunk(join(map(ungzip, parts)))_
//!
//! This pseudocode uses the following subroutines:
//!
//!   * _chunk_ encodes a sequence of proto messages into the chunked format.
//!   * _unchunk_ decodes a chunked stream into a sequence of messages.
//!   * _partition_ divides a stream of bytes into multiple parts.
//!   * _join_ sequentially combines multiple byte streams into one.
//!   * _gzip_ encodes a byte stream into the gzip format.
//!   * _ungzip_ decodes a byte stream from the gzip format.
//!
//! [chunked]: crate::chunked

/// Encodes the given iterator over protobuf messages into the gzchunked format.
///
/// This is a streaming encoder that lazily encodes the data and can be used to
/// effectively process megabytes of data. The function returns an iterator that
/// yields parts of the gzchunked file.
///
/// This variant uses default encoding settings. To customize them, one can use
/// the [`encode_with_opts`] function.
///
/// [`encode_with_opts`]: fn.encode_with_opts.html
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
/// use std::io::Write as _;
///
/// let items = vec! {
///     String::from("foo"),
///     String::from("bar"),
///     String::from("baz"),
/// };
///
/// let chunks = rrg::gzchunked::encode(items.into_iter());
/// for (idx, chunk) in chunks.enumerate() {
///     let mut file = File::open(format!("output.gzc.{}", idx)).unwrap();
///     file.write_all(chunk.unwrap().as_slice());
/// }
/// ```
pub fn encode<I>(iter: I) -> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    encode_with_opts(iter, EncodeOpts::default())
}

/// Encodes the given iterator over protobuf messages into the gzchunked format.
///
/// This is a variant of the [`encode`] function that allows customization of
/// encoding parameters. Refer to its documentation for more details.
///
/// [`encode`]: fn.encode.html
pub fn encode_with_opts<I>(iter: I, opts: EncodeOpts) -> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    Encode::with_opts(iter, opts)
}

/// Decodes an iterator over gzchunked file parts into a stream of messages.
///
/// This is a streaming decoder that performs the decoding in a lazy way and can
/// be used to effectively process megabytes of data.
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let paths = ["foo.gzc.1", "foo.gzc.2", "foo.gzc.3"];
/// let files = paths.iter().map(|path| File::open(path).unwrap());
///
/// for (idx, msg) in rrg::gzchunked::decode(files).enumerate() {
///     let msg: String = msg.unwrap();
///     println!("item #{}: {:?}", idx, msg);
/// }
/// ```
pub fn decode<I, M>(iter: I) -> impl Iterator<Item=std::io::Result<M>>
where
    I: Iterator,
    I::Item: std::io::Read,
    M: prost::Message + Default,
{
    let parts = iter.map(flate2::read::GzDecoder::new);
    crate::chunked::decode(crate::io::IterReader::new(parts))
}

/// A type describing compression level of a gzchunked output stream.
#[derive(Clone, Copy, Debug)]
pub struct Compression(flate2::Compression);

impl Compression {

    /// Creates a new compression descriptor at the specified level.
    ///
    /// The lower the number the worse the compression (with 0 meaning no
    /// compression at all).
    pub fn new(level: u32) -> Compression {
        Compression(flate2::Compression::new(level))
    }

    /// Creates a new compression descriptor with disabled compression.
    pub fn none() -> Compression {
        Compression(flate2::Compression::none())
    }

    /// Creates a new compression descriptor with highest compression.
    pub fn best() -> Compression {
        Compression(flate2::Compression::best())
    }
}

impl Default for Compression {

    fn default() -> Compression {
        Compression(flate2::Compression::new(5))
    }
}

/// Options and flags that configure encoding into the gzchuned format.
#[derive(Clone, Copy, Debug)]
pub struct EncodeOpts {
    /// Compression level used for the gzip encoding.
    pub compression: Compression,
    /// A rough file size limit for parts of the output file.
    pub part_size: u64,
}

impl Default for EncodeOpts {

    fn default() -> EncodeOpts {
        EncodeOpts {
            compression: Compression::default(),
            part_size: 1 * 1024 * 1024, // 1 MiB.
        }
    }
}

/// Streaming encoder for the gzchunked format.
///
/// It implements the `Iterator` trait, lazily polling the underlying iterator
/// over Protocol Buffers messages as more parts are needed.
///
/// Instances of this type can be constructed using the [`encode_with_opts`] or
/// [`encode`] function.
///
/// [`encode`]: fn.encode.html
/// [`encode_with_opts`]: fn.encode_with_opts.html
pub struct Encode<I> {
    chunked: crate::chunked::Encode<I>,
    opts: EncodeOpts,
}

impl<I> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    /// Creates a new encoder instance with the specified options.
    fn with_opts(iter: I, opts: EncodeOpts) -> Encode<I> {
        Encode {
            chunked: crate::chunked::encode(iter),
            opts: opts,
        }
    }

    /// Obtains the next part of the output file (if available).
    fn next_part(&mut self) -> std::io::Result<Option<Vec<u8>>> {
        use crate::io::copy_until;

        let compression = self.opts.compression.0;
        let part_size = self.opts.part_size;

        let mut encoder = flate2::write::GzEncoder::new(vec!(), compression);
        let len = copy_until(&mut self.chunked, &mut encoder, |_, encoder| {
            encoder.get_ref().len() as u64 >= part_size
        })?;

        if len == 0 {
            Ok(None)
        } else {
            Ok(Some(encoder.finish()?))
        }
    }
}

impl<I> Iterator for Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    type Item = std::io::Result<Vec<u8>>;

    fn next(&mut self) -> Option<std::io::Result<Vec<u8>>> {
        self.next_part().transpose()
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_encode_with_empty_iter() {
        let mut iter = encode(std::iter::empty::<()>())
            .map(Result::unwrap);

        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_decode_with_empty_iter() {
        let mut iter = decode::<_, ()>(std::iter::empty::<&[u8]>())
            .map(Result::unwrap);

        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_single_item_iter() {
        let chunks = encode(std::iter::once(String::from("foo")))
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(String::from("foo")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_multiple_items_iter() {
        let data = vec! {
            String::from("foo"),
            String::from("bar"),
            String::from("baz"),
        };

        let chunks = encode(data.into_iter())
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(String::from("foo")));
        assert_eq!(iter.next(), Some(String::from("bar")));
        assert_eq!(iter.next(), Some(String::from("baz")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_empty_items_iter() {
        let data = vec!((), (), ());

        let chunks = encode(data.into_iter())
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_many_items_iter() {
        let sample = rand::random::<[u8; 32]>().to_vec();
        let items = std::iter::repeat(sample.clone()).take(32 * 1024);

        let opts = EncodeOpts {
            compression: Compression::default(),
            part_size: 4 * 1024,
        };

        let chunks = encode_with_opts(items, opts)
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode::<_, Vec<u8>>(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert!(iter.all(|item| item == sample));
    }

    #[test]
    fn test_encode_and_decode_with_no_compression() {
        let sample = rand::random::<[u8; 32]>().to_vec();
        let items = std::iter::repeat(sample.clone()).take(32 * 1024);

        let opts = EncodeOpts {
            compression: Compression::none(),
            part_size: 4 * 1024,
        };

        let chunks = encode_with_opts(items, opts)
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode::<_, Vec<u8>>(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert!(iter.all(|item| item == sample));
    }

    #[test]
    fn test_encode_and_decode_with_best_compression() {
        let sample = rand::random::<[u8; 32]>().to_vec();
        let items = std::iter::repeat(sample.clone()).take(32 * 1024);

        let opts = EncodeOpts {
            compression: Compression::best(),
            part_size: 4 * 1024,
        };

        let chunks = encode_with_opts(items, opts)
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode::<_, Vec<u8>>(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert!(iter.all(|item| item == sample));
    }
}
<|endoftext|>


<commit_before>
use std::marker::PhantomData;
use crypto_stuff::*;

#[derive(Copy, Clone)]
pub enum Token {E, S, Dhee, Dhes, Dhse, Dhss, Empty}

pub trait HandshakePattern {
    fn name(out: &mut [u8]) -> usize;
    fn get(premsg_pattern_i: &mut [Token],
           premsg_pattern_r: &mut [Token],
           msg_patterns: &mut [[Token; 8]; 5]);
}

#[derive(Debug)]
pub enum NoiseError {DecryptError}

struct SymmetricState<C: Cipher, H: Hash> {
    cipherstate : CipherState<C>,
    has_key : bool,
    h : [u8; MAXHASHLEN], /* Change once Rust has trait-associated consts */
    ck: [u8; MAXHASHLEN], /* Change once Rust has trait-associated consts */
    wtf : PhantomData<H>, /* So rust thinks I'm using H, this is ugly */
}

pub struct HandshakeState<P: HandshakePattern, D: Dh, C: Cipher, H: Hash, R: Random> {
    symmetricstate: SymmetricState<C, H>,
    s: Option<D>,
    e: Option<D>,
    rs: Option<[u8; DHLEN]>,
    re: Option<[u8; DHLEN]>,
    my_turn_to_send : bool,
    msg_index: usize,
    messages : [[Token; 8]; 5],
    initiator: bool,
    rng : R,
    wtf : PhantomData<P>, /* So rust thinks I'm using P, this is ugly */
}

impl <C: Cipher, H: Hash> SymmetricState<C, H> {

    fn new(handshake_name: &[u8]) -> SymmetricState<C, H> {
        let mut hname = [0u8; MAXHASHLEN];
        if handshake_name.len() <= H::hash_len() {
            copy_memory(handshake_name, &mut hname);
        } else {
            let mut hasher = H::new();
            hasher.input(handshake_name);
            hasher.result(&mut hname);
        }
        SymmetricState{
            cipherstate: CipherState::new(&[0u8; CIPHERKEYLEN], 0),
            has_key : false,
            h: hname,
            ck : hname,
            wtf: PhantomData::<H>
        }
    }

    fn mix_key(&mut self, data: &[u8]) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], data, &mut hkdf_output.0, &mut hkdf_output.1);
        copy_memory(&hkdf_output.0, &mut self.ck);
        self.cipherstate = CipherState::new(&hkdf_output.1[..CIPHERKEYLEN], 0);
        self.has_key = true;
    }

    fn mix_hash(&mut self, data: &[u8]) {
        let mut hasher = H::new();
        hasher.input(&self.h[..H::hash_len()]);
        hasher.input(data);
        hasher.result(&mut self.h);
    }

    fn encrypt_and_hash(&mut self, plaintext: &[u8], out: &mut [u8]) -> usize {
        if self.has_key {
            self.cipherstate.encrypt_with_ad(&self.h[..H::hash_len()], plaintext, out);
            self.mix_hash(&out[..plaintext.len() + TAGLEN]);
            return plaintext.len() + TAGLEN;
        } else {
            copy_memory(plaintext, out);
            self.mix_hash(plaintext);
            return plaintext.len();
        }
    }

    fn decrypt_and_hash(&mut self, data: &[u8], out: &mut [u8]) -> bool {
        if self.has_key {
            if !self.cipherstate.decrypt_with_ad(&self.h[..H::hash_len()], data, out) {
                return false;
            }
        }
        else {
            copy_memory(data, out);
        }
        self.mix_hash(data);
        true
    }

    fn split(&self, initiator: bool) -> (CipherState<C>, CipherState<C>) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], &[0u8; 0], &mut hkdf_output.0, &mut hkdf_output.1);
        let c1 = CipherState::<C>::new(&hkdf_output.0[..CIPHERKEYLEN], 0);
        let c2 = CipherState::<C>::new(&hkdf_output.1[..CIPHERKEYLEN], 0);
        if initiator { (c1, c2) } else { (c2, c1) }
    }

}

impl <P: HandshakePattern, D: Dh, C: Cipher, H: Hash, R: Random> HandshakeState<P, D, C, H, R> {

    pub fn new(initiator: bool,
               new_s : Option<D>,
               new_e : Option<D>,
               new_rs: Option<[u8; DHLEN]>,
               new_re: Option<[u8; DHLEN]> ) -> HandshakeState<P, D, C, H, R> {
        let mut handshake_name = [0u8; 128];
        let mut name_len = P::name(&mut handshake_name);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += D::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += C::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += H::name(&mut handshake_name[name_len..]);

        let mut symmetricstate = SymmetricState::new(&handshake_name[..name_len]);

        let mut premsg_pattern_i = [Token::Empty; 2];
        let mut premsg_pattern_r = [Token::Empty; 2];
        let mut msg_patterns = [[Token::Empty; 8]; 5];
        P::get(&mut premsg_pattern_i, &mut premsg_pattern_r, &mut msg_patterns);
        if initiator {
            for token in &premsg_pattern_i {
                match *token {
                    Token::S => symmetricstate.mix_hash(new_s.as_ref().unwrap().pubkey()),
                    Token::E => symmetricstate.mix_hash(new_e.as_ref().unwrap().pubkey()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
            for token in &premsg_pattern_r {
                match *token {
                    Token::S => symmetricstate.mix_hash(&new_rs.unwrap()),
                    Token::E => symmetricstate.mix_hash(&new_re.unwrap()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
        } else {
            for token in &premsg_pattern_i {
                match *token {
                    Token::S => symmetricstate.mix_hash(&new_rs.unwrap()),
                    Token::E => symmetricstate.mix_hash(&new_re.unwrap()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
            for token in &premsg_pattern_r {
                match *token {
                    Token::S => symmetricstate.mix_hash(new_s.as_ref().unwrap().pubkey()),
                    Token::E => symmetricstate.mix_hash(new_e.as_ref().unwrap().pubkey()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
        }

        HandshakeState{
            symmetricstate: symmetricstate,
            s: new_s, e: new_e, rs: new_rs, re: new_re,
            my_turn_to_send: initiator,
            msg_index: 0,
            messages: msg_patterns,
            initiator: initiator,
            rng: R::new(),
            wtf: PhantomData::<P>}
    }

    pub fn write_message(&mut self,
                         payload: &[u8],
                         message: &mut [u8]) -> (usize, Option<(CipherState<C>, CipherState<C>)>) {
        assert!(self.my_turn_to_send);
        let tokens = self.messages[self.msg_index];
        let mut last = false;
        if let Token::Empty = self.messages[self.msg_index+1][0] {
            last = true;
        }
        self.msg_index += 1;

        let mut byte_index = 0;
        for token in &tokens {
            match *token {
                Token::E => {
                    self.e = Some(D::generate(&mut self.rng));
                    byte_index += self.symmetricstate.encrypt_and_hash(
                        &self.e.as_ref().unwrap().pubkey(), &mut message[byte_index..]);
                },
                Token::S => byte_index += self.symmetricstate.encrypt_and_hash(
                                &self.s.as_ref().unwrap().pubkey(), &mut message[byte_index..]),
                Token::Dhee => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhes => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Dhse => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhss => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Empty => break
            }
        }
        self.my_turn_to_send = false;
        byte_index += self.symmetricstate.encrypt_and_hash(payload, &mut message[byte_index..]);
        match last {
            true => (byte_index, Some(self.symmetricstate.split(self.initiator))),
            false => (byte_index, None)
        }
    }

    pub fn read_message(&mut self,
                        message: &[u8],
                        payload: &mut [u8]) ->
                            Result<(usize, Option<(CipherState<C>, CipherState<C>)>),
                                    NoiseError> {
        assert!(self.my_turn_to_send == false);
        let tokens = self.messages[self.msg_index];
        let mut last = false;
        if let Token::Empty = self.messages[self.msg_index+1][0] {
            last = true;
        }
        self.msg_index += 1;

        let mut ptr = message;
        for token in &tokens {
            match *token {
                Token::E | Token::S => {
                    let data = match self.symmetricstate.has_key {
                        true =>  {
                            let temp = &ptr[..DHLEN + TAGLEN];
                            ptr = &ptr[DHLEN + TAGLEN..];
                            temp
                        }
                        false => {
                            let temp = &ptr[..DHLEN];
                            ptr = &ptr[DHLEN..];
                            temp
                        }
                    };
                    let mut pubkey = [0u8; DHLEN];
                    if !self.symmetricstate.decrypt_and_hash(data, &mut pubkey) {
                        return Err(NoiseError::DecryptError);
                    }
                    match *token {
                        Token::E => self.re = Some(pubkey),
                        Token::S => self.rs = Some(pubkey),
                        _ => unreachable!(),
                    }
                },
                Token::Dhee => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhes => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhse => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Dhss => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Empty => break
            }
        }
        let payload_len = match self.symmetricstate.has_key {
            true => ptr.len() - TAGLEN,
            false => ptr.len()
        };
        if !self.symmetricstate.decrypt_and_hash(ptr, payload) {
            return Err(NoiseError::DecryptError);
        }
        self.my_turn_to_send = true;
        match last {
            true => Ok( (payload_len, Some(self.symmetricstate.split(self.initiator)) ) ),
            false => Ok( (payload_len, None) )
        }
    }

}

<commit_msg>Add MAXMSGLEN checks<commit_after>
use std::marker::PhantomData;
use crypto_stuff::*;

pub const MAXMSGLEN : usize = 65535;

#[derive(Copy, Clone)]
pub enum Token {E, S, Dhee, Dhes, Dhse, Dhss, Empty}

pub trait HandshakePattern {
    fn name(out: &mut [u8]) -> usize;
    fn get(premsg_pattern_i: &mut [Token],
           premsg_pattern_r: &mut [Token],
           msg_patterns: &mut [[Token; 8]; 5]);
}

#[derive(Debug)]
pub enum NoiseError {DecryptError}

struct SymmetricState<C: Cipher, H: Hash> {
    cipherstate : CipherState<C>,
    has_key : bool,
    h : [u8; MAXHASHLEN], /* Change once Rust has trait-associated consts */
    ck: [u8; MAXHASHLEN], /* Change once Rust has trait-associated consts */
    wtf : PhantomData<H>, /* So rust thinks I'm using H, this is ugly */
}

pub struct HandshakeState<P: HandshakePattern, D: Dh, C: Cipher, H: Hash, R: Random> {
    symmetricstate: SymmetricState<C, H>,
    s: Option<D>,
    e: Option<D>,
    rs: Option<[u8; DHLEN]>,
    re: Option<[u8; DHLEN]>,
    my_turn_to_send : bool,
    msg_index: usize,
    messages : [[Token; 8]; 5],
    initiator: bool,
    rng : R,
    wtf : PhantomData<P>, /* So rust thinks I'm using P, this is ugly */
}

impl <C: Cipher, H: Hash> SymmetricState<C, H> {

    fn new(handshake_name: &[u8]) -> SymmetricState<C, H> {
        let mut hname = [0u8; MAXHASHLEN];
        if handshake_name.len() <= H::hash_len() {
            copy_memory(handshake_name, &mut hname);
        } else {
            let mut hasher = H::new();
            hasher.input(handshake_name);
            hasher.result(&mut hname);
        }
        SymmetricState{
            cipherstate: CipherState::new(&[0u8; CIPHERKEYLEN], 0),
            has_key : false,
            h: hname,
            ck : hname,
            wtf: PhantomData::<H>
        }
    }

    fn mix_key(&mut self, data: &[u8]) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], data, &mut hkdf_output.0, &mut hkdf_output.1);
        copy_memory(&hkdf_output.0, &mut self.ck);
        self.cipherstate = CipherState::new(&hkdf_output.1[..CIPHERKEYLEN], 0);
        self.has_key = true;
    }

    fn mix_hash(&mut self, data: &[u8]) {
        let mut hasher = H::new();
        hasher.input(&self.h[..H::hash_len()]);
        hasher.input(data);
        hasher.result(&mut self.h);
    }

    fn encrypt_and_hash(&mut self, plaintext: &[u8], out: &mut [u8]) -> usize {
        if self.has_key {
            self.cipherstate.encrypt_with_ad(&self.h[..H::hash_len()], plaintext, out);
            self.mix_hash(&out[..plaintext.len() + TAGLEN]);
            return plaintext.len() + TAGLEN;
        } else {
            copy_memory(plaintext, out);
            self.mix_hash(plaintext);
            return plaintext.len();
        }
    }

    fn decrypt_and_hash(&mut self, data: &[u8], out: &mut [u8]) -> bool {
        if self.has_key {
            if !self.cipherstate.decrypt_with_ad(&self.h[..H::hash_len()], data, out) {
                return false;
            }
        }
        else {
            copy_memory(data, out);
        }
        self.mix_hash(data);
        true
    }

    fn split(&self, initiator: bool) -> (CipherState<C>, CipherState<C>) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], &[0u8; 0], &mut hkdf_output.0, &mut hkdf_output.1);
        let c1 = CipherState::<C>::new(&hkdf_output.0[..CIPHERKEYLEN], 0);
        let c2 = CipherState::<C>::new(&hkdf_output.1[..CIPHERKEYLEN], 0);
        if initiator { (c1, c2) } else { (c2, c1) }
    }

}

impl <P: HandshakePattern, D: Dh, C: Cipher, H: Hash, R: Random> HandshakeState<P, D, C, H, R> {

    pub fn new(initiator: bool,
               new_s : Option<D>,
               new_e : Option<D>,
               new_rs: Option<[u8; DHLEN]>,
               new_re: Option<[u8; DHLEN]> ) -> HandshakeState<P, D, C, H, R> {
        let mut handshake_name = [0u8; 128];
        let mut name_len = P::name(&mut handshake_name);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += D::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += C::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += H::name(&mut handshake_name[name_len..]);

        let mut symmetricstate = SymmetricState::new(&handshake_name[..name_len]);

        let mut premsg_pattern_i = [Token::Empty; 2];
        let mut premsg_pattern_r = [Token::Empty; 2];
        let mut msg_patterns = [[Token::Empty; 8]; 5];
        P::get(&mut premsg_pattern_i, &mut premsg_pattern_r, &mut msg_patterns);
        if initiator {
            for token in &premsg_pattern_i {
                match *token {
                    Token::S => symmetricstate.mix_hash(new_s.as_ref().unwrap().pubkey()),
                    Token::E => symmetricstate.mix_hash(new_e.as_ref().unwrap().pubkey()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
            for token in &premsg_pattern_r {
                match *token {
                    Token::S => symmetricstate.mix_hash(&new_rs.unwrap()),
                    Token::E => symmetricstate.mix_hash(&new_re.unwrap()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
        } else {
            for token in &premsg_pattern_i {
                match *token {
                    Token::S => symmetricstate.mix_hash(&new_rs.unwrap()),
                    Token::E => symmetricstate.mix_hash(&new_re.unwrap()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
            for token in &premsg_pattern_r {
                match *token {
                    Token::S => symmetricstate.mix_hash(new_s.as_ref().unwrap().pubkey()),
                    Token::E => symmetricstate.mix_hash(new_e.as_ref().unwrap().pubkey()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
        }

        HandshakeState{
            symmetricstate: symmetricstate,
            s: new_s, e: new_e, rs: new_rs, re: new_re,
            my_turn_to_send: initiator,
            msg_index: 0,
            messages: msg_patterns,
            initiator: initiator,
            rng: R::new(),
            wtf: PhantomData::<P>}
    }

    pub fn write_message(&mut self,
                         payload: &[u8],
                         message: &mut [u8]) -> (usize, Option<(CipherState<C>, CipherState<C>)>) {
        assert!(self.my_turn_to_send);
        let tokens = self.messages[self.msg_index];
        let mut last = false;
        if let Token::Empty = self.messages[self.msg_index+1][0] {
            last = true;
        }
        self.msg_index += 1;

        let mut byte_index = 0;
        for token in &tokens {
            match *token {
                Token::E => {
                    self.e = Some(D::generate(&mut self.rng));
                    byte_index += self.symmetricstate.encrypt_and_hash(
                        &self.e.as_ref().unwrap().pubkey(), &mut message[byte_index..]);
                },
                Token::S => byte_index += self.symmetricstate.encrypt_and_hash(
                                &self.s.as_ref().unwrap().pubkey(), &mut message[byte_index..]),
                Token::Dhee => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhes => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Dhse => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhss => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Empty => break
            }
        }
        self.my_turn_to_send = false;
        byte_index += self.symmetricstate.encrypt_and_hash(payload, &mut message[byte_index..]);
        assert!(byte_index <= MAXMSGLEN);
        match last {
            true => (byte_index, Some(self.symmetricstate.split(self.initiator))),
            false => (byte_index, None)
        }
    }

    pub fn read_message(&mut self,
                        message: &[u8],
                        payload: &mut [u8]) ->
                            Result<(usize, Option<(CipherState<C>, CipherState<C>)>),
                                    NoiseError> {
        assert!(self.my_turn_to_send == false);
        assert!(message.len() <= MAXMSGLEN);

        let tokens = self.messages[self.msg_index];
        let mut last = false;
        if let Token::Empty = self.messages[self.msg_index+1][0] {
            last = true;
        }
        self.msg_index += 1;

        let mut ptr = message;
        for token in &tokens {
            match *token {
                Token::E | Token::S => {
                    let data = match self.symmetricstate.has_key {
                        true =>  {
                            let temp = &ptr[..DHLEN + TAGLEN];
                            ptr = &ptr[DHLEN + TAGLEN..];
                            temp
                        }
                        false => {
                            let temp = &ptr[..DHLEN];
                            ptr = &ptr[DHLEN..];
                            temp
                        }
                    };
                    let mut pubkey = [0u8; DHLEN];
                    if !self.symmetricstate.decrypt_and_hash(data, &mut pubkey) {
                        return Err(NoiseError::DecryptError);
                    }
                    match *token {
                        Token::E => self.re = Some(pubkey),
                        Token::S => self.rs = Some(pubkey),
                        _ => unreachable!(),
                    }
                },
                Token::Dhee => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhes => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhse => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Dhss => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Empty => break
            }
        }
        let payload_len = match self.symmetricstate.has_key {
            true => ptr.len() - TAGLEN,
            false => ptr.len()
        };
        if !self.symmetricstate.decrypt_and_hash(ptr, payload) {
            return Err(NoiseError::DecryptError);
        }
        self.my_turn_to_send = true;
        match last {
            true => Ok( (payload_len, Some(self.symmetricstate.split(self.initiator)) ) ),
            false => Ok( (payload_len, None) )
        }
    }

}

<|endoftext|>