<commit_before>extern crate bit_vec;
extern crate ssz;

use bit_vec::BitVec;

use std::default;

/// A BooleanBitfield represents a set of booleans compactly stored as a vector of bits.
/// The BooleanBitfield is given a fixed size during construction. Reads outside of the current size return an out-of-bounds error. Writes outside of the current size expand the size of the set.
#[derive(Debug, Clone, PartialEq)]
pub struct BooleanBitfield(BitVec);

/// Error represents some reason a request against a bitfield was not satisfied
#[derive(Debug, PartialEq)]
pub enum Error {
    /// OutOfBounds refers to indexing into a bitfield where no bits exist; returns the illegal index and the current size of the bitfield, respectively
    OutOfBounds(usize, usize),
}

impl BooleanBitfield {
    /// Create a new bitfield.
    pub fn new() -> Self {
        Default::default()
    }

    pub fn with_capacity(initial_len: usize) -> Self {
        Self::from_elem(initial_len, false)
    }

    /// Create a new bitfield with the given length `initial_len` and all values set to `bit`.
    pub fn from_elem(inital_len: usize, bit: bool) -> Self {
        Self {
            0: BitVec::from_elem(inital_len, bit),
        }
    }

    /// Create a new bitfield using the supplied `bytes` as input
    pub fn from_bytes(bytes: &[u8]) -> Self {
        Self {
            0: BitVec::from_bytes(bytes),
        }
    }

    /// Read the value of a bit.
    ///
    /// If the index is in bounds, then result is Ok(value) where value is `true` if the bit is 1 and `false` if the bit is 0.
    /// If the index is out of bounds, we return an error to that extent.
    pub fn get(&self, i: usize) -> Result<bool, Error> {
        match self.0.get(i) {
            Some(value) => Ok(value),
            None => Err(Error::OutOfBounds(i, self.0.len())),
        }
    }

    /// Set the value of a bit.
    ///
    /// If the index is out of bounds, we expand the size of the underlying set to include the new index.
    /// Returns the previous value if there was one.
    pub fn set(&mut self, i: usize, value: bool) -> Option<bool> {
        let previous = match self.get(i) {
            Ok(previous) => Some(previous),
            Err(Error::OutOfBounds(_, len)) => {
                let new_len = i - len + 1;
                self.0.grow(new_len, false);
                None
            }
        };
        self.0.set(i, value);
        previous
    }

    /// Returns the index of the highest set bit. Some(n) if some bit is set, None otherwise.
    pub fn highest_set_bit(&self) -> Option<usize> {
        self.0.iter().rposition(|bit| bit)
    }

    /// Returns the number of bits in this bitfield.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Returns the number of bytes required to represent this bitfield.
    pub fn num_bytes(&self) -> usize {
        self.to_bytes().len()
    }

    /// Returns the number of `1` bits in the bitfield
    pub fn num_set_bits(&self) -> usize {
        self.0.iter().filter(|&bit| bit).count()
    }

    /// Returns a vector of bytes representing the bitfield
    /// Note that this returns the bit layout of the underlying implementation in the `bit-vec` crate.
    pub fn to_bytes(&self) -> Vec<u8> {
        self.0.to_bytes()
    }
}

impl default::Default for BooleanBitfield {
    /// default provides the "empty" bitfield
    /// Note: the empty bitfield is set to the `0` byte.
    fn default() -> Self {
        Self::from_elem(8, false)
    }
}

impl ssz::Encodable for BooleanBitfield {
    // ssz_append encodes Self according to the `ssz` spec.
    fn ssz_append(&self, s: &mut ssz::SszStream) {
        s.append_vec(&self.to_bytes())
    }
}

impl ssz::Decodable for BooleanBitfield {
    fn ssz_decode(bytes: &[u8], index: usize) -> Result<(Self, usize), ssz::DecodeError> {
        let len = ssz::decode::decode_length(bytes, index, ssz::LENGTH_BYTES)?;
        if (ssz::LENGTH_BYTES + len) > bytes.len() {
            return Err(ssz::DecodeError::TooShort);
        }

        if len == 0 {
            Ok((BooleanBitfield::new(), index + ssz::LENGTH_BYTES))
        } else {
            let bytes = &bytes[(index + 4)..(index + len + 4)];

            let count = len * 8;
            let mut field = BooleanBitfield::with_capacity(count);
            for (byte_index, byte) in bytes.iter().enumerate() {
                for i in 0..8 {
                    let bit = byte & (128 >> i);
                    if bit != 0 {
                        field.set(8 * byte_index + i, true);
                    }
                }
            }

            let index = index + ssz::LENGTH_BYTES + len;
            Ok((field, index))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ssz::{ssz_encode, Decodable, SszStream};

    #[test]
    fn test_new_bitfield() {
        let mut field = BooleanBitfield::new();
        let original_len = field.len();

        for i in 0..100 {
            if i < original_len {
                assert!(!field.get(i).unwrap());
            } else {
                assert!(field.get(i).is_err());
            }
            let previous = field.set(i, true);
            if i < original_len {
                assert!(!previous.unwrap());
            } else {
                assert!(previous.is_none());
            }
        }
    }

    #[test]
    fn test_empty_bitfield() {
        let mut field = BooleanBitfield::from_elem(0, false);
        let original_len = field.len();

        assert_eq!(original_len, 0);

        for i in 0..100 {
            if i < original_len {
                assert!(!field.get(i).unwrap());
            } else {
                assert!(field.get(i).is_err());
            }
            let previous = field.set(i, true);
            if i < original_len {
                assert!(!previous.unwrap());
            } else {
                assert!(previous.is_none());
            }
        }

        assert_eq!(field.len(), 100);
        assert_eq!(field.num_set_bits(), 100);
    }

    const INPUT: &[u8] = &[0b0000_0010, 0b0000_0010];

    #[test]
    fn test_get_from_bitfield() {
        let field = BooleanBitfield::from_bytes(INPUT);
        let unset = field.get(0).unwrap();
        assert!(!unset);
        let set = field.get(6).unwrap();
        assert!(set);
        let set = field.get(14).unwrap();
        assert!(set);
    }

    #[test]
    fn test_set_for_bitfield() {
        let mut field = BooleanBitfield::from_bytes(INPUT);
        let previous = field.set(10, true).unwrap();
        assert!(!previous);
        let previous = field.get(10).unwrap();
        assert!(previous);
        let previous = field.set(6, false).unwrap();
        assert!(previous);
        let previous = field.get(6).unwrap();
        assert!(!previous);
    }

    #[test]
    fn test_highest_set_bit() {
        let field = BooleanBitfield::from_bytes(INPUT);
        assert_eq!(field.highest_set_bit().unwrap(), 14);

        let field = BooleanBitfield::from_bytes(&[0b0000_0011]);
        assert_eq!(field.highest_set_bit().unwrap(), 7);

        let field = BooleanBitfield::new();
        assert_eq!(field.highest_set_bit(), None);
    }

    #[test]
    fn test_len() {
        let field = BooleanBitfield::from_bytes(INPUT);
        assert_eq!(field.len(), 16);

        let field = BooleanBitfield::new();
        assert_eq!(field.len(), 8);
    }

    #[test]
    fn test_num_set_bits() {
        let field = BooleanBitfield::from_bytes(INPUT);
        assert_eq!(field.num_set_bits(), 2);

        let field = BooleanBitfield::new();
        assert_eq!(field.num_set_bits(), 0);
    }

    #[test]
    fn test_to_bytes() {
        let field = BooleanBitfield::from_bytes(INPUT);
        assert_eq!(field.to_bytes(), INPUT);

        let field = BooleanBitfield::new();
        assert_eq!(field.to_bytes(), vec![0]);
    }

    #[test]
    fn test_out_of_bounds() {
        let mut field = BooleanBitfield::from_bytes(INPUT);

        let out_of_bounds_index = field.len();
        assert!(field.set(out_of_bounds_index, true).is_none());
        assert!(field.len() == out_of_bounds_index + 1);
        assert!(field.get(out_of_bounds_index).unwrap());

        for i in 0..100 {
            if i <= out_of_bounds_index {
                assert!(field.set(i, true).is_some());
            } else {
                assert!(field.set(i, true).is_none());
            }
        }
    }

    #[test]
    fn test_grows_with_false() {
        let input_all_set: &[u8] = &[0b1111_1111, 0b1111_1111];
        let mut field = BooleanBitfield::from_bytes(input_all_set);

        // Define `a` and `b`, where both are out of bounds and `b` is greater than `a`.
        let a = field.len();
        let b = a + 1;

        // Ensure `a` is out-of-bounds for test integrity.
        assert!(field.get(a).is_err());

        // Set `b` to `true`. Also, for test integrity, ensure it was previously out-of-bounds.
        assert!(field.set(b, true).is_none());

        // Ensure that `a` wasn't also set to `true` during the grow.
        assert_eq!(field.get(a), Ok(false));
        assert_eq!(field.get(b), Ok(true));
    }

    #[test]
    fn test_num_bytes() {
        let field = BooleanBitfield::from_bytes(INPUT);
        assert_eq!(field.num_bytes(), 2);

        let field = BooleanBitfield::from_elem(2, true);
        assert_eq!(field.num_bytes(), 1);

        let field = BooleanBitfield::from_elem(13, true);
        assert_eq!(field.num_bytes(), 2);
    }

    #[test]
    fn test_ssz_encode() {
        let field = create_test_bitfield();

        let mut stream = SszStream::new();
        stream.append(&field);
        assert_eq!(stream.drain(), vec![0, 0, 0, 2, 225, 192]);

        let field = BooleanBitfield::from_elem(18, true);
        let mut stream = SszStream::new();
        stream.append(&field);
        assert_eq!(stream.drain(), vec![0, 0, 0, 3, 255, 255, 192]);
    }

    fn create_test_bitfield() -> BooleanBitfield {
        let count = 2 * 8;
        let mut field = BooleanBitfield::with_capacity(count);

        let indices = &[0, 1, 2, 7, 8, 9];
        for &i in indices {
            field.set(i, true);
        }
        field
    }

    #[test]
    fn test_ssz_decode() {
        let encoded = vec![0, 0, 0, 2, 225, 192];
        let (field, _): (BooleanBitfield, usize) = ssz::decode_ssz(&encoded, 0).unwrap();
        let expected = create_test_bitfield();
        assert_eq!(field, expected);

        let encoded = vec![0, 0, 0, 3, 255, 255, 3];
        let (field, _): (BooleanBitfield, usize) = ssz::decode_ssz(&encoded, 0).unwrap();
        let expected = BooleanBitfield::from_bytes(&[255, 255, 3]);
        assert_eq!(field, expected);
    }

    #[test]
    fn test_ssz_round_trip() {
        let original = BooleanBitfield::from_bytes(&vec![18; 12][..]);
        let ssz = ssz_encode(&original);
        let (decoded, _) = BooleanBitfield::ssz_decode(&ssz, 0).unwrap();
        assert_eq!(original, decoded);
    }
}
<commit_msg>Adds custom `std::cmp::PartialEq` impl<commit_after>extern crate bit_vec;
extern crate ssz;

use bit_vec::BitVec;

use std::cmp;
use std::default;

/// A BooleanBitfield represents a set of booleans compactly stored as a vector of bits.
/// The BooleanBitfield is given a fixed size during construction. Reads outside of the current size return an out-of-bounds error. Writes outside of the current size expand the size of the set.
#[derive(Debug, Clone)]
pub struct BooleanBitfield(BitVec);

/// Error represents some reason a request against a bitfield was not satisfied
#[derive(Debug, PartialEq)]
pub enum Error {
    /// OutOfBounds refers to indexing into a bitfield where no bits exist; returns the illegal index and the current size of the bitfield, respectively
    OutOfBounds(usize, usize),
}

impl BooleanBitfield {
    /// Create a new bitfield.
    pub fn new() -> Self {
        Default::default()
    }

    pub fn with_capacity(initial_len: usize) -> Self {
        Self::from_elem(initial_len, false)
    }

    /// Create a new bitfield with the given length `initial_len` and all values set to `bit`.
    pub fn from_elem(inital_len: usize, bit: bool) -> Self {
        Self {
            0: BitVec::from_elem(inital_len, bit),
        }
    }

    /// Create a new bitfield using the supplied `bytes` as input
    pub fn from_bytes(bytes: &[u8]) -> Self {
        Self {
            0: BitVec::from_bytes(bytes),
        }
    }

    /// Read the value of a bit.
    ///
    /// If the index is in bounds, then result is Ok(value) where value is `true` if the bit is 1 and `false` if the bit is 0.
    /// If the index is out of bounds, we return an error to that extent.
    pub fn get(&self, i: usize) -> Result<bool, Error> {
        match self.0.get(i) {
            Some(value) => Ok(value),
            None => Err(Error::OutOfBounds(i, self.0.len())),
        }
    }

    /// Set the value of a bit.
    ///
    /// If the index is out of bounds, we expand the size of the underlying set to include the new index.
    /// Returns the previous value if there was one.
    pub fn set(&mut self, i: usize, value: bool) -> Option<bool> {
        let previous = match self.get(i) {
            Ok(previous) => Some(previous),
            Err(Error::OutOfBounds(_, len)) => {
                let new_len = i - len + 1;
                self.0.grow(new_len, false);
                None
            }
        };
        self.0.set(i, value);
        previous
    }

    /// Returns the index of the highest set bit. Some(n) if some bit is set, None otherwise.
    pub fn highest_set_bit(&self) -> Option<usize> {
        self.0.iter().rposition(|bit| bit)
    }

    /// Returns the number of bits in this bitfield.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Returns the number of bytes required to represent this bitfield.
    pub fn num_bytes(&self) -> usize {
        self.to_bytes().len()
    }

    /// Returns the number of `1` bits in the bitfield
    pub fn num_set_bits(&self) -> usize {
        self.0.iter().filter(|&bit| bit).count()
    }

    /// Returns a vector of bytes representing the bitfield
    /// Note that this returns the bit layout of the underlying implementation in the `bit-vec` crate.
    pub fn to_bytes(&self) -> Vec<u8> {
        self.0.to_bytes()
    }
}

impl default::Default for BooleanBitfield {
    /// default provides the "empty" bitfield
    /// Note: the empty bitfield is set to the `0` byte.
    fn default() -> Self {
        Self::from_elem(8, false)
    }
}

impl cmp::PartialEq for BooleanBitfield {
    /// Determines equality by comparing the `ssz` encoding of the two candidates.
    /// This method ensures that the presence of high-order (empty) bits in the highest byte do not exclude equality when they are in fact representing the same information.
    fn eq(&self, other: &Self) -> bool {
        ssz::ssz_encode(self) == ssz::ssz_encode(other)
    }
}

impl ssz::Encodable for BooleanBitfield {
    // ssz_append encodes Self according to the `ssz` spec.
    fn ssz_append(&self, s: &mut ssz::SszStream) {
        s.append_vec(&self.to_bytes())
    }
}

impl ssz::Decodable for BooleanBitfield {
    fn ssz_decode(bytes: &[u8], index: usize) -> Result<(Self, usize), ssz::DecodeError> {
        let len = ssz::decode::decode_length(bytes, index, ssz::LENGTH_BYTES)?;
        if (ssz::LENGTH_BYTES + len) > bytes.len() {
            return Err(ssz::DecodeError::TooShort);
        }

        if len == 0 {
            Ok((BooleanBitfield::new(), index + ssz::LENGTH_BYTES))
        } else {
            let bytes = &bytes[(index + 4)..(index + len + 4)];

            let count = len * 8;
            let mut field = BooleanBitfield::with_capacity(count);
            for (byte_index, byte) in bytes.iter().enumerate() {
                for i in 0..8 {
                    let bit = byte & (128 >> i);
                    if bit != 0 {
                        field.set(8 * byte_index + i, true);
                    }
                }
            }

            let index = index + ssz::LENGTH_BYTES + len;
            Ok((field, index))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ssz::{ssz_encode, Decodable, SszStream};

    #[test]
    fn test_new_bitfield() {
        let mut field = BooleanBitfield::new();
        let original_len = field.len();

        for i in 0..100 {
            if i < original_len {
                assert!(!field.get(i).unwrap());
            } else {
                assert!(field.get(i).is_err());
            }
            let previous = field.set(i, true);
            if i < original_len {
                assert!(!previous.unwrap());
            } else {
                assert!(previous.is_none());
            }
        }
    }

    #[test]
    fn test_empty_bitfield() {
        let mut field = BooleanBitfield::from_elem(0, false);
        let original_len = field.len();

        assert_eq!(original_len, 0);

        for i in 0..100 {
            if i < original_len {
                assert!(!field.get(i).unwrap());
            } else {
                assert!(field.get(i).is_err());
            }
            let previous = field.set(i, true);
            if i < original_len {
                assert!(!previous.unwrap());
            } else {
                assert!(previous.is_none());
            }
        }

        assert_eq!(field.len(), 100);
        assert_eq!(field.num_set_bits(), 100);
    }

    const INPUT: &[u8] = &[0b0000_0010, 0b0000_0010];

    #[test]
    fn test_get_from_bitfield() {
        let field = BooleanBitfield::from_bytes(INPUT);
        let unset = field.get(0).unwrap();
        assert!(!unset);
        let set = field.get(6).unwrap();
        assert!(set);
        let set = field.get(14).unwrap();
        assert!(set);
    }

    #[test]
    fn test_set_for_bitfield() {
        let mut field = BooleanBitfield::from_bytes(INPUT);
        let previous = field.set(10, true).unwrap();
        assert!(!previous);
        let previous = field.get(10).unwrap();
        assert!(previous);
        let previous = field.set(6, false).unwrap();
        assert!(previous);
        let previous = field.get(6).unwrap();
        assert!(!previous);
    }

    #[test]
    fn test_highest_set_bit() {
        let field = BooleanBitfield::from_bytes(INPUT);
        assert_eq!(field.highest_set_bit().unwrap(), 14);

        let field = BooleanBitfield::from_bytes(&[0b0000_0011]);
        assert_eq!(field.highest_set_bit().unwrap(), 7);

        let field = BooleanBitfield::new();
        assert_eq!(field.highest_set_bit(), None);
    }

    #[test]
    fn test_len() {
        let field = BooleanBitfield::from_bytes(INPUT);
        assert_eq!(field.len(), 16);

        let field = BooleanBitfield::new();
        assert_eq!(field.len(), 8);
    }

    #[test]
    fn test_num_set_bits() {
        let field = BooleanBitfield::from_bytes(INPUT);
        assert_eq!(field.num_set_bits(), 2);

        let field = BooleanBitfield::new();
        assert_eq!(field.num_set_bits(), 0);
    }

    #[test]
    fn test_to_bytes() {
        let field = BooleanBitfield::from_bytes(INPUT);
        assert_eq!(field.to_bytes(), INPUT);

        let field = BooleanBitfield::new();
        assert_eq!(field.to_bytes(), vec![0]);
    }

    #[test]
    fn test_out_of_bounds() {
        let mut field = BooleanBitfield::from_bytes(INPUT);

        let out_of_bounds_index = field.len();
        assert!(field.set(out_of_bounds_index, true).is_none());
        assert!(field.len() == out_of_bounds_index + 1);
        assert!(field.get(out_of_bounds_index).unwrap());

        for i in 0..100 {
            if i <= out_of_bounds_index {
                assert!(field.set(i, true).is_some());
            } else {
                assert!(field.set(i, true).is_none());
            }
        }
    }

    #[test]
    fn test_grows_with_false() {
        let input_all_set: &[u8] = &[0b1111_1111, 0b1111_1111];
        let mut field = BooleanBitfield::from_bytes(input_all_set);

        // Define `a` and `b`, where both are out of bounds and `b` is greater than `a`.
        let a = field.len();
        let b = a + 1;

        // Ensure `a` is out-of-bounds for test integrity.
        assert!(field.get(a).is_err());

        // Set `b` to `true`. Also, for test integrity, ensure it was previously out-of-bounds.
        assert!(field.set(b, true).is_none());

        // Ensure that `a` wasn't also set to `true` during the grow.
        assert_eq!(field.get(a), Ok(false));
        assert_eq!(field.get(b), Ok(true));
    }

    #[test]
    fn test_num_bytes() {
        let field = BooleanBitfield::from_bytes(INPUT);
        assert_eq!(field.num_bytes(), 2);

        let field = BooleanBitfield::from_elem(2, true);
        assert_eq!(field.num_bytes(), 1);

        let field = BooleanBitfield::from_elem(13, true);
        assert_eq!(field.num_bytes(), 2);
    }

    #[test]
    fn test_ssz_encode() {
        let field = create_test_bitfield();

        let mut stream = SszStream::new();
        stream.append(&field);
        assert_eq!(stream.drain(), vec![0, 0, 0, 2, 225, 192]);

        let field = BooleanBitfield::from_elem(18, true);
        let mut stream = SszStream::new();
        stream.append(&field);
        assert_eq!(stream.drain(), vec![0, 0, 0, 3, 255, 255, 192]);
    }

    fn create_test_bitfield() -> BooleanBitfield {
        let count = 2 * 8;
        let mut field = BooleanBitfield::with_capacity(count);

        let indices = &[0, 1, 2, 7, 8, 9];
        for &i in indices {
            field.set(i, true);
        }
        field
    }

    #[test]
    fn test_ssz_decode() {
        let encoded = vec![0, 0, 0, 2, 225, 192];
        let (field, _): (BooleanBitfield, usize) = ssz::decode_ssz(&encoded, 0).unwrap();
        let expected = create_test_bitfield();
        assert_eq!(field, expected);

        let encoded = vec![0, 0, 0, 3, 255, 255, 3];
        let (field, _): (BooleanBitfield, usize) = ssz::decode_ssz(&encoded, 0).unwrap();
        let expected = BooleanBitfield::from_bytes(&[255, 255, 3]);
        assert_eq!(field, expected);
    }

    #[test]
    fn test_ssz_round_trip() {
        let original = BooleanBitfield::from_bytes(&vec![18; 12][..]);
        let ssz = ssz_encode(&original);
        let (decoded, _) = BooleanBitfield::ssz_decode(&ssz, 0).unwrap();
        assert_eq!(original, decoded);
    }
}
<|endoftext|>


<commit_before>// Support code for rustc's built in test runner generator. Currently,
// none of this is meant for users. It is intended to support the
// simplest interface possible for representing and running tests
// while providing a base that other test frameworks may build off of.

import sort = sort::ivector;
import generic_os::getenv;

export test_name;
export test_fn;
export test_desc;
export test_main;
export test_result;
export test_opts;
export tr_ok;
export tr_failed;
export tr_ignored;
export run_tests_console;
export run_tests_console_;
export run_test;
export filter_tests;
export parse_opts;
export test_to_task;
export default_test_to_task;
export configure_test_task;

// The name of a test. By convention this follows the rules for rust
// paths, i.e it should be a series of identifiers seperated by double
// colons. This way if some test runner wants to arrange the tests
// heirarchically it may.
type test_name = str;

// A function that runs a test. If the function returns successfully,
// the test succeeds; if the function fails then the test fails. We
// may need to come up with a more clever definition of test in order
// to support isolation of tests into tasks.
type test_fn = fn() ;

// The definition of a single test. A test runner will run a list of
// these.
type test_desc = {name: test_name, fn: test_fn, ignore: bool};

// The default console test runner. It accepts the command line
// arguments and a vector of test_descs (generated at compile time).
fn test_main(args: &vec[str], tests: &test_desc[]) {
    let ivec_args =
        { let iargs = ~[]; for arg: str  in args { iargs += ~[arg] } iargs };
    check (ivec::is_not_empty(ivec_args));
    let opts =
        alt parse_opts(ivec_args) {
          either::left(o) { o }
          either::right(m) { fail m }
        };
    if !run_tests_console(opts, tests) { fail "Some tests failed"; }
}

type test_opts = {filter: option::t[str], run_ignored: bool};

type opt_res = either::t[test_opts, str];

// Parses command line arguments into test options
fn parse_opts(args: &str[]) : ivec::is_not_empty(args) -> opt_res {

    // FIXME (#649): Shouldn't have to check here
    check (ivec::is_not_empty(args));
    let args_ = ivec::tail(args);
    let opts = ~[getopts::optflag("ignored")];
    let match =
        alt getopts::getopts_ivec(args_, opts) {
          getopts::success(m) { m }
          getopts::failure(f) { ret either::right(getopts::fail_str(f)) }
        };

    let filter =
        if vec::len(match.free) > 0u {
            option::some(match.free.(0))
        } else { option::none };

    let run_ignored = getopts::opt_present(match, "ignored");

    let test_opts = {filter: filter, run_ignored: run_ignored};

    ret either::left(test_opts);
}

tag test_result { tr_ok; tr_failed; tr_ignored; }

// To get isolation and concurrency tests have to be run in their own tasks.
// In cases where test functions and closures it is not ok to just dump them
// into a task and run them, so this transformation gives the caller a chance
// to create the test task.
type test_to_task = fn(&fn()) -> task ;

// A simple console test runner
fn run_tests_console(opts: &test_opts, tests: &test_desc[]) -> bool {
    run_tests_console_(opts, tests, default_test_to_task)
}

fn run_tests_console_(opts: &test_opts, tests: &test_desc[],
                      to_task: &test_to_task) -> bool {

    let filtered_tests = filter_tests(opts, tests);

    let out = io::stdout();

    let total = ivec::len(filtered_tests);
    out.write_line(#fmt("\nrunning %u tests", total));

    let futures = ~[];

    let passed = 0u;
    let failed = 0u;
    let ignored = 0u;

    let failures = ~[];

    // It's tempting to just spawn all the tests at once but that doesn't
    // provide a great user experience because you might sit waiting for the
    // result of a particular test for an unusually long amount of time.
    let concurrency = get_concurrency();
    log #fmt("using %u test tasks", concurrency);
    let run_idx = 0u;
    let wait_idx = 0u;

    while wait_idx < total {
        while ivec::len(futures) < concurrency && run_idx < total {
            futures += ~[run_test(filtered_tests.(run_idx), to_task)];
            run_idx += 1u;
        }

        let future = futures.(0);
        out.write_str(#fmt("running %s ... ", future.test.name));
        let result = future.wait();
        alt result {
          tr_ok. {
            passed += 1u;
            write_ok(out, concurrency);
            out.write_line("");
          }
          tr_failed. {
            failed += 1u;
            write_failed(out, concurrency);
            out.write_line("");
            failures += ~[future.test];
          }
          tr_ignored. {
            ignored += 1u;
            write_ignored(out, concurrency);
            out.write_line("");
          }
        }
        futures = ivec::slice(futures, 1u, ivec::len(futures));
        wait_idx += 1u;
    }

    assert (passed + failed + ignored == total);
    let success = failed == 0u;

    if !success {
        out.write_line("\nfailures:");
        for test: test_desc  in failures {
            out.write_line(#fmt("    %s", test.name));
        }
    }

    out.write_str(#fmt("\nresult: "));
    if success {
        write_ok(out, concurrency);
    } else { write_failed(out, concurrency); }
    out.write_str(#fmt(". %u passed; %u failed; %u ignored\n\n", passed,
                       failed, ignored));

    ret success;

    fn write_ok(out: &io::writer, concurrency: uint) {
        write_pretty(out, "ok", term::color_green, concurrency);
    }

    fn write_failed(out: &io::writer, concurrency: uint) {
        write_pretty(out, "FAILED", term::color_red, concurrency);
    }

    fn write_ignored(out: &io::writer, concurrency: uint) {
        write_pretty(out, "ignored", term::color_yellow, concurrency);
    }

    fn write_pretty(out: &io::writer, word: &str, color: u8,
                    concurrency: uint) {
        // In the presence of concurrency, outputing control characters
        // can cause some crazy artifacting
        if concurrency == 1u && term::color_supported() {
            term::fg(out.get_buf_writer(), color);
        }
        out.write_str(word);
        if concurrency == 1u && term::color_supported() {
            term::reset(out.get_buf_writer());
        }
    }
}

fn get_concurrency() -> uint {
    alt getenv("RUST_THREADS") {
      option::some(t) {
        let threads = uint::parse_buf(str::bytes(t), 10u);
        threads > 0u ? threads : 1u
      }
      option::none. { 1u }
    }
}

fn filter_tests(opts: &test_opts, tests: &test_desc[]) -> test_desc[] {
    let filtered = tests;

    // Remove tests that don't match the test filter
    filtered =
        if option::is_none(opts.filter) {
            filtered
        } else {
            let filter_str =
                alt opts.filter {
                  option::some(f) { f }
                  option::none. { "" }
                };

            let filter =
                bind fn (test: &test_desc, filter_str: str) ->
                        option::t[test_desc] {
                         if str::find(test.name, filter_str) >= 0 {
                             ret option::some(test);
                         } else { ret option::none; }
                     }(_, filter_str);

            ivec::filter_map(filter, filtered)
        };

    // Maybe pull out the ignored test and unignore them
    filtered =
        if !opts.run_ignored {
            filtered
        } else {
            let filter =
                fn (test: &test_desc) -> option::t[test_desc] {
                    if test.ignore {
                        ret option::some({name: test.name,
                                          fn: test.fn,
                                          ignore: false});
                    } else { ret option::none; }
                };

            ivec::filter_map(filter, filtered)
        };

    // Sort the tests alphabetically
    filtered =
        {
            fn lteq(t1: &test_desc, t2: &test_desc) -> bool {
                str::lteq(t1.name, t2.name)
            }
            sort::merge_sort(lteq, filtered)
        };

    ret filtered;
}

type test_future =
    {test: test_desc, fnref: @fn() , wait: fn() -> test_result };

fn run_test(test: &test_desc, to_task: &test_to_task) -> test_future {
    // FIXME: Because of the unsafe way we're passing the test function
    // to the test task, we need to make sure we keep a reference to that
    // function around for longer than the lifetime of the task. To that end
    // we keep the function boxed in the test future.
    let fnref = @test.fn;
    if !test.ignore {
        let test_task = to_task(*fnref);
        ret {test: test,
             fnref: fnref,
             wait:
                 bind fn (test_task: &task) -> test_result {
                          alt task::join(test_task) {
                            task::tr_success. { tr_ok }
                            task::tr_failure. { tr_failed }
                          }
                      }(test_task)};
    } else {
        ret {test: test,
             fnref: fnref,
             wait: fn () -> test_result { tr_ignored }};
    }
}

native "rust" mod rustrt {
    fn hack_allow_leaks();
}

// We need to run our tests in another task in order to trap test failures.
// But, at least currently, functions can't be used as spawn arguments so
// we've got to treat our test functions as unsafe pointers.  This function
// only works with functions that don't contain closures.
fn default_test_to_task(f: &fn()) -> task {
    fn run_task(fptr: *mutable fn() ) {
        configure_test_task();
        // Run the test
        (*fptr)()
    }
    let fptr = ptr::addr_of(f);
    ret spawn run_task(fptr);
}

// Call from within a test task to make sure it's set up correctly
fn configure_test_task() {
    // If this task fails we don't want that failure to propagate to the
    // test runner or else we couldn't keep running tests
    task::unsupervise();

    // FIXME (236): Hack supreme - unwinding doesn't work yet so if this
    // task fails memory will not be freed correctly. This turns off the
    // sanity checks in the runtime's memory region for the task, so that
    // the test runner can continue.
    rustrt::hack_allow_leaks();
}

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C .. 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
<commit_msg>Refactor std::test<commit_after>// Support code for rustc's built in test runner generator. Currently,
// none of this is meant for users. It is intended to support the
// simplest interface possible for representing and running tests
// while providing a base that other test frameworks may build off of.

import sort = sort::ivector;
import generic_os::getenv;

export test_name;
export test_fn;
export test_desc;
export test_main;
export test_result;
export test_opts;
export tr_ok;
export tr_failed;
export tr_ignored;
export run_tests_console;
export run_tests_console_;
export run_test;
export filter_tests;
export parse_opts;
export test_to_task;
export default_test_to_task;
export configure_test_task;

// The name of a test. By convention this follows the rules for rust
// paths, i.e it should be a series of identifiers seperated by double
// colons. This way if some test runner wants to arrange the tests
// heirarchically it may.
type test_name = str;

// A function that runs a test. If the function returns successfully,
// the test succeeds; if the function fails then the test fails. We
// may need to come up with a more clever definition of test in order
// to support isolation of tests into tasks.
type test_fn = fn() ;

// The definition of a single test. A test runner will run a list of
// these.
type test_desc = {name: test_name, fn: test_fn, ignore: bool};

// The default console test runner. It accepts the command line
// arguments and a vector of test_descs (generated at compile time).
fn test_main(args: &vec[str], tests: &test_desc[]) {
    let ivec_args =
        { let iargs = ~[]; for arg: str  in args { iargs += ~[arg] } iargs };
    check (ivec::is_not_empty(ivec_args));
    let opts =
        alt parse_opts(ivec_args) {
          either::left(o) { o }
          either::right(m) { fail m }
        };
    if !run_tests_console(opts, tests) { fail "Some tests failed"; }
}

type test_opts = {filter: option::t[str], run_ignored: bool};

type opt_res = either::t[test_opts, str];

// Parses command line arguments into test options
fn parse_opts(args: &str[]) : ivec::is_not_empty(args) -> opt_res {

    // FIXME (#649): Shouldn't have to check here
    check (ivec::is_not_empty(args));
    let args_ = ivec::tail(args);
    let opts = ~[getopts::optflag("ignored")];
    let match =
        alt getopts::getopts_ivec(args_, opts) {
          getopts::success(m) { m }
          getopts::failure(f) { ret either::right(getopts::fail_str(f)) }
        };

    let filter =
        if vec::len(match.free) > 0u {
            option::some(match.free.(0))
        } else { option::none };

    let run_ignored = getopts::opt_present(match, "ignored");

    let test_opts = {filter: filter, run_ignored: run_ignored};

    ret either::left(test_opts);
}

tag test_result { tr_ok; tr_failed; tr_ignored; }

// To get isolation and concurrency tests have to be run in their own tasks.
// In cases where test functions and closures it is not ok to just dump them
// into a task and run them, so this transformation gives the caller a chance
// to create the test task.
type test_to_task = fn(&fn()) -> task ;

// A simple console test runner
fn run_tests_console(opts: &test_opts, tests: &test_desc[]) -> bool {
    run_tests_console_(opts, tests, default_test_to_task)
}

fn run_tests_console_(opts: &test_opts, tests: &test_desc[],
                      to_task: &test_to_task) -> bool {

    type test_state = @{
        out: io::writer,
        use_color: bool,
        mutable total: uint,
        mutable passed: uint,
        mutable failed: uint,
        mutable ignored: uint,
        mutable failures: test_desc[]
    };

    fn callback(event: testevent, st: test_state) {
        alt event {
          te_filtered(filtered_tests) {
            st.total = ivec::len(filtered_tests);
            st.out.write_line(#fmt("\nrunning %u tests", st.total));
          }
          te_result(test, result) {
            st.out.write_str(#fmt("test %s ... ", test.name));
            alt result {
              tr_ok. {
                st.passed += 1u;
                write_ok(st.out, st.use_color);
                st.out.write_line("");
              }
              tr_failed. {
                st.failed += 1u;
                write_failed(st.out, st.use_color);
                st.out.write_line("");
                st.failures += ~[test];
              }
              tr_ignored. {
                st.ignored += 1u;
                write_ignored(st.out, st.use_color);
                st.out.write_line("");
              }
            }
          }
        }
    }

    let st = @{
        out: io::stdout(),
        use_color: use_color(),
        mutable total: 0u,
        mutable passed: 0u,
        mutable failed: 0u,
        mutable ignored: 0u,
        mutable failures: ~[]
    };

    run_tests(opts, tests, to_task,
              bind callback(_, st));

    assert st.passed + st.failed + st.ignored == st.total;
    let success = st.failed == 0u;

    if !success {
        st.out.write_line("\nfailures:");
        for test: test_desc in st.failures {
            let testname = test.name; // Satisfy alias analysis
            st.out.write_line(#fmt("    %s", testname));
        }
    }

    st.out.write_str(#fmt("\nresult: "));
    if success {
        write_ok(st.out, st.use_color);
    } else { write_failed(st.out, st.use_color); }
    st.out.write_str(#fmt(". %u passed; %u failed; %u ignored\n\n",
                       st.passed, st.failed, st.ignored));

    ret success;

    fn write_ok(out: &io::writer, use_color: bool) {
        write_pretty(out, "ok", term::color_green, use_color);
    }

    fn write_failed(out: &io::writer, use_color: bool) {
        write_pretty(out, "FAILED", term::color_red, use_color);
    }

    fn write_ignored(out: &io::writer, use_color: bool) {
        write_pretty(out, "ignored", term::color_yellow, use_color);
    }

    fn write_pretty(out: &io::writer, word: &str, color: u8,
                    use_color: bool) {
        if use_color && term::color_supported() {
            term::fg(out.get_buf_writer(), color);
        }
        out.write_str(word);
        if use_color && term::color_supported() {
            term::reset(out.get_buf_writer());
        }
    }
}

fn use_color() -> bool {
    ret get_concurrency() == 1u;
}

tag testevent {
    te_filtered(test_desc[]);
    te_result(test_desc, test_result);
}

fn run_tests(opts: &test_opts, tests: &test_desc[],
             to_task: &test_to_task, callback: fn(testevent)) {

    let filtered_tests = filter_tests(opts, tests);

    callback(te_filtered(filtered_tests));

    // It's tempting to just spawn all the tests at once but that doesn't
    // provide a great user experience because you might sit waiting for the
    // result of a particular test for an unusually long amount of time.
    let concurrency = get_concurrency();
    log #fmt("using %u test tasks", concurrency);
    let total = ivec::len(filtered_tests);
    let run_idx = 0u;
    let wait_idx = 0u;
    let futures = ~[];

    while wait_idx < total {
        while ivec::len(futures) < concurrency && run_idx < total {
            futures += ~[run_test(filtered_tests.(run_idx), to_task)];
            run_idx += 1u;
        }

        let future = futures.(0);
        let result = future.wait();
        callback(te_result(future.test, result));
        futures = ivec::slice(futures, 1u, ivec::len(futures));
        wait_idx += 1u;
    }
}

fn get_concurrency() -> uint {
    alt getenv("RUST_THREADS") {
      option::some(t) {
        let threads = uint::parse_buf(str::bytes(t), 10u);
        threads > 0u ? threads : 1u
      }
      option::none. { 1u }
    }
}

fn filter_tests(opts: &test_opts, tests: &test_desc[]) -> test_desc[] {
    let filtered = tests;

    // Remove tests that don't match the test filter
    filtered =
        if option::is_none(opts.filter) {
            filtered
        } else {
            let filter_str =
                alt opts.filter {
                  option::some(f) { f }
                  option::none. { "" }
                };

            let filter =
                bind fn (test: &test_desc, filter_str: str) ->
                        option::t[test_desc] {
                         if str::find(test.name, filter_str) >= 0 {
                             ret option::some(test);
                         } else { ret option::none; }
                     }(_, filter_str);

            ivec::filter_map(filter, filtered)
        };

    // Maybe pull out the ignored test and unignore them
    filtered =
        if !opts.run_ignored {
            filtered
        } else {
            let filter =
                fn (test: &test_desc) -> option::t[test_desc] {
                    if test.ignore {
                        ret option::some({name: test.name,
                                          fn: test.fn,
                                          ignore: false});
                    } else { ret option::none; }
                };

            ivec::filter_map(filter, filtered)
        };

    // Sort the tests alphabetically
    filtered =
        {
            fn lteq(t1: &test_desc, t2: &test_desc) -> bool {
                str::lteq(t1.name, t2.name)
            }
            sort::merge_sort(lteq, filtered)
        };

    ret filtered;
}

type test_future =
    {test: test_desc, fnref: @fn() , wait: fn() -> test_result };

fn run_test(test: &test_desc, to_task: &test_to_task) -> test_future {
    // FIXME: Because of the unsafe way we're passing the test function
    // to the test task, we need to make sure we keep a reference to that
    // function around for longer than the lifetime of the task. To that end
    // we keep the function boxed in the test future.
    let fnref = @test.fn;
    if !test.ignore {
        let test_task = to_task(*fnref);
        ret {test: test,
             fnref: fnref,
             wait:
                 bind fn (test_task: &task) -> test_result {
                          alt task::join(test_task) {
                            task::tr_success. { tr_ok }
                            task::tr_failure. { tr_failed }
                          }
                      }(test_task)};
    } else {
        ret {test: test,
             fnref: fnref,
             wait: fn () -> test_result { tr_ignored }};
    }
}

native "rust" mod rustrt {
    fn hack_allow_leaks();
}

// We need to run our tests in another task in order to trap test failures.
// But, at least currently, functions can't be used as spawn arguments so
// we've got to treat our test functions as unsafe pointers.  This function
// only works with functions that don't contain closures.
fn default_test_to_task(f: &fn()) -> task {
    fn run_task(fptr: *mutable fn() ) {
        configure_test_task();
        // Run the test
        (*fptr)()
    }
    let fptr = ptr::addr_of(f);
    ret spawn run_task(fptr);
}

// Call from within a test task to make sure it's set up correctly
fn configure_test_task() {
    // If this task fails we don't want that failure to propagate to the
    // test runner or else we couldn't keep running tests
    task::unsupervise();

    // FIXME (236): Hack supreme - unwinding doesn't work yet so if this
    // task fails memory will not be freed correctly. This turns off the
    // sanity checks in the runtime's memory region for the task, so that
    // the test runner can continue.
    rustrt::hack_allow_leaks();
}

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C .. 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
<|endoftext|>


<commit_before>use std::{io, time};

use actix_codec::{AsyncRead, AsyncWrite, Framed};
use bytes::Bytes;
use futures::future::{ok, Either};
use futures::{Async, Future, Poll, Sink, Stream};

use crate::error::PayloadError;
use crate::h1;
use crate::message::{RequestHead, ResponseHead};
use crate::payload::{Payload, PayloadStream};

use super::connection::{ConnectionLifetime, ConnectionType, IoConnection};
use super::error::{ConnectError, SendRequestError};
use super::pool::Acquired;
use crate::body::{BodySize, MessageBody};

pub(crate) fn send_request<T, B>(
    io: T,
    head: RequestHead,
    body: B,
    created: time::Instant,
    pool: Option<Acquired<T>>,
) -> impl Future<Item = (ResponseHead, Payload), Error = SendRequestError>
where
    T: AsyncRead + AsyncWrite + 'static,
    B: MessageBody,
{
    let io = H1Connection {
        created,
        pool,
        io: Some(io),
    };

    let len = body.length();

    // create Framed and send reqest
    Framed::new(io, h1::ClientCodec::default())
        .send((head, len).into())
        .from_err()
        // send request body
        .and_then(move |framed| match body.length() {
            BodySize::None | BodySize::Empty | BodySize::Sized(0) => {
                Either::A(ok(framed))
            }
            _ => Either::B(SendBody::new(body, framed)),
        })
        // read response and init read body
        .and_then(|framed| {
            framed
                .into_future()
                .map_err(|(e, _)| SendRequestError::from(e))
                .and_then(|(item, framed)| {
                    if let Some(res) = item {
                        match framed.get_codec().message_type() {
                            h1::MessageType::None => {
                                let force_close = !framed.get_codec().keepalive();
                                release_connection(framed, force_close);
                                Ok((res, Payload::None))
                            }
                            _ => {
                                let pl: PayloadStream = Box::new(PlStream::new(framed));
                                Ok((res, pl.into()))
                            }
                        }
                    } else {
                        Err(ConnectError::Disconnected.into())
                    }
                })
        })
}

pub(crate) fn open_tunnel<T>(
    io: T,
    head: RequestHead,
) -> impl Future<Item = (ResponseHead, Framed<T, h1::ClientCodec>), Error = SendRequestError>
where
    T: AsyncRead + AsyncWrite + 'static,
{
    // create Framed and send reqest
    Framed::new(io, h1::ClientCodec::default())
        .send((head, BodySize::None).into())
        .from_err()
        // read response
        .and_then(|framed| {
            framed
                .into_future()
                .map_err(|(e, _)| SendRequestError::from(e))
                .and_then(|(head, framed)| {
                    if let Some(head) = head {
                        Ok((head, framed))
                    } else {
                        Err(SendRequestError::from(ConnectError::Disconnected))
                    }
                })
        })
}

#[doc(hidden)]
/// HTTP client connection
pub struct H1Connection<T> {
    io: Option<T>,
    created: time::Instant,
    pool: Option<Acquired<T>>,
}

impl<T: AsyncRead + AsyncWrite + 'static> ConnectionLifetime for H1Connection<T> {
    /// Close connection
    fn close(&mut self) {
        if let Some(mut pool) = self.pool.take() {
            if let Some(io) = self.io.take() {
                pool.close(IoConnection::new(
                    ConnectionType::H1(io),
                    self.created,
                    None,
                ));
            }
        }
    }

    /// Release this connection to the connection pool
    fn release(&mut self) {
        if let Some(mut pool) = self.pool.take() {
            if let Some(io) = self.io.take() {
                pool.release(IoConnection::new(
                    ConnectionType::H1(io),
                    self.created,
                    None,
                ));
            }
        }
    }
}

impl<T: AsyncRead + AsyncWrite + 'static> io::Read for H1Connection<T> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.io.as_mut().unwrap().read(buf)
    }
}

impl<T: AsyncRead + AsyncWrite + 'static> AsyncRead for H1Connection<T> {}

impl<T: AsyncRead + AsyncWrite + 'static> io::Write for H1Connection<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.io.as_mut().unwrap().write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.io.as_mut().unwrap().flush()
    }
}

impl<T: AsyncRead + AsyncWrite + 'static> AsyncWrite for H1Connection<T> {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        self.io.as_mut().unwrap().shutdown()
    }
}

/// Future responsible for sending request body to the peer
pub(crate) struct SendBody<I, B> {
    body: Option<B>,
    framed: Option<Framed<I, h1::ClientCodec>>,
    flushed: bool,
}

impl<I, B> SendBody<I, B>
where
    I: AsyncRead + AsyncWrite + 'static,
    B: MessageBody,
{
    pub(crate) fn new(body: B, framed: Framed<I, h1::ClientCodec>) -> Self {
        SendBody {
            body: Some(body),
            framed: Some(framed),
            flushed: true,
        }
    }
}

impl<I, B> Future for SendBody<I, B>
where
    I: ConnectionLifetime,
    B: MessageBody,
{
    type Item = Framed<I, h1::ClientCodec>;
    type Error = SendRequestError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut body_ready = true;
        loop {
            while body_ready
                && self.body.is_some()
                && !self.framed.as_ref().unwrap().is_write_buf_full()
            {
                match self.body.as_mut().unwrap().poll_next()? {
                    Async::Ready(item) => {
                        // check if body is done
                        if item.is_none() {
                            let _ = self.body.take();
                        }
                        self.flushed = false;
                        self.framed
                            .as_mut()
                            .unwrap()
                            .force_send(h1::Message::Chunk(item))?;
                        break;
                    }
                    Async::NotReady => body_ready = false,
                }
            }

            if !self.flushed {
                match self.framed.as_mut().unwrap().poll_complete()? {
                    Async::Ready(_) => {
                        self.flushed = true;
                        continue;
                    }
                    Async::NotReady => return Ok(Async::NotReady),
                }
            }

            if self.body.is_none() {
                return Ok(Async::Ready(self.framed.take().unwrap()));
            }
            return Ok(Async::NotReady);
        }
    }
}

pub(crate) struct PlStream<Io> {
    framed: Option<Framed<Io, h1::ClientPayloadCodec>>,
}

impl<Io: ConnectionLifetime> PlStream<Io> {
    fn new(framed: Framed<Io, h1::ClientCodec>) -> Self {
        PlStream {
            framed: Some(framed.map_codec(|codec| codec.into_payload_codec())),
        }
    }
}

impl<Io: ConnectionLifetime> Stream for PlStream<Io> {
    type Item = Bytes;
    type Error = PayloadError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        match self.framed.as_mut().unwrap().poll()? {
            Async::NotReady => Ok(Async::NotReady),
            Async::Ready(Some(chunk)) => {
                if let Some(chunk) = chunk {
                    Ok(Async::Ready(Some(chunk)))
                } else {
                    let framed = self.framed.take().unwrap();
                    let force_close = framed.get_codec().keepalive();
                    release_connection(framed, force_close);
                    Ok(Async::Ready(None))
                }
            }
            Async::Ready(None) => Ok(Async::Ready(None)),
        }
    }
}

fn release_connection<T, U>(framed: Framed<T, U>, force_close: bool)
where
    T: ConnectionLifetime,
{
    let mut parts = framed.into_parts();
    if !force_close && parts.read_buf.is_empty() && parts.write_buf.is_empty() {
        parts.io.release()
    } else {
        parts.io.close()
    }
}
<commit_msg>set host header for http1 connections<commit_after>use std::io::Write;
use std::{io, time};

use actix_codec::{AsyncRead, AsyncWrite, Framed};
use bytes::{BufMut, Bytes, BytesMut};
use futures::future::{ok, Either};
use futures::{Async, Future, Poll, Sink, Stream};

use crate::error::PayloadError;
use crate::h1;
use crate::http::header::{IntoHeaderValue, HOST};
use crate::message::{RequestHead, ResponseHead};
use crate::payload::{Payload, PayloadStream};

use super::connection::{ConnectionLifetime, ConnectionType, IoConnection};
use super::error::{ConnectError, SendRequestError};
use super::pool::Acquired;
use crate::body::{BodySize, MessageBody};

pub(crate) fn send_request<T, B>(
    io: T,
    mut head: RequestHead,
    body: B,
    created: time::Instant,
    pool: Option<Acquired<T>>,
) -> impl Future<Item = (ResponseHead, Payload), Error = SendRequestError>
where
    T: AsyncRead + AsyncWrite + 'static,
    B: MessageBody,
{
    // set request host header
    if !head.headers.contains_key(HOST) {
        if let Some(host) = head.uri.host() {
            let mut wrt = BytesMut::with_capacity(host.len() + 5).writer();

            let _ = match head.uri.port_u16() {
                None | Some(80) | Some(443) => write!(wrt, "{}", host),
                Some(port) => write!(wrt, "{}:{}", host, port),
            };

            match wrt.get_mut().take().freeze().try_into() {
                Ok(value) => {
                    head.headers.insert(HOST, value);
                }
                Err(e) => {
                    log::error!("Can not set HOST header {}", e);
                }
            }
        }
    }

    let io = H1Connection {
        created,
        pool,
        io: Some(io),
    };

    let len = body.length();

    // create Framed and send reqest
    Framed::new(io, h1::ClientCodec::default())
        .send((head, len).into())
        .from_err()
        // send request body
        .and_then(move |framed| match body.length() {
            BodySize::None | BodySize::Empty | BodySize::Sized(0) => {
                Either::A(ok(framed))
            }
            _ => Either::B(SendBody::new(body, framed)),
        })
        // read response and init read body
        .and_then(|framed| {
            framed
                .into_future()
                .map_err(|(e, _)| SendRequestError::from(e))
                .and_then(|(item, framed)| {
                    if let Some(res) = item {
                        match framed.get_codec().message_type() {
                            h1::MessageType::None => {
                                let force_close = !framed.get_codec().keepalive();
                                release_connection(framed, force_close);
                                Ok((res, Payload::None))
                            }
                            _ => {
                                let pl: PayloadStream = Box::new(PlStream::new(framed));
                                Ok((res, pl.into()))
                            }
                        }
                    } else {
                        Err(ConnectError::Disconnected.into())
                    }
                })
        })
}

pub(crate) fn open_tunnel<T>(
    io: T,
    head: RequestHead,
) -> impl Future<Item = (ResponseHead, Framed<T, h1::ClientCodec>), Error = SendRequestError>
where
    T: AsyncRead + AsyncWrite + 'static,
{
    // create Framed and send reqest
    Framed::new(io, h1::ClientCodec::default())
        .send((head, BodySize::None).into())
        .from_err()
        // read response
        .and_then(|framed| {
            framed
                .into_future()
                .map_err(|(e, _)| SendRequestError::from(e))
                .and_then(|(head, framed)| {
                    if let Some(head) = head {
                        Ok((head, framed))
                    } else {
                        Err(SendRequestError::from(ConnectError::Disconnected))
                    }
                })
        })
}

#[doc(hidden)]
/// HTTP client connection
pub struct H1Connection<T> {
    io: Option<T>,
    created: time::Instant,
    pool: Option<Acquired<T>>,
}

impl<T: AsyncRead + AsyncWrite + 'static> ConnectionLifetime for H1Connection<T> {
    /// Close connection
    fn close(&mut self) {
        if let Some(mut pool) = self.pool.take() {
            if let Some(io) = self.io.take() {
                pool.close(IoConnection::new(
                    ConnectionType::H1(io),
                    self.created,
                    None,
                ));
            }
        }
    }

    /// Release this connection to the connection pool
    fn release(&mut self) {
        if let Some(mut pool) = self.pool.take() {
            if let Some(io) = self.io.take() {
                pool.release(IoConnection::new(
                    ConnectionType::H1(io),
                    self.created,
                    None,
                ));
            }
        }
    }
}

impl<T: AsyncRead + AsyncWrite + 'static> io::Read for H1Connection<T> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.io.as_mut().unwrap().read(buf)
    }
}

impl<T: AsyncRead + AsyncWrite + 'static> AsyncRead for H1Connection<T> {}

impl<T: AsyncRead + AsyncWrite + 'static> io::Write for H1Connection<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.io.as_mut().unwrap().write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.io.as_mut().unwrap().flush()
    }
}

impl<T: AsyncRead + AsyncWrite + 'static> AsyncWrite for H1Connection<T> {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        self.io.as_mut().unwrap().shutdown()
    }
}

/// Future responsible for sending request body to the peer
pub(crate) struct SendBody<I, B> {
    body: Option<B>,
    framed: Option<Framed<I, h1::ClientCodec>>,
    flushed: bool,
}

impl<I, B> SendBody<I, B>
where
    I: AsyncRead + AsyncWrite + 'static,
    B: MessageBody,
{
    pub(crate) fn new(body: B, framed: Framed<I, h1::ClientCodec>) -> Self {
        SendBody {
            body: Some(body),
            framed: Some(framed),
            flushed: true,
        }
    }
}

impl<I, B> Future for SendBody<I, B>
where
    I: ConnectionLifetime,
    B: MessageBody,
{
    type Item = Framed<I, h1::ClientCodec>;
    type Error = SendRequestError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut body_ready = true;
        loop {
            while body_ready
                && self.body.is_some()
                && !self.framed.as_ref().unwrap().is_write_buf_full()
            {
                match self.body.as_mut().unwrap().poll_next()? {
                    Async::Ready(item) => {
                        // check if body is done
                        if item.is_none() {
                            let _ = self.body.take();
                        }
                        self.flushed = false;
                        self.framed
                            .as_mut()
                            .unwrap()
                            .force_send(h1::Message::Chunk(item))?;
                        break;
                    }
                    Async::NotReady => body_ready = false,
                }
            }

            if !self.flushed {
                match self.framed.as_mut().unwrap().poll_complete()? {
                    Async::Ready(_) => {
                        self.flushed = true;
                        continue;
                    }
                    Async::NotReady => return Ok(Async::NotReady),
                }
            }

            if self.body.is_none() {
                return Ok(Async::Ready(self.framed.take().unwrap()));
            }
            return Ok(Async::NotReady);
        }
    }
}

pub(crate) struct PlStream<Io> {
    framed: Option<Framed<Io, h1::ClientPayloadCodec>>,
}

impl<Io: ConnectionLifetime> PlStream<Io> {
    fn new(framed: Framed<Io, h1::ClientCodec>) -> Self {
        PlStream {
            framed: Some(framed.map_codec(|codec| codec.into_payload_codec())),
        }
    }
}

impl<Io: ConnectionLifetime> Stream for PlStream<Io> {
    type Item = Bytes;
    type Error = PayloadError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        match self.framed.as_mut().unwrap().poll()? {
            Async::NotReady => Ok(Async::NotReady),
            Async::Ready(Some(chunk)) => {
                if let Some(chunk) = chunk {
                    Ok(Async::Ready(Some(chunk)))
                } else {
                    let framed = self.framed.take().unwrap();
                    let force_close = framed.get_codec().keepalive();
                    release_connection(framed, force_close);
                    Ok(Async::Ready(None))
                }
            }
            Async::Ready(None) => Ok(Async::Ready(None)),
        }
    }
}

fn release_connection<T, U>(framed: Framed<T, U>, force_close: bool)
where
    T: ConnectionLifetime,
{
    let mut parts = framed.into_parts();
    if !force_close && parts.read_buf.is_empty() && parts.write_buf.is_empty() {
        parts.io.release()
    } else {
        parts.io.close()
    }
}
<|endoftext|>


<commit_before>//! Protocol messages
//!
//! All protocol messages are serialized by serde according to the bincode data format. For
//! transmission on the wire, `tokio_io::codec::length_delimited` is used to prepend protocol
//! messages by an additional length header, thereby creating frames. It is crucial to
//!
//! The main type `Message` and the types in its fields are dump containers, which are not
//! responsible for any protocol logic (except for syntactic validation including validation of
//! `PublicKey` and `SecretKey` fields). Consequently, all fields of `Message` and all fields of
//! its contained types such as `Header` and `Payload` are public.

use secp256k1::key::{PublicKey, SecretKey};
use secp256k1::Signature;
use ::{SessionId, PeerIndex, SymmetricKey, SequenceNum};
use field::Fp;

/// A protocol message
///
/// Protocol messages consist of a header and a payload.
#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct Message {
    pub header: Header,
    pub payload: Payload,
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct Header {
    pub session_id: SessionId, // just for consistency checks
    pub peer_index: PeerIndex,
    pub sequence_num: SequenceNum, // just for consistency checks
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub enum Payload {
    KeyExchange(KeyExchange),
    DcExponential(DcExponential),
    DcMain(DcMain),
    Blame(Blame),
    Confirm(Confirm),
    Reveal(Reveal),
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct KeyExchange {
    pub ke_pk: PublicKey,
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct DcExponential {
    pub commitment: [u8; 32],
    pub dc_exp: Vec<Fp>,
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct DcMain {
    pub ok: bool,
    pub dc_xor: Vec<Vec<u8>>,
    pub ke_pk: PublicKey,
    pub extension: Extension,
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub enum Extension {
    None,
    DcAddSecp256k1Scalar(/* TODO */),
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct Blame {
    pub ke_sk: SecretKey,
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct Confirm {
    pub data: Vec<u8>,
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct Reveal {
    pub keys: Vec<(PeerIndex, SymmetricKey)>,
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;
    use tokio_io::codec::length_delimited;
    use tokio_serde_bincode::{ReadBincode, WriteBincode};
    use futures::sink::Sink;
    use futures::{Future, Stream};

    use super::*;
    use secp256k1::key::SecretKey;

    #[cfg(test)]
    fn roundtrip_serde_bincode(payload: Payload) {
        let msg1 = Message {
            header: dummy_header(),
            payload: payload,
        }.add_signature();

        // Write message
        let pipe = Cursor::new(Vec::new());
        let write = WriteBincode::new(length_delimited::FramedWrite::new(pipe));
        let write = write.send(msg1.clone()).wait().unwrap();

        // Read message again
        let pipe = Cursor::new(write.into_inner().into_inner().into_inner());
        let read = ReadBincode::<_ ,Message>::new(length_delimited::FramedRead::new(pipe));
        let msg2 = read.wait().next().unwrap().unwrap();

        assert_eq!(msg1, msg2);
    }

    #[test]
    fn roundtrip_keyexchange() {
        let slice: [u8; 32] = [0x4f; 32];
        let sk = SecretKey::from_slice(&::SECP256K1, &slice).unwrap();
        let ke_pk = PublicKey::from_secret_key(&::SECP256K1, &sk).unwrap();
        assert!(ke_pk.is_valid());

        let payload = Payload::KeyExchange(KeyExchange {
            ke_pk: ke_pk,
        });

        roundtrip_serde_bincode(payload);
    }

    #[test]
    fn roundtrip_dcexponential() {
        let payload = Payload::DcExponential(DcExponential {
            commitment: [9; 32],
            dc_exp: vec![Fp::from_u127(Fp::prime()), Fp::from_u127(0), Fp::from_u127(656)],
        });

        roundtrip_serde_bincode(payload);
    }

    #[cfg(test)]
    fn dummy_header() -> Header {
        Header {
            peer_index: 2,
            session_id: SessionId([56; 32]),
            sequence_num: 14,
        }
    }

    #[cfg(test)]
    fn sign(msg: Message) -> SignedMessage {
        // Create secret key, public key, message digest and signature
        let slice: [u8; 32] = [0xab; 32];
        let sk = SecretKey::from_slice(&::SECP256K1, &slice).unwrap();
        // TODO woha, this is ugly
        // maybe https://docs.rs/tokio-io/0.1/tokio_io/codec/trait.Decoder.html
        let digest = ::secp256k1::Message::from_slice(msg.serialize()).unwrap();
        let sig = ::SECP256K1.sign(&digest, &sk).unwrap();
        SignedMessage {
            message: msg,
            signature: sig,
        }
    }
}
<commit_msg>Remove obsolete tests<commit_after>//! Protocol messages
//!
//! All protocol messages are serialized by serde according to the bincode data format. For
//! transmission on the wire, `tokio_io::codec::length_delimited` is used to prepend protocol
//! messages by an additional length header, thereby creating frames. It is crucial to
//!
//! The main type `Message` and the types in its fields are dump containers, which are not
//! responsible for any protocol logic (except for syntactic validation including validation of
//! `PublicKey` and `SecretKey` fields). Consequently, all fields of `Message` and all fields of
//! its contained types such as `Header` and `Payload` are public.

use secp256k1::key::{PublicKey, SecretKey};
use secp256k1::Signature;
use ::{SessionId, PeerIndex, SymmetricKey, SequenceNum};
use field::Fp;

/// A protocol message
///
/// Protocol messages consist of a header and a payload.
#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct Message {
    pub header: Header,
    pub payload: Payload,
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct Header {
    pub session_id: SessionId, // just for consistency checks
    pub peer_index: PeerIndex,
    pub sequence_num: SequenceNum, // just for consistency checks
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub enum Payload {
    KeyExchange(KeyExchange),
    DcExponential(DcExponential),
    DcMain(DcMain),
    Blame(Blame),
    Confirm(Confirm),
    Reveal(Reveal),
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct KeyExchange {
    pub ke_pk: PublicKey,
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct DcExponential {
    pub commitment: [u8; 32],
    pub dc_exp: Vec<Fp>,
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct DcMain {
    pub ok: bool,
    pub dc_xor: Vec<Vec<u8>>,
    pub ke_pk: PublicKey,
    pub extension: Extension,
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub enum Extension {
    None,
    DcAddSecp256k1Scalar(/* TODO */),
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct Blame {
    pub ke_sk: SecretKey,
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct Confirm {
    pub data: Vec<u8>,
}

#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct Reveal {
    pub keys: Vec<(PeerIndex, SymmetricKey)>,
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;
    use tokio_io::codec::length_delimited;
    use futures::sink::Sink;
    use futures::{Future, Stream};

    use super::*;
    use secp256k1::key::SecretKey;

    #[test]
    fn roundtrip_keyexchange() {
        let slice: [u8; 32] = [0x4f; 32];
        let sk = SecretKey::from_slice(&::SECP256K1, &slice).unwrap();
        let ke_pk = PublicKey::from_secret_key(&::SECP256K1, &sk).unwrap();
        assert!(ke_pk.is_valid());

        let payload = Payload::KeyExchange(KeyExchange {
            ke_pk: ke_pk,
        });

        roundtrip_serde_bincode(payload);
    }

    #[cfg(test)]
    fn roundtrip_serde_bincode(payload: Payload) {
        // TODO
    }
}
<|endoftext|>


<commit_before>//! Cluster metadata.

use std::ffi::CStr;
use std::slice;

use rdkafka_sys as rdsys;
use rdkafka_sys::types::*;

use crate::error::IsError;

/// Broker metadata information.
pub struct MetadataBroker(RDKafkaMetadataBroker);

impl MetadataBroker {
    /// Returns the id of the broker.
    pub fn id(&self) -> i32 {
        self.0.id
    }

    /// Returns the host name of the broker.
    pub fn host(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.host)
                .to_str()
                .expect("Broker host is not a valid UTF-8 string")
        }
    }

    /// Returns the port of the broker.
    pub fn port(&self) -> i32 {
        self.0.port
    }
}

/// Partition metadata information.
pub struct MetadataPartition(RDKafkaMetadataPartition);

impl MetadataPartition {
    /// Returns the id of the partition.
    pub fn id(&self) -> i32 {
        self.0.id
    }

    /// Returns the broker id of the leader broker for the partition.
    pub fn leader(&self) -> i32 {
        self.0.leader
    }

    // TODO: return result?
    /// Returns the metadata error for the partition, or None if there is no error.
    pub fn error(&self) -> Option<RDKafkaRespErr> {
        if self.0.err.is_error() {
            Some(self.0.err)
        } else {
            None
        }
    }

    /// Returns the broker ids of the replicas.
    pub fn replicas(&self) -> &[i32] {
        unsafe { slice::from_raw_parts(self.0.replicas, self.0.replica_cnt as usize) }
    }

    /// Returns the broker ids of the in sync replicas.
    pub fn isr(&self) -> &[i32] {
        unsafe { slice::from_raw_parts(self.0.isrs, self.0.isr_cnt as usize) }
    }
}

/// Topic metadata information.
pub struct MetadataTopic(RDKafkaMetadataTopic);

impl MetadataTopic {
    /// Returns the name of the topic.
    pub fn name(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.topic)
                .to_str()
                .expect("Topic name is not a valid UTF-8 string")
        }
    }

    /// Returns the partition metadata information for all the partitions.
    pub fn partitions(&self) -> &[MetadataPartition] {
        unsafe {
            slice::from_raw_parts(
                self.0.partitions as *const MetadataPartition,
                self.0.partition_cnt as usize,
            )
        }
    }

    /// Returns the metadata error, or None if there was no error.
    pub fn error(&self) -> Option<RDKafkaRespErr> {
        if self.0.err.is_error() {
            Some(self.0.err)
        } else {
            None
        }
    }
}

/// Metadata container. This structure wraps the metadata pointer returned by rdkafka-sys,
/// and deallocates all the native resources when dropped.
pub struct Metadata(*const RDKafkaMetadata);

impl Metadata {
    /// Creates a new Metadata container given a pointer to the native rdkafka-sys metadata.
    pub(crate) unsafe fn from_ptr(ptr: *const RDKafkaMetadata) -> Metadata {
        Metadata(ptr)
    }

    /// Returns the id of the broker originating this metadata.
    pub fn orig_broker_id(&self) -> i32 {
        unsafe { (*self.0).orig_broker_id }
    }

    /// Returns the hostname of the broker originating this metadata.
    pub fn orig_broker_name(&self) -> &str {
        unsafe {
            CStr::from_ptr((*self.0).orig_broker_name)
                .to_str()
                .expect("Broker name is not a valid UTF-8 string")
        }
    }

    /// Returns the metadata information for all the brokers in the cluster.
    pub fn brokers(&self) -> &[MetadataBroker] {
        unsafe {
            slice::from_raw_parts(
                (*self.0).brokers as *const MetadataBroker,
                (*self.0).broker_cnt as usize,
            )
        }
    }

    /// Returns the metadata information for all the topics in the cluster.
    pub fn topics(&self) -> &[MetadataTopic] {
        unsafe {
            slice::from_raw_parts(
                (*self.0).topics as *const MetadataTopic,
                (*self.0).topic_cnt as usize,
            )
        }
    }
}

impl Drop for Metadata {
    fn drop(&mut self) {
        unsafe { rdsys::rd_kafka_metadata_destroy(self.0) };
    }
}
<commit_msg>Implement Send and Sync for Metadata<commit_after>//! Cluster metadata.

use std::ffi::CStr;
use std::slice;

use rdkafka_sys as rdsys;
use rdkafka_sys::types::*;

use crate::error::IsError;

/// Broker metadata information.
pub struct MetadataBroker(RDKafkaMetadataBroker);

impl MetadataBroker {
    /// Returns the id of the broker.
    pub fn id(&self) -> i32 {
        self.0.id
    }

    /// Returns the host name of the broker.
    pub fn host(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.host)
                .to_str()
                .expect("Broker host is not a valid UTF-8 string")
        }
    }

    /// Returns the port of the broker.
    pub fn port(&self) -> i32 {
        self.0.port
    }
}

/// Partition metadata information.
pub struct MetadataPartition(RDKafkaMetadataPartition);

impl MetadataPartition {
    /// Returns the id of the partition.
    pub fn id(&self) -> i32 {
        self.0.id
    }

    /// Returns the broker id of the leader broker for the partition.
    pub fn leader(&self) -> i32 {
        self.0.leader
    }

    // TODO: return result?
    /// Returns the metadata error for the partition, or None if there is no error.
    pub fn error(&self) -> Option<RDKafkaRespErr> {
        if self.0.err.is_error() {
            Some(self.0.err)
        } else {
            None
        }
    }

    /// Returns the broker ids of the replicas.
    pub fn replicas(&self) -> &[i32] {
        unsafe { slice::from_raw_parts(self.0.replicas, self.0.replica_cnt as usize) }
    }

    /// Returns the broker ids of the in sync replicas.
    pub fn isr(&self) -> &[i32] {
        unsafe { slice::from_raw_parts(self.0.isrs, self.0.isr_cnt as usize) }
    }
}

/// Topic metadata information.
pub struct MetadataTopic(RDKafkaMetadataTopic);

impl MetadataTopic {
    /// Returns the name of the topic.
    pub fn name(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.topic)
                .to_str()
                .expect("Topic name is not a valid UTF-8 string")
        }
    }

    /// Returns the partition metadata information for all the partitions.
    pub fn partitions(&self) -> &[MetadataPartition] {
        unsafe {
            slice::from_raw_parts(
                self.0.partitions as *const MetadataPartition,
                self.0.partition_cnt as usize,
            )
        }
    }

    /// Returns the metadata error, or None if there was no error.
    pub fn error(&self) -> Option<RDKafkaRespErr> {
        if self.0.err.is_error() {
            Some(self.0.err)
        } else {
            None
        }
    }
}

/// Metadata container. This structure wraps the metadata pointer returned by rdkafka-sys,
/// and deallocates all the native resources when dropped.
pub struct Metadata(*const RDKafkaMetadata);

impl Metadata {
    /// Creates a new Metadata container given a pointer to the native rdkafka-sys metadata.
    pub(crate) unsafe fn from_ptr(ptr: *const RDKafkaMetadata) -> Metadata {
        Metadata(ptr)
    }

    /// Returns the id of the broker originating this metadata.
    pub fn orig_broker_id(&self) -> i32 {
        unsafe { (*self.0).orig_broker_id }
    }

    /// Returns the hostname of the broker originating this metadata.
    pub fn orig_broker_name(&self) -> &str {
        unsafe {
            CStr::from_ptr((*self.0).orig_broker_name)
                .to_str()
                .expect("Broker name is not a valid UTF-8 string")
        }
    }

    /// Returns the metadata information for all the brokers in the cluster.
    pub fn brokers(&self) -> &[MetadataBroker] {
        unsafe {
            slice::from_raw_parts(
                (*self.0).brokers as *const MetadataBroker,
                (*self.0).broker_cnt as usize,
            )
        }
    }

    /// Returns the metadata information for all the topics in the cluster.
    pub fn topics(&self) -> &[MetadataTopic] {
        unsafe {
            slice::from_raw_parts(
                (*self.0).topics as *const MetadataTopic,
                (*self.0).topic_cnt as usize,
            )
        }
    }
}

impl Drop for Metadata {
    fn drop(&mut self) {
        unsafe { rdsys::rd_kafka_metadata_destroy(self.0) };
    }
}

unsafe impl Send for Metadata {}
unsafe impl Sync for Metadata {}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::BufReader;
use std::io::Read;
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::str;

use rustc_serialize::json::{self, Json};
use rustc_serialize::Decodable;

use git2::{Oid, Repository};

#[derive(RustcDecodable, RustcEncodable)]
pub struct TreeConfigPaths {
    pub index_path: String,
    pub files_path: String,
    pub git_path: Option<String>,
    pub git_blame_path: Option<String>,
    pub objdir_path: String,
    pub hg_root: Option<String>,
    pub github_repo: Option<String>,
}

pub struct GitData {
    pub repo: Repository,
    pub blame_repo: Option<Repository>,

    pub blame_map: HashMap<Oid, Oid>, // Maps repo OID to blame_repo OID.
    pub hg_map: HashMap<Oid, String>, // Maps repo OID to Hg rev.

    pub mailmap: Mailmap,
}

pub struct TreeConfig {
    pub paths: TreeConfigPaths,
    pub git: Option<GitData>,
}

pub struct Config {
    pub trees: BTreeMap<String, TreeConfig>,
    pub mozsearch_path: String,
}

pub fn get_git(tree_config: &TreeConfig) -> Result<&GitData, &'static str> {
    match &tree_config.git {
        &Some(ref git) => Ok(git),
        &None => Err("History data unavailable"),
    }
}

pub fn get_git_path(tree_config: &TreeConfig) -> Result<&str, &'static str> {
    match &tree_config.paths.git_path {
        &Some(ref git_path) => Ok(git_path),
        &None => Err("History data unavailable"),
    }
}

pub fn get_hg_root(tree_config: &TreeConfig) -> String {
    // For temporary backwards compatibility, produce the m-c root if
    // there isn't one specified. We can remove this once all relevant
    // deployed config.json files have an explicit hg root, and make
    // this return an Option<&str> instead.
    match tree_config.paths.hg_root {
        Some(ref hg_root) => hg_root.clone(),
        None => String::from("https://hg.mozilla.org/mozilla-central"),
    }
}

pub fn get_git_commit_link(tree_config: &TreeConfig, commit_id: &str) -> String {
    // For backwards compatibility, produce the gecko-dev link if there isn't
    // one specified. We can remove this if we want as described above for
    // get_hg_root and such.
    let repo = match tree_config.paths.github_repo {
        Some(ref repo) => &**repo,
        None => "https://github.com/mozilla/gecko-dev",
    };

    format!("{}/commit/{}", repo, commit_id)
}

fn index_blame(_repo: &Repository, blame_repo: &Repository) -> (HashMap<Oid, Oid>, HashMap<Oid, String>) {
    let mut walk = blame_repo.revwalk().unwrap();
    walk.push_head().unwrap();

    let mut blame_map = HashMap::new();
    let mut hg_map = HashMap::new();
    for r in walk {
        let oid = r.unwrap();
        let commit = blame_repo.find_commit(oid).unwrap();

        let msg = commit.message().unwrap();
        let pieces = msg.split_whitespace().collect::<Vec<_>>();

        let orig_oid = Oid::from_str(pieces[1]).unwrap();
        blame_map.insert(orig_oid, commit.id());

        if pieces.len() > 2 {
            let hg_id = pieces[3].to_owned();
            hg_map.insert(orig_oid, hg_id);
        }
    }

    (blame_map, hg_map)
}

pub fn load(config_path: &str, need_indexes: bool) -> Config {
    let config_file = File::open(config_path).unwrap();
    let mut reader = BufReader::new(&config_file);
    let mut input = String::new();
    reader.read_to_string(&mut input).unwrap();
    let config = Json::from_str(&input).unwrap();

    let mut obj = config.as_object().unwrap().clone();

    let mozsearch_json = obj.remove("mozsearch_path").unwrap();
    let mozsearch = mozsearch_json.as_string().unwrap();

    let trees_obj = obj.get("trees").unwrap().as_object().unwrap().clone();

    let mut trees = BTreeMap::new();
    for (tree_name, tree_config) in trees_obj {
        let mut decoder = json::Decoder::new(tree_config);
        let paths = TreeConfigPaths::decode(&mut decoder).unwrap();

        let git = match (&paths.git_path, &paths.git_blame_path) {
            (&Some(ref git_path), &Some(ref git_blame_path)) => {
                let repo = Repository::open(&git_path).unwrap();
                let mailmap = Mailmap::load(&repo);

                let blame_repo = Repository::open(&git_blame_path).unwrap();
                let (blame_map, hg_map) = if need_indexes {
                    index_blame(&repo, &blame_repo)
                } else {
                    (HashMap::new(), HashMap::new())
                };

                Some(GitData {
                    repo: repo,
                    blame_repo: Some(blame_repo),
                    blame_map: blame_map,
                    hg_map: hg_map,
                    mailmap: mailmap,
                })
            },
            (&Some(ref git_path), &None) => {
                let repo = Repository::open(&git_path).unwrap();
                let mailmap = Mailmap::load(&repo);

                Some(GitData {
                    repo: repo,
                    blame_repo: None,
                    blame_map: HashMap::new(),
                    hg_map: HashMap::new(),
                    mailmap: mailmap,
                })
            },
            _ => None,
        };

        trees.insert(tree_name, TreeConfig {
            paths: paths,
            git: git,
        });
    }

    Config { trees: trees, mozsearch_path: mozsearch.to_owned() }
}

#[derive(Hash, Eq, PartialEq, Debug)]
struct MailmapKey(Option<String>, Option<String>);

/// Mapping from names and emails to replace to the real names and emails for
/// these authors.
pub struct Mailmap {
    /// Map from old name and email to real name and email
    entries: HashMap<MailmapKey, MailmapKey>,
}

impl Mailmap {
    // Look up an entry in the mailmap, and return the real name and email.
    pub fn lookup<'a>(&'a self, name: &'a str, email: &'a str) -> (&'a str, &'a str) {
        // Unfortunately, we need to actually own our key strings due to type
        // matching & the keys being tuple-structs. I doubt this will have any
        // meaningful perf impact.

        // Try to look up with both name & email.
        let mut key = MailmapKey(Some(name.to_owned()), Some(email.to_owned()));
        if let Some(&MailmapKey(ref new_name, ref new_email)) = self.entries.get(&key) {
            return (
                new_name.as_ref().map(String::as_str).unwrap_or(name),
                new_email.as_ref().map(String::as_str).unwrap_or(email),
            )
        }

        // Try looking up only by email.
        key.0 = None;
        if let Some(&MailmapKey(ref new_name, ref new_email)) = self.entries.get(&key) {
            return (
                new_name.as_ref().map(String::as_str).unwrap_or(name),
                new_email.as_ref().map(String::as_str).unwrap_or(email),
            )
        }

        // Not in the mailmap, return it as-is.
        (name, email)
    }

    /// Load the Mailmap for the given repository.
    pub fn load(repo: &Repository) -> Self {
        // Repo may not have a mailmap file, in which case we can just generate
        // an empty one.
        Mailmap::try_load(repo).unwrap_or_else(|| Mailmap { entries: HashMap::new() })
    }

    fn parse_line(mut line: &str) -> Option<(MailmapKey, MailmapKey)> {
        fn nonempty(s: &str) -> Option<String> {
            if s.is_empty() { None } else { Some(s.to_owned()) }
        }

        // Remove text after a '#' comment from the line.
        line = line.split('#').next().unwrap();

        // name_a is the optional string before the first email.
        let idx = line.find('<')?;
        let name_a = nonempty(line[..idx].trim());
        line = &line[idx + 1..];

        // email_a is the required string until the end of the email block.
        let idx = line.find('>')?;
        let email_a = line[..idx].trim().to_owned();
        line = &line[idx + 1..];

        // name_b and email_b are optional. name_b requires email_b.
        let (name_b, email_b) = if let Some(idx) = line.find('<') {
            let name_b = nonempty(line[..idx].trim());
            line = &line[idx + 1..];

            let idx = line.find('>')?;
            let email_b = line[..idx].trim().to_owned();
            line = &line[idx + 1..];

            (name_b, Some(email_b))
        } else {
            (None, None)
        };

        // If we have junk at the end of the line - ignore it.
        if !line.trim().is_empty() {
            return None;
        }

        // Determine which format was being used, and build up our old and new
        // mailmap keys.
        let old;
        let new;
        if let Some(email_b) = email_b {
            new = MailmapKey(name_a, Some(email_a));
            old = MailmapKey(name_b, Some(email_b));
        } else {
            new = MailmapKey(name_a, None);
            old = MailmapKey(None, Some(email_a));
        }

        Some((old, new))
    }

    fn try_load(repo: &Repository) -> Option<Self> {
        // Get current mailmap from the repository.
        let obj = repo.revparse_single("HEAD:.mailmap").ok()?;
        let blob = obj.peel_to_blob().ok()?;
        let data = str::from_utf8(blob.content()).ok()?;

        // Parse each entry in turn
        let mut entries = HashMap::new();
        for line in data.lines() {
            if let Some((old, new)) = Mailmap::parse_line(line) {
                entries.insert(old, new);
            }
        }

        Some(Mailmap { entries })
    }
}
<commit_msg>Bug 1507923 - Read the .git-blame-ignore-revs file into a BlameIgnoreList<commit_after>use std::fs::File;
use std::io::BufReader;
use std::io::Read;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::str;

use rustc_serialize::json::{self, Json};
use rustc_serialize::Decodable;

use git2::{Oid, Repository};

#[derive(RustcDecodable, RustcEncodable)]
pub struct TreeConfigPaths {
    pub index_path: String,
    pub files_path: String,
    pub git_path: Option<String>,
    pub git_blame_path: Option<String>,
    pub objdir_path: String,
    pub hg_root: Option<String>,
    pub github_repo: Option<String>,
}

pub struct GitData {
    pub repo: Repository,
    pub blame_repo: Option<Repository>,

    pub blame_map: HashMap<Oid, Oid>, // Maps repo OID to blame_repo OID.
    pub hg_map: HashMap<Oid, String>, // Maps repo OID to Hg rev.

    pub mailmap: Mailmap,
    /// Revs that we want to skip over during blame computation
    pub blame_ignore: BlameIgnoreList,
}

pub struct TreeConfig {
    pub paths: TreeConfigPaths,
    pub git: Option<GitData>,
}

pub struct Config {
    pub trees: BTreeMap<String, TreeConfig>,
    pub mozsearch_path: String,
}

pub fn get_git(tree_config: &TreeConfig) -> Result<&GitData, &'static str> {
    match &tree_config.git {
        &Some(ref git) => Ok(git),
        &None => Err("History data unavailable"),
    }
}

pub fn get_git_path(tree_config: &TreeConfig) -> Result<&str, &'static str> {
    match &tree_config.paths.git_path {
        &Some(ref git_path) => Ok(git_path),
        &None => Err("History data unavailable"),
    }
}

pub fn get_hg_root(tree_config: &TreeConfig) -> String {
    // For temporary backwards compatibility, produce the m-c root if
    // there isn't one specified. We can remove this once all relevant
    // deployed config.json files have an explicit hg root, and make
    // this return an Option<&str> instead.
    match tree_config.paths.hg_root {
        Some(ref hg_root) => hg_root.clone(),
        None => String::from("https://hg.mozilla.org/mozilla-central"),
    }
}

pub fn get_git_commit_link(tree_config: &TreeConfig, commit_id: &str) -> String {
    // For backwards compatibility, produce the gecko-dev link if there isn't
    // one specified. We can remove this if we want as described above for
    // get_hg_root and such.
    let repo = match tree_config.paths.github_repo {
        Some(ref repo) => &**repo,
        None => "https://github.com/mozilla/gecko-dev",
    };

    format!("{}/commit/{}", repo, commit_id)
}

fn index_blame(_repo: &Repository, blame_repo: &Repository) -> (HashMap<Oid, Oid>, HashMap<Oid, String>) {
    let mut walk = blame_repo.revwalk().unwrap();
    walk.push_head().unwrap();

    let mut blame_map = HashMap::new();
    let mut hg_map = HashMap::new();
    for r in walk {
        let oid = r.unwrap();
        let commit = blame_repo.find_commit(oid).unwrap();

        let msg = commit.message().unwrap();
        let pieces = msg.split_whitespace().collect::<Vec<_>>();

        let orig_oid = Oid::from_str(pieces[1]).unwrap();
        blame_map.insert(orig_oid, commit.id());

        if pieces.len() > 2 {
            let hg_id = pieces[3].to_owned();
            hg_map.insert(orig_oid, hg_id);
        }
    }

    (blame_map, hg_map)
}

pub fn load(config_path: &str, need_indexes: bool) -> Config {
    let config_file = File::open(config_path).unwrap();
    let mut reader = BufReader::new(&config_file);
    let mut input = String::new();
    reader.read_to_string(&mut input).unwrap();
    let config = Json::from_str(&input).unwrap();

    let mut obj = config.as_object().unwrap().clone();

    let mozsearch_json = obj.remove("mozsearch_path").unwrap();
    let mozsearch = mozsearch_json.as_string().unwrap();

    let trees_obj = obj.get("trees").unwrap().as_object().unwrap().clone();

    let mut trees = BTreeMap::new();
    for (tree_name, tree_config) in trees_obj {
        let mut decoder = json::Decoder::new(tree_config);
        let paths = TreeConfigPaths::decode(&mut decoder).unwrap();

        let git = match (&paths.git_path, &paths.git_blame_path) {
            (&Some(ref git_path), &Some(ref git_blame_path)) => {
                let repo = Repository::open(&git_path).unwrap();
                let mailmap = Mailmap::load(&repo);
                let blame_ignore = BlameIgnoreList::load(&repo);

                let blame_repo = Repository::open(&git_blame_path).unwrap();
                let (blame_map, hg_map) = if need_indexes {
                    index_blame(&repo, &blame_repo)
                } else {
                    (HashMap::new(), HashMap::new())
                };

                Some(GitData {
                    repo: repo,
                    blame_repo: Some(blame_repo),
                    blame_map: blame_map,
                    hg_map: hg_map,
                    mailmap: mailmap,
                    blame_ignore: blame_ignore,
                })
            },
            (&Some(ref git_path), &None) => {
                let repo = Repository::open(&git_path).unwrap();
                let mailmap = Mailmap::load(&repo);
                let blame_ignore = BlameIgnoreList::load(&repo);

                Some(GitData {
                    repo: repo,
                    blame_repo: None,
                    blame_map: HashMap::new(),
                    hg_map: HashMap::new(),
                    mailmap: mailmap,
                    blame_ignore: blame_ignore,
                })
            },
            _ => None,
        };

        trees.insert(tree_name, TreeConfig {
            paths: paths,
            git: git,
        });
    }

    Config { trees: trees, mozsearch_path: mozsearch.to_owned() }
}

#[derive(Hash, Eq, PartialEq, Debug)]
struct MailmapKey(Option<String>, Option<String>);

/// Mapping from names and emails to replace to the real names and emails for
/// these authors.
pub struct Mailmap {
    /// Map from old name and email to real name and email
    entries: HashMap<MailmapKey, MailmapKey>,
}

impl Mailmap {
    // Look up an entry in the mailmap, and return the real name and email.
    pub fn lookup<'a>(&'a self, name: &'a str, email: &'a str) -> (&'a str, &'a str) {
        // Unfortunately, we need to actually own our key strings due to type
        // matching & the keys being tuple-structs. I doubt this will have any
        // meaningful perf impact.

        // Try to look up with both name & email.
        let mut key = MailmapKey(Some(name.to_owned()), Some(email.to_owned()));
        if let Some(&MailmapKey(ref new_name, ref new_email)) = self.entries.get(&key) {
            return (
                new_name.as_ref().map(String::as_str).unwrap_or(name),
                new_email.as_ref().map(String::as_str).unwrap_or(email),
            )
        }

        // Try looking up only by email.
        key.0 = None;
        if let Some(&MailmapKey(ref new_name, ref new_email)) = self.entries.get(&key) {
            return (
                new_name.as_ref().map(String::as_str).unwrap_or(name),
                new_email.as_ref().map(String::as_str).unwrap_or(email),
            )
        }

        // Not in the mailmap, return it as-is.
        (name, email)
    }

    /// Load the Mailmap for the given repository.
    pub fn load(repo: &Repository) -> Self {
        // Repo may not have a mailmap file, in which case we can just generate
        // an empty one.
        Mailmap::try_load(repo).unwrap_or_else(|| Mailmap { entries: HashMap::new() })
    }

    fn parse_line(mut line: &str) -> Option<(MailmapKey, MailmapKey)> {
        fn nonempty(s: &str) -> Option<String> {
            if s.is_empty() { None } else { Some(s.to_owned()) }
        }

        // Remove text after a '#' comment from the line.
        line = line.split('#').next().unwrap();

        // name_a is the optional string before the first email.
        let idx = line.find('<')?;
        let name_a = nonempty(line[..idx].trim());
        line = &line[idx + 1..];

        // email_a is the required string until the end of the email block.
        let idx = line.find('>')?;
        let email_a = line[..idx].trim().to_owned();
        line = &line[idx + 1..];

        // name_b and email_b are optional. name_b requires email_b.
        let (name_b, email_b) = if let Some(idx) = line.find('<') {
            let name_b = nonempty(line[..idx].trim());
            line = &line[idx + 1..];

            let idx = line.find('>')?;
            let email_b = line[..idx].trim().to_owned();
            line = &line[idx + 1..];

            (name_b, Some(email_b))
        } else {
            (None, None)
        };

        // If we have junk at the end of the line - ignore it.
        if !line.trim().is_empty() {
            return None;
        }

        // Determine which format was being used, and build up our old and new
        // mailmap keys.
        let old;
        let new;
        if let Some(email_b) = email_b {
            new = MailmapKey(name_a, Some(email_a));
            old = MailmapKey(name_b, Some(email_b));
        } else {
            new = MailmapKey(name_a, None);
            old = MailmapKey(None, Some(email_a));
        }

        Some((old, new))
    }

    fn try_load(repo: &Repository) -> Option<Self> {
        // Get current mailmap from the repository.
        let obj = repo.revparse_single("HEAD:.mailmap").ok()?;
        let blob = obj.peel_to_blob().ok()?;
        let data = str::from_utf8(blob.content()).ok()?;

        // Parse each entry in turn
        let mut entries = HashMap::new();
        for line in data.lines() {
            if let Some((old, new)) = Mailmap::parse_line(line) {
                entries.insert(old, new);
            }
        }

        Some(Mailmap { entries })
    }
}

#[derive(Default)]
pub struct BlameIgnoreList {
    entries: HashSet<String>,
}

impl BlameIgnoreList {
    pub fn load(repo: &Repository) -> Self {
        // Produce an empty list if we fail to load anything
        BlameIgnoreList::try_load(repo).unwrap_or_default()
    }

    fn try_load(repo: &Repository) -> Option<Self> {
        let obj = repo.revparse_single("HEAD:.git-blame-ignore-revs").ok()?;
        let blob = obj.peel_to_blob().ok()?;
        let data = str::from_utf8(blob.content()).ok()?;

        let mut entries = HashSet::new();
        for line in data.lines() {
            let trimmed = line.split('#').next().unwrap().trim();
            // I guess we could also verify these are actually revisions but
            // that will eat CPU cycles for not much real benefit
            if !trimmed.is_empty() {
                entries.insert(trimmed.to_owned());
            }
        }

        Some(BlameIgnoreList { entries })
    }

    pub fn should_ignore(&self, rev: &str) -> bool {
        self.entries.contains(rev)
    }
}

impl GitData {
    pub fn should_ignore_for_blame(&self, rev: &str) -> bool {
        // TODO: we might want to pull the commit message and check for
        // special annotations like "#skip-blame" or backouts as well.
        // For now just check the list.
        self.blame_ignore.should_ignore(rev)
    }
}
<|endoftext|>