<commit_before>// Copyright (c) 2016 Fedor Gogolev <knsd@knsd.net>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::io::{Read, BufRead, Error as IoError, ErrorKind};
use std::string::{FromUtf8Error};
use std::collections::{HashMap};
use std::cell::{RefCell};
use std::rc::{Rc};

use num::{Zero};
use num::bigint::{BigInt, ToBigInt, Sign};
use byteorder::{ReadBytesExt, LittleEndian, BigEndian, Error as ByteorderError};
use from_ascii::{FromAscii, ParseIntError, ParseFloatError};

use string::{unescape, Error as UnescapeError};
use value::{Value};

use opcodes::*;

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        EmptyMarker
        StackTooSmall
        EmptyStack
        InvalidValueOnStack
        InvalidGetValue
        InvalidPutValue
        NotImplemented

        Read(err: ByteorderError) {
            from()
        }
        Io(err: IoError) {
            from()
        }
        UnknownOpcode(opcode: u8) {}

        InvalidInt {
            from(ParseIntError)
        }
        InvalidLong
        InvalidFloat {
            from(ParseFloatError)
        }

        InvalidString {
            from(FromUtf8Error)
        }
        UnescapeError(err: UnescapeError) {
            from()
        }

        InvalidProto(proto: u8)
        NegativeLength {}

        #[doc(hidden)]
        __Nonexhaustive
    }
}

#[derive(Debug, PartialEq)]
pub enum BooleanOrInt {
    Boolean(bool),
    Int(i64),
}

macro_rules! rc {
    ($term: expr) => (Rc::new(RefCell::new($term)))
}

fn read_exact<R>(rd: &mut R, mut buf: &mut [u8]) -> Result<(), IoError> where R: Read {
    while !buf.is_empty() {
        match rd.read(buf) {
            Ok(0) => break,
            Ok(n) => { let tmp = buf; buf = &mut tmp[n..]; }
            Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    if !buf.is_empty() {
        Err(IoError::new(ErrorKind::Other,
                       "failed to fill whole buffer"))
    } else {
        Ok(())
    }
}

fn read_until_newline<R>(rd: &mut R) -> Result<Vec<u8>, Error> where R: Read + BufRead {
    let mut buf = Vec::new();
    try!(rd.read_until('\n' as u8, &mut buf));

    // Skip last symbol — \n
    match buf.split_last() {
        Some((&b'\n', init)) => Ok(init.to_vec()),
        _ => Err(Error::InvalidString),
    }
}

fn read_decimal_int<R>(rd: &mut R) -> Result<BooleanOrInt, Error> where R: Read + BufRead {
    let s = try!(read_until_newline(rd));
    let val = match &s[..] {
        b"00" => BooleanOrInt::Boolean(false),
        b"01" => BooleanOrInt::Boolean(true),
        _ => BooleanOrInt::Int(try!(i64::from_ascii(&s)))
    };
    Ok(val)
}

fn read_decimal_long<R>(rd: &mut R) -> Result<BigInt, Error> where R: Read + BufRead {
    let s = try!(read_until_newline(rd));
    let init = match s.split_last() {
        None => return Err(Error::InvalidString),
        Some((&b'L', init)) => init,
        Some(_) => &s[..],
    };

    match BigInt::parse_bytes(&init, 10) {
        Some(i) => Ok(i),
        None => Err(Error::InvalidLong)
    }
}

fn read_long<R>(rd: &mut R, length: usize) -> Result<BigInt, Error> where R: Read + BufRead {
    let mut buf = vec![0; length];
    try!(read_exact(rd, buf.as_mut()));

    let mut n = BigInt::from_bytes_le(Sign::Plus, &buf);

    let last = match buf.last_mut() {
        None => return Err(Error::InvalidLong),
        Some(last) => last,
    };

    if *last > 127 {
        n = n - (1.to_bigint().unwrap() << (length * 8))
    }

    Ok(n)
}

pub struct Machine {
    stack: Vec<Value>,
    memo: HashMap<usize, Value>,
    marker: Option<usize>,
}

impl Machine {
    pub fn new() -> Self {
        Machine {
            stack: Vec::new(),
            memo: HashMap::new(),
            marker: None,
        }
    }

    fn split_off(&mut self) -> Result<Vec<Value>, Error> {
        let at = match self.marker {
            None => return Err(Error::EmptyMarker),
            Some(marker) => marker,
        };

        if at > self.stack.len() {
            return Err(Error::StackTooSmall);
        }

        Ok(self.stack.split_off(at))
    }

    fn pop(&mut self) -> Result<Value, Error> {
        match self.stack.pop() {
            None => Err(Error::EmptyStack),
            Some(value) => Ok(value),
        }
    }

    fn handle_get(&mut self, i: usize) -> Result<(), Error> {
        let value = match self.memo.get(&i) {
            None => return Err(Error::InvalidGetValue),
            Some(ref v) => (*v).clone(),
        };
        self.stack.push(value);
        Ok(())
    }

    fn handle_put(&mut self, i: usize) -> Result<(), Error> {
        let value = match self.stack.last() {
            None => return Err(Error::EmptyStack),
            Some(ref v) => (*v).clone(),
        };
        self.memo.insert(i, value);
        Ok(())
    }

    pub fn execute<R>(&mut self, rd: &mut R) -> Result<bool, Error> where R: Read + BufRead {
        macro_rules! ensure_not_negative {
            ($n: expr) => ({
                if $n < Zero::zero() {
                    return Err(Error::NegativeLength)
                }
            })
        }

        match try!(rd.read_u8()) {
            PROTO => {
                let version = try!(rd.read_u8());
                if version < 2 {
                    return Err(Error::InvalidProto(version))
                }
            },
            STOP => return Ok(true),

            INT => {
                self.stack.push(match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Boolean(v) => Value::Bool(v),
                    BooleanOrInt::Int(v) => Value::Long(BigInt::from(v)),
                })
            },
            BININT => self.stack.push(Value::Int(try!(rd.read_i32::<LittleEndian>()) as isize)),
            BININT1 => self.stack.push(Value::Int(try!(rd.read_u8()) as isize)),
            BININT2 => self.stack.push(Value::Int(try!(rd.read_u16::<LittleEndian>()) as isize)),
            LONG => self.stack.push(Value::Long(BigInt::from(try!(read_decimal_long(rd))))),
            LONG1 => {
                let length = try!(rd.read_u8());
                self.stack.push(Value::Long(BigInt::from(try!(read_long(rd, length as usize)))))
            }
            LONG4 => {
                let length = try!(rd.read_i32::<LittleEndian>());
                self.stack.push(Value::Long(BigInt::from(try!(read_long(rd, length as usize)))))
            }

            STRING => self.stack.push(Value::String(try!(unescape(&try!(read_until_newline(rd)), false)))),
            BINSTRING => {
                let length = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(length);

                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, &mut buf));
                self.stack.push(Value::String(buf))
            },
            SHORT_BINSTRING => {
                let length = try!(rd.read_u8());
                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, &mut buf));
                self.stack.push(Value::String(buf))
            },

            NONE => self.stack.push(Value::None),
            NEWTRUE => self.stack.push(Value::Bool(true)),
            NEWFALSE => self.stack.push(Value::Bool(false)),

            UNICODE => {
                let buf = try!(unescape(&try!(read_until_newline(rd)), true));
                self.stack.push(Value::Unicode(try!(String::from_utf8(buf))))
            },
            BINUNICODE => {
                let length = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(length);
                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, buf.as_mut()));
                self.stack.push(Value::Unicode(try!(String::from_utf8(buf))))
            },

            FLOAT => {
                let s = try!(read_until_newline(rd));
                self.stack.push(Value::Float(try!(f64::from_ascii(&s))))
            },
            BINFLOAT => {
                self.stack.push(Value::Float(try!(rd.read_f64::<BigEndian>())))
            },

            EMPTY_LIST => {
                self.stack.push(Value::List(rc!(Vec::new())))
            },
            APPEND => {
                let v = try!(self.pop());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::List(ref mut list)) => (*list.borrow_mut()).push(v),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            APPENDS => {
                let values = try!(self.split_off());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::List(ref mut list)) => (*list.borrow_mut()).extend(values),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            LIST => {
                let values = try!(self.split_off());
                self.stack.push(Value::List(rc!(values)));
            },

            EMPTY_TUPLE => self.stack.push(Value::Tuple(rc!(Vec::new()))),
            TUPLE => {
                let values = try!(self.split_off());
                self.stack.push(Value::Tuple(rc!(values)));
            },
            TUPLE1 => {
                let v1 = try!(self.pop());
                self.stack.push(Value::Tuple(rc!(vec![v1])))
            },
            TUPLE2 => {
                let v1 = try!(self.pop());
                let v2 = try!(self.pop());
                self.stack.push(Value::Tuple(rc!(vec![v1, v2])))
            },
            TUPLE3 => {
                let v1 = try!(self.pop());
                let v2 = try!(self.pop());
                let v3 = try!(self.pop());
                self.stack.push(Value::Tuple(rc!(vec![v1, v2, v3])))
            }

            EMPTY_DICT => self.stack.push(Value::Dict(rc!(Vec::new()))),
            DICT => {
                let mut values = try!(self.split_off());
                let mut dict = Vec::new();

                for i in 0 .. values.len() / 2 { // TODO: Check panic
                    let key = values.remove(2 * i);
                    let value = values.remove(2 * i + 1);
                    dict.push((key, value));
                }
                self.stack.push(Value::Dict(rc!(dict)));
            },
            SETITEM => {
                let value = try!(self.pop());
                let key = try!(self.pop());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::Dict(ref mut dict)) => (*dict.borrow_mut()).push((key, value)),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            SETITEMS => {
                let mut values = try!(self.split_off());

                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::Dict(ref mut dict_ref)) => {
                        for i in 0 .. values.len() / 2 { // TODO: Check panic
                                let key = values.remove(2 * i);
                                let value = values.remove(2 * i + 1);
                                (*dict_ref.borrow_mut()).push((key, value));
                            }
                    },
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },

            POP => {
                try!(self.pop());
            },
            DUP => {
                let value = match self.stack.last() {
                    None => return Err(Error::EmptyStack),
                    Some(ref v) => (*v).clone(),
                };
                self.stack.push(value)
            },
            MARK => {
                self.marker = Some(self.stack.len())
            },
            POP_MARK => {
                try!(self.split_off());
            },

            GET => {
                let n = match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Int(n) => n,
                    BooleanOrInt::Boolean(false) => 0,
                    BooleanOrInt::Boolean(true) => 1,
                };
                ensure_not_negative!(n);
                try!(self.handle_get(n as usize))
            }
            BINGET => {
                try!(self.handle_get(try!(rd.read_u8()) as usize))
            }
            LONG_BINGET => {
                let n = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(n);
                try!(self.handle_get(n as usize))
            }

            PUT => {
                let n = match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Int(n) => n,
                    BooleanOrInt::Boolean(false) => 0,
                    BooleanOrInt::Boolean(true) => 1,
                };
                ensure_not_negative!(n);
                try!(self.handle_put(n as usize))
            }
            BINPUT => {
                try!(self.handle_put(try!(rd.read_u8()) as usize))
            }
            LONG_BINPUT => {
                let n = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(n);
                try!(self.handle_put(n as usize))
            }

            c => return Err(Error::UnknownOpcode(c)),
        }
        Ok(false)
    }
}

pub fn unpickle<R>(rd: &mut R) -> Result<Value, Error> where R: Read + BufRead {
    let mut machine = Machine::new();
    loop {
        if try!(machine.execute(rd)) {
            break
        }
    }
    Ok(try!(machine.pop()))
}

#[cfg(test)]
mod tests {
    use std::io::{Cursor};

    use num::{FromPrimitive};

    use super::{unpickle};
    use super::super::value::{Value};

    macro_rules! t {
        ($buffer: expr, $pat:pat, $result:expr) => ({
            match unpickle(&mut Cursor::new(&$buffer[..])) {
                Ok($pat) => $result,
                other => {
                    println!("ERROR {:?}", other);
                    assert!(false)
                },
            }
        })
    }

    macro_rules! e {
        ($buffer: expr, $pat:pat) => ({
            match unpickle(&mut Cursor::new(&$buffer[..])) {
                Err($pat) => (),
                other => {
                    println!("ERROR {:?}", other);
                    assert!(false)
                },
            }
        })
    }

    macro_rules! n {
        ($x: expr) => ({FromPrimitive::from_isize($x).unwrap()})
    }
}<commit_msg>Add int test<commit_after>// Copyright (c) 2016 Fedor Gogolev <knsd@knsd.net>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::io::{Read, BufRead, Error as IoError, ErrorKind};
use std::string::{FromUtf8Error};
use std::collections::{HashMap};
use std::cell::{RefCell};
use std::rc::{Rc};

use num::{Zero};
use num::bigint::{BigInt, ToBigInt, Sign};
use byteorder::{ReadBytesExt, LittleEndian, BigEndian, Error as ByteorderError};
use from_ascii::{FromAscii, ParseIntError, ParseFloatError};

use string::{unescape, Error as UnescapeError};
use value::{Value};

use opcodes::*;

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        EmptyMarker
        StackTooSmall
        EmptyStack
        InvalidValueOnStack
        InvalidGetValue
        InvalidPutValue
        NotImplemented

        Read(err: ByteorderError) {
            from()
        }
        Io(err: IoError) {
            from()
        }
        UnknownOpcode(opcode: u8) {}

        InvalidInt {
            from(ParseIntError)
        }
        InvalidLong
        InvalidFloat {
            from(ParseFloatError)
        }

        InvalidString {
            from(FromUtf8Error)
        }
        UnescapeError(err: UnescapeError) {
            from()
        }

        InvalidProto(proto: u8)
        NegativeLength {}

        #[doc(hidden)]
        __Nonexhaustive
    }
}

#[derive(Debug, PartialEq)]
pub enum BooleanOrInt {
    Boolean(bool),
    Int(i64),
}

macro_rules! rc {
    ($term: expr) => (Rc::new(RefCell::new($term)))
}

fn read_exact<R>(rd: &mut R, mut buf: &mut [u8]) -> Result<(), IoError> where R: Read {
    while !buf.is_empty() {
        match rd.read(buf) {
            Ok(0) => break,
            Ok(n) => { let tmp = buf; buf = &mut tmp[n..]; }
            Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    if !buf.is_empty() {
        Err(IoError::new(ErrorKind::Other,
                       "failed to fill whole buffer"))
    } else {
        Ok(())
    }
}

fn read_until_newline<R>(rd: &mut R) -> Result<Vec<u8>, Error> where R: Read + BufRead {
    let mut buf = Vec::new();
    try!(rd.read_until('\n' as u8, &mut buf));

    // Skip last symbol — \n
    match buf.split_last() {
        Some((&b'\n', init)) => Ok(init.to_vec()),
        _ => Err(Error::InvalidString),
    }
}

fn read_decimal_int<R>(rd: &mut R) -> Result<BooleanOrInt, Error> where R: Read + BufRead {
    let s = try!(read_until_newline(rd));
    let val = match &s[..] {
        b"00" => BooleanOrInt::Boolean(false),
        b"01" => BooleanOrInt::Boolean(true),
        _ => BooleanOrInt::Int(try!(i64::from_ascii(&s)))
    };
    Ok(val)
}

fn read_decimal_long<R>(rd: &mut R) -> Result<BigInt, Error> where R: Read + BufRead {
    let s = try!(read_until_newline(rd));
    let init = match s.split_last() {
        None => return Err(Error::InvalidString),
        Some((&b'L', init)) => init,
        Some(_) => &s[..],
    };

    match BigInt::parse_bytes(&init, 10) {
        Some(i) => Ok(i),
        None => Err(Error::InvalidLong)
    }
}

fn read_long<R>(rd: &mut R, length: usize) -> Result<BigInt, Error> where R: Read + BufRead {
    let mut buf = vec![0; length];
    try!(read_exact(rd, buf.as_mut()));

    let mut n = BigInt::from_bytes_le(Sign::Plus, &buf);

    let last = match buf.last_mut() {
        None => return Err(Error::InvalidLong),
        Some(last) => last,
    };

    if *last > 127 {
        n = n - (1.to_bigint().unwrap() << (length * 8))
    }

    Ok(n)
}

pub struct Machine {
    stack: Vec<Value>,
    memo: HashMap<usize, Value>,
    marker: Option<usize>,
}

impl Machine {
    pub fn new() -> Self {
        Machine {
            stack: Vec::new(),
            memo: HashMap::new(),
            marker: None,
        }
    }

    fn split_off(&mut self) -> Result<Vec<Value>, Error> {
        let at = match self.marker {
            None => return Err(Error::EmptyMarker),
            Some(marker) => marker,
        };

        if at > self.stack.len() {
            return Err(Error::StackTooSmall);
        }

        Ok(self.stack.split_off(at))
    }

    fn pop(&mut self) -> Result<Value, Error> {
        match self.stack.pop() {
            None => Err(Error::EmptyStack),
            Some(value) => Ok(value),
        }
    }

    fn handle_get(&mut self, i: usize) -> Result<(), Error> {
        let value = match self.memo.get(&i) {
            None => return Err(Error::InvalidGetValue),
            Some(ref v) => (*v).clone(),
        };
        self.stack.push(value);
        Ok(())
    }

    fn handle_put(&mut self, i: usize) -> Result<(), Error> {
        let value = match self.stack.last() {
            None => return Err(Error::EmptyStack),
            Some(ref v) => (*v).clone(),
        };
        self.memo.insert(i, value);
        Ok(())
    }

    pub fn execute<R>(&mut self, rd: &mut R) -> Result<bool, Error> where R: Read + BufRead {
        macro_rules! ensure_not_negative {
            ($n: expr) => ({
                if $n < Zero::zero() {
                    return Err(Error::NegativeLength)
                }
            })
        }

        match try!(rd.read_u8()) {
            PROTO => {
                let version = try!(rd.read_u8());
                if version < 2 {
                    return Err(Error::InvalidProto(version))
                }
            },
            STOP => return Ok(true),

            INT => {
                self.stack.push(match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Boolean(v) => Value::Bool(v),
                    BooleanOrInt::Int(v) => Value::Long(BigInt::from(v)), // FIXME: or int?
                })
            },
            BININT => self.stack.push(Value::Int(try!(rd.read_i32::<LittleEndian>()) as isize)),
            BININT1 => self.stack.push(Value::Int(try!(rd.read_u8()) as isize)),
            BININT2 => self.stack.push(Value::Int(try!(rd.read_u16::<LittleEndian>()) as isize)),
            LONG => self.stack.push(Value::Long(BigInt::from(try!(read_decimal_long(rd))))),
            LONG1 => {
                let length = try!(rd.read_u8());
                self.stack.push(Value::Long(BigInt::from(try!(read_long(rd, length as usize)))))
            }
            LONG4 => {
                let length = try!(rd.read_i32::<LittleEndian>());
                self.stack.push(Value::Long(BigInt::from(try!(read_long(rd, length as usize)))))
            }

            STRING => self.stack.push(Value::String(try!(unescape(&try!(read_until_newline(rd)), false)))),
            BINSTRING => {
                let length = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(length);

                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, &mut buf));
                self.stack.push(Value::String(buf))
            },
            SHORT_BINSTRING => {
                let length = try!(rd.read_u8());
                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, &mut buf));
                self.stack.push(Value::String(buf))
            },

            NONE => self.stack.push(Value::None),
            NEWTRUE => self.stack.push(Value::Bool(true)),
            NEWFALSE => self.stack.push(Value::Bool(false)),

            UNICODE => {
                let buf = try!(unescape(&try!(read_until_newline(rd)), true));
                self.stack.push(Value::Unicode(try!(String::from_utf8(buf))))
            },
            BINUNICODE => {
                let length = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(length);
                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, buf.as_mut()));
                self.stack.push(Value::Unicode(try!(String::from_utf8(buf))))
            },

            FLOAT => {
                let s = try!(read_until_newline(rd));
                self.stack.push(Value::Float(try!(f64::from_ascii(&s))))
            },
            BINFLOAT => {
                self.stack.push(Value::Float(try!(rd.read_f64::<BigEndian>())))
            },

            EMPTY_LIST => {
                self.stack.push(Value::List(rc!(Vec::new())))
            },
            APPEND => {
                let v = try!(self.pop());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::List(ref mut list)) => (*list.borrow_mut()).push(v),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            APPENDS => {
                let values = try!(self.split_off());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::List(ref mut list)) => (*list.borrow_mut()).extend(values),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            LIST => {
                let values = try!(self.split_off());
                self.stack.push(Value::List(rc!(values)));
            },

            EMPTY_TUPLE => self.stack.push(Value::Tuple(rc!(Vec::new()))),
            TUPLE => {
                let values = try!(self.split_off());
                self.stack.push(Value::Tuple(rc!(values)));
            },
            TUPLE1 => {
                let v1 = try!(self.pop());
                self.stack.push(Value::Tuple(rc!(vec![v1])))
            },
            TUPLE2 => {
                let v1 = try!(self.pop());
                let v2 = try!(self.pop());
                self.stack.push(Value::Tuple(rc!(vec![v1, v2])))
            },
            TUPLE3 => {
                let v1 = try!(self.pop());
                let v2 = try!(self.pop());
                let v3 = try!(self.pop());
                self.stack.push(Value::Tuple(rc!(vec![v1, v2, v3])))
            }

            EMPTY_DICT => self.stack.push(Value::Dict(rc!(Vec::new()))),
            DICT => {
                let mut values = try!(self.split_off());
                let mut dict = Vec::new();

                for i in 0 .. values.len() / 2 { // TODO: Check panic
                    let key = values.remove(2 * i);
                    let value = values.remove(2 * i + 1);
                    dict.push((key, value));
                }
                self.stack.push(Value::Dict(rc!(dict)));
            },
            SETITEM => {
                let value = try!(self.pop());
                let key = try!(self.pop());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::Dict(ref mut dict)) => (*dict.borrow_mut()).push((key, value)),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            SETITEMS => {
                let mut values = try!(self.split_off());

                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::Dict(ref mut dict_ref)) => {
                        for i in 0 .. values.len() / 2 { // TODO: Check panic
                                let key = values.remove(2 * i);
                                let value = values.remove(2 * i + 1);
                                (*dict_ref.borrow_mut()).push((key, value));
                            }
                    },
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },

            POP => {
                try!(self.pop());
            },
            DUP => {
                let value = match self.stack.last() {
                    None => return Err(Error::EmptyStack),
                    Some(ref v) => (*v).clone(),
                };
                self.stack.push(value)
            },
            MARK => {
                self.marker = Some(self.stack.len())
            },
            POP_MARK => {
                try!(self.split_off());
            },

            GET => {
                let n = match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Int(n) => n,
                    BooleanOrInt::Boolean(false) => 0,
                    BooleanOrInt::Boolean(true) => 1,
                };
                ensure_not_negative!(n);
                try!(self.handle_get(n as usize))
            }
            BINGET => {
                try!(self.handle_get(try!(rd.read_u8()) as usize))
            }
            LONG_BINGET => {
                let n = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(n);
                try!(self.handle_get(n as usize))
            }

            PUT => {
                let n = match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Int(n) => n,
                    BooleanOrInt::Boolean(false) => 0,
                    BooleanOrInt::Boolean(true) => 1,
                };
                ensure_not_negative!(n);
                try!(self.handle_put(n as usize))
            }
            BINPUT => {
                try!(self.handle_put(try!(rd.read_u8()) as usize))
            }
            LONG_BINPUT => {
                let n = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(n);
                try!(self.handle_put(n as usize))
            }

            c => return Err(Error::UnknownOpcode(c)),
        }
        Ok(false)
    }
}

pub fn unpickle<R>(rd: &mut R) -> Result<Value, Error> where R: Read + BufRead {
    let mut machine = Machine::new();
    loop {
        if try!(machine.execute(rd)) {
            break
        }
    }
    Ok(try!(machine.pop()))
}

#[cfg(test)]
mod tests {
    use std::io::{Cursor};

    use num::{FromPrimitive};

    use super::{unpickle};
    use super::super::value::{Value};

    macro_rules! t {
        ($buffer: expr, $pat:pat, $result:expr) => ({
            match unpickle(&mut Cursor::new(&$buffer[..])) {
                Ok($pat) => $result,
                other => {
                    println!("ERROR {:?}", other);
                    assert!(false)
                },
            }
        })
    }

    macro_rules! e {
        ($buffer: expr, $pat:pat) => ({
            match unpickle(&mut Cursor::new(&$buffer[..])) {
                Err($pat) => (),
                other => {
                    println!("ERROR {:?}", other);
                    assert!(false)
                },
            }
        })
    }

    macro_rules! n {
        ($x: expr) => ({FromPrimitive::from_isize($x).unwrap()})
    }

    #[test]
    fn test_int() {
        t!(b"I1\n.", Value::Long(n), assert_eq!(n, n!(1)));
        t!(b"K\x01.", Value::Int(n), assert_eq!(n, 1));
        t!(b"\x80\x02K\x01.", Value::Int(n), assert_eq!(n, 1));
    }
}<|endoftext|>


<commit_before>extern crate cranelift_codegen;
extern crate cranelift_entity;
extern crate cranelift_frontend;
extern crate cranelift_module;
extern crate cranelift_simplejit;

use cranelift_codegen::ir::*;
use cranelift_codegen::settings::*;
use cranelift_codegen::Context;
use cranelift_entity::EntityRef;
use cranelift_frontend::*;
use cranelift_module::*;
use cranelift_simplejit::*;

#[test]
fn error_on_incompatible_sig_in_declare_function() {
    let mut module: Module<SimpleJITBackend> = Module::new(SimpleJITBuilder::new());
    let mut sig = Signature {
        params: vec![AbiParam::new(types::I64)],
        returns: vec![],
        call_conv: CallConv::SystemV,
        argument_bytes: None,
    };
    module
        .declare_function("abc", Linkage::Local, &sig)
        .unwrap();
    sig.params[0] = AbiParam::new(types::I32);
    module
        .declare_function("abc", Linkage::Local, &sig)
        .err()
        .unwrap(); // Make sure this is an error
}

fn define_simple_function(module: &mut Module<SimpleJITBackend>) -> FuncId {
    let sig = Signature {
        params: vec![],
        returns: vec![],
        call_conv: CallConv::SystemV,
        argument_bytes: None,
    };

    let func_id = module
        .declare_function("abc", Linkage::Local, &sig)
        .unwrap();

    let mut ctx = Context::new();
    ctx.func = Function::with_name_signature(ExternalName::user(0, func_id.index() as u32), sig);
    let mut func_ctx = FunctionBuilderContext::new();
    {
        let mut bcx: FunctionBuilder<Variable> = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
        let ebb = bcx.create_ebb();
        bcx.switch_to_block(ebb);
        bcx.ins().return_(&[]);
    }

    module.define_function(func_id, &mut ctx).unwrap();

    func_id
}

#[test]
#[should_panic(expected = "function can't be finalized twice")]
fn panic_on_double_finalize() {
    let mut module: Module<SimpleJITBackend> = Module::new(SimpleJITBuilder::new());

    let func_id = define_simple_function(&mut module);
    module.finalize_function(func_id);
    module.finalize_function(func_id);
}

#[test]
#[should_panic(expected = "Result::unwrap()` on an `Err` value: DuplicateDefinition(\"abc\")")]
fn panic_on_define_after_finalize() {
    let mut module: Module<SimpleJITBackend> = Module::new(SimpleJITBuilder::new());

    let func_id = define_simple_function(&mut module);
    module.finalize_function(func_id);
    define_simple_function(&mut module);
}
<commit_msg>Update a testcase for API changes.<commit_after>extern crate cranelift_codegen;
extern crate cranelift_entity;
extern crate cranelift_frontend;
extern crate cranelift_module;
extern crate cranelift_simplejit;

use cranelift_codegen::ir::*;
use cranelift_codegen::settings::*;
use cranelift_codegen::Context;
use cranelift_entity::EntityRef;
use cranelift_frontend::*;
use cranelift_module::*;
use cranelift_simplejit::*;

#[test]
fn error_on_incompatible_sig_in_declare_function() {
    let mut module: Module<SimpleJITBackend> = Module::new(SimpleJITBuilder::new());
    let mut sig = Signature {
        params: vec![AbiParam::new(types::I64)],
        returns: vec![],
        call_conv: CallConv::SystemV,
        argument_bytes: None,
    };
    module
        .declare_function("abc", Linkage::Local, &sig)
        .unwrap();
    sig.params[0] = AbiParam::new(types::I32);
    module
        .declare_function("abc", Linkage::Local, &sig)
        .err()
        .unwrap(); // Make sure this is an error
}

fn define_simple_function(module: &mut Module<SimpleJITBackend>) -> FuncId {
    let sig = Signature {
        params: vec![],
        returns: vec![],
        call_conv: CallConv::SystemV,
        argument_bytes: None,
    };

    let func_id = module
        .declare_function("abc", Linkage::Local, &sig)
        .unwrap();

    let mut ctx = Context::new();
    ctx.func = Function::with_name_signature(ExternalName::user(0, func_id.index() as u32), sig);
    let mut func_ctx = FunctionBuilderContext::new();
    {
        let mut bcx: FunctionBuilder = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
        let ebb = bcx.create_ebb();
        bcx.switch_to_block(ebb);
        bcx.ins().return_(&[]);
    }

    module.define_function(func_id, &mut ctx).unwrap();

    func_id
}

#[test]
#[should_panic(expected = "function can't be finalized twice")]
fn panic_on_double_finalize() {
    let mut module: Module<SimpleJITBackend> = Module::new(SimpleJITBuilder::new());

    let func_id = define_simple_function(&mut module);
    module.finalize_function(func_id);
    module.finalize_function(func_id);
}

#[test]
#[should_panic(expected = "Result::unwrap()` on an `Err` value: DuplicateDefinition(\"abc\")")]
fn panic_on_define_after_finalize() {
    let mut module: Module<SimpleJITBackend> = Module::new(SimpleJITBuilder::new());

    let func_id = define_simple_function(&mut module);
    module.finalize_function(func_id);
    define_simple_function(&mut module);
}
<|endoftext|>


<commit_before>//! Clients for high level interactions with TUF repositories.

use std::collections::{HashSet, VecDeque};

use Result;
use error::Error;
use interchange::DataInterchange;
use metadata::{MetadataVersion, RootMetadata, Role, MetadataPath, TargetPath};
use repository::Repository;
use tuf::Tuf;

/// A client that interacts with TUF repositories.
pub struct Client<D, L, R>
where
    D: DataInterchange,
    L: Repository<D>,
    R: Repository<D>,
{
    tuf: Tuf<D>,
    config: Config,
    local: L,
    remote: R,
}

impl<D, L, R> Client<D, L, R>
where
    D: DataInterchange,
    L: Repository<D>,
    R: Repository<D>,
{
    /// Create a new TUF client from the given `Tuf` (metadata storage) and local and remote
    /// repositories.
    pub fn new(tuf: Tuf<D>, config: Config, mut local: L, mut remote: R) -> Result<Self> {
        local.initialize()?;
        remote.initialize()?;

        Ok(Client {
            tuf: tuf,
            config: config,
            local: local,
            remote: remote,
        })
    }

    /// Update TUF metadata from the local repository.
    ///
    /// Returns `true` if an update occurred and `false` otherwise.
    pub fn update_local(&mut self) -> Result<bool> {
        let r = Self::update_root(
            &mut self.tuf,
            &mut self.local,
            &self.config.max_root_size,
            self.config.min_bytes_per_second,
        )?;
        let ts = match Self::update_timestamp(
            &mut self.tuf,
            &mut self.local,
            &self.config.max_timestamp_size,
            self.config.min_bytes_per_second,
        ) {
            Ok(b) => b,
            Err(e) => {
                warn!(
                    "Error updating timestamp metadata from local sources: {:?}",
                    e
                );
                false
            }
        };
        let sn = match Self::update_snapshot(
            &mut self.tuf,
            &mut self.local,
            self.config.min_bytes_per_second,
        ) {
            Ok(b) => b,
            Err(e) => {
                warn!(
                    "Error updating snapshot metadata from local sources: {:?}",
                    e
                );
                false
            }
        };
        let ta = match Self::update_targets(
            &mut self.tuf,
            &mut self.local,
            self.config.min_bytes_per_second,
        ) {
            Ok(b) => b,
            Err(e) => {
                warn!(
                    "Error updating targets metadata from local sources: {:?}",
                    e
                );
                false
            }
        };

        let de = match Self::update_delegations(
            &mut self.tuf,
            &mut self.local,
            self.config.min_bytes_per_second,
        ) {
            Ok(b) => b,
            Err(e) => {
                warn!(
                    "Error updating delegation metadata from local sources: {:?}",
                    e
                );
                // TODO this might be untrue because of a partial update
                false
            }
        };

        Ok(r || ts || sn || ta || de)
    }

    /// Update TUF metadata from the remote repository.
    ///
    /// Returns `true` if an update occurred and `false` otherwise.
    pub fn update_remote(&mut self) -> Result<bool> {
        let r = Self::update_root(
            &mut self.tuf,
            &mut self.remote,
            &self.config.max_root_size,
            self.config.min_bytes_per_second,
        )?;
        let ts = Self::update_timestamp(
            &mut self.tuf,
            &mut self.remote,
            &self.config.max_timestamp_size,
            self.config.min_bytes_per_second,
        )?;
        let sn = Self::update_snapshot(
            &mut self.tuf,
            &mut self.remote,
            self.config.min_bytes_per_second,
        )?;
        let ta = Self::update_targets(
            &mut self.tuf,
            &mut self.remote,
            self.config.min_bytes_per_second,
        )?;
        let de = Self::update_delegations(
            &mut self.tuf,
            &mut self.remote,
            self.config.min_bytes_per_second,
        )?;

        Ok(r || ts || sn || ta || de)
    }

    /// Returns `true` if an update occurred and `false` otherwise.
    fn update_root<T>(
        tuf: &mut Tuf<D>,
        repo: &mut T,
        max_root_size: &Option<usize>,
        min_bytes_per_second: u32,
    ) -> Result<bool>
    where
        T: Repository<D>,
    {
        let latest_root = repo.fetch_metadata(
            &Role::Root,
            &MetadataPath::from_role(&Role::Root),
            &MetadataVersion::None,
            max_root_size,
            min_bytes_per_second,
            None,
        )?;
        let latest_version = D::deserialize::<RootMetadata>(latest_root.signed())?
            .version();

        if latest_version < tuf.root().version() {
            return Err(Error::VerificationFailure(format!(
                "Latest root version is lower than current root version: {} < {}",
                latest_version,
                tuf.root().version()
            )));
        } else if latest_version == tuf.root().version() {
            return Ok(false);
        }

        let err_msg = "TUF claimed no update occurred when one should have. \
                       This is a programming error. Please report this as a bug.";

        for i in (tuf.root().version() + 1)..latest_version {
            let signed = repo.fetch_metadata(
                &Role::Root,
                &MetadataPath::from_role(&Role::Root),
                &MetadataVersion::Number(i),
                max_root_size,
                min_bytes_per_second,
                None,
            )?;
            if !tuf.update_root(signed)? {
                error!("{}", err_msg);
                return Err(Error::Programming(err_msg.into()));
            }
        }

        if !tuf.update_root(latest_root)? {
            error!("{}", err_msg);
            return Err(Error::Programming(err_msg.into()));
        }
        Ok(true)
    }

    /// Returns `true` if an update occurred and `false` otherwise.
    fn update_timestamp<T>(
        tuf: &mut Tuf<D>,
        repo: &mut T,
        max_timestamp_size: &Option<usize>,
        min_bytes_per_second: u32,
    ) -> Result<bool>
    where
        T: Repository<D>,
    {
        let ts = repo.fetch_metadata(
            &Role::Timestamp,
            &MetadataPath::from_role(&Role::Timestamp),
            &MetadataVersion::None,
            max_timestamp_size,
            min_bytes_per_second,
            None,
        )?;
        tuf.update_timestamp(ts)
    }

    /// Returns `true` if an update occurred and `false` otherwise.
    fn update_snapshot<T>(tuf: &mut Tuf<D>, repo: &mut T, min_bytes_per_second: u32) -> Result<bool>
    where
        T: Repository<D>,
    {
        let snapshot_description = match tuf.timestamp() {
            Some(ts) => {
                match ts.meta().get(&MetadataPath::from_role(&Role::Snapshot)) {
                    Some(d) => Ok(d),
                    None => Err(Error::VerificationFailure(
                        "Timestamp metadata did not contain a description of the \
                                current snapshot metadata."
                            .into(),
                    )),
                }
            }
            None => Err(Error::MissingMetadata(Role::Timestamp)),
        }?
            .clone();

        if snapshot_description.version() <= tuf.snapshot().map(|s| s.version()).unwrap_or(0) {
            return Ok(false);
        }

        let snap = repo.fetch_metadata(
            &Role::Snapshot,
            &MetadataPath::from_role(&Role::Snapshot),
            &MetadataVersion::None,
            &None,
            min_bytes_per_second,
            None,
        )?;
        tuf.update_snapshot(snap)
    }

    /// Returns `true` if an update occurred and `false` otherwise.
    fn update_targets<T>(tuf: &mut Tuf<D>, repo: &mut T, min_bytes_per_second: u32) -> Result<bool>
    where
        T: Repository<D>,
    {
        let targets_description = match tuf.snapshot() {
            Some(sn) => {
                match sn.meta().get(&MetadataPath::from_role(&Role::Targets)) {
                    Some(d) => Ok(d),
                    None => Err(Error::VerificationFailure(
                        "Snapshot metadata did not contain a description of the \
                                current targets metadata."
                            .into(),
                    )),
                }
            }
            None => Err(Error::MissingMetadata(Role::Snapshot)),
        }?
            .clone();

        if targets_description.version() <= tuf.targets().map(|t| t.version()).unwrap_or(0) {
            return Ok(false);
        }

        let targets = repo.fetch_metadata(
            &Role::Targets,
            &MetadataPath::from_role(&Role::Targets),
            &MetadataVersion::None,
            &None,
            min_bytes_per_second,
            None,
        )?;
        tuf.update_targets(targets)
    }

    /// Returns `true` if an update occurred and `false` otherwise.
    fn update_delegations<T>(
        tuf: &mut Tuf<D>,
        repo: &mut T,
        min_bytes_per_second: u32,
    ) -> Result<bool>
    where
        T: Repository<D>,
    {
        let _ = match tuf.snapshot() {
            Some(s) => s,
            None => return Err(Error::MissingMetadata(Role::Snapshot)),
        }.clone();
        let targets = match tuf.targets() {
            Some(t) => t,
            None => return Err(Error::MissingMetadata(Role::Targets)),
        }.clone();
        let delegations = match targets.delegations() {
            Some(d) => d,
            None => return Ok(false),
        }.clone();

        let mut visited = HashSet::new();
        let mut to_visit = VecDeque::new();

        for role in delegations.roles().iter().map(|r| r.role()) {
            let _ = to_visit.push_back(role.clone());
        }

        let mut updated = false;
        while let Some(role) = to_visit.pop_front() {
            if visited.contains(&role) {
                continue;
            }
            let _ = visited.insert(role.clone());

            let delegation = match repo.fetch_metadata(
                &Role::Targets,
                &role,
                &MetadataVersion::None,
                &None,
                min_bytes_per_second,
                None,
            ) {
                Ok(d) => d,
                Err(e) => {
                    warn!("Failed to fetuch delegation {:?}: {:?}", role, e);
                    continue;
                }
            };

            match tuf.update_delegation(&role, delegation) {
                Ok(u) => updated |= u,
                Err(e) => {
                    warn!("Failed to update delegation {:?}: {:?}", role, e);
                    continue;
                }
            };

            if let Some(ds) = tuf.delegations().get(&role).and_then(|t| t.delegations()) {
                for d in ds.roles() {
                    let _ = to_visit.push_back(d.role().clone());
                }
            }
        }

        Ok(updated)
    }

    /// Fetch a target from the remote repo and write it to the local repo.
    pub fn fetch_target(&mut self, target: &TargetPath) -> Result<()> {
        let target_description = self.tuf.target_description(target)?;
        let read = self.remote.fetch_target(
            target,
            &target_description,
            self.config.min_bytes_per_second,
        )?;
        self.local.store_target(read, target)
    }
}

/// Configuration for a TUF `Client`.
#[derive(Debug)]
pub struct Config {
    max_root_size: Option<usize>,
    max_timestamp_size: Option<usize>,
    min_bytes_per_second: u32,
}

impl Config {
    /// Initialize a `ConfigBuilder` with the default values.
    pub fn build() -> ConfigBuilder {
        ConfigBuilder::default()
    }

    /// Return the optional maximum root metadata size.
    pub fn max_root_size(&self) -> &Option<usize> {
        &self.max_root_size
    }

    /// Return the optional maximum timestamp metadata size.
    pub fn max_timestamp_size(&self) -> &Option<usize> {
        &self.max_timestamp_size
    }

    /// The minimum bytes per second for a read to be considered good.
    pub fn min_bytes_per_second(&self) -> u32 {
        self.min_bytes_per_second
    }
}

/// Helper for building and validating a TUF `Config`.
#[derive(Debug, PartialEq)]
pub struct ConfigBuilder {
    max_root_size: Option<usize>,
    max_timestamp_size: Option<usize>,
    min_bytes_per_second: u32,
}

impl ConfigBuilder {
    /// Validate this builder return a `Config` if validation succeeds.
    pub fn finish(self) -> Result<Config> {
        Ok(Config {
            max_root_size: self.max_root_size,
            max_timestamp_size: self.max_timestamp_size,
            min_bytes_per_second: self.min_bytes_per_second,
        })
    }

    /// Set the optional maximum download size for root metadata.
    pub fn max_root_size(mut self, max: Option<usize>) -> Self {
        self.max_root_size = max;
        self
    }

    /// Set the optional maximum download size for timestamp metadata.
    pub fn max_timestamp_size(mut self, max: Option<usize>) -> Self {
        self.max_timestamp_size = max;
        self
    }

    /// Set the minimum bytes per second for a read to be considered good.
    pub fn min_bytes_per_second(mut self, min: u32) -> Self {
        self.min_bytes_per_second = min;
        self
    }
}

impl Default for ConfigBuilder {
    /// ```
    /// use tuf::client::ConfigBuilder;
    ///
    /// let default = ConfigBuilder::default();
    /// let config = ConfigBuilder::default()
    ///     .max_root_size(Some(1024 * 1024))
    ///     .max_timestamp_size(Some(32 * 1024))
    ///     .min_bytes_per_second(4096);
    /// assert_eq!(config, default);
    /// assert!(default.finish().is_ok())
    /// ```
    fn default() -> Self {
        ConfigBuilder {
            max_root_size: Some(1024 * 1024),
            max_timestamp_size: Some(32 * 1024),
            min_bytes_per_second: 4096,
        }
    }
}
<commit_msg>remove delegation-auto-updater<commit_after>//! Clients for high level interactions with TUF repositories.

use Result;
use error::Error;
use interchange::DataInterchange;
use metadata::{MetadataVersion, RootMetadata, Role, MetadataPath, TargetPath};
use repository::Repository;
use tuf::Tuf;

/// A client that interacts with TUF repositories.
pub struct Client<D, L, R>
where
    D: DataInterchange,
    L: Repository<D>,
    R: Repository<D>,
{
    tuf: Tuf<D>,
    config: Config,
    local: L,
    remote: R,
}

impl<D, L, R> Client<D, L, R>
where
    D: DataInterchange,
    L: Repository<D>,
    R: Repository<D>,
{
    /// Create a new TUF client from the given `Tuf` (metadata storage) and local and remote
    /// repositories.
    pub fn new(tuf: Tuf<D>, config: Config, mut local: L, mut remote: R) -> Result<Self> {
        local.initialize()?;
        remote.initialize()?;

        Ok(Client {
            tuf: tuf,
            config: config,
            local: local,
            remote: remote,
        })
    }

    /// Update TUF metadata from the local repository.
    ///
    /// Returns `true` if an update occurred and `false` otherwise.
    pub fn update_local(&mut self) -> Result<bool> {
        let r = Self::update_root(
            &mut self.tuf,
            &mut self.local,
            &self.config.max_root_size,
            self.config.min_bytes_per_second,
        )?;
        let ts = match Self::update_timestamp(
            &mut self.tuf,
            &mut self.local,
            &self.config.max_timestamp_size,
            self.config.min_bytes_per_second,
        ) {
            Ok(b) => b,
            Err(e) => {
                warn!(
                    "Error updating timestamp metadata from local sources: {:?}",
                    e
                );
                false
            }
        };
        let sn = match Self::update_snapshot(
            &mut self.tuf,
            &mut self.local,
            self.config.min_bytes_per_second,
        ) {
            Ok(b) => b,
            Err(e) => {
                warn!(
                    "Error updating snapshot metadata from local sources: {:?}",
                    e
                );
                false
            }
        };
        let ta = match Self::update_targets(
            &mut self.tuf,
            &mut self.local,
            self.config.min_bytes_per_second,
        ) {
            Ok(b) => b,
            Err(e) => {
                warn!(
                    "Error updating targets metadata from local sources: {:?}",
                    e
                );
                false
            }
        };

        Ok(r || ts || sn || ta)
    }

    /// Update TUF metadata from the remote repository.
    ///
    /// Returns `true` if an update occurred and `false` otherwise.
    pub fn update_remote(&mut self) -> Result<bool> {
        let r = Self::update_root(
            &mut self.tuf,
            &mut self.remote,
            &self.config.max_root_size,
            self.config.min_bytes_per_second,
        )?;
        let ts = Self::update_timestamp(
            &mut self.tuf,
            &mut self.remote,
            &self.config.max_timestamp_size,
            self.config.min_bytes_per_second,
        )?;
        let sn = Self::update_snapshot(
            &mut self.tuf,
            &mut self.remote,
            self.config.min_bytes_per_second,
        )?;
        let ta = Self::update_targets(
            &mut self.tuf,
            &mut self.remote,
            self.config.min_bytes_per_second,
        )?;

        Ok(r || ts || sn || ta)
    }

    /// Returns `true` if an update occurred and `false` otherwise.
    fn update_root<T>(
        tuf: &mut Tuf<D>,
        repo: &mut T,
        max_root_size: &Option<usize>,
        min_bytes_per_second: u32,
    ) -> Result<bool>
    where
        T: Repository<D>,
    {
        let latest_root = repo.fetch_metadata(
            &Role::Root,
            &MetadataPath::from_role(&Role::Root),
            &MetadataVersion::None,
            max_root_size,
            min_bytes_per_second,
            None,
        )?;
        let latest_version = D::deserialize::<RootMetadata>(latest_root.signed())?
            .version();

        if latest_version < tuf.root().version() {
            return Err(Error::VerificationFailure(format!(
                "Latest root version is lower than current root version: {} < {}",
                latest_version,
                tuf.root().version()
            )));
        } else if latest_version == tuf.root().version() {
            return Ok(false);
        }

        let err_msg = "TUF claimed no update occurred when one should have. \
                       This is a programming error. Please report this as a bug.";

        for i in (tuf.root().version() + 1)..latest_version {
            let signed = repo.fetch_metadata(
                &Role::Root,
                &MetadataPath::from_role(&Role::Root),
                &MetadataVersion::Number(i),
                max_root_size,
                min_bytes_per_second,
                None,
            )?;
            if !tuf.update_root(signed)? {
                error!("{}", err_msg);
                return Err(Error::Programming(err_msg.into()));
            }
        }

        if !tuf.update_root(latest_root)? {
            error!("{}", err_msg);
            return Err(Error::Programming(err_msg.into()));
        }
        Ok(true)
    }

    /// Returns `true` if an update occurred and `false` otherwise.
    fn update_timestamp<T>(
        tuf: &mut Tuf<D>,
        repo: &mut T,
        max_timestamp_size: &Option<usize>,
        min_bytes_per_second: u32,
    ) -> Result<bool>
    where
        T: Repository<D>,
    {
        let ts = repo.fetch_metadata(
            &Role::Timestamp,
            &MetadataPath::from_role(&Role::Timestamp),
            &MetadataVersion::None,
            max_timestamp_size,
            min_bytes_per_second,
            None,
        )?;
        tuf.update_timestamp(ts)
    }

    /// Returns `true` if an update occurred and `false` otherwise.
    fn update_snapshot<T>(tuf: &mut Tuf<D>, repo: &mut T, min_bytes_per_second: u32) -> Result<bool>
    where
        T: Repository<D>,
    {
        let snapshot_description = match tuf.timestamp() {
            Some(ts) => {
                match ts.meta().get(&MetadataPath::from_role(&Role::Snapshot)) {
                    Some(d) => Ok(d),
                    None => Err(Error::VerificationFailure(
                        "Timestamp metadata did not contain a description of the \
                                current snapshot metadata."
                            .into(),
                    )),
                }
            }
            None => Err(Error::MissingMetadata(Role::Timestamp)),
        }?
            .clone();

        if snapshot_description.version() <= tuf.snapshot().map(|s| s.version()).unwrap_or(0) {
            return Ok(false);
        }

        let snap = repo.fetch_metadata(
            &Role::Snapshot,
            &MetadataPath::from_role(&Role::Snapshot),
            &MetadataVersion::None,
            &None,
            min_bytes_per_second,
            None,
        )?;
        tuf.update_snapshot(snap)
    }

    /// Returns `true` if an update occurred and `false` otherwise.
    fn update_targets<T>(tuf: &mut Tuf<D>, repo: &mut T, min_bytes_per_second: u32) -> Result<bool>
    where
        T: Repository<D>,
    {
        let targets_description = match tuf.snapshot() {
            Some(sn) => {
                match sn.meta().get(&MetadataPath::from_role(&Role::Targets)) {
                    Some(d) => Ok(d),
                    None => Err(Error::VerificationFailure(
                        "Snapshot metadata did not contain a description of the \
                                current targets metadata."
                            .into(),
                    )),
                }
            }
            None => Err(Error::MissingMetadata(Role::Snapshot)),
        }?
            .clone();

        if targets_description.version() <= tuf.targets().map(|t| t.version()).unwrap_or(0) {
            return Ok(false);
        }

        let targets = repo.fetch_metadata(
            &Role::Targets,
            &MetadataPath::from_role(&Role::Targets),
            &MetadataVersion::None,
            &None,
            min_bytes_per_second,
            None,
        )?;
        tuf.update_targets(targets)
    }

    /// Fetch a target from the remote repo and write it to the local repo.
    pub fn fetch_target(&mut self, target: &TargetPath) -> Result<()> {
        let target_description = self.tuf.target_description(target)?;
        let read = self.remote.fetch_target(
            target,
            &target_description,
            self.config.min_bytes_per_second,
        )?;
        self.local.store_target(read, target)
    }
}

/// Configuration for a TUF `Client`.
#[derive(Debug)]
pub struct Config {
    max_root_size: Option<usize>,
    max_timestamp_size: Option<usize>,
    min_bytes_per_second: u32,
}

impl Config {
    /// Initialize a `ConfigBuilder` with the default values.
    pub fn build() -> ConfigBuilder {
        ConfigBuilder::default()
    }

    /// Return the optional maximum root metadata size.
    pub fn max_root_size(&self) -> &Option<usize> {
        &self.max_root_size
    }

    /// Return the optional maximum timestamp metadata size.
    pub fn max_timestamp_size(&self) -> &Option<usize> {
        &self.max_timestamp_size
    }

    /// The minimum bytes per second for a read to be considered good.
    pub fn min_bytes_per_second(&self) -> u32 {
        self.min_bytes_per_second
    }
}

/// Helper for building and validating a TUF `Config`.
#[derive(Debug, PartialEq)]
pub struct ConfigBuilder {
    max_root_size: Option<usize>,
    max_timestamp_size: Option<usize>,
    min_bytes_per_second: u32,
}

impl ConfigBuilder {
    /// Validate this builder return a `Config` if validation succeeds.
    pub fn finish(self) -> Result<Config> {
        Ok(Config {
            max_root_size: self.max_root_size,
            max_timestamp_size: self.max_timestamp_size,
            min_bytes_per_second: self.min_bytes_per_second,
        })
    }

    /// Set the optional maximum download size for root metadata.
    pub fn max_root_size(mut self, max: Option<usize>) -> Self {
        self.max_root_size = max;
        self
    }

    /// Set the optional maximum download size for timestamp metadata.
    pub fn max_timestamp_size(mut self, max: Option<usize>) -> Self {
        self.max_timestamp_size = max;
        self
    }

    /// Set the minimum bytes per second for a read to be considered good.
    pub fn min_bytes_per_second(mut self, min: u32) -> Self {
        self.min_bytes_per_second = min;
        self
    }
}

impl Default for ConfigBuilder {
    /// ```
    /// use tuf::client::ConfigBuilder;
    ///
    /// let default = ConfigBuilder::default();
    /// let config = ConfigBuilder::default()
    ///     .max_root_size(Some(1024 * 1024))
    ///     .max_timestamp_size(Some(32 * 1024))
    ///     .min_bytes_per_second(4096);
    /// assert_eq!(config, default);
    /// assert!(default.finish().is_ok())
    /// ```
    fn default() -> Self {
        ConfigBuilder {
            max_root_size: Some(1024 * 1024),
            max_timestamp_size: Some(32 * 1024),
            min_bytes_per_second: 4096,
        }
    }
}
<|endoftext|>


<commit_before>/// User facing module of this library.
///
/// Fetching message from multiple (topic, partition) pair or producing messages to multiple
/// topics is not yet supported.
/// It should be added very soon.

use error::{Result, Error};
use utils;
use protocol;
use connection::KafkaConnection;
use codecs::{ToByte, FromByte};
use std::collections::HashMap;
use std::io::Cursor;
use std::io::Read;

const CLIENTID: &'static str = "kafka-rust";
const DEFAULT_TIMEOUT: i32 = 120; // seconds

/// Client struct. It keeps track of brokers and topic metadata
///
/// # Examples
///
/// ```no_run
/// let mut client = kafka::client::KafkaClient::new(&vec!("localhost:9092".to_string()));
/// let res = client.load_metadata_all();
/// ```
///
/// You will have to load metadata before making any other request.
#[derive(Default, Debug)]
pub struct KafkaClient {
    clientid: String,
    timeout: i32,
    hosts: Vec<String>,
    correlation: i32,
    conns: HashMap<String, KafkaConnection>,
    pub topic_partitions: HashMap<String, Vec<i32>>,
    topic_brokers: HashMap<String, String>
}

impl KafkaClient {
    /// Create a new instance of KafkaClient
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut client = kafka::client::KafkaClient::new(&vec!("localhost:9092".to_string()));
    /// ```
    pub fn new(hosts: Vec<String>) -> KafkaClient {
        KafkaClient { hosts: hosts, clientid: CLIENTID.to_string(),
                      timeout: DEFAULT_TIMEOUT, ..KafkaClient::default()}
    }

    fn get_conn(& mut self, host: &str) -> Result<KafkaConnection> {
        match self.conns.get(host) {
            Some (conn) => return conn.clone(),
            None => {}
        }
        // TODO
        // Keeping this out here since get is causing ownership issues
        // Will refactor once I know better
        self.conns.insert(host.to_string(),
                          try!(KafkaConnection::new(host, self.timeout)));
        self.get_conn(host)
    }

    fn next_id(&mut self) -> i32{
        self.correlation = (self.correlation + 1) % (1i32 << 30);
        self.correlation
    }

    /// Resets and loads metadata for all topics.
    pub fn load_metadata_all(&mut self) -> Result<()>{
        self.reset_metadata();
        self.load_metadata(vec!())
    }

    /// Reloads metadata for a list of supplied topics
    ///
    /// returns Result<(), error::Error>
    pub fn load_metadata (&mut self, topics: Vec<String>) -> Result<()>{
        let resp = try!(self.get_metadata(topics));

        let mut brokers: HashMap<i32, String> = HashMap::new();
        for broker in resp.brokers {
            brokers.insert(broker.nodeid, format!("{}:{}", broker.host, broker.port));
        }

        self.topic_brokers.clear();
        for topic in resp.topics {
            self.topic_partitions.insert(topic.topic.clone(), vec!());

            for partition in topic.partitions {
                match brokers.get(&partition.leader) {
                    Some(broker) => {
                        self.topic_partitions.get_mut(&topic.topic).unwrap().push(partition.id);
                        self.topic_brokers.insert(
                            format!("{}-{}", topic.topic, partition.id),
                            broker.clone());
                    },
                    None => {}
                }
            }
        }
        Ok(())
    }

    /// Clears metadata stored in the client. You must load metadata after this call if you want
    /// to use the client
    pub fn reset_metadata(&mut self) {
        self.topic_partitions.clear();
        self.topic_brokers.clear();
    }

    fn get_metadata(&mut self, topics: Vec<String>) -> Result<protocol::MetadataResponse> {
        let correlation = self.next_id();
        for host in self.hosts.to_vec() {
            let req = protocol::MetadataRequest::new(correlation, self.clientid.clone(), topics.to_vec());
            match self.get_conn(&host) {
                Ok(mut conn) => if self.send_request(&mut conn, req).is_ok() {
                    return self.get_response::<protocol::MetadataResponse>(&mut conn);
                },
                Err(_) => {}
            }
        }

        Err(Error::NoHostReachable)
    }

    /// Fetch offsets for a list of topics
    /// Not implemented as yet.
    pub fn fetch_offsets(&mut self, _topics: &Vec<String>) {
        // TODO - Implement method to fetch offsets for more than 1 topic

    }

    /// Fetch offset for a topic.
    /// It gets the latest offset only. Support for getting earliest will be added soon
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut client = kafka::client::KafkaClient::new(&vec!("localhost:9092".to_string()));
    /// let res = client.load_metadata_all();
    /// let offsets = client.fetch_topic_offset(&"my-topic".to_string());
    /// ```
    /// Returns a vector of (topic, partition offset data).
    /// PartitionOffset will contain parition and offset info Or Error code as returned by Kafka.
    pub fn fetch_topic_offset(&mut self, topic: String) -> Result<Vec<(String, Vec<utils::PartitionOffset>)>> {
        // Doing it like this because HashMap will not return borrow of self otherwise
        let partitions = self.topic_partitions
                             .get(&topic)
                             .unwrap_or(&vec!())
                             .clone();

        // Maybe TODO: Can this be simplified without complexifying?
        let mut brokers: HashMap<String, Vec<i32>> = HashMap:: new();
        for p in partitions {
            let key = format!("{}-{}", topic, p);
            match self.topic_brokers.get(&key) {
                Some(broker) => {
                    if !brokers.contains_key(broker) {brokers.insert(broker.clone(), Vec::new());}
                    brokers.get_mut(broker).unwrap().push(p);
                },
                None => {}
            }
        }

        let mut res: Vec<utils::PartitionOffset> = vec!();
        for (host, partitions) in brokers.iter() {
            let v = vec!(utils::TopicPartitions{
                topic: topic.clone(),
                partitions: partitions.to_vec()
                });
            for tpo in try!(self.fetch_offset(v, host.clone())) {
                res.push(utils::PartitionOffset{partition: tpo.partition, offset: tpo.offset});
            }
        }
        Ok(vec!((topic.clone(), res)))
    }

    fn fetch_offset(&mut self, topic_partitions: Vec<utils::TopicPartitions>, host: String)
                            -> Result<Vec<utils::TopicPartitionOffset>> {
        let correlation = self.next_id();
        let req = protocol::OffsetRequest::new_latest(topic_partitions, correlation, self.clientid.clone());

        let resp = try!(self.send_receive::<protocol::OffsetRequest, protocol::OffsetResponse>(&host, req));

        Ok(resp.get_offsets())

    }

    fn get_broker(&mut self, topic: &String, partition: &i32) -> Option<String> {
        let key = format!("{}-{}", topic, partition);
        match self.topic_brokers.get(&key) {
            Some(broker) => {
                Some(broker.clone())
            },
            None => None
        }
    }

    /// Fetch messages from Kafka
    ///
    /// It takes a single topic, parition and offset and return a vector of messages
    /// or error::Error
    /// You can figure out the appropriate partition and offset using client's
    /// client.topic_partitions and client.fetch_topic_offset(topic)
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut client = kafka::client::KafkaClient::new(&vec!("localhost:9092".to_string()));
    /// let res = client.load_metadata_all();
    /// let msgs = client.fetch_messages(&"my-topic".to_string(), 0, 0);
    /// ```
    pub fn fetch_messages(&mut self, topic: String, partition: i32, offset: i64) -> Result<Vec<utils::OffsetMessage>>{

        let host = self.get_broker(&topic, &partition).unwrap();

        let correlation = self.next_id();
        let req = protocol::FetchRequest::new_single(topic, partition, offset, correlation, self.clientid.clone());

        let resp = try!(self.send_receive::<protocol::FetchRequest, protocol::FetchResponse>(&host, req));
        Ok(resp.get_messages()
                .iter()
                .filter(|ref x| x.offset >= offset)
                .cloned()
                .collect()
            )
    }

    /// Send a message to Kafka
    ///
    /// You can figure out the appropriate partition and offset using client's
    /// client.topic_partitions and client.fetch_topic_offset(topic)
    ///
    /// `required_acks` - indicates how many acknowledgements the servers should receive before
    /// responding to the request. If it is 0 the server will not send any response
    /// (this is the only case where the server will not reply to a request).
    /// If it is 1, the server will wait the data is written to the local log before sending
    /// a response. If it is -1 the server will block until the message is committed by all
    /// in sync replicas before sending a response. For any number > 1 the server will block
    /// waiting for this number of acknowledgements to occur (but the server will never wait
    /// for more acknowledgements than there are in-sync replicas).
    ///
    /// `timeout` - This provides a maximum time in milliseconds the server can await the
    /// receipt of the number of acknowledgements in `required_acks`
    /// `message` - A single message as a vector of u8s
    ///
    /// # Example
    ///
    /// ```no_run
    /// let mut client = kafka::client::KafkaClient::new(&vec!("localhost:9092".to_string()));
    /// let res = client.load_metadata_all();
    /// let msgs = client.send_message(&"my-topic".to_string(), 0, 1,
    ///                  100, &"b".to_string().into_bytes());
    /// ```
    /// The return value will contain topic, partition, offset and error if any
    /// OR error:Error
    pub fn send_message(&mut self, topic: String, partition: i32, required_acks: i16,
                      timeout: i32, message: Vec<u8>) -> Result<Vec<utils::TopicPartitionOffset>> {

        let host = self.get_broker(&topic, &partition).unwrap();

        let correlation = self.next_id();
        let req = protocol::ProduceRequest::new_single(topic, partition, required_acks,
                                                       timeout, message, correlation,
                                                       self.clientid.clone());

        let resp = try!(self.send_receive
            ::<protocol::ProduceRequest, protocol::ProduceResponse>(&host, req));
        Ok(resp.get_response())

    }

    pub fn commit_offset(&mut self, group: String, topic: String,
                         partition: i32, offset: i64) -> Result<()>{
        let host = self.get_broker(&topic, &partition).unwrap();

        let correlation = self.next_id();

        let req = protocol::OffsetCommitRequest::new(group, topic, partition, offset,
                                                    "".to_string(), correlation, self.clientid.clone());

        try!(self.send_receive
            ::<protocol::OffsetCommitRequest, protocol::OffsetCommitResponse>(&host, req));

        Ok(())
    }

    pub fn fetch_group_topic_offset(&mut self, group: String, topic: String, partition: i32) -> Result<i64> {
        let host = self.get_broker(&topic, &partition).unwrap();

        let correlation = self.next_id();
        let req = protocol::OffsetFetchRequest::new(group, vec!(topic.clone()),
                                                    vec!(partition), correlation, self.clientid.clone());

        let resp = try!(self.send_receive
                    ::<protocol::OffsetFetchRequest, protocol::OffsetFetchResponse>(&host, req));

        Ok(resp.get_offset_partition(topic, partition))

    }

    fn send_receive<T: ToByte, V: FromByte>(&mut self, host: &str, req: T) -> Result<V::R> {
        let mut conn = try!(self.get_conn(&host));
        try!(self.send_request(&mut conn, req));
        self.get_response::<V>(&mut conn)
    }

    fn send_request<T: ToByte>(&self, conn: &mut KafkaConnection, request: T) -> Result<usize>{
        let mut buffer = vec!();
        try!(request.encode(&mut buffer));

        let mut s = vec!();
        try!((buffer.len() as i32).encode(&mut s));
        for byte in buffer.iter() { s.push(*byte); }

        conn.send(&s)
    }

    fn get_response<T: FromByte>(&self, conn:&mut KafkaConnection) -> Result<T::R>{
        let mut v: Vec<u8> = vec!();
        let _ = conn.read(4, &mut v);

        let size = try!(i32::decode_new(&mut Cursor::new(v)));

        let mut resp: Vec<u8> = vec!();
        let _ = try!(conn.read(size as u64, &mut resp));

        T::decode_new(&mut Cursor::new(resp))
    }

}
<commit_msg>All tests pass<commit_after>/// User facing module of this library.
///
/// Fetching message from multiple (topic, partition) pair or producing messages to multiple
/// topics is not yet supported.
/// It should be added very soon.

use error::{Result, Error};
use utils;
use protocol;
use connection::KafkaConnection;
use codecs::{ToByte, FromByte};
use std::collections::HashMap;
use std::io::Cursor;
use std::io::Read;

const CLIENTID: &'static str = "kafka-rust";
const DEFAULT_TIMEOUT: i32 = 120; // seconds

/// Client struct. It keeps track of brokers and topic metadata
///
/// # Examples
///
/// ```no_run
/// let mut client = kafka::client::KafkaClient::new(vec!("localhost:9092".to_string()));
/// let res = client.load_metadata_all();
/// ```
///
/// You will have to load metadata before making any other request.
#[derive(Default, Debug)]
pub struct KafkaClient {
    clientid: String,
    timeout: i32,
    hosts: Vec<String>,
    correlation: i32,
    conns: HashMap<String, KafkaConnection>,
    pub topic_partitions: HashMap<String, Vec<i32>>,
    topic_brokers: HashMap<String, String>
}

impl KafkaClient {
    /// Create a new instance of KafkaClient
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut client = kafka::client::KafkaClient::new(vec!("localhost:9092".to_string()));
    /// ```
    pub fn new(hosts: Vec<String>) -> KafkaClient {
        KafkaClient { hosts: hosts, clientid: CLIENTID.to_string(),
                      timeout: DEFAULT_TIMEOUT, ..KafkaClient::default()}
    }

    fn get_conn(& mut self, host: &str) -> Result<KafkaConnection> {
        match self.conns.get(host) {
            Some (conn) => return conn.clone(),
            None => {}
        }
        // TODO
        // Keeping this out here since get is causing ownership issues
        // Will refactor once I know better
        self.conns.insert(host.to_string(),
                          try!(KafkaConnection::new(host, self.timeout)));
        self.get_conn(host)
    }

    fn next_id(&mut self) -> i32{
        self.correlation = (self.correlation + 1) % (1i32 << 30);
        self.correlation
    }

    /// Resets and loads metadata for all topics.
    pub fn load_metadata_all(&mut self) -> Result<()>{
        self.reset_metadata();
        self.load_metadata(vec!())
    }

    /// Reloads metadata for a list of supplied topics
    ///
    /// returns Result<(), error::Error>
    pub fn load_metadata (&mut self, topics: Vec<String>) -> Result<()>{
        let resp = try!(self.get_metadata(topics));

        let mut brokers: HashMap<i32, String> = HashMap::new();
        for broker in resp.brokers {
            brokers.insert(broker.nodeid, format!("{}:{}", broker.host, broker.port));
        }

        self.topic_brokers.clear();
        for topic in resp.topics {
            self.topic_partitions.insert(topic.topic.clone(), vec!());

            for partition in topic.partitions {
                match brokers.get(&partition.leader) {
                    Some(broker) => {
                        self.topic_partitions.get_mut(&topic.topic).unwrap().push(partition.id);
                        self.topic_brokers.insert(
                            format!("{}-{}", topic.topic, partition.id),
                            broker.clone());
                    },
                    None => {}
                }
            }
        }
        Ok(())
    }

    /// Clears metadata stored in the client. You must load metadata after this call if you want
    /// to use the client
    pub fn reset_metadata(&mut self) {
        self.topic_partitions.clear();
        self.topic_brokers.clear();
    }

    fn get_metadata(&mut self, topics: Vec<String>) -> Result<protocol::MetadataResponse> {
        let correlation = self.next_id();
        for host in self.hosts.to_vec() {
            let req = protocol::MetadataRequest::new(correlation, self.clientid.clone(), topics.to_vec());
            match self.get_conn(&host) {
                Ok(mut conn) => if self.send_request(&mut conn, req).is_ok() {
                    return self.get_response::<protocol::MetadataResponse>(&mut conn);
                },
                Err(_) => {}
            }
        }

        Err(Error::NoHostReachable)
    }

    /// Fetch offsets for a list of topics
    /// Not implemented as yet.
    pub fn fetch_offsets(&mut self, _topics: &Vec<String>) {
        // TODO - Implement method to fetch offsets for more than 1 topic

    }

    /// Fetch offset for a topic.
    /// It gets the latest offset only. Support for getting earliest will be added soon
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut client = kafka::client::KafkaClient::new(vec!("localhost:9092".to_string()));
    /// let res = client.load_metadata_all();
    /// let offsets = client.fetch_topic_offset("my-topic".to_string());
    /// ```
    /// Returns a vector of (topic, partition offset data).
    /// PartitionOffset will contain parition and offset info Or Error code as returned by Kafka.
    pub fn fetch_topic_offset(&mut self, topic: String) -> Result<Vec<(String, Vec<utils::PartitionOffset>)>> {
        // Doing it like this because HashMap will not return borrow of self otherwise
        let partitions = self.topic_partitions
                             .get(&topic)
                             .unwrap_or(&vec!())
                             .clone();

        // Maybe TODO: Can this be simplified without complexifying?
        let mut brokers: HashMap<String, Vec<i32>> = HashMap:: new();
        for p in partitions {
            let key = format!("{}-{}", topic, p);
            match self.topic_brokers.get(&key) {
                Some(broker) => {
                    if !brokers.contains_key(broker) {brokers.insert(broker.clone(), Vec::new());}
                    brokers.get_mut(broker).unwrap().push(p);
                },
                None => {}
            }
        }

        let mut res: Vec<utils::PartitionOffset> = vec!();
        for (host, partitions) in brokers.iter() {
            let v = vec!(utils::TopicPartitions{
                topic: topic.clone(),
                partitions: partitions.to_vec()
                });
            for tpo in try!(self.fetch_offset(v, host.clone())) {
                res.push(utils::PartitionOffset{partition: tpo.partition, offset: tpo.offset});
            }
        }
        Ok(vec!((topic.clone(), res)))
    }

    fn fetch_offset(&mut self, topic_partitions: Vec<utils::TopicPartitions>, host: String)
                            -> Result<Vec<utils::TopicPartitionOffset>> {
        let correlation = self.next_id();
        let req = protocol::OffsetRequest::new_latest(topic_partitions, correlation, self.clientid.clone());

        let resp = try!(self.send_receive::<protocol::OffsetRequest, protocol::OffsetResponse>(&host, req));

        Ok(resp.get_offsets())

    }

    fn get_broker(&mut self, topic: &String, partition: &i32) -> Option<String> {
        let key = format!("{}-{}", topic, partition);
        match self.topic_brokers.get(&key) {
            Some(broker) => {
                Some(broker.clone())
            },
            None => None
        }
    }

    /// Fetch messages from Kafka
    ///
    /// It takes a single topic, parition and offset and return a vector of messages
    /// or error::Error
    /// You can figure out the appropriate partition and offset using client's
    /// client.topic_partitions and client.fetch_topic_offset(topic)
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut client = kafka::client::KafkaClient::new(vec!("localhost:9092".to_string()));
    /// let res = client.load_metadata_all();
    /// let msgs = client.fetch_messages("my-topic".to_string(), 0, 0);
    /// ```
    pub fn fetch_messages(&mut self, topic: String, partition: i32, offset: i64) -> Result<Vec<utils::OffsetMessage>>{

        let host = self.get_broker(&topic, &partition).unwrap();

        let correlation = self.next_id();
        let req = protocol::FetchRequest::new_single(topic, partition, offset, correlation, self.clientid.clone());

        let resp = try!(self.send_receive::<protocol::FetchRequest, protocol::FetchResponse>(&host, req));
        Ok(resp.get_messages()
                .iter()
                .filter(|ref x| x.offset >= offset)
                .cloned()
                .collect()
            )
    }

    /// Send a message to Kafka
    ///
    /// You can figure out the appropriate partition and offset using client's
    /// client.topic_partitions and client.fetch_topic_offset(topic)
    ///
    /// `required_acks` - indicates how many acknowledgements the servers should receive before
    /// responding to the request. If it is 0 the server will not send any response
    /// (this is the only case where the server will not reply to a request).
    /// If it is 1, the server will wait the data is written to the local log before sending
    /// a response. If it is -1 the server will block until the message is committed by all
    /// in sync replicas before sending a response. For any number > 1 the server will block
    /// waiting for this number of acknowledgements to occur (but the server will never wait
    /// for more acknowledgements than there are in-sync replicas).
    ///
    /// `timeout` - This provides a maximum time in milliseconds the server can await the
    /// receipt of the number of acknowledgements in `required_acks`
    /// `message` - A single message as a vector of u8s
    ///
    /// # Example
    ///
    /// ```no_run
    /// let mut client = kafka::client::KafkaClient::new(vec!("localhost:9092".to_string()));
    /// let res = client.load_metadata_all();
    /// let msgs = client.send_message("my-topic".to_string(), 0, 1,
    ///                  100, "b".to_string().into_bytes());
    /// ```
    /// The return value will contain topic, partition, offset and error if any
    /// OR error:Error
    pub fn send_message(&mut self, topic: String, partition: i32, required_acks: i16,
                      timeout: i32, message: Vec<u8>) -> Result<Vec<utils::TopicPartitionOffset>> {

        let host = self.get_broker(&topic, &partition).unwrap();

        let correlation = self.next_id();
        let req = protocol::ProduceRequest::new_single(topic, partition, required_acks,
                                                       timeout, message, correlation,
                                                       self.clientid.clone());

        let resp = try!(self.send_receive
            ::<protocol::ProduceRequest, protocol::ProduceResponse>(&host, req));
        Ok(resp.get_response())

    }

    pub fn commit_offset(&mut self, group: String, topic: String,
                         partition: i32, offset: i64) -> Result<()>{
        let host = self.get_broker(&topic, &partition).unwrap();

        let correlation = self.next_id();

        let req = protocol::OffsetCommitRequest::new(group, topic, partition, offset,
                                                    "".to_string(), correlation, self.clientid.clone());

        try!(self.send_receive
            ::<protocol::OffsetCommitRequest, protocol::OffsetCommitResponse>(&host, req));

        Ok(())
    }

    pub fn fetch_group_topic_offset(&mut self, group: String, topic: String, partition: i32) -> Result<i64> {
        let host = self.get_broker(&topic, &partition).unwrap();

        let correlation = self.next_id();
        let req = protocol::OffsetFetchRequest::new(group, vec!(topic.clone()),
                                                    vec!(partition), correlation, self.clientid.clone());

        let resp = try!(self.send_receive
                    ::<protocol::OffsetFetchRequest, protocol::OffsetFetchResponse>(&host, req));

        Ok(resp.get_offset_partition(topic, partition))

    }

    fn send_receive<T: ToByte, V: FromByte>(&mut self, host: &str, req: T) -> Result<V::R> {
        let mut conn = try!(self.get_conn(&host));
        try!(self.send_request(&mut conn, req));
        self.get_response::<V>(&mut conn)
    }

    fn send_request<T: ToByte>(&self, conn: &mut KafkaConnection, request: T) -> Result<usize>{
        let mut buffer = vec!();
        try!(request.encode(&mut buffer));

        let mut s = vec!();
        try!((buffer.len() as i32).encode(&mut s));
        for byte in buffer.iter() { s.push(*byte); }

        conn.send(&s)
    }

    fn get_response<T: FromByte>(&self, conn:&mut KafkaConnection) -> Result<T::R>{
        let mut v: Vec<u8> = vec!();
        let _ = conn.read(4, &mut v);

        let size = try!(i32::decode_new(&mut Cursor::new(v)));

        let mut resp: Vec<u8> = vec!();
        let _ = try!(conn.read(size as u64, &mut resp));

        T::decode_new(&mut Cursor::new(resp))
    }

}
<|endoftext|>


<commit_before>use palette;
use palette::pixel::RgbPixel;

pub type Srgb = [f32; 3];

fn str_to_srgb(color: &str) -> [f32; 3] {
    let rgba = palette::named::from_str(color)
        .expect(&format!("unknown color: {}", color))
        .to_rgba();
    palette::pixel::RgbPixel::from_rgba(rgba.0, rgba.1, rgba.2, 1.0)
}

lazy_static! {
    pub static ref BLACK:  Srgb = str_to_srgb("black");
    pub static ref BROWN:  Srgb = str_to_srgb("saddlebrown");
    pub static ref BLUE:   Srgb = str_to_srgb("blue");
    pub static ref RED:    Srgb = str_to_srgb("red");
    pub static ref GREEN:  Srgb = str_to_srgb("green");
    pub static ref YELLOW: Srgb = str_to_srgb("yellow");
    pub static ref WHITE:  Srgb = str_to_srgb("white");
}
<commit_msg>fix(colors): web hex values ARE sRGB so we need to convert to linear RGB to get fragment shader to output correct color<commit_after>use palette;
use palette::pixel::RgbPixel;

pub type Srgb = [f32; 3]; // TODO rename rgb

fn normalize(color: [f32; 3]) -> [f32; 3] {
    [color[0] / 255.0, color[1] / 255.0, color[2] / 255.0]
}

fn srgb_to_rgb(color: [f32; 3]) -> [f32; 3] {
    let srgb = palette::pixel::Srgb::new(color[0], color[1], color[2]);
    let rgb = palette::Rgb::from(srgb);
    [rgb.red, rgb.green, rgb.blue]
}

fn str_to_rgb(color: &str) -> Srgb {
    let rgba = palette::named::from_str(color)
        .expect(&format!("unknown color: {}", color))
        .to_rgba();
    srgb_to_rgb([rgba.0, rgba.1, rgba.2])
}

lazy_static! {
    pub static ref BLACK:      Srgb = str_to_rgb("black");
    pub static ref CYAN:       Srgb = str_to_rgb("cyan");
    pub static ref BROWN:      Srgb = str_to_rgb("saddlebrown");
    pub static ref BLUE:       Srgb = str_to_rgb("blue");
    pub static ref RED:        Srgb = str_to_rgb("red");
    pub static ref GREEN:      Srgb = str_to_rgb("green");
    pub static ref YELLOW:     Srgb = str_to_rgb("yellow");
    pub static ref WHITE:      Srgb = str_to_rgb("white");

    pub static ref YOGA_TEAL:  Srgb = srgb_to_rgb(normalize([151.0, 220.0, 207.0]));
    pub static ref YOGA_GRAY:  Srgb = srgb_to_rgb(normalize([48.0,   56.0,  70.0]));
}
<|endoftext|>


<commit_before>use rustc_serialize::json::{Json, ToJson};
use rustc_serialize::{Encodable, Encoder};
use std::collections::BTreeMap;
use std::error::{Error, FromError};
use std::num::ToPrimitive;

use error::{WebDriverResult, WebDriverError, ErrorStatus};

pub static ELEMENT_KEY: &'static str = "element-6066-11e4-a52e-4f735466cecf";

#[derive(RustcEncodable, PartialEq, Debug)]
pub struct Date(u64);

impl Date {
    pub fn new(timestamp: u64) -> Date {
        Date(timestamp)
    }
}

impl ToJson for Date {
    fn to_json(&self) -> Json {
        let &Date(x) = self;
        x.to_json()
    }
}

#[derive(PartialEq, Clone, Debug)]
pub enum Nullable<T: ToJson> {
    Value(T),
    Null
}

impl<T: ToJson> Nullable<T> {
     pub fn is_null(&self) -> bool {
        match *self {
            Nullable::Value(_) => false,
            Nullable::Null => true
        }
    }

     pub fn is_value(&self) -> bool {
        match *self {
            Nullable::Value(_) => true,
            Nullable::Null => false
        }
    }
}

impl<T: ToJson> Nullable<T> {
    //This is not very pretty
    pub fn from_json<F: FnOnce(&Json) -> WebDriverResult<T>>(value: &Json, f: F) -> WebDriverResult<Nullable<T>> {
        if value.is_null() {
            Ok(Nullable::Null)
        } else {
            Ok(Nullable::Value(try!(f(value))))
        }
    }
}

impl<T: ToJson> ToJson for Nullable<T> {
    fn to_json(&self) -> Json {
        match *self {
            Nullable::Value(ref x) => x.to_json(),
            Nullable::Null => Json::Null
        }
    }
}

impl<T: ToJson> Encodable for Nullable<T> {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        match *self {
            Nullable::Value(ref x) => x.to_json().encode(s),
            Nullable::Null => s.emit_option_none()
        }
    }
}

#[derive(PartialEq)]
pub struct WebElement {
    pub id: String
}

impl WebElement {
    pub fn new(id: String) -> WebElement {
        WebElement {
            id: id
        }
    }

    pub fn from_json(data: &Json) -> WebDriverResult<WebElement> {
        let object = try_opt!(data.as_object(),
                              ErrorStatus::InvalidArgument,
                              "Could not convert webelement to object");
        let id_value = try_opt!(object.get(ELEMENT_KEY),
                                ErrorStatus::InvalidArgument,
                                "Could not find webelement key");

        let id = try_opt!(id_value.as_string(),
                          ErrorStatus::InvalidArgument,
                          "Could not convert web element to string").to_string();

        Ok(WebElement::new(id))
    }
}

impl ToJson for WebElement {
    fn to_json(&self) -> Json {
        let mut data = BTreeMap::new();
        data.insert(ELEMENT_KEY.to_string(), self.id.to_json());
        Json::Object(data)
    }
}

#[derive(PartialEq)]
pub enum FrameId {
    Short(u16),
    Element(WebElement),
    Null
}

impl FrameId {
    pub fn from_json(data: &Json) -> WebDriverResult<FrameId> {
        match data {
            &Json::U64(x) => {
                let id = try_opt!(x.to_u16(),
                                  ErrorStatus::NoSuchFrame,
                                  "frame id out of range");
                Ok(FrameId::Short(id))
            },
            &Json::Null => Ok(FrameId::Null),
            &Json::String(ref x) => Ok(FrameId::Element(WebElement::new(x.clone()))),
            _ => Err(WebDriverError::new(ErrorStatus::NoSuchFrame,
                                         "frame id has unexpected type"))
        }
    }
}

impl ToJson for FrameId {
    fn to_json(&self) -> Json {
        match *self {
            FrameId::Short(x) => {
                Json::U64(x as u64)
            },
            FrameId::Element(ref x) => {
                Json::String(x.id.clone())
            },
            FrameId::Null => {
                Json::Null
            }
        }
    }
}

#[derive(PartialEq)]
pub enum LocatorStrategy {
    CSSSelector,
    LinkText,
    PartialLinkText,
    XPath
}

impl LocatorStrategy {
    pub fn from_json(body: &Json) -> WebDriverResult<LocatorStrategy> {
        match try_opt!(body.as_string(),
                       ErrorStatus::InvalidArgument,
                       "Cound not convert strategy to string") {
            "css selector" => Ok(LocatorStrategy::CSSSelector),
            "link text" => Ok(LocatorStrategy::LinkText),
            "partial link text" => Ok(LocatorStrategy::PartialLinkText),
            "xpath" => Ok(LocatorStrategy::XPath),
            x => Err(WebDriverError::new(ErrorStatus::InvalidArgument,
                                         &format!("Unknown locator strategy {}", x)[..]))
        }
    }
}

impl ToJson for LocatorStrategy {
    fn to_json(&self) -> Json {
        Json::String(match *self {
            LocatorStrategy::CSSSelector => "css selector",
            LocatorStrategy::LinkText => "link text",
            LocatorStrategy::PartialLinkText => "partial link text",
            LocatorStrategy::XPath => "xpath"
        }.to_string())
    }
}
<commit_msg>Fix parsing of FrameIds with Elements<commit_after>use rustc_serialize::json::{Json, ToJson};
use rustc_serialize::{Encodable, Encoder};
use std::collections::BTreeMap;
use std::error::{Error, FromError};
use std::num::ToPrimitive;

use error::{WebDriverResult, WebDriverError, ErrorStatus};

pub static ELEMENT_KEY: &'static str = "element-6066-11e4-a52e-4f735466cecf";

#[derive(RustcEncodable, PartialEq, Debug)]
pub struct Date(u64);

impl Date {
    pub fn new(timestamp: u64) -> Date {
        Date(timestamp)
    }
}

impl ToJson for Date {
    fn to_json(&self) -> Json {
        let &Date(x) = self;
        x.to_json()
    }
}

#[derive(PartialEq, Clone, Debug)]
pub enum Nullable<T: ToJson> {
    Value(T),
    Null
}

impl<T: ToJson> Nullable<T> {
     pub fn is_null(&self) -> bool {
        match *self {
            Nullable::Value(_) => false,
            Nullable::Null => true
        }
    }

     pub fn is_value(&self) -> bool {
        match *self {
            Nullable::Value(_) => true,
            Nullable::Null => false
        }
    }
}

impl<T: ToJson> Nullable<T> {
    //This is not very pretty
    pub fn from_json<F: FnOnce(&Json) -> WebDriverResult<T>>(value: &Json, f: F) -> WebDriverResult<Nullable<T>> {
        if value.is_null() {
            Ok(Nullable::Null)
        } else {
            Ok(Nullable::Value(try!(f(value))))
        }
    }
}

impl<T: ToJson> ToJson for Nullable<T> {
    fn to_json(&self) -> Json {
        match *self {
            Nullable::Value(ref x) => x.to_json(),
            Nullable::Null => Json::Null
        }
    }
}

impl<T: ToJson> Encodable for Nullable<T> {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        match *self {
            Nullable::Value(ref x) => x.to_json().encode(s),
            Nullable::Null => s.emit_option_none()
        }
    }
}

#[derive(PartialEq)]
pub struct WebElement {
    pub id: String
}

impl WebElement {
    pub fn new(id: String) -> WebElement {
        WebElement {
            id: id
        }
    }

    pub fn from_json(data: &Json) -> WebDriverResult<WebElement> {
        let object = try_opt!(data.as_object(),
                              ErrorStatus::InvalidArgument,
                              "Could not convert webelement to object");
        let id_value = try_opt!(object.get(ELEMENT_KEY),
                                ErrorStatus::InvalidArgument,
                                "Could not find webelement key");

        let id = try_opt!(id_value.as_string(),
                          ErrorStatus::InvalidArgument,
                          "Could not convert web element to string").to_string();

        Ok(WebElement::new(id))
    }
}

impl ToJson for WebElement {
    fn to_json(&self) -> Json {
        let mut data = BTreeMap::new();
        data.insert(ELEMENT_KEY.to_string(), self.id.to_json());
        Json::Object(data)
    }
}

#[derive(PartialEq)]
pub enum FrameId {
    Short(u16),
    Element(WebElement),
    Null
}

impl FrameId {
    pub fn from_json(data: &Json) -> WebDriverResult<FrameId> {
        match data {
            &Json::U64(x) => {
                let id = try_opt!(x.to_u16(),
                                  ErrorStatus::NoSuchFrame,
                                  "frame id out of range");
                Ok(FrameId::Short(id))
            },
            &Json::Null => Ok(FrameId::Null),
            &Json::Object(_) => Ok(FrameId::Element(
                try!(WebElement::from_json(data)))),
            _ => Err(WebDriverError::new(ErrorStatus::NoSuchFrame,
                                         "frame id has unexpected type"))
        }
    }
}

impl ToJson for FrameId {
    fn to_json(&self) -> Json {
        match *self {
            FrameId::Short(x) => {
                Json::U64(x as u64)
            },
            FrameId::Element(ref x) => {
                Json::String(x.id.clone())
            },
            FrameId::Null => {
                Json::Null
            }
        }
    }
}

#[derive(PartialEq)]
pub enum LocatorStrategy {
    CSSSelector,
    LinkText,
    PartialLinkText,
    XPath
}

impl LocatorStrategy {
    pub fn from_json(body: &Json) -> WebDriverResult<LocatorStrategy> {
        match try_opt!(body.as_string(),
                       ErrorStatus::InvalidArgument,
                       "Cound not convert strategy to string") {
            "css selector" => Ok(LocatorStrategy::CSSSelector),
            "link text" => Ok(LocatorStrategy::LinkText),
            "partial link text" => Ok(LocatorStrategy::PartialLinkText),
            "xpath" => Ok(LocatorStrategy::XPath),
            x => Err(WebDriverError::new(ErrorStatus::InvalidArgument,
                                         &format!("Unknown locator strategy {}", x)[..]))
        }
    }
}

impl ToJson for LocatorStrategy {
    fn to_json(&self) -> Json {
        Json::String(match *self {
            LocatorStrategy::CSSSelector => "css selector",
            LocatorStrategy::LinkText => "link text",
            LocatorStrategy::PartialLinkText => "partial link text",
            LocatorStrategy::XPath => "xpath"
        }.to_string())
    }
}
<|endoftext|>