<commit_before>#![feature(libc)]

extern crate libc;
use std::ffi::CString;
use std::slice;
use std::ptr;

mod ffi;

#[derive(Debug, Copy, Clone)]
pub struct Env {
    env: ffi::Voidptr
}

#[derive(Debug, Copy, Clone)]
pub struct Db {
    db: ffi::Voidptr
}

// XXX: Use Arc
unsafe impl Sync for Db {}
unsafe impl Send for Db {}
unsafe impl Sync for Env {}
unsafe impl Send for Env {}

impl Env {
    pub fn new() -> Env {
        Env {env: unsafe{ffi::sp_env()}}
    }

    pub fn destroy(self) {
        unsafe{ffi::sp_destroy(self.env)};
    }

    pub fn db(&mut self, dbname: &str) {
        // XXX: is dbname copied?
        let dbname = CString::new(dbname).unwrap();
        unsafe {ffi::setstring(self.env, "db\0".as_bytes(), dbname.as_bytes(), 0) };
    }

    pub fn open(&mut self) {
        unsafe {ffi::sp_open(self.env)};
    }

    /*
    fn get_object(&mut self, attr: &str) -> Voidptr {
        let attr = CString::new(attr).unwrap();
        unsafe {ffi::sp_getobject(self.env, attr.as_ptr() as *const u8)}
    }
    */

    pub fn get_db(&self, dbname: &str) -> Option<Db> {
        let dbstr = "db.".to_string() + dbname;
        let attr = CString::new(&dbstr[..]).unwrap();
        unsafe {
            let obj = ffi::sp_getobject(self.env, attr.as_ptr() as *const u8);
            if obj.is_null() {
                None
            }
            else {
                Some(Db {db: obj})
            }
        }
    }

    pub fn setattr(&mut self, key: &str, val: &str) {
        let key = CString::new(key).unwrap();
        let val = CString::new(val).unwrap();
        unsafe {ffi::setstring(self.env, key.as_bytes(), val.as_bytes(), 0) };
    }

    pub fn setintattr(&mut self, key: &str, val: i64) {
        let key = CString::new(key).unwrap();
        unsafe {ffi::sp_setint(self.env, key.as_ptr(), val) };
    }
}

pub struct Object<'a> {
    obj: ffi::Voidptr,
    key: &'a [u8]
}

impl<'a> Object<'a> {
    pub fn get_value<'b>(&'b mut self) -> Option<&'b[u8]> {
        unsafe {
            let mut sz = 0;
            let valptr = ffi::sp_getstring(self.obj, "value\0".as_ptr(), &mut sz);

            // XXX: what if we stored an empty value?
            if valptr.is_null() {
                return None;
            }

            let s = slice::from_raw_parts(valptr as *const u8, sz as usize);
            return Some(s);
        }
    }
}

impl<'a> Drop for Object<'a> {
    fn drop(&mut self) {
        unsafe {ffi::sp_destroy(self.obj);}
    }
}

pub struct Cursor {
    obj: ffi::Voidptr,
}

pub struct CursorObject<'a> {
    obj: ffi::Voidptr,
    cur: &'a Cursor
}

impl<'a> CursorObject<'a> {
    pub fn get_key<'b>(&'b mut self) -> Option<&'b[u8]> {
        unsafe {
            let mut sz = 0;
            let keyptr = ffi::sp_getstring(self.obj, "key\0".as_ptr(), &mut sz);

            if keyptr.is_null() {
                return None;
            }

            let s = slice::from_raw_parts(keyptr as *const u8, sz as usize);
            return Some(s);
        }
    }

    pub fn get_value<'b>(&'b mut self) -> Option<&'b[u8]> {
        unsafe {
            let mut sz = 0;
            let valptr = ffi::sp_getstring(self.obj, "value\0".as_ptr(), &mut sz);

            // XXX: what if we stored an empty value?
            if valptr.is_null() {
                return None;
            }

            let s = slice::from_raw_parts(valptr as *const u8, sz as usize);
            return Some(s);
        }
    }
}

impl<'a> Drop for CursorObject<'a> {
    fn drop(&mut self) {
        unsafe {ffi::sp_destroy(self.obj);}
    }
}

impl Drop for Cursor {
    fn drop(&mut self) {
        unsafe {ffi::sp_destroy(self.obj);}
    }
}

impl Cursor {
    fn next<'a>(&'a self) -> Option<CursorObject<'a>> {
        let res = unsafe { ffi::sp_get(self.obj, ptr::null_mut()) };
        if res.is_null() {
            None
        }
        else {
            Some(CursorObject{obj: res, cur: &self})
        }
    }
}

impl Db {
    pub fn set(&mut self, key: &[u8], value: &[u8]) {
        unsafe {
            //ffi::sp_set_kv(self.db, key.as_ptr(), key.len() as i32, value.as_ptr(), value.len() as i32);
            let obj = ffi::sp_object(self.db);
            assert!(!obj.is_null());
            ffi::setkey(obj, key);
            ffi::setvalue(obj, value);
            ffi::sp_set(self.db, obj);
            ffi::sp_destroy(obj);
        }
    }

    pub fn iter_all(&mut self) -> Cursor {
        unsafe {
            let obj = ffi::sp_object(self.db);
            assert!(!obj.is_null());
            ffi::setstring(obj, "order\0".as_bytes(), ">=\0".as_bytes(), 0);
            let cursor = ffi::sp_cursor(self.db, obj);
            assert!(!cursor.is_null());
            Cursor {obj: cursor}
        }
    }

    pub fn get<'a>(&self, key: &'a [u8]) -> Option<Object<'a>> {
        assert!(key.len() > 0);
        unsafe {
            let pattern = ffi::sp_object(self.db);
            if pattern.is_null() {
                // XXX: Return Err
                return None;
            }
            ffi::setkey(pattern, key);
            let res = ffi::sp_get(self.db, pattern);
            if res.is_null() {
                ffi::sp_destroy(pattern);
                return None;
            }
            return Some(Object{obj: res, key: key});
        }
    }
}
<commit_msg>Make Cursor#next() public<commit_after>#![feature(libc)]

extern crate libc;
use std::ffi::CString;
use std::slice;
use std::ptr;

mod ffi;

#[derive(Debug, Copy, Clone)]
pub struct Env {
    env: ffi::Voidptr
}

#[derive(Debug, Copy, Clone)]
pub struct Db {
    db: ffi::Voidptr
}

// XXX: Use Arc
unsafe impl Sync for Db {}
unsafe impl Send for Db {}
unsafe impl Sync for Env {}
unsafe impl Send for Env {}

impl Env {
    pub fn new() -> Env {
        Env {env: unsafe{ffi::sp_env()}}
    }

    pub fn destroy(self) {
        unsafe{ffi::sp_destroy(self.env)};
    }

    pub fn db(&mut self, dbname: &str) {
        // XXX: is dbname copied?
        let dbname = CString::new(dbname).unwrap();
        unsafe {ffi::setstring(self.env, "db\0".as_bytes(), dbname.as_bytes(), 0) };
    }

    pub fn open(&mut self) {
        unsafe {ffi::sp_open(self.env)};
    }

    /*
    fn get_object(&mut self, attr: &str) -> Voidptr {
        let attr = CString::new(attr).unwrap();
        unsafe {ffi::sp_getobject(self.env, attr.as_ptr() as *const u8)}
    }
    */

    pub fn get_db(&self, dbname: &str) -> Option<Db> {
        let dbstr = "db.".to_string() + dbname;
        let attr = CString::new(&dbstr[..]).unwrap();
        unsafe {
            let obj = ffi::sp_getobject(self.env, attr.as_ptr() as *const u8);
            if obj.is_null() {
                None
            }
            else {
                Some(Db {db: obj})
            }
        }
    }

    pub fn setattr(&mut self, key: &str, val: &str) {
        let key = CString::new(key).unwrap();
        let val = CString::new(val).unwrap();
        unsafe {ffi::setstring(self.env, key.as_bytes(), val.as_bytes(), 0) };
    }

    pub fn setintattr(&mut self, key: &str, val: i64) {
        let key = CString::new(key).unwrap();
        unsafe {ffi::sp_setint(self.env, key.as_ptr(), val) };
    }
}

pub struct Object<'a> {
    obj: ffi::Voidptr,
    key: &'a [u8]
}

impl<'a> Object<'a> {
    pub fn get_value<'b>(&'b mut self) -> Option<&'b[u8]> {
        unsafe {
            let mut sz = 0;
            let valptr = ffi::sp_getstring(self.obj, "value\0".as_ptr(), &mut sz);

            // XXX: what if we stored an empty value?
            if valptr.is_null() {
                return None;
            }

            let s = slice::from_raw_parts(valptr as *const u8, sz as usize);
            return Some(s);
        }
    }
}

impl<'a> Drop for Object<'a> {
    fn drop(&mut self) {
        unsafe {ffi::sp_destroy(self.obj);}
    }
}

pub struct Cursor {
    obj: ffi::Voidptr,
}

pub struct CursorObject<'a> {
    obj: ffi::Voidptr,
    cur: &'a Cursor
}

impl<'a> CursorObject<'a> {
    pub fn get_key<'b>(&'b mut self) -> Option<&'b[u8]> {
        unsafe {
            let mut sz = 0;
            let keyptr = ffi::sp_getstring(self.obj, "key\0".as_ptr(), &mut sz);

            if keyptr.is_null() {
                return None;
            }

            let s = slice::from_raw_parts(keyptr as *const u8, sz as usize);
            return Some(s);
        }
    }

    pub fn get_value<'b>(&'b mut self) -> Option<&'b[u8]> {
        unsafe {
            let mut sz = 0;
            let valptr = ffi::sp_getstring(self.obj, "value\0".as_ptr(), &mut sz);

            // XXX: what if we stored an empty value?
            if valptr.is_null() {
                return None;
            }

            let s = slice::from_raw_parts(valptr as *const u8, sz as usize);
            return Some(s);
        }
    }
}

impl<'a> Drop for CursorObject<'a> {
    fn drop(&mut self) {
        unsafe {ffi::sp_destroy(self.obj);}
    }
}

impl Drop for Cursor {
    fn drop(&mut self) {
        unsafe {ffi::sp_destroy(self.obj);}
    }
}

impl Cursor {
    pub fn next<'a>(&'a self) -> Option<CursorObject<'a>> {
        let res = unsafe { ffi::sp_get(self.obj, ptr::null_mut()) };
        if res.is_null() {
            None
        }
        else {
            Some(CursorObject{obj: res, cur: &self})
        }
    }
}

impl Db {
    pub fn set(&mut self, key: &[u8], value: &[u8]) {
        unsafe {
            //ffi::sp_set_kv(self.db, key.as_ptr(), key.len() as i32, value.as_ptr(), value.len() as i32);
            let obj = ffi::sp_object(self.db);
            assert!(!obj.is_null());
            ffi::setkey(obj, key);
            ffi::setvalue(obj, value);
            ffi::sp_set(self.db, obj);
            ffi::sp_destroy(obj);
        }
    }

    pub fn iter_all(&mut self) -> Cursor {
        unsafe {
            let obj = ffi::sp_object(self.db);
            assert!(!obj.is_null());
            ffi::setstring(obj, "order\0".as_bytes(), ">=\0".as_bytes(), 0);
            let cursor = ffi::sp_cursor(self.db, obj);
            assert!(!cursor.is_null());
            Cursor {obj: cursor}
        }
    }

    pub fn get<'a>(&self, key: &'a [u8]) -> Option<Object<'a>> {
        assert!(key.len() > 0);
        unsafe {
            let pattern = ffi::sp_object(self.db);
            if pattern.is_null() {
                // XXX: Return Err
                return None;
            }
            ffi::setkey(pattern, key);
            let res = ffi::sp_get(self.db, pattern);
            if res.is_null() {
                ffi::sp_destroy(pattern);
                return None;
            }
            return Some(Object{obj: res, key: key});
        }
    }
}
<|endoftext|>


<commit_before>pub mod quad;

pub type Float = f32;
pub type Point = [Float; 3];

/// Generic grid type.
pub trait Grid {
    /// Cell coordinate on the grid.
    type Coordinate: Copy + Debug + PartialEq;
    /// Direciton on the grid.
    type Direction: Copy + Debug + PartialEq;
    /// Get the center point of a cell.
    fn get_cell_center(&self, Self::Coordinate) -> Point;
    /// Get the closest coordinate of a point.
    fn get_coordinate(&self, &Point) -> Self::Coordinate;
    /// Get the cell neighbor in a direction.
    fn get_neighbour(&self, Self::Coordinate, Self::Direction) -> Self::Coordinate;
    /// Fold over all neighbours.
    fn fold_neighbours<U, F: Fn(U, Self::Coordinate, Self::Direction) -> U>(&self, Self::Coordinate, F, U) -> U;
    /// Fold over all cells in a radius.
    fn fold_in_radius<U, F: Fn(U, Self::Coordinate) -> U>(&self, Self::Coordinate, Float, F, U) -> U;
}

/// Position on the grid.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Position<G: Grid> {
    /// Coordinate.
    pub coord: G::Coordinate,
    /// Direction.
    pub dir: G::Direction,
}
<commit_msg>Added the Debug import<commit_after>pub mod quad;

use std::fmt::Debug;

/// Scalar type to be used in `Point`.
pub type Float = f32;
/// A point in physical space.
pub type Point = [Float; 3];

/// Generic grid type.
pub trait Grid {
    /// Cell coordinate on the grid.
    type Coordinate: Copy + Debug + PartialEq;
    /// Direciton on the grid.
    type Direction: Copy + Debug + PartialEq;
    /// Get the center point of a cell.
    fn get_cell_center(&self, Self::Coordinate) -> Point;
    /// Get the closest coordinate of a point.
    fn get_coordinate(&self, &Point) -> Self::Coordinate;
    /// Get the cell neighbor in a direction.
    fn get_neighbour(&self, Self::Coordinate, Self::Direction) -> Self::Coordinate;
    /// Fold over all neighbours.
    fn fold_neighbours<U, F: Fn(U, Self::Coordinate, Self::Direction) -> U>(&self, Self::Coordinate, F, U) -> U;
    /// Fold over all cells in a radius.
    fn fold_in_radius<U, F: Fn(U, Self::Coordinate) -> U>(&self, Self::Coordinate, Float, F, U) -> U;
}

/// Position on the grid.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Position<G: Grid> {
    /// Coordinate.
    pub coord: G::Coordinate,
    /// Direction.
    pub dir: G::Direction,
}
<|endoftext|>


<commit_before>#![feature(slicing_syntax)]

extern crate lzf;
use std::str;
use lzf::decompress;
use std::io::MemReader;
use std::io;

mod version {
    pub const SUPPORTED_MINIMUM : u32 = 1;
    pub const SUPPORTED_MAXIMUM : u32 = 7;
}

mod constants {
    pub const RDB_6BITLEN : u8 = 0;
    pub const RDB_14BITLEN : u8 = 1;
    pub const RDB_ENCVAL : u8 = 3;
    pub const RDB_MAGIC : &'static str = "REDIS";
}

mod op_codes {
    pub const RESIZEDB : u8 = 251;
    pub const EXPIRETIME_MS : u8 = 252;
    pub const EXPIRETIME : u8 = 253;
    pub const SELECTDB   : u8 = 254;
    pub const EOF : u8 = 255;
}

mod types {
    pub const STRING : u8 = 0;
    pub const LIST : u8 = 1;
    pub const SET : u8 = 2;
    pub const ZSET : u8 = 3;
    pub const HASH : u8 = 4;
    pub const HASH_ZIPMAP : u8 = 9;
    pub const LIST_ZIPLIST : u8 = 10;
    pub const SET_INTSET : u8 = 11;
    pub const ZSET_ZIPLIST : u8 = 12;
    pub const HASH_ZIPLIST : u8 = 13;
}

mod encoding {
    pub const INT8 : u32 = 0;
    pub const INT16 : u32 = 1;
    pub const INT32 : u32 = 2;
    pub const LZF : u32 = 3;
}

#[derive(Show)]
pub enum DataType {
    String(Vec<u8>),
    Number(i64),
    ListOfTypes(Vec<DataType>),
    HashOfTypes(Vec<DataType>),
    SortedSetOfTypes(Vec<DataType>),
    Intset(Vec<i64>),
    List(Vec<Vec<u8>>),
    Set(Vec<Vec<u8>>),
    SortedSet(Vec<(f64,Vec<u8>)>),
    Hash(Vec<Vec<u8>>),
    Unknown
}

pub fn parse<R: Reader>(input: &mut R) {
    assert!(verify_magic(input));
    assert!(verify_version(input));

    let mut out = io::stdout();
    loop {
        let next_op = input.read_byte().unwrap();

        match next_op {
            op_codes::SELECTDB => {
                let db = read_length(input);
                println!("SELECTDB: {}", db);
            },
            op_codes::EOF => {
                println!("EOF");
                let chksum = input.read_to_end().unwrap();
                println!("checksum: {}", chksum);
                break ;
            },
            op_codes::EXPIRETIME_MS => {
                let expiretime_ms = input.read_le_u64().unwrap();
                println!("EXPIRETIME_MS: {}", expiretime_ms);
            },
            op_codes::EXPIRETIME => {
                let expiretime = input.read_be_u32().unwrap();
                println!("EXPIRETIME: {}", expiretime);
            },
            op_codes::RESIZEDB => {
                let db_size = read_length(input);
                let expires_size = read_length(input);
                println!("DB Size: {}, Expires Size: {}",
                         db_size, expires_size);
            }
            _ => {
                let key = read_blob(input);
                let _ = out.write(key[]);
                let _ = out.write_str(": ");
                let _ = out.flush();

                match read_type(next_op, input) {
                    DataType::String(t) => {
                        let _ = out.write(t[]);
                        let _ = out.write_str("\n");
                    },
                    DataType::Number(t) => { println!("{}", t) },
                    DataType::ListOfTypes(t) => { println!("{}", t) },
                    DataType::Intset(t) => { println!("{}", t) },
                    DataType::Hash(t) => {
                        for val in t.iter() {
                            let _ = out.write(val[]);
                            let _ = out.write_str(", ");
                        }
                        let _ = out.write_str("\n");
                    },
                    DataType::HashOfTypes(t) => { println!("{}", t) },
                    _ => {}
                }
            }
        }

    }

    return;

}

pub fn verify_magic<R: Reader>(input: &mut R) -> bool {
    let magic = input.read_exact(5).unwrap();

    // Meeeeeh.
    magic[0] == constants::RDB_MAGIC.as_bytes()[0] &&
        magic[1] == constants::RDB_MAGIC.as_bytes()[1] &&
        magic[2] == constants::RDB_MAGIC.as_bytes()[2] &&
        magic[3] == constants::RDB_MAGIC.as_bytes()[3] &&
        magic[4] == constants::RDB_MAGIC.as_bytes()[4]
}

pub fn verify_version<R: Reader>(input: &mut R) -> bool {
    let version = input.read_exact(4).unwrap();

    let version = (version[0]-48) as u32 * 1000 +
        (version[1]-48) as u32 * 100 +
        (version[2]-48) as u32 * 10 +
        (version[3]-48) as u32;

    version >= version::SUPPORTED_MINIMUM &&
        version <= version::SUPPORTED_MAXIMUM
}

fn read_linked_list<T: Reader>(input: &mut T) -> Vec<Vec<u8>> {
    let mut len = read_length(input);
    let mut list = vec![];

    while len > 0 {
        let blob = read_blob(input);
        list.push(blob);
        len -= 1;
    }

    list
}

fn read_sorted_set<T: Reader>(input: &mut T) -> Vec<(f64,Vec<u8>)> {
    let mut set = vec![];
    let mut set_items = read_length(input);

    while set_items > 0 {
        let val = read_blob(input);
        let score_length = input.read_byte().unwrap();
        let score = input.read_exact(score_length as uint).unwrap();
        let score = unsafe{str::from_utf8_unchecked(score[])}.
            parse::<f64>().unwrap();

        set.push((score, val));

        set_items -= 1;
    }

    set
}

fn read_hash<T: Reader>(input: &mut T) -> Vec<Vec<u8>> {
    let mut hash = vec![];
    let mut hash_items = read_length(input);
    hash_items = 2*hash_items;

    while hash_items > 0 {
        let val = read_blob(input);
        hash.push(val);

        hash_items -= 1;
    }

    hash
}

fn read_list_ziplist<T: Reader>(input: &mut T) -> Vec<DataType> {
    let mut list = vec![];

    let ziplist = read_blob(input);

    let mut reader = MemReader::new(ziplist);

    let _zlbytes = reader.read_le_u32().unwrap();
    let _zltail = reader.read_le_u32().unwrap();
    let zllen = reader.read_le_u16().unwrap();

    for _ in range(0, zllen) {
        // 1. 1 or 5 bytes length of previous entry
        match reader.read_byte().unwrap() {
            254 => {
                let _ = reader.read_exact(4).unwrap();
            },
            _ => {}
        }

        let mut length : u64;
        let mut number_value : i64;

        // 2. Read flag or number value
        let flag = reader.read_byte().unwrap();

        match (flag & 0xC0) >> 6 {
            0 => { length = (flag & 0x3F) as u64 },
            1 => {
                let next_byte = reader.read_byte().unwrap();
                length = (((flag & 0x3F) as u64) << 8) | next_byte as u64;
            },
            2 => {
                length = reader.read_be_u32().unwrap() as u64;
            },
            _ => {
                match (flag & 0xF0) >> 4 {
                    0xC => {
                        number_value = reader.read_le_i16().unwrap() as i64 },
                    0xD => {
                        number_value = reader.read_le_i32().unwrap() as i64 },
                    0xE => {
                        number_value = reader.read_le_i64().unwrap() as i64 },
                    0xF => {
                        match flag & 0xF {
                            0 => {
                                let bytes = reader.read_exact(3).unwrap();
                                number_value = ((bytes[0] as i64) << 16) ^
                                    ((bytes[1] as i64) << 8) ^
                                    (bytes[2] as i64);
                            },
                            0xE => {
                                number_value = reader.read_byte().unwrap() as i64 },
                            _ => { number_value = (flag & 0xF) as i64 - 1; }
                        }
                    },
                    _ => {
                        println!("Flag not handled: {}", flag);
                        continue;
                    }

                }

                list.push(DataType::Number(number_value));
                continue;
            }
        }

        // 3. Read value
        let rawval = reader.read_exact(length as uint).unwrap();
        list.push(DataType::String(rawval));
    }
    assert!(reader.read_byte().unwrap() == 0xFF);

    list
}

fn read_hash_zipmap<T: Reader>(input: &mut T) -> Vec<Vec<u8>> {
    let zipmap = read_blob(input);

    let mut reader = MemReader::new(zipmap);

    let zmlen = reader.read_byte().unwrap();

    let mut length;
    let mut hash;
    if zmlen <= 254 {
        length = zmlen as uint;
        hash = Vec::with_capacity(length);
    } else {
        length = -1;
        hash = Vec::with_capacity(255);
    }

    loop {
        let next_byte = reader.read_byte().unwrap();

        if next_byte == 0xFF {
            break; // End of list.
        }

        let mut elem_len;
        match next_byte {
            253 => { elem_len = reader.read_le_u32().unwrap() },
            254 | 255 => { panic!("Invalid length value in zipmap: {}", next_byte) },
            _ => { elem_len = next_byte as u32 }
        }

        let key = reader.read_exact(elem_len as uint).unwrap();
        hash.push(key);

        let next_byte = reader.read_byte().unwrap();
        match next_byte {
            253 => { elem_len = reader.read_le_u32().unwrap() },
            254 | 255 => { panic!("Invalid length value in zipmap: {}", next_byte) },
            _ => { elem_len = next_byte as u32 }
        }

        let _free = reader.read_byte().unwrap();
        let value = reader.read_exact(elem_len as uint).unwrap();
        hash.push(value);

        if length > 0 {
            length -= 1;
        }

        if length == 0 {
            assert!(reader.read_byte().unwrap() == 0xFF);
            break;
        }
    }

    hash
}

fn read_set_intset<T: Reader>(input: &mut T) -> Vec<i64> {
    let mut set = vec![];

    let intset = read_blob(input);

    let mut reader = MemReader::new(intset);
    let byte_size = reader.read_le_u32().unwrap();
    let intset_length = reader.read_le_u32().unwrap();

    for _ in range(0, intset_length) {
        let val = match byte_size {
            2 => reader.read_le_i16().unwrap() as i64,
            4 => reader.read_le_i32().unwrap() as i64,
            8 => reader.read_le_i64().unwrap(),
            _ => panic!("unhandled byte size in intset: {}", byte_size)
        };

        set.push(val);
    }

    set
}

fn read_type<T: Reader>(value_type: u8, input: &mut T) -> DataType {
    match value_type {
        types::STRING => {
            DataType::String(read_blob(input))
        },
        types::LIST => {
            DataType::List(read_linked_list(input))
        },
        types::SET => {
            DataType::Set(read_linked_list(input))
        },
        types::ZSET => {
            DataType::SortedSet(read_sorted_set(input))
        },
        types::HASH => {
            DataType::Hash(read_hash(input))
        },
        types::HASH_ZIPMAP => {
            DataType::Hash(read_hash_zipmap(input))
        },
        types::LIST_ZIPLIST => {
            DataType::ListOfTypes(read_list_ziplist(input))
        },
        types::SET_INTSET => {
            DataType::Intset(read_set_intset(input))
        },
        types::ZSET_ZIPLIST => {
            DataType::SortedSetOfTypes(read_list_ziplist(input))
        },
        types::HASH_ZIPLIST => {
            DataType::ListOfTypes(read_list_ziplist(input))
        },
        _ => { panic!("Value Type not implemented: {}", value_type) }
    }
}

pub fn read_length_with_encoding<T: Reader>(input: &mut T) -> (u32, bool) {
    let mut length;
    let mut is_encoded = false;

    let enc_type = input.read_byte().unwrap();

    match (enc_type & 0xC0) >> 6 {
        constants::RDB_ENCVAL => {
            is_encoded = true;
            length = (enc_type & 0x3F) as u32;
        },
        constants::RDB_6BITLEN => {
            length = (enc_type & 0x3F) as u32;
        },
        constants::RDB_14BITLEN => {
            let next_byte = input.read_byte().unwrap();
            length = (((enc_type & 0x3F) as u32) <<8) | next_byte as u32;
        },
        _ => {
            length = input.read_be_u32().unwrap();
        }
    }

    (length, is_encoded)
}

pub fn read_length<T: Reader>(input: &mut T) -> u32 {
    let (length, _) = read_length_with_encoding(input);
    length
}

fn int_to_vec(number: i32) -> Vec<u8> {
    let number = number.to_string();
    let mut result = Vec::with_capacity(number.len());
    for &c in number.as_bytes().iter() {
        result.push(c);
    }
    result
}

pub fn read_blob<T: Reader>(input: &mut T) -> Vec<u8> {
    let (length, is_encoded) = read_length_with_encoding(input);

    if is_encoded {
        match length {
            encoding::INT8 => { int_to_vec(input.read_i8().unwrap() as i32) },
            encoding::INT16 => { int_to_vec(input.read_le_i16().unwrap() as i32) },
            encoding::INT32 => { int_to_vec(input.read_le_i32().unwrap() as i32) },
            encoding::LZF => {
                let compressed_length = read_length(input);
                let real_length = read_length(input);
                let data = input.read_exact(compressed_length as uint).unwrap();
                lzf::decompress(data.as_slice(), real_length as uint).unwrap()
            },
            _ => { panic!("Unknown encoding: {}", length) }
        }
    } else {
        input.read_exact(length as uint).unwrap()
    }
}
<commit_msg>Refactor zipmap entry reading into own method<commit_after>#![feature(slicing_syntax)]

extern crate lzf;
use std::str;
use lzf::decompress;
use std::io::MemReader;
use std::io;

mod version {
    pub const SUPPORTED_MINIMUM : u32 = 1;
    pub const SUPPORTED_MAXIMUM : u32 = 7;
}

mod constants {
    pub const RDB_6BITLEN : u8 = 0;
    pub const RDB_14BITLEN : u8 = 1;
    pub const RDB_ENCVAL : u8 = 3;
    pub const RDB_MAGIC : &'static str = "REDIS";
}

mod op_codes {
    pub const RESIZEDB : u8 = 251;
    pub const EXPIRETIME_MS : u8 = 252;
    pub const EXPIRETIME : u8 = 253;
    pub const SELECTDB   : u8 = 254;
    pub const EOF : u8 = 255;
}

mod types {
    pub const STRING : u8 = 0;
    pub const LIST : u8 = 1;
    pub const SET : u8 = 2;
    pub const ZSET : u8 = 3;
    pub const HASH : u8 = 4;
    pub const HASH_ZIPMAP : u8 = 9;
    pub const LIST_ZIPLIST : u8 = 10;
    pub const SET_INTSET : u8 = 11;
    pub const ZSET_ZIPLIST : u8 = 12;
    pub const HASH_ZIPLIST : u8 = 13;
}

mod encoding {
    pub const INT8 : u32 = 0;
    pub const INT16 : u32 = 1;
    pub const INT32 : u32 = 2;
    pub const LZF : u32 = 3;
}

#[derive(Show)]
pub enum DataType {
    String(Vec<u8>),
    Number(i64),
    ListOfTypes(Vec<DataType>),
    HashOfTypes(Vec<DataType>),
    SortedSetOfTypes(Vec<DataType>),
    Intset(Vec<i64>),
    List(Vec<Vec<u8>>),
    Set(Vec<Vec<u8>>),
    SortedSet(Vec<(f64,Vec<u8>)>),
    Hash(Vec<Vec<u8>>),
    Unknown
}

pub fn parse<R: Reader>(input: &mut R) {
    assert!(verify_magic(input));
    assert!(verify_version(input));

    let mut out = io::stdout();
    loop {
        let next_op = input.read_byte().unwrap();

        match next_op {
            op_codes::SELECTDB => {
                let db = read_length(input);
                println!("SELECTDB: {}", db);
            },
            op_codes::EOF => {
                println!("EOF");
                let chksum = input.read_to_end().unwrap();
                println!("checksum: {}", chksum);
                break ;
            },
            op_codes::EXPIRETIME_MS => {
                let expiretime_ms = input.read_le_u64().unwrap();
                println!("EXPIRETIME_MS: {}", expiretime_ms);
            },
            op_codes::EXPIRETIME => {
                let expiretime = input.read_be_u32().unwrap();
                println!("EXPIRETIME: {}", expiretime);
            },
            op_codes::RESIZEDB => {
                let db_size = read_length(input);
                let expires_size = read_length(input);
                println!("DB Size: {}, Expires Size: {}",
                         db_size, expires_size);
            }
            _ => {
                let key = read_blob(input);
                let _ = out.write(key[]);
                let _ = out.write_str(": ");
                let _ = out.flush();

                match read_type(next_op, input) {
                    DataType::String(t) => {
                        let _ = out.write(t[]);
                        let _ = out.write_str("\n");
                    },
                    DataType::Number(t) => { println!("{}", t) },
                    DataType::ListOfTypes(t) => { println!("{}", t) },
                    DataType::Intset(t) => { println!("{}", t) },
                    DataType::Hash(t) => {
                        for val in t.iter() {
                            let _ = out.write(val[]);
                            let _ = out.write_str(", ");
                        }
                        let _ = out.write_str("\n");
                    },
                    DataType::HashOfTypes(t) => { println!("{}", t) },
                    _ => {}
                }
            }
        }

    }

    return;

}

pub fn verify_magic<R: Reader>(input: &mut R) -> bool {
    let magic = input.read_exact(5).unwrap();

    // Meeeeeh.
    magic[0] == constants::RDB_MAGIC.as_bytes()[0] &&
        magic[1] == constants::RDB_MAGIC.as_bytes()[1] &&
        magic[2] == constants::RDB_MAGIC.as_bytes()[2] &&
        magic[3] == constants::RDB_MAGIC.as_bytes()[3] &&
        magic[4] == constants::RDB_MAGIC.as_bytes()[4]
}

pub fn verify_version<R: Reader>(input: &mut R) -> bool {
    let version = input.read_exact(4).unwrap();

    let version = (version[0]-48) as u32 * 1000 +
        (version[1]-48) as u32 * 100 +
        (version[2]-48) as u32 * 10 +
        (version[3]-48) as u32;

    version >= version::SUPPORTED_MINIMUM &&
        version <= version::SUPPORTED_MAXIMUM
}

fn read_linked_list<T: Reader>(input: &mut T) -> Vec<Vec<u8>> {
    let mut len = read_length(input);
    let mut list = vec![];

    while len > 0 {
        let blob = read_blob(input);
        list.push(blob);
        len -= 1;
    }

    list
}

fn read_sorted_set<T: Reader>(input: &mut T) -> Vec<(f64,Vec<u8>)> {
    let mut set = vec![];
    let mut set_items = read_length(input);

    while set_items > 0 {
        let val = read_blob(input);
        let score_length = input.read_byte().unwrap();
        let score = input.read_exact(score_length as uint).unwrap();
        let score = unsafe{str::from_utf8_unchecked(score[])}.
            parse::<f64>().unwrap();

        set.push((score, val));

        set_items -= 1;
    }

    set
}

fn read_hash<T: Reader>(input: &mut T) -> Vec<Vec<u8>> {
    let mut hash = vec![];
    let mut hash_items = read_length(input);
    hash_items = 2*hash_items;

    while hash_items > 0 {
        let val = read_blob(input);
        hash.push(val);

        hash_items -= 1;
    }

    hash
}

fn read_list_ziplist<T: Reader>(input: &mut T) -> Vec<DataType> {
    let mut list = vec![];

    let ziplist = read_blob(input);

    let mut reader = MemReader::new(ziplist);

    let _zlbytes = reader.read_le_u32().unwrap();
    let _zltail = reader.read_le_u32().unwrap();
    let zllen = reader.read_le_u16().unwrap();

    for _ in range(0, zllen) {
        // 1. 1 or 5 bytes length of previous entry
        match reader.read_byte().unwrap() {
            254 => {
                let _ = reader.read_exact(4).unwrap();
            },
            _ => {}
        }

        let mut length : u64;
        let mut number_value : i64;

        // 2. Read flag or number value
        let flag = reader.read_byte().unwrap();

        match (flag & 0xC0) >> 6 {
            0 => { length = (flag & 0x3F) as u64 },
            1 => {
                let next_byte = reader.read_byte().unwrap();
                length = (((flag & 0x3F) as u64) << 8) | next_byte as u64;
            },
            2 => {
                length = reader.read_be_u32().unwrap() as u64;
            },
            _ => {
                match (flag & 0xF0) >> 4 {
                    0xC => {
                        number_value = reader.read_le_i16().unwrap() as i64 },
                    0xD => {
                        number_value = reader.read_le_i32().unwrap() as i64 },
                    0xE => {
                        number_value = reader.read_le_i64().unwrap() as i64 },
                    0xF => {
                        match flag & 0xF {
                            0 => {
                                let bytes = reader.read_exact(3).unwrap();
                                number_value = ((bytes[0] as i64) << 16) ^
                                    ((bytes[1] as i64) << 8) ^
                                    (bytes[2] as i64);
                            },
                            0xE => {
                                number_value = reader.read_byte().unwrap() as i64 },
                            _ => { number_value = (flag & 0xF) as i64 - 1; }
                        }
                    },
                    _ => {
                        println!("Flag not handled: {}", flag);
                        continue;
                    }

                }

                list.push(DataType::Number(number_value));
                continue;
            }
        }

        // 3. Read value
        let rawval = reader.read_exact(length as uint).unwrap();
        list.push(DataType::String(rawval));
    }
    assert!(reader.read_byte().unwrap() == 0xFF);

    list
}

fn read_zipmap_entry<R: Reader>(next_byte: u8, input: &mut R) -> Vec<u8> {
    let mut elem_len;
    match next_byte {
        253 => { elem_len = input.read_le_u32().unwrap() },
        254 | 255 => {
            panic!("Invalid length value in zipmap: {}", next_byte)
        },
        _ => { elem_len = next_byte as u32 }
    }

    input.read_exact(elem_len as uint).unwrap()
}

fn read_hash_zipmap<T: Reader>(input: &mut T) -> Vec<Vec<u8>> {
    let zipmap = read_blob(input);

    let mut reader = MemReader::new(zipmap);

    let zmlen = reader.read_byte().unwrap();

    let mut length;
    let mut hash;
    if zmlen <= 254 {
        length = zmlen as uint;
        hash = Vec::with_capacity(length);
    } else {
        length = -1;
        hash = Vec::with_capacity(255);
    }

    loop {
        let next_byte = reader.read_byte().unwrap();

        if next_byte == 0xFF {
            break; // End of list.
        }

        let key = read_zipmap_entry(next_byte, &mut reader);
        hash.push(key);

        let next_byte = reader.read_byte().unwrap();
        let _free = reader.read_byte().unwrap();
        let value = read_zipmap_entry(next_byte, &mut reader);
        hash.push(value);

        if length > 0 {
            length -= 1;
        }

        if length == 0 {
            assert!(reader.read_byte().unwrap() == 0xFF);
            break;
        }
    }

    hash
}

fn read_set_intset<T: Reader>(input: &mut T) -> Vec<i64> {
    let mut set = vec![];

    let intset = read_blob(input);

    let mut reader = MemReader::new(intset);
    let byte_size = reader.read_le_u32().unwrap();
    let intset_length = reader.read_le_u32().unwrap();

    for _ in range(0, intset_length) {
        let val = match byte_size {
            2 => reader.read_le_i16().unwrap() as i64,
            4 => reader.read_le_i32().unwrap() as i64,
            8 => reader.read_le_i64().unwrap(),
            _ => panic!("unhandled byte size in intset: {}", byte_size)
        };

        set.push(val);
    }

    set
}

fn read_type<T: Reader>(value_type: u8, input: &mut T) -> DataType {
    match value_type {
        types::STRING => {
            DataType::String(read_blob(input))
        },
        types::LIST => {
            DataType::List(read_linked_list(input))
        },
        types::SET => {
            DataType::Set(read_linked_list(input))
        },
        types::ZSET => {
            DataType::SortedSet(read_sorted_set(input))
        },
        types::HASH => {
            DataType::Hash(read_hash(input))
        },
        types::HASH_ZIPMAP => {
            DataType::Hash(read_hash_zipmap(input))
        },
        types::LIST_ZIPLIST => {
            DataType::ListOfTypes(read_list_ziplist(input))
        },
        types::SET_INTSET => {
            DataType::Intset(read_set_intset(input))
        },
        types::ZSET_ZIPLIST => {
            DataType::SortedSetOfTypes(read_list_ziplist(input))
        },
        types::HASH_ZIPLIST => {
            DataType::ListOfTypes(read_list_ziplist(input))
        },
        _ => { panic!("Value Type not implemented: {}", value_type) }
    }
}

pub fn read_length_with_encoding<T: Reader>(input: &mut T) -> (u32, bool) {
    let mut length;
    let mut is_encoded = false;

    let enc_type = input.read_byte().unwrap();

    match (enc_type & 0xC0) >> 6 {
        constants::RDB_ENCVAL => {
            is_encoded = true;
            length = (enc_type & 0x3F) as u32;
        },
        constants::RDB_6BITLEN => {
            length = (enc_type & 0x3F) as u32;
        },
        constants::RDB_14BITLEN => {
            let next_byte = input.read_byte().unwrap();
            length = (((enc_type & 0x3F) as u32) <<8) | next_byte as u32;
        },
        _ => {
            length = input.read_be_u32().unwrap();
        }
    }

    (length, is_encoded)
}

pub fn read_length<T: Reader>(input: &mut T) -> u32 {
    let (length, _) = read_length_with_encoding(input);
    length
}

fn int_to_vec(number: i32) -> Vec<u8> {
    let number = number.to_string();
    let mut result = Vec::with_capacity(number.len());
    for &c in number.as_bytes().iter() {
        result.push(c);
    }
    result
}

pub fn read_blob<T: Reader>(input: &mut T) -> Vec<u8> {
    let (length, is_encoded) = read_length_with_encoding(input);

    if is_encoded {
        match length {
            encoding::INT8 => { int_to_vec(input.read_i8().unwrap() as i32) },
            encoding::INT16 => { int_to_vec(input.read_le_i16().unwrap() as i32) },
            encoding::INT32 => { int_to_vec(input.read_le_i32().unwrap() as i32) },
            encoding::LZF => {
                let compressed_length = read_length(input);
                let real_length = read_length(input);
                let data = input.read_exact(compressed_length as uint).unwrap();
                lzf::decompress(data.as_slice(), real_length as uint).unwrap()
            },
            _ => { panic!("Unknown encoding: {}", length) }
        }
    } else {
        input.read_exact(length as uint).unwrap()
    }
}
<|endoftext|>


<commit_before>use std::io;

/// Main bitmap structure
pub struct Bitmap {
    width:  i32,
    height: i32,
    pixels: Vec<u8>,
}

impl Bitmap {
    /// Create a new blank (white) bitmap of a specified size.
    pub fn new(width: i32, height: i32) -> Bitmap {
        Bitmap {
            width:  width,
            height: height,
            // Create a vector to store the pixels in, ensuring that it is padded to a multiple of 4
            // bytes of each row.
            pixels: Vec::from_elem((height * (width * 3 + width % 4)) as uint, 0xFFu8)
        }
    }

    /// Set a pixel at (x, y) to a specified color (r, g, b).
    pub fn set_pixel(&mut self, x: i32, y: i32, color: (u8, u8, u8)) {
        // Calculate the byte offset for x
        let i = ((self.height - y - 1) * (self.width * 3) + x * 3) as uint;

        let (r, g, b) = color;
        // Note: Pixel order for bitmaps is (blue, green, red)
        *self.pixels.get_mut(i + 0) = b;
        *self.pixels.get_mut(i + 1) = g;
        *self.pixels.get_mut(i + 2) = r;
    }

    /// Write the stored data to a file with given filename.
    pub fn write_to_file(&self, filename: &str) -> io::IoResult<()> {
        const FILE_HEADER_SIZE: u32 = 14;
        const BMP_INFO_SIZE: u32 = 40;
        const TOTAL_HEADER_SIZE: u32 = FILE_HEADER_SIZE + BMP_INFO_SIZE;

        let image_size = (self.height * self.width*3 + self.height * (self.width % 4)) as u32;
        let file_size = image_size + TOTAL_HEADER_SIZE;

        // Bitmap file header
        let file_header: [u8, ..FILE_HEADER_SIZE as uint] = [
            'B' as u8, 'M' as u8,
            file_size as u8, (file_size>>8) as u8, (file_size>>16) as u8, (file_size>>24) as u8,
            0, 0,
            0, 0,
            TOTAL_HEADER_SIZE as u8, 0, 0, 0
        ];
        // Bitmap information header
        let info_header: [u8, ..BMP_INFO_SIZE as uint] = [
            BMP_INFO_SIZE as u8, 0, 0, 0,
            self.width as u8, (self.width>>8) as u8, (self.width>>16) as u8, (self.width>>24) as u8,
            self.height as u8, (self.height>>8) as u8, (self.height>>16) as u8, (self.height>>24) as u8,
            1, 0,
            24, 0,
            0, 0, 0, 0,
            image_size as u8, (image_size>>8) as u8, (image_size>>16) as u8, (image_size>>24) as u8,
            72, 0, 0, 0,
            72, 0, 0, 0,
            0, 0, 0, 0,
            0, 0, 0, 0
        ];

        // Set up the file writer
        let mut file = io::File::create(&Path::new(filename));

        // Write the bitmap headers to file
        try!(file.write(file_header));
        try!(file.write(info_header));

        // Write data to file
        file.write(self.pixels.as_slice())
    }

    /// Get the width of the bitmap.
    pub fn width(&self) -> i32 {
        self.width
    }

    /// Get the height of the bitmap.
    pub fn height(&self) -> i32 {
        self.height
    }
}
<commit_msg>Update to new version of Rust<commit_after>use std::io;

/// Main bitmap structure
pub struct Bitmap {
    width:  i32,
    height: i32,
    pixels: Vec<u8>,
}

impl Bitmap {
    /// Create a new blank (white) bitmap of a specified size.
    pub fn new(width: i32, height: i32) -> Bitmap {
        Bitmap {
            width: width,
            height: height,
            // Create a vector to store the pixels in, ensuring that it is padded to a multiple of 4
            // bytes of each row.
            pixels: Vec::from_elem((height * (width * 3 + width % 4)) as uint, 0xFFu8)
        }
    }

    /// Set a pixel at (x, y) to a specified color (r, g, b).
    pub fn set_pixel(&mut self, x: i32, y: i32, color: (u8, u8, u8)) {
        // Calculate the byte offset for x
        let i = ((self.height - y - 1) * (self.width * 3) + x * 3) as uint;

        let (r, g, b) = color;
        // Note: Pixel order for bitmaps is (blue, green, red)
        self.pixels[i + 0] = b;
        self.pixels[i + 1] = g;
        self.pixels[i + 2] = r;
    }

    /// Write the stored data to a file with given filename.
    pub fn write_to_file(&self, filename: &str) -> io::IoResult<()> {
        const FILE_HEADER_SIZE: u32 = 14;
        const BMP_INFO_SIZE: u32 = 40;
        const TOTAL_HEADER_SIZE: u32 = FILE_HEADER_SIZE + BMP_INFO_SIZE;

        let image_size = (self.height * self.width*3 + self.height * (self.width % 4)) as u32;
        let file_size = image_size + TOTAL_HEADER_SIZE;

        // Bitmap file header
        let file_header: [u8, ..FILE_HEADER_SIZE as uint] = [
            'B' as u8, 'M' as u8,
            file_size as u8, (file_size>>8) as u8, (file_size>>16) as u8, (file_size>>24) as u8,
            0, 0,
            0, 0,
            TOTAL_HEADER_SIZE as u8, 0, 0, 0
        ];
        // Bitmap information header
        let info_header: [u8, ..BMP_INFO_SIZE as uint] = [
            BMP_INFO_SIZE as u8, 0, 0, 0,
            self.width as u8, (self.width>>8) as u8, (self.width>>16) as u8, (self.width>>24) as u8,
            self.height as u8, (self.height>>8) as u8, (self.height>>16) as u8, (self.height>>24) as u8,
            1, 0,
            24, 0,
            0, 0, 0, 0,
            image_size as u8, (image_size>>8) as u8, (image_size>>16) as u8, (image_size>>24) as u8,
            72, 0, 0, 0,
            72, 0, 0, 0,
            0, 0, 0, 0,
            0, 0, 0, 0
        ];

        // Set up the file writer
        let mut file = io::File::create(&Path::new(filename));

        // Write the bitmap headers to file
        try!(file.write(file_header));
        try!(file.write(info_header));

        // Write data to file
        file.write(self.pixels.as_slice())
    }

    /// Get the width of the bitmap.
    pub fn width(&self) -> i32 {
        self.width
    }

    /// Get the height of the bitmap.
    pub fn height(&self) -> i32 {
        self.height
    }
}
<|endoftext|>


<commit_before>#![feature(custom_derive)]
#![feature(no_std)]
#![feature(core)]
#![no_std]

#![crate_name = "elfloader"]
#![crate_type = "lib"]

#[macro_use]
extern crate core;
#[macro_use]
extern crate klogger;

#[cfg(test)]
#[macro_use]
extern crate std;

pub mod elf;
use core::prelude::*;
use core::fmt;
use core::mem::{transmute, size_of};
use core::slice;

pub type PAddr = u64;
pub type VAddr = usize;

/// Abstract representation of a loadable ELF binary.
pub struct ElfBinary {
    name: &'static str,
    region: &'static [u8],
    header: &'static elf::FileHeader,
}

impl fmt::Debug for ElfBinary {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} {}", self.name, self.header)
    }
}

/// Verify that memory region starts with a correct ELF magic.
fn valid_elf_magic(region: &'static [u8]) -> bool {
    region[0] == elf::ELFMAG0 &&
    region[1] == elf::ELFMAG1 &&
    region[2] == elf::ELFMAG2 &&
    region[3] == elf::ELFMAG3
}

/// Implement this for ELF loading.
pub trait ElfLoader {
    /// Allocates a virtual region of size amount of bytes.
    fn allocate(&mut self, base: VAddr, size: usize, flags: elf::ProgFlag);

    /// Copies the region into the base.
    fn load(&mut self, base: VAddr, region: &'static [u8]);
}

impl ElfBinary {

    /// Create a new ElfBinary.
    /// Makes sure that the provided region has valid ELF magic byte sequence
    /// and is big enough to contain at least the ELF file header
    /// otherwise it will return None.
    pub fn new(name: &'static str, region: &'static [u8]) -> Option<ElfBinary> {
        if region.len() >= size_of::<elf::FileHeader>() && valid_elf_magic(region) {
            let header: &elf::FileHeader = unsafe { transmute(&region[0]) };
            return Some(ElfBinary { name: name, region: region, header: header });
        }

        None
    }

    /// Print the program headers.
    pub fn print_program_headers(&self) {
        for p in self.program_headers() {
            log!("pheader = {}", p);
        }
    }

    /// Create a slice of the program headers.
    fn program_headers(&self) -> &[elf::ProgramHeader] {
        let correct_header_size = self.header.phentsize as usize == size_of::<elf::ProgramHeader>();
        let pheader_region_size = self.header.phoff as usize + self.header.phnum as usize * self.header.phentsize as usize;
        let big_enough_region = self.region.len() >= pheader_region_size;

        if !correct_header_size || !big_enough_region {
            return &[];
        }

        let pheaders: &[elf::ProgramHeader] = unsafe {
            core::slice::from_raw_parts(
                transmute(&self.region[self.header.phoff as usize]),
                self.header.phnum as usize)
        };
        return pheaders;
    }

    /// Can we load the binary on our platform?
    fn can_load(&self) -> bool {
        let correct_class = self.header.ident.class == elf::ELFCLASS64;
        let correct_elfversion = self.header.ident.version == elf::EV_CURRENT;
        let correct_data = self.header.ident.data == elf::ELFDATA2LSB;
        let correct_osabi = self.header.ident.osabi == elf::ELFOSABI_SYSV || self.header.ident.osabi == elf::ELFOSABI_LINUX;
        let correct_type = self.header.elftype == elf::ET_EXEC || self.header.elftype == elf::ET_DYN;
        let correct_machine = self.header.machine == elf::EM_X86_64;

        correct_class && correct_data && correct_elfversion && correct_machine && correct_osabi && correct_type
    }

    fn load_header(&self, p: &elf::ProgramHeader, loader: &mut ElfLoader) {
        let big_enough_region = self.region.len() >= (p.offset + p.filesz) as usize;
        if !big_enough_region {
            log!("Unable to load {}", p);
            return;
        }

        loader.allocate(p.vaddr, p.memsz as usize, p.flags);
        let segment: &'static [u8] = unsafe {
            core::slice::from_raw_parts(
                transmute(&self.region[p.offset as usize]), p.filesz as usize)
        };
        loader.load(p.vaddr, segment);
    }

    pub fn load(&self, loader: &mut ElfLoader) {
        for p in self.program_headers() {
            let x = match p.progtype {
                elf::PT_LOAD => self.load_header(p, loader),
                _ => ()
            };
        }
    }

}
<commit_msg>Update elfloader features.<commit_after>#![feature(no_std, core, core_prelude, core_slice_ext, custom_derive)]
#![no_std]

#![crate_name = "elfloader"]
#![crate_type = "lib"]

#[macro_use]
extern crate core;
#[macro_use]
extern crate klogger;

#[cfg(test)]
#[macro_use]
extern crate std;

pub mod elf;
use core::prelude::*;
use core::fmt;
use core::mem::{transmute, size_of};

pub type PAddr = u64;
pub type VAddr = usize;

/// Abstract representation of a loadable ELF binary.
pub struct ElfBinary {
    name: &'static str,
    region: &'static [u8],
    header: &'static elf::FileHeader,
}

impl fmt::Debug for ElfBinary {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} {}", self.name, self.header)
    }
}

/// Verify that memory region starts with a correct ELF magic.
fn valid_elf_magic(region: &'static [u8]) -> bool {
    region[0] == elf::ELFMAG0 &&
    region[1] == elf::ELFMAG1 &&
    region[2] == elf::ELFMAG2 &&
    region[3] == elf::ELFMAG3
}

/// Implement this for ELF loading.
pub trait ElfLoader {
    /// Allocates a virtual region of size amount of bytes.
    fn allocate(&mut self, base: VAddr, size: usize, flags: elf::ProgFlag);

    /// Copies the region into the base.
    fn load(&mut self, base: VAddr, region: &'static [u8]);
}

impl ElfBinary {

    /// Create a new ElfBinary.
    /// Makes sure that the provided region has valid ELF magic byte sequence
    /// and is big enough to contain at least the ELF file header
    /// otherwise it will return None.
    pub fn new(name: &'static str, region: &'static [u8]) -> Option<ElfBinary> {
        if region.len() >= size_of::<elf::FileHeader>() && valid_elf_magic(region) {
            let header: &elf::FileHeader = unsafe { transmute(&region[0]) };
            return Some(ElfBinary { name: name, region: region, header: header });
        }

        None
    }

    /// Print the program headers.
    pub fn print_program_headers(&self) {
        for p in self.program_headers() {
            log!("pheader = {}", p);
        }
    }

    /// Create a slice of the program headers.
    fn program_headers(&self) -> &[elf::ProgramHeader] {
        let correct_header_size = self.header.phentsize as usize == size_of::<elf::ProgramHeader>();
        let pheader_region_size = self.header.phoff as usize + self.header.phnum as usize * self.header.phentsize as usize;
        let big_enough_region = self.region.len() >= pheader_region_size;

        if !correct_header_size || !big_enough_region {
            return &[];
        }

        let pheaders: &[elf::ProgramHeader] = unsafe {
            core::slice::from_raw_parts(
                transmute(&self.region[self.header.phoff as usize]),
                self.header.phnum as usize)
        };
        return pheaders;
    }

    /// Can we load the binary on our platform?
    fn can_load(&self) -> bool {
        let correct_class = self.header.ident.class == elf::ELFCLASS64;
        let correct_elfversion = self.header.ident.version == elf::EV_CURRENT;
        let correct_data = self.header.ident.data == elf::ELFDATA2LSB;
        let correct_osabi = self.header.ident.osabi == elf::ELFOSABI_SYSV || self.header.ident.osabi == elf::ELFOSABI_LINUX;
        let correct_type = self.header.elftype == elf::ET_EXEC || self.header.elftype == elf::ET_DYN;
        let correct_machine = self.header.machine == elf::EM_X86_64;

        correct_class && correct_data && correct_elfversion && correct_machine && correct_osabi && correct_type
    }

    fn load_header(&self, p: &elf::ProgramHeader, loader: &mut ElfLoader) {
        let big_enough_region = self.region.len() >= (p.offset + p.filesz) as usize;
        if !big_enough_region {
            log!("Unable to load {}", p);
            return;
        }

        loader.allocate(p.vaddr, p.memsz as usize, p.flags);
        let segment: &'static [u8] = unsafe {
            core::slice::from_raw_parts(
                transmute(&self.region[p.offset as usize]), p.filesz as usize)
        };
        loader.load(p.vaddr, segment);
    }

    pub fn load(&self, loader: &mut ElfLoader) {
        for p in self.program_headers() {
            let x = match p.progtype {
                elf::PT_LOAD => self.load_header(p, loader),
                _ => ()
            };
        }
    }

}
<|endoftext|>


<commit_before>#![crate_name = "raft"]
#![crate_type="lib"]

#![allow(unstable)]
extern crate "rustc-serialize" as rustc_serialize;

use std::io::net::ip::SocketAddr;
use std::io::net::udp::UdpSocket;
use std::io::timer::Timer;
use std::time::Duration;
use std::thread::Thread;
use std::rand::{thread_rng, Rng};
use std::sync::mpsc::{channel, Sender, Receiver};
use std::str;
use rustc_serialize::{json, Encodable, Decodable};
use RemoteProcedureCall::{AppendEntries, RequestVote};

// The maximum size of the read buffer.
const BUFFER_SIZE: usize = 4096;
const HEARTBEAT_MIN: i64 = 150;
const HEARTBEAT_MAX: i64 = 300;

/// The Raft Distributed Consensus Algorithm requires two RPC calls to be available:
///
///   * `append_entries` which is used as both a heartbeat (with no payload) and the primary
///     interface for requests.
///   * `request_vote` which is used by candidates during campaigns to obtain a vote.
///
pub trait Raft<T: Encodable + Decodable + Send + Clone> {
    /// Returns (term, success)
    fn append_entries(term: usize, leader_id: usize, prev_log_index: usize,
                      prev_log_term: usize, entries: Vec<T>,
                      leader_commit: usize) -> (usize, bool);

    /// Returns (term, voteGranted)
    fn request_vote(term: usize, candidate_id: usize, last_log_index: usize,
                    last_log_term: usize) -> (usize, bool);
}

/// A `RaftNode` acts as a replicated state machine. The server's role in the cluster depends on it's
/// own status. It will maintain both volatile state (which can be safely lost) and persistent
/// state (which must be carefully stored and kept safe).
pub struct RaftNode<T: Encodable + Decodable + Send + Clone> {
    // Raft related.
    state: NodeState,
    persistent_state: PersistentState<T>,
    volatile_state: VolatileState,
    // Auxilary Data.
    // TODO: This should probably be split off.
    // All nodes need to know this otherwise they can't effectively lead or hold elections.
    known_nodes: Vec<SocketAddr>,
    leader_node: Option<SocketAddr>,
    own_socket: SocketAddr,
    heartbeat: Timer,
}

/// The implementation of the RaftNode. In most use cases, creating a `RaftNode` should just be
/// done via `::new()`.
///
/// ```
/// use raft::RaftNode;
/// use std::io::net::ip::SocketAddr;
/// use std::io::net::ip::IpAddr::Ipv4Addr;
///
/// let my_node = RaftNode::<String>::new(
///     // This node's address.
///     SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11111 },
///     // The list of nodes.
///     vec![
///     SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11112 },
///     SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11113 }
///     ]);
/// ```
impl<T: Encodable + Decodable + Send + Clone> RaftNode<T> {
    /// Creates a new RaftNode with the neighbors specified.
    pub fn new(own_socket: SocketAddr, node_list: Vec<SocketAddr>) -> RaftNode<T> {
        RaftNode {
            state: NodeState::Follower,
            persistent_state: PersistentState {
                current_term: 0, // TODO: Double Check.
                voted_for: 0,    // TODO: Better type?
                log: Vec::<T>::new(),
            },
            volatile_state: VolatileState {
                commit_index: 0,
                last_applied: 0,
            },
            known_nodes: node_list,
            leader_node: None,
            own_socket: own_socket,
            // Blank timer for now.
            heartbeat: Timer::new().unwrap(), // If this fails we're in trouble.
        }
    }
    /// Spins up a thread that listens and handles RPCs.
    pub fn spinup(mut self) -> (Sender<T>, Receiver<T>) {
        let (client_in, node_out) = channel::<T>();
        let (node_in, client_out) = channel::<T>();
        // Need to clone the SocketAddr for lifetime reasons.
        Thread::spawn(move || {
            // Setup the socket, make it not block.
            let mut socket = UdpSocket::bind(self.own_socket)
                .unwrap(); // TODO: Can we do better?
            socket.set_read_timeout(Some(0));
            // Start up a RNG.
            let mut rng = thread_rng();
            self.heartbeat.oneshot(Duration::milliseconds(rng.gen_range::<i64>(HEARTBEAT_MIN, HEARTBEAT_MAX)));
            // We need a read buffer.
            let mut read_buffer = [0; BUFFER_SIZE];
            loop {
                // Read anything from the socket and handle it.
                match socket.recv_from(&mut read_buffer) {
                    Ok((num_read, source)) => {
                        println!("Got something from socket");
                        let string_slice = str::from_utf8(read_buffer.slice_to_mut(num_read))
                            .unwrap(); // TODO: Can we do better?
                        let mut rpc = json::decode::<RemoteProcedureCall<T>>(string_slice)
                            .unwrap(); // TODO: Can we do better?
                        match rpc {
                            AppendEntries { entries: mut entries, .. } => {
                                // TODO: Log data properly.
                                for entry in entries.iter_mut() {
                                    node_in.send(entry.clone());
                                }
                            },
                            RequestVote { .. } => {
                                // TODO: Do something.
                            },
                        }
                    },
                    // TODO: Maybe need to handle this better.
                    Err(_) => (),
                }
                // If there is something in the reciever we need to handle it.
                match (node_out.try_recv(), &self.state) {
                    (Ok(entry), &NodeState::Leader(_)) => {
                        println!("Sending things!");
                        let rpc = AppendEntries {
                            term: 0,
                            leader_id: 0,
                            prev_log_index: 0,
                            entries: vec![entry],
                            leader_commit: 0,
                        };
                        let encoded = json::encode::<RemoteProcedureCall<T>>(&rpc);
                        for node in self.known_nodes.iter() {
                            socket.send_to(encoded.as_bytes(), *node)
                                .unwrap(); // TODO: Can we do better?
                        }
                    },
                    // Need to deal with election issues.
                    (Ok(entry), _) => {
                        println!("Got entry, not leader");
                    },
                    // TODO: Maybe need to handle this better.
                    (Err(_), _) => (),
                }
            }
        });
        (client_in, client_out)
    }
}

/// The RPC calls required by the Raft protocol.
impl<T: Encodable + Decodable + Send + Clone> Raft<T> for RaftNode<T> {
    /// Returns (term, success)
    fn append_entries(term: usize, leader_id: usize, prev_log_index: usize,
                      prev_log_term: usize, entries: Vec<T>,
                      leader_commit: usize) -> (usize, bool) {
        (0, false) // TODO: Implement
    }
    /// Returns (term, voteGranted)
    fn request_vote(term: usize, candidate_id: usize, last_log_index: usize,
                    last_log_term: usize) -> (usize, bool) {
        (0, false) // TODO: Implement
    }
}

/// Data interchange format for RPC calls.
#[derive(RustcEncodable, RustcDecodable, Show, Clone)]
pub enum RemoteProcedureCall<T> {
    AppendEntries {
        term: usize,
        leader_id: usize,
        prev_log_index: usize,
        entries: Vec<T>,
        leader_commit: usize,
    },
    RequestVote {
        term: usize,
        candidate_id: usize,
        last_log_index: usize,
        last_log_term: usize,
    },
}

/// Nodes can either be:
///
///   * A `Follower`, which replicates AppendEntries requests and votes for it's leader.
///   * A `Leader`, which leads the cluster by serving incoming requests, ensuring data is
///     replicated, and issuing heartbeats..
///   * A `Candidate`, which campaigns in an election and may become a `Leader` (if it gets enough
///     votes) or a `Follower`, if it hears from a `Leader`.
pub enum NodeState {
    Follower,
    Leader(LeaderState),
    Candidate,
}

/// Persistent state
/// **Must be updated to stable storage before RPC response.**
pub struct PersistentState<T: Encodable + Decodable + Send + Clone> {
    current_term: usize,
    voted_for: usize, // Better way? Can we use a IpAddr?
    log: Vec<T>,
}

/// Volatile state
pub struct VolatileState {
    commit_index: usize,
    last_applied: usize
}

/// Leader Only
/// **Reinitialized after election.**
pub struct LeaderState {
    next_index: Vec<usize>,
    match_index: Vec<usize>
}

#[cfg(test)]
mod tests {
    use std::io::net::ip::SocketAddr;
    use std::io::net::udp::UdpSocket;
    use std::io::net::ip::IpAddr::Ipv4Addr;
    use std::thread::Thread;
    use rustc_serialize::json;
    use std::sync::mpsc::{channel, Sender, Receiver};
    use std::str;
    use super::*;

    #[test]
    fn basic_test() {
        let send_addr = SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11111 };
        let recieve_addr = SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11112 };
        // Create the node.
        let send_node = RaftNode::<String>::new(
            send_addr, // The node's address.
            vec![
            recieve_addr, // The neighbor.
            ]);
        let recieve_node = RaftNode::<String>::new(
            recieve_addr, // The node's address.
            vec![
            send_addr, // The neighbor.
            ]);
        // Start up the node.
        let (_, log_reciever) = recieve_node.spinup();
        let (log_sender, _) = send_node.spinup();
        // Make a test send to that port.
        let test_value = "foo".to_string();
        println!("About to send");
        log_sender.send(test_value.clone()).unwrap();
        println!("After send!");
        // Get the result.
        let event = log_reciever.recv().unwrap();
        assert_eq!(event, "foo".to_string());
    }

}
<commit_msg>API changes<commit_after>#![crate_name = "raft"]
#![crate_type="lib"]

#![allow(unstable)]
extern crate "rustc-serialize" as rustc_serialize;

use std::io::net::ip::SocketAddr;
use std::io::net::udp::UdpSocket;
use std::io::timer::Timer;
use std::time::Duration;
use std::thread::Thread;
use std::rand::{thread_rng, Rng};
use std::sync::mpsc::{channel, Sender, Receiver};
use std::str;
use rustc_serialize::{json, Encodable, Decodable};
use std::collections::HashMap;

use RemoteProcedureCall::{AppendEntries, RequestVote};
use NodeState::{Follower, Leader, Candidate};

// The maximum size of the read buffer.
const BUFFER_SIZE: usize = 4096;
const HEARTBEAT_MIN: i64 = 150;
const HEARTBEAT_MAX: i64 = 300;

/// The Raft Distributed Consensus Algorithm requires two RPC calls to be available:
///
///   * `append_entries` which is used as both a heartbeat (with no payload) and the primary
///     interface for requests.
///   * `request_vote` which is used by candidates during campaigns to obtain a vote.
///
pub trait Raft<T: Encodable + Decodable + Send + Clone> {
    /// Returns (term, success)
    fn append_entries(term: u64, leader_id: u64, prev_log_index: u64,
                      prev_log_term: u64, entries: Vec<T>,
                      leader_commit: u64) -> (u64, bool);

    /// Returns (term, voteGranted)
    fn request_vote(term: u64, candidate_id: u64, last_log_index: u64,
                    last_log_term: u64) -> (u64, bool);
}

/// A `RaftNode` acts as a replicated state machine. The server's role in the cluster depends on it's
/// own status. It will maintain both volatile state (which can be safely lost) and persistent
/// state (which must be carefully stored and kept safe).
pub struct RaftNode<T: Encodable + Decodable + Send + Clone> {
    // Raft related.
    state: NodeState,
    persistent_state: PersistentState<T>,
    volatile_state: VolatileState,
    // Auxilary Data.
    // TODO: This should probably be split off.
    // All nodes need to know this otherwise they can't effectively lead or hold elections.
    leader_node: Option<SocketAddr>,
    own_id: u64,
    nodes: HashMap<u64, SocketAddr>,
    heartbeat: Timer,
}

/// The implementation of the RaftNode. In most use cases, creating a `RaftNode` should just be
/// done via `::new()`.
///
/// ```
/// use raft::RaftNode;
/// use std::io::net::ip::SocketAddr;
/// use std::io::net::ip::IpAddr::Ipv4Addr;
///
/// let my_node = RaftNode::<String>::new(
///     // This node's address.
///     SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11111 },
///     // The list of nodes.
///     vec![
///     SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11112 },
///     SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11113 }
///     ]);
/// ```
impl<T: Encodable + Decodable + Send + Clone> RaftNode<T> {
    /// Creates a new RaftNode with the neighbors specified. `id` should be a valid index into
    /// `nodes`. The idea is that you can use the same `nodes` on all of the clients and only vary
    /// `id`.
    pub fn new(id: u64, nodes: HashMap<u64, SocketAddr>) -> RaftNode<T> {
        // TODO: Check index.
        RaftNode {
            state: Follower,
            persistent_state: PersistentState {
                current_term: 0, // TODO: Double Check.
                voted_for: 0,    // TODO: Better type?
                log: Vec::<T>::new(),
            },
            volatile_state: VolatileState {
                commit_index: 0,
                last_applied: 0,
            },
            leader_node: None,
            own_id: id,
            nodes: nodes,
            // Blank timer for now.
            heartbeat: Timer::new().unwrap(), // If this fails we're in trouble.
        }
    }
    /// Spins up a thread that listens and handles RPCs.
    pub fn spinup(mut self) -> (Sender<T>, Receiver<T>) {
        let (client_in, node_out) = channel::<T>();
        let (node_in, client_out) = channel::<T>();
        // Need to clone the SocketAddr for lifetime reasons.
        Thread::spawn(move || {
            // Setup the socket, make it not block.
            let own_socket_addr = self.nodes.get(&mut self.own_id)
                .unwrap(); // TODO: Can we do better?
            let mut socket = UdpSocket::bind(*own_socket_addr)
                .unwrap(); // TODO: Can we do better?
            socket.set_read_timeout(Some(0));
            // Start up a RNG.
            let mut rng = thread_rng();
            self.heartbeat.oneshot(Duration::milliseconds(rng.gen_range::<i64>(HEARTBEAT_MIN, HEARTBEAT_MAX)));
            // We need a read buffer.
            let mut read_buffer = [0; BUFFER_SIZE];
            loop {
                // Read anything from the socket and handle it.
                match socket.recv_from(&mut read_buffer) {
                    Ok((num_read, source)) => {
                        println!("Got something from socket");
                        let string_slice = str::from_utf8(read_buffer.slice_to_mut(num_read))
                            .unwrap(); // TODO: Can we do better?
                        let mut rpc = json::decode::<RemoteProcedureCall<T>>(string_slice)
                            .unwrap(); // TODO: Can we do better?
                        match rpc {
                            AppendEntries { entries: mut entries, .. } => {
                                // TODO: Log data properly.
                                for entry in entries.iter_mut() {
                                    node_in.send(entry.clone());
                                }
                            },
                            RequestVote { .. } => {
                                // TODO: Do something.
                            },
                        }
                    },
                    // TODO: Maybe need to handle this better.
                    Err(_) => (),
                }
                // If there is something in the reciever we need to handle it.
                match (node_out.try_recv(), &self.state) {
                    (Ok(entry), &NodeState::Leader(_)) => {
                        println!("Sending things!");
                        let rpc = AppendEntries {
                            term: 0,
                            leader_id: 0,
                            prev_log_index: 0,
                            entries: vec![entry],
                            leader_commit: 0,
                        };
                        let encoded = json::encode::<RemoteProcedureCall<T>>(&rpc)
                            .unwrap(); // TODO: Can we do better?
                        for (id, node) in self.nodes.iter() {
                            println!("Sending to {} on {}", id, node);
                            socket.send_to(encoded.as_bytes(), *node)
                                .unwrap(); // TODO: Can we do better?
                        }
                    },
                    // Need to deal with election issues.
                    (Ok(entry), _) => {
                        println!("Got entry, not leader");
                    },
                    // TODO: Maybe need to handle this better.
                    (Err(_), _) => (),
                }
            }
        });
        (client_in, client_out)
    }
}

/// The RPC calls required by the Raft protocol.
impl<T: Encodable + Decodable + Send + Clone> Raft<T> for RaftNode<T> {
    /// Returns (term, success)
    fn append_entries(term: u64, leader_id: u64, prev_log_index: u64,
                      prev_log_term: u64, entries: Vec<T>,
                      leader_commit: u64) -> (u64, bool) {
        (0, false) // TODO: Implement
    }
    /// Returns (term, voteGranted)
    fn request_vote(term: u64, candidate_id: u64, last_log_index: u64,
                    last_log_term: u64) -> (u64, bool) {
        (0, false) // TODO: Implement
    }
}

/// Data interchange format for RPC calls.
#[derive(RustcEncodable, RustcDecodable, Show, Clone)]
pub enum RemoteProcedureCall<T> {
    AppendEntries {
        term: u64,
        leader_id: u64,
        prev_log_index: u64,
        entries: Vec<T>,
        leader_commit: u64,
    },
    RequestVote {
        term: u64,
        candidate_id: u64,
        last_log_index: u64,
        last_log_term: u64,
    },
}

/// Nodes can either be:
///
///   * A `Follower`, which replicates AppendEntries requests and votes for it's leader.
///   * A `Leader`, which leads the cluster by serving incoming requests, ensuring data is
///     replicated, and issuing heartbeats..
///   * A `Candidate`, which campaigns in an election and may become a `Leader` (if it gets enough
///     votes) or a `Follower`, if it hears from a `Leader`.
pub enum NodeState {
    Follower,
    Leader(LeaderState),
    Candidate,
}

/// Persistent state
/// **Must be updated to stable storage before RPC response.**
pub struct PersistentState<T: Encodable + Decodable + Send + Clone> {
    current_term: u64,
    voted_for: u64, // Better way? Can we use a IpAddr?
    log: Vec<T>,
}

/// Volatile state
pub struct VolatileState {
    commit_index: u64,
    last_applied: u64
}

/// Leader Only
/// **Reinitialized after election.**
pub struct LeaderState {
    next_index: Vec<u64>,
    match_index: Vec<u64>
}

#[cfg(test)]
mod tests {
    use std::io::net::ip::SocketAddr;
    use std::io::net::udp::UdpSocket;
    use std::io::net::ip::IpAddr::Ipv4Addr;
    use std::thread::Thread;
    use std::collections::HashMap;
    use rustc_serialize::json;
    use std::sync::mpsc::{channel, Sender, Receiver};
    use std::str;
    use super::*;

    #[test]
    fn basic_test() {
        let mut nodes = HashMap::new();
        nodes.insert(1, SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11111 });
        nodes.insert(2, SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11112 });
        // Create the nodes.
        let send_node = RaftNode::<String>::new(1, nodes.clone());
        let recieve_node = RaftNode::<String>::new(2, nodes.clone());
        // Start up the node.
        let (_, log_reciever) = recieve_node.spinup();
        let (log_sender, _) = send_node.spinup();
        // Make a test send to that port.
        let test_value = "foo".to_string();
        log_sender.send(test_value.clone()).unwrap();
        // Get the result.
        let event = log_reciever.recv().unwrap();
        assert_eq!(event, "foo".to_string());
    }

}
<|endoftext|>


<commit_before>#![recursion_limit = "1024"]

mod async_rt;
mod backend;
mod codec;
mod dealer;
mod endpoint;
mod error;
mod fair_queue;
mod message;
mod r#pub;
mod pull;
mod push;
mod rep;
mod req;
mod router;
mod sub;
mod task_handle;
mod transport;
pub mod util;

#[doc(hidden)]
pub mod __async_rt {
    //! DO NOT USE! PRIVATE IMPLEMENTATION, EXPOSED ONLY FOR INTEGRATION TESTS.
    pub use super::async_rt::*;
}

pub use crate::dealer::*;
pub use crate::endpoint::{Endpoint, Host, Transport, TryIntoEndpoint};
pub use crate::error::{ZmqError, ZmqResult};
pub use crate::pull::*;
pub use crate::push::*;
pub use crate::r#pub::*;
pub use crate::rep::*;
pub use crate::req::*;
pub use crate::router::*;
pub use crate::sub::*;
pub use message::*;

use crate::codec::*;
use crate::transport::AcceptStopHandle;
use util::PeerIdentity;

#[macro_use]
extern crate enum_primitive_derive;

use async_trait::async_trait;
use asynchronous_codec::FramedWrite;
use futures::channel::mpsc;
use futures::FutureExt;
use num_traits::ToPrimitive;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fmt::{Debug, Display};
use std::sync::Arc;

#[allow(clippy::upper_case_acronyms)]
#[derive(Clone, Copy, Debug, PartialEq, Primitive)]
pub enum SocketType {
    PAIR = 0,
    PUB = 1,
    SUB = 2,
    REQ = 3,
    REP = 4,
    DEALER = 5,
    ROUTER = 6,
    PULL = 7,
    PUSH = 8,
    XPUB = 9,
    XSUB = 10,
    STREAM = 11,
}

impl TryFrom<&str> for SocketType {
    type Error = ZmqError;

    fn try_from(s: &str) -> Result<Self, ZmqError> {
        Ok(match s {
            "PAIR" => SocketType::PAIR,
            "PUB" => SocketType::PUB,
            "SUB" => SocketType::SUB,
            "REQ" => SocketType::REQ,
            "REP" => SocketType::REP,
            "DEALER" => SocketType::DEALER,
            "ROUTER" => SocketType::ROUTER,
            "PULL" => SocketType::PULL,
            "PUSH" => SocketType::PUSH,
            "XPUB" => SocketType::XPUB,
            "XSUB" => SocketType::XSUB,
            "STREAM" => SocketType::STREAM,
            _ => return Err(ZmqError::Other("Unknown socket type")),
        })
    }
}

impl Display for SocketType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SocketType::PAIR => write!(f, "PAIR"),
            SocketType::PUB => write!(f, "PUB"),
            SocketType::SUB => write!(f, "SUB"),
            SocketType::REQ => write!(f, "REQ"),
            SocketType::REP => write!(f, "REP"),
            SocketType::DEALER => write!(f, "DEALER"),
            SocketType::ROUTER => write!(f, "ROUTER"),
            SocketType::PULL => write!(f, "PULL"),
            SocketType::PUSH => write!(f, "PUSH"),
            SocketType::XPUB => write!(f, "XPUB"),
            SocketType::XSUB => write!(f, "XSUB"),
            SocketType::STREAM => write!(f, "STREAM"),
        }
    }
}

#[derive(Debug)]
pub enum SocketEvent {
    Connected(Endpoint, PeerIdentity),
    ConnectDelayed,
    ConnectRetried,
    Listening(Endpoint),
    Accepted(Endpoint, PeerIdentity),
    AcceptFailed(ZmqError),
    Closed,
    CloseFailed,
    Disconnected(PeerIdentity),
}

pub struct SocketOptions {
    pub(crate) peer_id: Option<PeerIdentity>,
}

impl SocketOptions {
    pub fn peer_identity(&mut self, peer_id: PeerIdentity) -> &mut Self {
        self.peer_id = Some(peer_id);
        self
    }
}

impl Default for SocketOptions {
    fn default() -> Self {
        Self { peer_id: None }
    }
}

#[async_trait]
pub trait MultiPeerBackend: SocketBackend {
    /// This should not be public..
    /// Find a better way of doing this

    async fn peer_connected(self: Arc<Self>, peer_id: &PeerIdentity, io: FramedIo);
    fn peer_disconnected(&self, peer_id: &PeerIdentity);
}

pub trait SocketBackend: Send + Sync {
    fn socket_type(&self) -> SocketType;
    fn socket_options(&self) -> &SocketOptions;
    fn shutdown(&self);
    fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>>;
}

#[async_trait]
pub trait SocketRecv {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage>;
}

#[async_trait]
pub trait SocketSend {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()>;
}

/// Marker trait that express the fact that only certain types of sockets might be used
/// in [proxy] function as a capture parameter
pub trait CaptureSocket: SocketSend {}

#[async_trait]
pub trait Socket: Sized + Send {
    fn new() -> Self {
        Self::with_options(SocketOptions::default())
    }

    fn with_options(options: SocketOptions) -> Self;

    fn backend(&self) -> Arc<dyn MultiPeerBackend>;

    /// Binds to the endpoint and starts a coroutine to accept new connections
    /// on it.
    ///
    /// Returns the endpoint resolved to the exact bound location if applicable
    /// (port # resolved, for example).
    async fn bind(&mut self, endpoint: &str) -> ZmqResult<Endpoint> {
        let endpoint = endpoint.try_into()?;

        let cloned_backend = self.backend();
        let cback = move |result| {
            let cloned_backend = cloned_backend.clone();
            async move {
                let result = match result {
                    Ok((socket, endpoint)) => {
                        match util::peer_connected(socket, cloned_backend.clone()).await {
                            Ok(peer_id) => Ok((endpoint, peer_id)),
                            Err(e) => Err(e),
                        }
                    }
                    Err(e) => Err(e),
                };
                match result {
                    Ok((endpoint, peer_id)) => {
                        if let Some(monitor) = cloned_backend.monitor().lock().as_mut() {
                            let _ = monitor.try_send(SocketEvent::Accepted(endpoint, peer_id));
                        }
                    }
                    Err(e) => {
                        if let Some(monitor) = cloned_backend.monitor().lock().as_mut() {
                            let _ = monitor.try_send(SocketEvent::AcceptFailed(e));
                        }
                    }
                }
            }
        };

        let (endpoint, stop_handle) = transport::begin_accept(endpoint, cback).await?;

        if let Some(monitor) = self.backend().monitor().lock().as_mut() {
            let _ = monitor.try_send(SocketEvent::Listening(endpoint.clone()));
        }

        self.binds().insert(endpoint.clone(), stop_handle);
        Ok(endpoint)
    }

    fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle>;

    /// Unbinds the endpoint, blocking until the associated endpoint is no
    /// longer in use
    ///
    /// # Errors
    /// May give a `ZmqError::NoSuchBind` if `endpoint` isn't bound. May also
    /// give any other zmq errors encountered when attempting to disconnect
    async fn unbind(&mut self, endpoint: Endpoint) -> ZmqResult<()> {
        let stop_handle = self.binds().remove(&endpoint);
        let stop_handle = stop_handle.ok_or(ZmqError::NoSuchBind(endpoint))?;
        stop_handle.0.shutdown().await
    }

    /// Unbinds all bound endpoints, blocking until finished.
    async fn unbind_all(&mut self) -> Vec<ZmqError> {
        let mut errs = Vec::new();
        let endpoints: Vec<_> = self
            .binds()
            .iter()
            .map(|(endpoint, _)| endpoint.clone())
            .collect();
        for endpoint in endpoints {
            if let Err(err) = self.unbind(endpoint).await {
                errs.push(err);
            }
        }
        errs
    }

    /// Connects to the given endpoint.
    async fn connect(&mut self, endpoint: &str) -> ZmqResult<()> {
        let backend = self.backend();
        let endpoint = endpoint.try_into()?;

        let result = match util::connect_forever(endpoint).await {
            Ok((socket, endpoint)) => match util::peer_connected(socket, backend).await {
                Ok(peer_id) => Ok((endpoint, peer_id)),
                Err(e) => Err(e),
            },
            Err(e) => Err(e),
        };
        match result {
            Ok((endpoint, peer_id)) => {
                if let Some(monitor) = self.backend().monitor().lock().as_mut() {
                    let _ = monitor.try_send(SocketEvent::Connected(endpoint, peer_id));
                }
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    /// Creates and setups new socket monitor
    ///
    /// Subsequent calls to this method each create a new monitor channel.
    /// Sender side of previous one is dropped.
    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent>;

    // TODO: async fn connections(&self) -> ?

    /// Disconnects from the given endpoint, blocking until finished.
    ///
    /// # Errors
    /// May give a `ZmqError::NoSuchConnection` if `endpoint` isn't connected.
    /// May also give any other zmq errors encountered when attempting to
    /// disconnect.
    // TODO: async fn disconnect(&mut self, endpoint: impl TryIntoEndpoint + 'async_trait) ->
    // ZmqResult<()>;

    /// Disconnects all connecttions, blocking until finished.
    // TODO: async fn disconnect_all(&mut self) -> ZmqResult<()>;

    /// Closes the socket, blocking until all associated binds are closed.
    /// This is equivalent to `drop()`, but with the benefit of blocking until
    /// resources are released, and getting any underlying errors.
    ///
    /// Returns any encountered errors.
    // TODO: Call disconnect_all() when added
    async fn close(mut self) -> Vec<ZmqError> {
        // self.disconnect_all().await?;
        self.unbind_all().await
    }
}

pub async fn proxy<Frontend: SocketSend + SocketRecv, Backend: SocketSend + SocketRecv>(
    mut frontend: Frontend,
    mut backend: Backend,
    mut capture: Option<Box<dyn CaptureSocket>>,
) -> ZmqResult<()> {
    loop {
        futures::select! {
            frontend_mess = frontend.recv().fuse() => {
                match frontend_mess {
                    Ok(message) => {
                        if let Some(capture) = &mut capture {
                            capture.send(message.clone()).await?;
                        }
                        backend.send(message).await?;
                    }
                    Err(_) => {
                        todo!()
                    }
                }
            },
            backend_mess = backend.recv().fuse() => {
                match backend_mess {
                    Ok(message) => {
                        if let Some(capture) = &mut capture {
                            capture.send(message.clone()).await?;
                        }
                        frontend.send(message).await?;
                    }
                    Err(_) => {
                        todo!()
                    }
                }
            }
        };
    }
}

pub mod prelude {
    //! Re-exports important traits. Consider glob-importing.

    pub use crate::{Socket, SocketRecv, SocketSend, TryIntoEndpoint};
}
<commit_msg>Derive Default to make clippy happy<commit_after>#![recursion_limit = "1024"]

mod async_rt;
mod backend;
mod codec;
mod dealer;
mod endpoint;
mod error;
mod fair_queue;
mod message;
mod r#pub;
mod pull;
mod push;
mod rep;
mod req;
mod router;
mod sub;
mod task_handle;
mod transport;
pub mod util;

#[doc(hidden)]
pub mod __async_rt {
    //! DO NOT USE! PRIVATE IMPLEMENTATION, EXPOSED ONLY FOR INTEGRATION TESTS.
    pub use super::async_rt::*;
}

pub use crate::dealer::*;
pub use crate::endpoint::{Endpoint, Host, Transport, TryIntoEndpoint};
pub use crate::error::{ZmqError, ZmqResult};
pub use crate::pull::*;
pub use crate::push::*;
pub use crate::r#pub::*;
pub use crate::rep::*;
pub use crate::req::*;
pub use crate::router::*;
pub use crate::sub::*;
pub use message::*;

use crate::codec::*;
use crate::transport::AcceptStopHandle;
use util::PeerIdentity;

#[macro_use]
extern crate enum_primitive_derive;

use async_trait::async_trait;
use asynchronous_codec::FramedWrite;
use futures::channel::mpsc;
use futures::FutureExt;
use num_traits::ToPrimitive;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fmt::{Debug, Display};
use std::sync::Arc;

#[allow(clippy::upper_case_acronyms)]
#[derive(Clone, Copy, Debug, PartialEq, Primitive)]
pub enum SocketType {
    PAIR = 0,
    PUB = 1,
    SUB = 2,
    REQ = 3,
    REP = 4,
    DEALER = 5,
    ROUTER = 6,
    PULL = 7,
    PUSH = 8,
    XPUB = 9,
    XSUB = 10,
    STREAM = 11,
}

impl TryFrom<&str> for SocketType {
    type Error = ZmqError;

    fn try_from(s: &str) -> Result<Self, ZmqError> {
        Ok(match s {
            "PAIR" => SocketType::PAIR,
            "PUB" => SocketType::PUB,
            "SUB" => SocketType::SUB,
            "REQ" => SocketType::REQ,
            "REP" => SocketType::REP,
            "DEALER" => SocketType::DEALER,
            "ROUTER" => SocketType::ROUTER,
            "PULL" => SocketType::PULL,
            "PUSH" => SocketType::PUSH,
            "XPUB" => SocketType::XPUB,
            "XSUB" => SocketType::XSUB,
            "STREAM" => SocketType::STREAM,
            _ => return Err(ZmqError::Other("Unknown socket type")),
        })
    }
}

impl Display for SocketType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SocketType::PAIR => write!(f, "PAIR"),
            SocketType::PUB => write!(f, "PUB"),
            SocketType::SUB => write!(f, "SUB"),
            SocketType::REQ => write!(f, "REQ"),
            SocketType::REP => write!(f, "REP"),
            SocketType::DEALER => write!(f, "DEALER"),
            SocketType::ROUTER => write!(f, "ROUTER"),
            SocketType::PULL => write!(f, "PULL"),
            SocketType::PUSH => write!(f, "PUSH"),
            SocketType::XPUB => write!(f, "XPUB"),
            SocketType::XSUB => write!(f, "XSUB"),
            SocketType::STREAM => write!(f, "STREAM"),
        }
    }
}

#[derive(Debug)]
pub enum SocketEvent {
    Connected(Endpoint, PeerIdentity),
    ConnectDelayed,
    ConnectRetried,
    Listening(Endpoint),
    Accepted(Endpoint, PeerIdentity),
    AcceptFailed(ZmqError),
    Closed,
    CloseFailed,
    Disconnected(PeerIdentity),
}

#[derive(Default)]
pub struct SocketOptions {
    pub(crate) peer_id: Option<PeerIdentity>,
}

impl SocketOptions {
    pub fn peer_identity(&mut self, peer_id: PeerIdentity) -> &mut Self {
        self.peer_id = Some(peer_id);
        self
    }
}

#[async_trait]
pub trait MultiPeerBackend: SocketBackend {
    /// This should not be public..
    /// Find a better way of doing this

    async fn peer_connected(self: Arc<Self>, peer_id: &PeerIdentity, io: FramedIo);
    fn peer_disconnected(&self, peer_id: &PeerIdentity);
}

pub trait SocketBackend: Send + Sync {
    fn socket_type(&self) -> SocketType;
    fn socket_options(&self) -> &SocketOptions;
    fn shutdown(&self);
    fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>>;
}

#[async_trait]
pub trait SocketRecv {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage>;
}

#[async_trait]
pub trait SocketSend {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()>;
}

/// Marker trait that express the fact that only certain types of sockets might be used
/// in [proxy] function as a capture parameter
pub trait CaptureSocket: SocketSend {}

#[async_trait]
pub trait Socket: Sized + Send {
    fn new() -> Self {
        Self::with_options(SocketOptions::default())
    }

    fn with_options(options: SocketOptions) -> Self;

    fn backend(&self) -> Arc<dyn MultiPeerBackend>;

    /// Binds to the endpoint and starts a coroutine to accept new connections
    /// on it.
    ///
    /// Returns the endpoint resolved to the exact bound location if applicable
    /// (port # resolved, for example).
    async fn bind(&mut self, endpoint: &str) -> ZmqResult<Endpoint> {
        let endpoint = endpoint.try_into()?;

        let cloned_backend = self.backend();
        let cback = move |result| {
            let cloned_backend = cloned_backend.clone();
            async move {
                let result = match result {
                    Ok((socket, endpoint)) => {
                        match util::peer_connected(socket, cloned_backend.clone()).await {
                            Ok(peer_id) => Ok((endpoint, peer_id)),
                            Err(e) => Err(e),
                        }
                    }
                    Err(e) => Err(e),
                };
                match result {
                    Ok((endpoint, peer_id)) => {
                        if let Some(monitor) = cloned_backend.monitor().lock().as_mut() {
                            let _ = monitor.try_send(SocketEvent::Accepted(endpoint, peer_id));
                        }
                    }
                    Err(e) => {
                        if let Some(monitor) = cloned_backend.monitor().lock().as_mut() {
                            let _ = monitor.try_send(SocketEvent::AcceptFailed(e));
                        }
                    }
                }
            }
        };

        let (endpoint, stop_handle) = transport::begin_accept(endpoint, cback).await?;

        if let Some(monitor) = self.backend().monitor().lock().as_mut() {
            let _ = monitor.try_send(SocketEvent::Listening(endpoint.clone()));
        }

        self.binds().insert(endpoint.clone(), stop_handle);
        Ok(endpoint)
    }

    fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle>;

    /// Unbinds the endpoint, blocking until the associated endpoint is no
    /// longer in use
    ///
    /// # Errors
    /// May give a `ZmqError::NoSuchBind` if `endpoint` isn't bound. May also
    /// give any other zmq errors encountered when attempting to disconnect
    async fn unbind(&mut self, endpoint: Endpoint) -> ZmqResult<()> {
        let stop_handle = self.binds().remove(&endpoint);
        let stop_handle = stop_handle.ok_or(ZmqError::NoSuchBind(endpoint))?;
        stop_handle.0.shutdown().await
    }

    /// Unbinds all bound endpoints, blocking until finished.
    async fn unbind_all(&mut self) -> Vec<ZmqError> {
        let mut errs = Vec::new();
        let endpoints: Vec<_> = self
            .binds()
            .iter()
            .map(|(endpoint, _)| endpoint.clone())
            .collect();
        for endpoint in endpoints {
            if let Err(err) = self.unbind(endpoint).await {
                errs.push(err);
            }
        }
        errs
    }

    /// Connects to the given endpoint.
    async fn connect(&mut self, endpoint: &str) -> ZmqResult<()> {
        let backend = self.backend();
        let endpoint = endpoint.try_into()?;

        let result = match util::connect_forever(endpoint).await {
            Ok((socket, endpoint)) => match util::peer_connected(socket, backend).await {
                Ok(peer_id) => Ok((endpoint, peer_id)),
                Err(e) => Err(e),
            },
            Err(e) => Err(e),
        };
        match result {
            Ok((endpoint, peer_id)) => {
                if let Some(monitor) = self.backend().monitor().lock().as_mut() {
                    let _ = monitor.try_send(SocketEvent::Connected(endpoint, peer_id));
                }
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    /// Creates and setups new socket monitor
    ///
    /// Subsequent calls to this method each create a new monitor channel.
    /// Sender side of previous one is dropped.
    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent>;

    // TODO: async fn connections(&self) -> ?

    /// Disconnects from the given endpoint, blocking until finished.
    ///
    /// # Errors
    /// May give a `ZmqError::NoSuchConnection` if `endpoint` isn't connected.
    /// May also give any other zmq errors encountered when attempting to
    /// disconnect.
    // TODO: async fn disconnect(&mut self, endpoint: impl TryIntoEndpoint + 'async_trait) ->
    // ZmqResult<()>;

    /// Disconnects all connecttions, blocking until finished.
    // TODO: async fn disconnect_all(&mut self) -> ZmqResult<()>;

    /// Closes the socket, blocking until all associated binds are closed.
    /// This is equivalent to `drop()`, but with the benefit of blocking until
    /// resources are released, and getting any underlying errors.
    ///
    /// Returns any encountered errors.
    // TODO: Call disconnect_all() when added
    async fn close(mut self) -> Vec<ZmqError> {
        // self.disconnect_all().await?;
        self.unbind_all().await
    }
}

pub async fn proxy<Frontend: SocketSend + SocketRecv, Backend: SocketSend + SocketRecv>(
    mut frontend: Frontend,
    mut backend: Backend,
    mut capture: Option<Box<dyn CaptureSocket>>,
) -> ZmqResult<()> {
    loop {
        futures::select! {
            frontend_mess = frontend.recv().fuse() => {
                match frontend_mess {
                    Ok(message) => {
                        if let Some(capture) = &mut capture {
                            capture.send(message.clone()).await?;
                        }
                        backend.send(message).await?;
                    }
                    Err(_) => {
                        todo!()
                    }
                }
            },
            backend_mess = backend.recv().fuse() => {
                match backend_mess {
                    Ok(message) => {
                        if let Some(capture) = &mut capture {
                            capture.send(message.clone()).await?;
                        }
                        frontend.send(message).await?;
                    }
                    Err(_) => {
                        todo!()
                    }
                }
            }
        };
    }
}

pub mod prelude {
    //! Re-exports important traits. Consider glob-importing.

    pub use crate::{Socket, SocketRecv, SocketSend, TryIntoEndpoint};
}
<|endoftext|>