<commit_before>extern crate hyper;
extern crate libc;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
extern crate log4rs;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;

use hyper::Client;
use hyper::client::response::Response;
use hyper::header::{ Headers, ContentType };
use hyper::mime::{ Mime, TopLevel, SubLevel };
use hyper::status::StatusCode;
use libc::{ uint8_t, size_t };
use log::LogLevelFilter;
use log4rs::append::file::FileAppender;
use log4rs::encode::pattern::PatternEncoder;
use log4rs::config::{ Appender, Config, Logger, Root };
use serde::Serialize;
use std::collections::HashMap;
use std::env::var;
use std::io::Read;
use std::slice;
use std::sync::{ Once, ONCE_INIT, Mutex };

const BASE_URL: &'static str = "http://localhost:8080/k2/ctapi/";

static OK: i8 = 0;
static ERR_INVALID: i8 = -1;
static ERR_HOST: i8 = -127;

static INIT: Once = ONCE_INIT;

lazy_static! {
    static ref MAP: Mutex<HashMap<u16, u16>> = Mutex::new(HashMap::new());
}

#[derive(Serialize)]
struct Empty();

#[derive(Serialize)]
struct RequestData {
    ctn: u16,
    dad: u8,
    sad: u8,
    lenc: usize,
    command: Vec<u8>,
    lenr: usize
}

#[allow(non_snake_case)]
#[derive(Deserialize)]
struct ResponseData {
    dad: u8,
    sad: u8,
    lenr: usize,
    response: Vec<u8>,
    responseCode: i8
}

macro_rules! post_request {
    ($path:expr) => (post_request($path, &Empty{}));
    ($path:expr, $data:expr) => (post_request($path, $data));
}

#[no_mangle]
#[allow(non_snake_case, unused_must_use)]
pub extern fn CT_init(ctn: u16, pn: u16) -> i8 {
    init_logging();

    debug!("CT_init( ctn: {}; pn: {} )", ctn, pn);

    // Do we know this CTN?
    if MAP.lock().unwrap().contains_key(&ctn) {
        return ERR_INVALID
    }

    // Build the request URL
    let endpoint = "ct_init".to_string();
    let path = endpoint + "/" + &ctn.to_string() + "/" + &pn.to_string();

    debug!("Request path: {}", path);

    // Perform the request
    let mut response = post_request!(&path);

    match response.status {
        StatusCode::Ok => {
            debug!("Response received!");

            // Cast server response
            let mut body = String::new();
            response.read_to_string(&mut body).unwrap();

            let status = body.parse::<i8>().unwrap();
            if status == OK {
                // Store CTN
                MAP.lock().unwrap().insert(ctn, pn);
            }

            debug!("Return status: {}", status);
            status
        },
        _ => {
            error!("Response not OK! Returning {}", ERR_HOST);
            debug!("Response: {:?}", response);
            ERR_HOST
        }
    }
}

#[no_mangle]
#[allow(non_snake_case)]
pub extern fn CT_data(ctn: u16, dad: *mut uint8_t, sad: *mut uint8_t, lenc: size_t, command: *const uint8_t, lenr: *mut size_t, response: *mut uint8_t) -> i8 {
    init_logging();

    debug!("CT_data for ctn: {}", ctn);
    unsafe {
        if !MAP.lock().unwrap().contains_key(&ctn) {
            return ERR_INVALID
        }

        let dad: &mut u8 = &mut *dad;

        let sad: &mut u8 = &mut *sad;

        let lenr: &mut size_t = &mut *lenr;

        let command = slice::from_raw_parts(command, lenc as usize);

        let response = slice::from_raw_parts_mut(response, *lenr);

        let requestData = RequestData {
            ctn: ctn,
            dad: *dad,
            sad: *sad,
            lenc: lenc,
            command: command.to_vec(),
            lenr: *lenr
        };

        debug!("Additonal request parameter: dad: {}, sad: {}, lenc: {}, command: {:?}, lenr: {}, response.len(): {}", dad, sad, lenc, command, lenr, response.len());

        let pn = MAP.lock().unwrap();
        let pn = pn.get(&ctn).unwrap();
        let endpoint = "ct_data".to_string();
        let path = endpoint + "/" + &ctn.to_string() + "/" + &pn.to_string();

        debug!("Request path: {}", path);

        let mut http_response = post_request(&path, &requestData);

        match http_response.status {
            StatusCode::Ok => {
                debug!("Response received!");

                // decode server response
                let mut body = String::new();
                http_response.read_to_string(&mut body).unwrap();
                let responseData: ResponseData = serde_json::from_str(&body).unwrap();

                if responseData.responseCode == OK {

                    debug!("Body data from response: dad: {}, sad: {}, lenr: {}", responseData.dad, responseData.sad, responseData.lenr);

                    *dad = responseData.dad;
                    *sad = responseData.sad;
                    *lenr = responseData.lenr;

                    debug!("Content of response array");
                    for (place, element) in response.iter_mut().zip(responseData.response.iter()) {
                        debug!("[{}] {}", place, element);
                        *place = *element;
                    }
                }
                debug!("Return status: {}", responseData.responseCode);
                return responseData.responseCode;
            },
            _ => {
                error!("Response not OK! Returning {}", ERR_HOST);
                debug!("Response: {:?}", response);
                ERR_HOST
            }
        }
    }
}

#[no_mangle]
#[allow(non_snake_case)]
pub extern fn CT_close(ctn: u16) -> i8 {
    init_logging();

    debug!("CT_close( ctn: {} )", ctn);

    // Do we know this CTN?
    if !MAP.lock().unwrap().contains_key(&ctn) {
        return ERR_INVALID
    }

    // Build the request URL
    let pn = MAP.lock().unwrap().get(&ctn).unwrap().clone();
    let endpoint = "ct_close".to_string();
    let path = endpoint + "/" + &ctn.to_string() + "/" + &pn.to_string();

    debug!("Request path: {}", path);

    // Perform the request
    let mut response = post_request!(&path);

    match response.status {
        StatusCode::Ok => {
            debug!("Response received!");

            // Cast server response
            let mut body = String::new();
            response.read_to_string(&mut body).unwrap();

            let status = body.parse::<i8>().unwrap();
            if status == OK {
                // Remove CTN
                MAP.lock().unwrap().remove(&ctn);
            }

            debug!("Return status: {}", status);
            status
        },
        _ => {
            error!("Response not OK! Returning {}", ERR_HOST);
            debug!("Response: {:?}", response);
            ERR_HOST
        }
    }
}

fn init_logging() {
   match var("K2_LOG_PATH") {
        Ok(path) => {
            INIT.call_once(|| {
                let file = FileAppender::builder()
                    .encoder(Box::new(PatternEncoder::new("{d} {l} {M}: {m}{n}")))
                    .build(path + "/" + &"k2_basecamp.log".to_string())
                    .unwrap();

                let config = Config::builder()
                    .appender(Appender::builder().build("file", Box::new(file)))
                    .logger(Logger::builder()
                        .appender("file")
                        .additive(false)
                        .build("k2_basecamp", LogLevelFilter::Debug))
                    .build(Root::builder().appender("file").build(LogLevelFilter::Error))
                    .unwrap();

                log4rs::init_config(config).unwrap();
            })
        },
        _ => ()
    }
}

fn post_request<T>(path: &str, payload: &T) -> Response
    where T: Serialize
{
    let base_url = var("K2_BASE_URL").unwrap_or(BASE_URL.to_string());
    let url = base_url + path;

    let client = Client::new();

    let mut headers = Headers::new();
    headers.set(
        ContentType(Mime(TopLevel::Application, SubLevel::Json, vec![]))
    );

    let body = serde_json::to_string(&payload).unwrap();

    return client.post(&url)
        .headers(headers)
        .body(&body[..])
        .send()
        .unwrap();
}
<commit_msg>add http response to error logging<commit_after>extern crate hyper;
extern crate libc;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
extern crate log4rs;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;

use hyper::Client;
use hyper::client::response::Response;
use hyper::header::{ Headers, ContentType };
use hyper::mime::{ Mime, TopLevel, SubLevel };
use hyper::status::StatusCode;
use libc::{ uint8_t, size_t };
use log::LogLevelFilter;
use log4rs::append::file::FileAppender;
use log4rs::encode::pattern::PatternEncoder;
use log4rs::config::{ Appender, Config, Logger, Root };
use serde::Serialize;
use std::collections::HashMap;
use std::env::var;
use std::io::Read;
use std::slice;
use std::sync::{ Once, ONCE_INIT, Mutex };

const BASE_URL: &'static str = "http://localhost:8080/k2/ctapi/";

static OK: i8 = 0;
static ERR_INVALID: i8 = -1;
static ERR_HOST: i8 = -127;

static INIT: Once = ONCE_INIT;

lazy_static! {
    static ref MAP: Mutex<HashMap<u16, u16>> = Mutex::new(HashMap::new());
}

#[derive(Serialize)]
struct Empty();

#[derive(Serialize)]
struct RequestData {
    ctn: u16,
    dad: u8,
    sad: u8,
    lenc: usize,
    command: Vec<u8>,
    lenr: usize
}

#[allow(non_snake_case)]
#[derive(Deserialize)]
struct ResponseData {
    dad: u8,
    sad: u8,
    lenr: usize,
    response: Vec<u8>,
    responseCode: i8
}

macro_rules! post_request {
    ($path:expr) => (post_request($path, &Empty{}));
    ($path:expr, $data:expr) => (post_request($path, $data));
}

#[no_mangle]
#[allow(non_snake_case, unused_must_use)]
pub extern fn CT_init(ctn: u16, pn: u16) -> i8 {
    init_logging();

    debug!("CT_init( ctn: {}; pn: {} )", ctn, pn);

    // Do we know this CTN?
    if MAP.lock().unwrap().contains_key(&ctn) {
        return ERR_INVALID
    }

    // Build the request URL
    let endpoint = "ct_init".to_string();
    let path = endpoint + "/" + &ctn.to_string() + "/" + &pn.to_string();

    debug!("Request path: {}", path);

    // Perform the request
    let mut response = post_request!(&path);

    match response.status {
        StatusCode::Ok => {
            debug!("Response received!");

            // Cast server response
            let mut body = String::new();
            response.read_to_string(&mut body).unwrap();

            let status = body.parse::<i8>().unwrap();
            if status == OK {
                // Store CTN
                MAP.lock().unwrap().insert(ctn, pn);
            }

            debug!("Return status: {}", status);
            status
        },
        _ => {
            error!("Response not OK! Returning {}", ERR_HOST);
            debug!("Response: {:?}", response);
            ERR_HOST
        }
    }
}

#[no_mangle]
#[allow(non_snake_case)]
pub extern fn CT_data(ctn: u16, dad: *mut uint8_t, sad: *mut uint8_t, lenc: size_t, command: *const uint8_t, lenr: *mut size_t, response: *mut uint8_t) -> i8 {
    init_logging();

    debug!("CT_data for ctn: {}", ctn);
    unsafe {
        if !MAP.lock().unwrap().contains_key(&ctn) {
            return ERR_INVALID
        }

        let dad: &mut u8 = &mut *dad;

        let sad: &mut u8 = &mut *sad;

        let lenr: &mut size_t = &mut *lenr;

        let command = slice::from_raw_parts(command, lenc as usize);

        let response = slice::from_raw_parts_mut(response, *lenr);

        let requestData = RequestData {
            ctn: ctn,
            dad: *dad,
            sad: *sad,
            lenc: lenc,
            command: command.to_vec(),
            lenr: *lenr
        };

        debug!("Additonal request parameter: dad: {}, sad: {}, lenc: {}, command: {:?}, lenr: {}, response.len(): {}", dad, sad, lenc, command, lenr, response.len());

        let pn = MAP.lock().unwrap();
        let pn = pn.get(&ctn).unwrap();
        let endpoint = "ct_data".to_string();
        let path = endpoint + "/" + &ctn.to_string() + "/" + &pn.to_string();

        debug!("Request path: {}", path);

        let mut http_response = post_request(&path, &requestData);

        match http_response.status {
            StatusCode::Ok => {
                debug!("Response received!");

                // decode server response
                let mut body = String::new();
                http_response.read_to_string(&mut body).unwrap();
                let responseData: ResponseData = serde_json::from_str(&body).unwrap();

                if responseData.responseCode == OK {

                    debug!("Body data from response: dad: {}, sad: {}, lenr: {}", responseData.dad, responseData.sad, responseData.lenr);

                    *dad = responseData.dad;
                    *sad = responseData.sad;
                    *lenr = responseData.lenr;

                    debug!("Content of response array");
                    for (place, element) in response.iter_mut().zip(responseData.response.iter()) {
                        debug!("[{}] {}", place, element);
                        *place = *element;
                    }
                }
                debug!("Return status: {}", responseData.responseCode);
                return responseData.responseCode;
            },
            _ => {
                error!("Response not OK! Returning {}", ERR_HOST);
                debug!("Response: {:?}", http_response);
                ERR_HOST
            }
        }
    }
}

#[no_mangle]
#[allow(non_snake_case)]
pub extern fn CT_close(ctn: u16) -> i8 {
    init_logging();

    debug!("CT_close( ctn: {} )", ctn);

    // Do we know this CTN?
    if !MAP.lock().unwrap().contains_key(&ctn) {
        return ERR_INVALID
    }

    // Build the request URL
    let pn = MAP.lock().unwrap().get(&ctn).unwrap().clone();
    let endpoint = "ct_close".to_string();
    let path = endpoint + "/" + &ctn.to_string() + "/" + &pn.to_string();

    debug!("Request path: {}", path);

    // Perform the request
    let mut response = post_request!(&path);

    match response.status {
        StatusCode::Ok => {
            debug!("Response received!");

            // Cast server response
            let mut body = String::new();
            response.read_to_string(&mut body).unwrap();

            let status = body.parse::<i8>().unwrap();
            if status == OK {
                // Remove CTN
                MAP.lock().unwrap().remove(&ctn);
            }

            debug!("Return status: {}", status);
            status
        },
        _ => {
            error!("Response not OK! Returning {}", ERR_HOST);
            debug!("Response: {:?}", response);
            ERR_HOST
        }
    }
}

fn init_logging() {
   match var("K2_LOG_PATH") {
        Ok(path) => {
            INIT.call_once(|| {
                let file = FileAppender::builder()
                    .encoder(Box::new(PatternEncoder::new("{d} {l} {M}: {m}{n}")))
                    .build(path + "/" + &"k2_basecamp.log".to_string())
                    .unwrap();

                let config = Config::builder()
                    .appender(Appender::builder().build("file", Box::new(file)))
                    .logger(Logger::builder()
                        .appender("file")
                        .additive(false)
                        .build("k2_basecamp", LogLevelFilter::Debug))
                    .build(Root::builder().appender("file").build(LogLevelFilter::Error))
                    .unwrap();

                log4rs::init_config(config).unwrap();
            })
        },
        _ => ()
    }
}

fn post_request<T>(path: &str, payload: &T) -> Response
    where T: Serialize
{
    let base_url = var("K2_BASE_URL").unwrap_or(BASE_URL.to_string());
    let url = base_url + path;

    let client = Client::new();

    let mut headers = Headers::new();
    headers.set(
        ContentType(Mime(TopLevel::Application, SubLevel::Json, vec![]))
    );

    let body = serde_json::to_string(&payload).unwrap();

    return client.post(&url)
        .headers(headers)
        .body(&body[..])
        .send()
        .unwrap();
}
<|endoftext|>


<commit_before>#![crate_name = "cql"]
//#![comment = "A Rust CQl binary protocol implementation"]
//#![license = "MIT/ASL2"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![feature(custom_derive, core, into_cow)]

pub use client::connect;
pub use def::Consistency;
pub use def::BatchType;
pub use def::CqlValue;
pub use def::CqlValue::CqlFloat;
pub use def::CqlValue::CqlVarchar;
pub use def::CQLList;
pub use def::CQLMap;
pub use def::CQLSet;
pub use def::Query::QueryStr;
pub use def::Query::QueryPrepared;
pub use def::OpcodeResponse;
pub use def::CqlResponseBody;
use std::borrow::IntoCow;
pub use def::RCResult;
pub use def::RCError;

#[macro_export]
macro_rules! try_bo(
    ($call: expr, $msg: expr) => {
        match $call {
            Ok(val) => val,
            Err(self::byteorder::Error::UnexpectedEOF) => return Err($crate::def::RCError::new(format!("{} -> {}", $msg, "Unexpected EOF"), $crate::def::RCErrorType::IOError)),
            Err(self::byteorder::Error::Io(ref err)) => {
            	use std::error::Error;
            	return Err($crate::def::RCError::new(format!("{} -> {}", $msg, err.description()), $crate::def::RCErrorType::IOError))
            }
        };
    }
);

#[macro_export]
macro_rules! try_io(
    ($call: expr, $msg: expr) => {
        match $call {
            Ok(val) => val,
            Err(ref err) => {
            	use std::error::Error;
            	return Err(RCError::new(format!("{} -> {}", $msg, err.description()), RCErrorType::IOError))
            }
        };
    }
);

#[macro_export]
macro_rules! try_rc(
    ($call: expr, $msg: expr) => {
        match $call {
            Ok(val) => val,
            Err(ref err) => return Err($crate::def::RCError::new(format!("{} -> {}", $msg, err.description()), $crate::def::RCErrorType::IOError))
        };
    }
);

macro_rules! try_rc_length(
    ($call: expr, $msg: expr) => {
        match $call {
            Ok(-1) => return Ok(None),
            Ok(val) => val,
            Err(ref err) => return Err($crate::def::RCError::new(format!("{} -> {}", $msg, err.description()), $crate::def::RCErrorType::IOError))
        };
    }
);

macro_rules! try_rc_noption(
    ($call: expr, $msg: expr) => {
        match $call {
            Ok(option) => match option {
                None => return Err($crate::def::RCError::new(format!("{} -> {}", $msg, "No data found (length == -1)"), $crate::def::RCErrorType::IOError)),
                Some(val) => val
            },
            Err(ref err) => return Err($crate::def::RCError::new(format!("{} -> {}", $msg, err.description()), $crate::def::RCErrorType::IOError))
        };
    }
);

macro_rules! sendstr_tuple_void(
    () => {
        ("".into_cow(), "".into_cow())
    }
);

mod def;
mod reader;
mod serialize;
pub mod client;
<commit_msg>Enable unstable library `convert`<commit_after>#![crate_name = "cql"]
//#![comment = "A Rust CQl binary protocol implementation"]
//#![license = "MIT/ASL2"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![feature(custom_derive, core, into_cow, convert)]

pub use client::connect;
pub use def::Consistency;
pub use def::BatchType;
pub use def::CqlValue;
pub use def::CqlValue::CqlFloat;
pub use def::CqlValue::CqlVarchar;
pub use def::CQLList;
pub use def::CQLMap;
pub use def::CQLSet;
pub use def::Query::QueryStr;
pub use def::Query::QueryPrepared;
pub use def::OpcodeResponse;
pub use def::CqlResponseBody;
use std::borrow::IntoCow;
pub use def::RCResult;
pub use def::RCError;

#[macro_export]
macro_rules! try_bo(
    ($call: expr, $msg: expr) => {
        match $call {
            Ok(val) => val,
            Err(self::byteorder::Error::UnexpectedEOF) => return Err($crate::def::RCError::new(format!("{} -> {}", $msg, "Unexpected EOF"), $crate::def::RCErrorType::IOError)),
            Err(self::byteorder::Error::Io(ref err)) => {
            	use std::error::Error;
            	return Err($crate::def::RCError::new(format!("{} -> {}", $msg, err.description()), $crate::def::RCErrorType::IOError))
            }
        };
    }
);

#[macro_export]
macro_rules! try_io(
    ($call: expr, $msg: expr) => {
        match $call {
            Ok(val) => val,
            Err(ref err) => {
            	use std::error::Error;
            	return Err(RCError::new(format!("{} -> {}", $msg, err.description()), RCErrorType::IOError))
            }
        };
    }
);

#[macro_export]
macro_rules! try_rc(
    ($call: expr, $msg: expr) => {
        match $call {
            Ok(val) => val,
            Err(ref err) => return Err($crate::def::RCError::new(format!("{} -> {}", $msg, err.description()), $crate::def::RCErrorType::IOError))
        };
    }
);

macro_rules! try_rc_length(
    ($call: expr, $msg: expr) => {
        match $call {
            Ok(-1) => return Ok(None),
            Ok(val) => val,
            Err(ref err) => return Err($crate::def::RCError::new(format!("{} -> {}", $msg, err.description()), $crate::def::RCErrorType::IOError))
        };
    }
);

macro_rules! try_rc_noption(
    ($call: expr, $msg: expr) => {
        match $call {
            Ok(option) => match option {
                None => return Err($crate::def::RCError::new(format!("{} -> {}", $msg, "No data found (length == -1)"), $crate::def::RCErrorType::IOError)),
                Some(val) => val
            },
            Err(ref err) => return Err($crate::def::RCError::new(format!("{} -> {}", $msg, err.description()), $crate::def::RCErrorType::IOError))
        };
    }
);

macro_rules! sendstr_tuple_void(
    () => {
        ("".into_cow(), "".into_cow())
    }
);

mod def;
mod reader;
mod serialize;
pub mod client;
<|endoftext|>


<commit_before>extern crate x264_sys as ffi;

use std::mem;
use ffi::x264::*;
use std::ptr::null;
use std::os::raw::c_int;

pub struct Picture {
    pic: x264_picture_t,
    height: i32, // to compute the slice dimension
    native: bool,
}

impl Picture {
    /*
    pub fn new() -> Picture {
        let mut pic = unsafe { mem::uninitialized() };

        unsafe { x264_picture_init(&mut pic as *mut x264_picture_t) };

        Picture { pic: pic }
    }
*/
    pub fn from_param(param: &Param) -> Result<Picture, &'static str> {
        let mut pic = unsafe { mem::uninitialized() };

        let ret = unsafe {
            x264_picture_alloc(&mut pic as *mut x264_picture_t,
                               param.par.i_csp,
                               param.par.i_width,
                               param.par.i_height)
        };
        if ret < 0 {
            Err("Allocation Failure")
        } else {
            Ok(Picture {
                   pic: pic,
                   height: param.par.i_height,
                   native: true,
               })
        }
    }

    pub fn as_slice<'a>(&'a self, plane: usize) -> Result<&'a [u8], &'static str> {
        if plane > self.pic.img.i_plane as usize {
            Err("Invalid Argument")
        } else {
            let size = self.height * self.pic.img.i_stride[plane];
            Ok(unsafe { std::slice::from_raw_parts(self.pic.img.plane[plane], size as usize) })
        }
    }

    pub fn as_slice_mut<'a>(&'a mut self, plane: usize) -> Result<&'a mut [u8], &'static str> {
        if plane > self.pic.img.i_plane as usize {
            Err("Invalid Argument")
        } else {
            let size = self.height * self.pic.img.i_stride[plane];
            Ok(unsafe { std::slice::from_raw_parts_mut(self.pic.img.plane[plane], size as usize) })
        }
    }

    pub fn set_timestamp(mut self, pts: i64) -> Picture {
        self.pic.i_pts = pts;
        self
    }
}

impl Drop for Picture {
    fn drop(&mut self) {
        if self.native {
            unsafe { x264_picture_clean(&mut self.pic as *mut x264_picture_t) };
        }
    }
}

// TODO: Provide a builder API instead?
pub struct Param {
    par: x264_param_t,
}

impl Param {
    pub fn new() -> Param {
        let mut par = unsafe { mem::uninitialized() };

        unsafe {
            x264_param_default(&mut par as *mut x264_param_t);
        }

        Param { par: par }
    }
    pub fn default_preset(tune: Option<&str>, preset: Option<&str>) -> Result<Param, &'static str> {
        let mut par = unsafe { mem::uninitialized() };

        match unsafe {
                  x264_param_default_preset(&mut par as *mut x264_param_t,
                                            tune.map_or_else(|| null(),
                                                             |v| v.as_ptr() as *const i8),
                                            preset.map_or_else(|| null(), |v| v.as_ptr()) as
                                            *const i8)
              } {
            -1 => Err("Invalid Argument"),
            0 => Ok(Param { par: par }),
            _ => Err("Unexpected"),
        }
    }
    pub fn apply_profile(mut self, profile: &str) -> Result<Param, &'static str> {
        match unsafe { x264_param_apply_profile(&mut self.par, profile.as_ptr() as *const i8) } {
            -1 => Err("Invalid Argument"),
            0 => Ok(self),
            _ => Err("Unexpected"),
        }
    }
    pub fn param_parse<'a>(mut self, name: &str, value: &str) -> Result<Param, &'static str> {
        match unsafe {
                  x264_param_parse(&mut self.par,
                                   name.as_ptr() as *const i8,
                                   value.as_ptr() as *const i8)
              } {
            -1 => Err("Invalid Argument"),
            0 => Ok(self),
            _ => Err("Unexpected"),
        }
    }
    pub fn set_dimension(mut self, height: usize, width: usize) -> Param {
        self.par.i_height = height as c_int;
        self.par.i_width = width as c_int;

        self
    }
}

// TODO: Expose a NAL abstraction
pub struct NalData {
    vec: Vec<u8>,
}

impl NalData {
    /*
     * x264 functions return x264_nal_t arrays that are valid only until another
     * function of that kind is called.
     *
     * Always copy the data over.
     *
     * TODO: Consider using Bytes as backing store.
     */
    fn from_nals(c_nals: *mut x264_nal_t, nb_nal: usize) -> NalData {
        let mut data = NalData { vec: Vec::new() };

        for i in 0..nb_nal {
            let nal = unsafe { Box::from_raw(c_nals.offset(i as isize)) };

            let payload =
                unsafe { std::slice::from_raw_parts(nal.p_payload, nal.i_payload as usize) };

            data.vec.extend_from_slice(payload);

            mem::forget(payload);
            mem::forget(nal);
        }

        data
    }

    pub fn as_bytes(&self) -> &[u8] {
        self.vec.as_slice()
    }
}

pub struct Encoder {
    enc: *mut x264_t,
}

impl Encoder {
    pub fn open(par: &mut Param) -> Result<Encoder, &'static str> {
        let enc = unsafe { x264_encoder_open(&mut par.par as *mut x264_param_t) };

        if enc.is_null() {
            Err("Out of Memory")
        } else {
            Ok(Encoder { enc: enc })
        }
    }

    pub fn get_headers(&mut self) -> Result<NalData, &'static str> {
        let mut nb_nal: c_int = 0;
        let mut c_nals: *mut x264_nal_t = unsafe { mem::uninitialized() };

        let bytes = unsafe {
            x264_encoder_headers(self.enc,
                                 &mut c_nals as *mut *mut x264_nal_t,
                                 &mut nb_nal as *mut c_int)
        };

        if bytes < 0 {
            Err("Encoding Headers Failed")
        } else {
            Ok(NalData::from_nals(c_nals, nb_nal as usize))
        }
    }

    pub fn encode<'a, P>(&mut self, pic: P) -> Result<Option<(NalData, i64, i64)>, &'static str>
        where P: Into<Option<&'a Picture>>
    {
        let mut pic_out: x264_picture_t = unsafe { mem::uninitialized() };
        let mut c_nals: *mut x264_nal_t = unsafe { mem::uninitialized() };
        let mut nb_nal: c_int = 0;
        let c_pic = pic.into().map_or_else(|| null(), |v| &v.pic as *const x264_picture_t);

        let ret = unsafe {
            x264_encoder_encode(self.enc,
                                &mut c_nals as *mut *mut x264_nal_t,
                                &mut nb_nal as *mut c_int,
                                c_pic as *mut x264_picture_t,
                                &mut pic_out as *mut x264_picture_t)
        };
        if ret < 0 {
            Err("Error encoding")
        } else {
            if nb_nal > 0 {
                let data = NalData::from_nals(c_nals, nb_nal as usize);
                Ok(Some((data, pic_out.i_pts, pic_out.i_dts)))
            } else {
                Ok(None)
            }
        }
    }

    pub fn delayed_frames(&self) -> bool {
        unsafe { x264_encoder_delayed_frames(self.enc) != 0 }
    }
}

impl Drop for Encoder {
    fn drop(&mut self) {
        unsafe { x264_encoder_close(self.enc) };
    }
}

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

    #[test]
    fn test_open() {
        let mut par = Param::new().set_dimension(640, 480);

        let mut enc = Encoder::open(&mut par).unwrap();

        let headers = enc.get_headers().unwrap();

        println!("Headers len {}", headers.as_bytes().len());
    }

    #[test]
    fn test_picture() {
        let par = Param::new().set_dimension(640, 480);
        {
            let mut pic = Picture::from_param(&par).unwrap();
            {
                let p = pic.as_slice_mut(0).unwrap();
                p[0] = 1;
            }
            let p = pic.as_slice(0).unwrap();

            assert_eq!(p[0], 1);
        }
    }

    #[test]
    fn test_encode() {
        let mut par = Param::new().set_dimension(640, 480);
        let mut enc = Encoder::open(&mut par).unwrap();
        let mut pic = Picture::from_param(&par).unwrap();

        let headers = enc.get_headers().unwrap();

        println!("Headers len {}", headers.as_bytes().len());

        for pts in 0..5 {
            pic = pic.set_timestamp(pts as i64);
            let ret = enc.encode(&pic).unwrap();
            match ret {
                Some((_, pts, dts)) => println!("Frame pts {}, dts {}", pts, dts),
                _ => (),
            }
        }

        while enc.delayed_frames() {
            let ret = enc.encode(None).unwrap();
            match ret {
                Some((_, pts, dts)) => println!("Frame pts {}, dts {}", pts, dts),
                _ => (),
            }
        }
    }
}
<commit_msg>param: Use the Into Option pattern here as well<commit_after>extern crate x264_sys as ffi;

use std::mem;
use ffi::x264::*;
use std::ptr::null;
use std::os::raw::c_int;

pub struct Picture {
    pic: x264_picture_t,
    height: i32, // to compute the slice dimension
    native: bool,
}

impl Picture {
    /*
    pub fn new() -> Picture {
        let mut pic = unsafe { mem::uninitialized() };

        unsafe { x264_picture_init(&mut pic as *mut x264_picture_t) };

        Picture { pic: pic }
    }
*/
    pub fn from_param(param: &Param) -> Result<Picture, &'static str> {
        let mut pic = unsafe { mem::uninitialized() };

        let ret = unsafe {
            x264_picture_alloc(&mut pic as *mut x264_picture_t,
                               param.par.i_csp,
                               param.par.i_width,
                               param.par.i_height)
        };
        if ret < 0 {
            Err("Allocation Failure")
        } else {
            Ok(Picture {
                   pic: pic,
                   height: param.par.i_height,
                   native: true,
               })
        }
    }

    pub fn as_slice<'a>(&'a self, plane: usize) -> Result<&'a [u8], &'static str> {
        if plane > self.pic.img.i_plane as usize {
            Err("Invalid Argument")
        } else {
            let size = self.height * self.pic.img.i_stride[plane];
            Ok(unsafe { std::slice::from_raw_parts(self.pic.img.plane[plane], size as usize) })
        }
    }

    pub fn as_slice_mut<'a>(&'a mut self, plane: usize) -> Result<&'a mut [u8], &'static str> {
        if plane > self.pic.img.i_plane as usize {
            Err("Invalid Argument")
        } else {
            let size = self.height * self.pic.img.i_stride[plane];
            Ok(unsafe { std::slice::from_raw_parts_mut(self.pic.img.plane[plane], size as usize) })
        }
    }

    pub fn set_timestamp(mut self, pts: i64) -> Picture {
        self.pic.i_pts = pts;
        self
    }
}

impl Drop for Picture {
    fn drop(&mut self) {
        if self.native {
            unsafe { x264_picture_clean(&mut self.pic as *mut x264_picture_t) };
        }
    }
}

// TODO: Provide a builder API instead?
pub struct Param {
    par: x264_param_t,
}

impl Param {
    pub fn new() -> Param {
        let mut par = unsafe { mem::uninitialized() };

        unsafe {
            x264_param_default(&mut par as *mut x264_param_t);
        }

        Param { par: par }
    }
    pub fn default_preset<'a, O>(tune: O, preset: O) -> Result<Param, &'static str>
        where O: Into<Option<&'a str>>
    {
        let mut par = unsafe { mem::uninitialized() };
        let c_tune = tune.into().map_or_else(|| null(), |v| v.as_ptr() as *const i8);
        let c_preset = preset.into().map_or_else(|| null(), |v| v.as_ptr() as *const i8);
        match unsafe {
                  x264_param_default_preset(&mut par as *mut x264_param_t, c_tune, c_preset)
              } {
            -1 => Err("Invalid Argument"),
            0 => Ok(Param { par: par }),
            _ => Err("Unexpected"),
        }
    }
    pub fn apply_profile(mut self, profile: &str) -> Result<Param, &'static str> {
        match unsafe { x264_param_apply_profile(&mut self.par, profile.as_ptr() as *const i8) } {
            -1 => Err("Invalid Argument"),
            0 => Ok(self),
            _ => Err("Unexpected"),
        }
    }
    pub fn param_parse<'a>(mut self, name: &str, value: &str) -> Result<Param, &'static str> {
        match unsafe {
                  x264_param_parse(&mut self.par,
                                   name.as_ptr() as *const i8,
                                   value.as_ptr() as *const i8)
              } {
            -1 => Err("Invalid Argument"),
            0 => Ok(self),
            _ => Err("Unexpected"),
        }
    }
    pub fn set_dimension(mut self, height: usize, width: usize) -> Param {
        self.par.i_height = height as c_int;
        self.par.i_width = width as c_int;

        self
    }
}

// TODO: Expose a NAL abstraction
pub struct NalData {
    vec: Vec<u8>,
}

impl NalData {
    /*
     * x264 functions return x264_nal_t arrays that are valid only until another
     * function of that kind is called.
     *
     * Always copy the data over.
     *
     * TODO: Consider using Bytes as backing store.
     */
    fn from_nals(c_nals: *mut x264_nal_t, nb_nal: usize) -> NalData {
        let mut data = NalData { vec: Vec::new() };

        for i in 0..nb_nal {
            let nal = unsafe { Box::from_raw(c_nals.offset(i as isize)) };

            let payload =
                unsafe { std::slice::from_raw_parts(nal.p_payload, nal.i_payload as usize) };

            data.vec.extend_from_slice(payload);

            mem::forget(payload);
            mem::forget(nal);
        }

        data
    }

    pub fn as_bytes(&self) -> &[u8] {
        self.vec.as_slice()
    }
}

pub struct Encoder {
    enc: *mut x264_t,
}

impl Encoder {
    pub fn open(par: &mut Param) -> Result<Encoder, &'static str> {
        let enc = unsafe { x264_encoder_open(&mut par.par as *mut x264_param_t) };

        if enc.is_null() {
            Err("Out of Memory")
        } else {
            Ok(Encoder { enc: enc })
        }
    }

    pub fn get_headers(&mut self) -> Result<NalData, &'static str> {
        let mut nb_nal: c_int = 0;
        let mut c_nals: *mut x264_nal_t = unsafe { mem::uninitialized() };

        let bytes = unsafe {
            x264_encoder_headers(self.enc,
                                 &mut c_nals as *mut *mut x264_nal_t,
                                 &mut nb_nal as *mut c_int)
        };

        if bytes < 0 {
            Err("Encoding Headers Failed")
        } else {
            Ok(NalData::from_nals(c_nals, nb_nal as usize))
        }
    }

    pub fn encode<'a, P>(&mut self, pic: P) -> Result<Option<(NalData, i64, i64)>, &'static str>
        where P: Into<Option<&'a Picture>>
    {
        let mut pic_out: x264_picture_t = unsafe { mem::uninitialized() };
        let mut c_nals: *mut x264_nal_t = unsafe { mem::uninitialized() };
        let mut nb_nal: c_int = 0;
        let c_pic = pic.into().map_or_else(|| null(), |v| &v.pic as *const x264_picture_t);

        let ret = unsafe {
            x264_encoder_encode(self.enc,
                                &mut c_nals as *mut *mut x264_nal_t,
                                &mut nb_nal as *mut c_int,
                                c_pic as *mut x264_picture_t,
                                &mut pic_out as *mut x264_picture_t)
        };
        if ret < 0 {
            Err("Error encoding")
        } else {
            if nb_nal > 0 {
                let data = NalData::from_nals(c_nals, nb_nal as usize);
                Ok(Some((data, pic_out.i_pts, pic_out.i_dts)))
            } else {
                Ok(None)
            }
        }
    }

    pub fn delayed_frames(&self) -> bool {
        unsafe { x264_encoder_delayed_frames(self.enc) != 0 }
    }
}

impl Drop for Encoder {
    fn drop(&mut self) {
        unsafe { x264_encoder_close(self.enc) };
    }
}

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

    #[test]
    fn test_open() {
        let mut par = Param::new().set_dimension(640, 480);

        let mut enc = Encoder::open(&mut par).unwrap();

        let headers = enc.get_headers().unwrap();

        println!("Headers len {}", headers.as_bytes().len());
    }

    #[test]
    fn test_picture() {
        let par = Param::new().set_dimension(640, 480);
        {
            let mut pic = Picture::from_param(&par).unwrap();
            {
                let p = pic.as_slice_mut(0).unwrap();
                p[0] = 1;
            }
            let p = pic.as_slice(0).unwrap();

            assert_eq!(p[0], 1);
        }
    }

    #[test]
    fn test_encode() {
        let mut par = Param::new().set_dimension(640, 480);
        let mut enc = Encoder::open(&mut par).unwrap();
        let mut pic = Picture::from_param(&par).unwrap();

        let headers = enc.get_headers().unwrap();

        println!("Headers len {}", headers.as_bytes().len());

        for pts in 0..5 {
            pic = pic.set_timestamp(pts as i64);
            let ret = enc.encode(&pic).unwrap();
            match ret {
                Some((_, pts, dts)) => println!("Frame pts {}, dts {}", pts, dts),
                _ => (),
            }
        }

        while enc.delayed_frames() {
            let ret = enc.encode(None).unwrap();
            match ret {
                Some((_, pts, dts)) => println!("Frame pts {}, dts {}", pts, dts),
                _ => (),
            }
        }
    }
}
<|endoftext|>


<commit_before>// Module for parsing ISO Base Media Format aka video/mp4 streams.

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

/// Basic ISO box structure.
pub struct BoxHeader {
    /// Four character box type
pub name: u32,
    /// Size of the box in bytes
pub size: u64,
    /// Offset to the start of the contained data (or header size).
pub offset: u64,
}

/// File type box 'ftyp'.
pub struct FileTypeBox {
    name: u32,
    size: u64,
    major_brand: u32,
    minor_version: u32,
    compatible_brands: Vec<u32>,
}

/// Movie header box 'mvhd'.
pub struct MovieHeaderBox {
    pub name: u32,
    pub size: u64,
    pub timescale: u32,
    pub duration: u64,
    // Ignore other fields.
}

pub struct TrackHeaderBox {
    pub name: u32,
    pub size: u64,
    pub track_id: u32,
    pub duration: u64,
    pub width: u32,
    pub height: u32,
}

extern crate byteorder;
use byteorder::{BigEndian, ReadBytesExt};
use std::io::{Error, ErrorKind, Read, Result, Seek, SeekFrom, Take};
use std::io::Cursor;

/// Parse a box out of a data buffer.
pub fn read_box_header<T: ReadBytesExt>(src: &mut T) -> Result<BoxHeader> {
    let tmp_size = try!(src.read_u32::<BigEndian>());
    let name = try!(src.read_u32::<BigEndian>());
    let size = match tmp_size {
        1 => try!(src.read_u64::<BigEndian>()),
        _ => tmp_size as u64,
    };
    assert!(size >= 8);
    if tmp_size == 1 {
        assert!(size >= 16);
    }
    let offset = match tmp_size {
        1 => 4 + 4 + 8,
        _ => 4 + 4,
    };
    assert!(offset <= size);
    Ok(BoxHeader{
      name: name,
      size: size,
      offset: offset,
    })
}

/// Parse the extra header fields for a full box.
fn read_fullbox_extra<T: ReadBytesExt>(src: &mut T) -> (u8, u32) {
    let version = src.read_u8().unwrap();
    let flags_a = src.read_u8().unwrap();
    let flags_b = src.read_u8().unwrap();
    let flags_c = src.read_u8().unwrap();
    (version, (flags_a as u32) << 16 |
              (flags_b as u32) <<  8 |
              (flags_c as u32))
}

/// Skip over the contents of a box.
pub fn skip_box_content<T: ReadBytesExt + Seek>
  (src: &mut T, header: &BoxHeader)
  -> std::io::Result<u64>
{
    src.seek(SeekFrom::Current((header.size - header.offset) as i64))
}

/// Helper to construct a Take over the contents of a box.
fn limit<'a, T: Read>(f: &'a mut T, h: &BoxHeader) -> Take<&'a mut T> {
    f.take(h.size - h.offset)
}

/// Helper to construct a Cursor over the contents of a box.
fn recurse<T: Read>(f: &mut T, h: &BoxHeader) -> Result<()> {
    println!("{} -- recursing", h);
    // FIXME: I couldn't figure out how to do this without copying.
    // We use Seek on the Read we return in skip_box_content, but
    // that trait isn't implemented for a Take like our limit()
    // returns. Slurping the buffer and wrapping it in a Cursor
    // functions as a work around.
    let buf: Vec<u8> = limit(f, &h)
        .bytes()
        .map(|u| u.unwrap())
        .collect();
    let mut content = Cursor::new(buf);
    loop {
        match read_box(&mut content) {
            Ok(_) => {},
            Err(e) => {
                println!("Error '{:?}' reading box", e.kind());
                return Err(e);
            },
        }
    }
    println!("{} -- end", h);
    Ok(())
}

/// Read the contents of a box, including sub boxes.
/// Right now it just prints the box value rather than
/// returning anything.
pub fn read_box<T: Read + Seek>(f: &mut T) -> Result<()> {
    read_box_header(f).and_then(|h| {
        match &(fourcc_to_string(h.name))[..] {
            "ftyp" => {
                let mut content = limit(f, &h);
                let ftyp = read_ftyp(&mut content, &h).unwrap();
                println!("{}", ftyp);
            },
            "moov" => try!(recurse(f, &h)),
            "mvhd" => {
                let mut content = limit(f, &h);
                let mvhd = read_mvhd(&mut content, &h).unwrap();
                println!("  {}", mvhd);
            },
            "trak" => try!(recurse(f, &h)),
            "tkhd" => {
                let mut content = limit(f, &h);
                let tkhd = read_tkhd(&mut content, &h).unwrap();
                println!("  {}", tkhd);
            },
            _ => {
                // Skip the contents of unknown chunks.
                println!("{} (skipped)", h);
                try!(skip_box_content(f, &h).and(Ok(())));
            },
        };
        Ok(()) // and_then needs a Result to return.
    })
}

/// Entry point for C language callers.
/// Take a buffer and call read_box() on it.
#[no_mangle]
pub unsafe extern fn read_box_from_buffer(buffer: *const u8, size: usize)
  -> bool {
    use std::slice;
    use std::thread;

    // Validate arguments from C.
    if buffer.is_null() || size < 8 {
        return false;
    }

    // Wrap the buffer we've been give in a slice.
    let b = slice::from_raw_parts(buffer, size);
    let mut c = Cursor::new(b);

    // Parse in a subthread.
    let task = thread::spawn(move || {
        read_box(&mut c).unwrap();
    });
    // Catch any panics.
    task.join().is_ok()
}

/// Parse an ftype box.
pub fn read_ftyp<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<FileTypeBox> {
    let major = src.read_u32::<BigEndian>().unwrap();
    let minor = src.read_u32::<BigEndian>().unwrap();
    let brand_count = (head.size - 8 - 8) /4;
    let mut brands = Vec::new();
    for _ in 0..brand_count {
        brands.push(src.read_u32::<BigEndian>().unwrap());
    }
    Some(FileTypeBox{
        name: head.name,
        size: head.size,
        major_brand: major,
        minor_version: minor,
        compatible_brands: brands,
    })
}

/// Parse an mvhd box.
pub fn read_mvhd<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<MovieHeaderBox> {
    let (version, _) = read_fullbox_extra(src);
    match version {
        1 => {
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 16];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        0 => {
            // 32 bit creation and modification times.
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 8];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        _ => panic!("invalid mhdr version"),
    }
    let timescale = src.read_u32::<BigEndian>().unwrap();
    let duration = match version {
        1 => src.read_u64::<BigEndian>().unwrap(),
        0 => src.read_u32::<BigEndian>().unwrap() as u64,
        _ => panic!("invalid mhdr version"),
    };
    // Skip remaining fields.
    let mut skip: Vec<u8> = vec![0; 80];
    let r = src.read(&mut skip).unwrap();
    assert!(r == skip.len());
    Some(MovieHeaderBox {
        name: head.name,
        size: head.size,
        timescale: timescale,
        duration: duration,
    })
}

/// Parse a tkhd box.
pub fn read_tkhd<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Result<TrackHeaderBox> {
    let (version, flags) = read_fullbox_extra(src);
    if flags & 0x1u32 == 0 || flags & 0x2u32 == 0 {
        return Err(Error::new(ErrorKind::Other, "Track is disabled"));
    }
    match version {
        1 => {
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 16];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        0 => {
            // 32 bit creation and modification times.
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 8];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        _ => panic!("invalid tkhd version"),
    }
    let track_id = src.read_u32::<BigEndian>().unwrap();
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    assert!(_reserved == 0);
    let duration = match version {
        1 => {
            src.read_u64::<BigEndian>().unwrap()
        },
        0 => src.read_u32::<BigEndian>().unwrap() as u64,
        _ => panic!("invalid tkhd version"),
    };
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    // Skip uninterested fields.
    let mut skip: Vec<u8> = vec![0; 44];
    let r = src.read(&mut skip).unwrap();
    assert!(r == skip.len());
    let width = src.read_u32::<BigEndian>().unwrap();
    let height = src.read_u32::<BigEndian>().unwrap();
    Ok(TrackHeaderBox {
       name: head.name,
       size: head.size,
       track_id: track_id,
       duration: duration,
       width: width,
       height: height,
    })
}

/// Convert the iso box type or other 4-character value to a string.
fn fourcc_to_string(name: u32) -> String {
    let u32_to_vec = |u| {
        vec!((u >> 24 & 0xffu32) as u8,
             (u >> 16 & 0xffu32) as u8,
             (u >>  8 & 0xffu32) as u8,
             (u & 0xffu32) as u8)
    };
    let name_bytes = u32_to_vec(name);
    String::from_utf8_lossy(&name_bytes).into_owned()
}

use std::fmt;
impl fmt::Display for BoxHeader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "'{}' {} bytes", fourcc_to_string(self.name), self.size)
    }
}

impl fmt::Display for FileTypeBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        let brand = fourcc_to_string(self.major_brand);
        write!(f, "'{}' {} bytes '{}' v{}", name, self.size,
            brand, self.minor_version)
    }
}

impl fmt::Display for MovieHeaderBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        write!(f, "'{}' {} bytes duration {}s", name, self.size,
            (self.duration as f64)/(self.timescale as f64))
    }
}

use std::u16;
impl fmt::Display for TrackHeaderBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        // Dimensions are 16.16 fixed-point.
        let base = u16::MAX as f64 + 1.0;
        let width = (self.width as f64) / base;
        let height = (self.height as f64) / base;
        write!(f, "'{}' {} bytes duration {} id {} {}x{}",
            name, self.size, self.duration, self.track_id,
            width, height)
    }
}

#[test]
fn test_read_box_header() {
    use std::io::Cursor;
    use std::io::Write;
    let mut test: Vec<u8> = vec![0, 0, 0, 8];  // minimal box length
    write!(&mut test, "test").unwrap(); // box type
    let mut stream = Cursor::new(test);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, 1952805748);
    assert_eq!(parsed.size, 8);
    println!("box {}", parsed);
}

#[test]
fn test_read_box_header_long() {
    use std::io::Cursor;
    let mut test: Vec<u8> = vec![0, 0, 0, 1]; // long box extension code
    test.extend("long".to_string().into_bytes()); // box type
    test.extend(vec![0, 0, 0, 0, 0, 0, 16, 0]); // 64 bit size
    // Skip generating box content.
    let mut stream = Cursor::new(test);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, 1819242087);
    assert_eq!(parsed.size, 4096);
    println!("box {}", parsed);
}

#[test]
fn test_read_ftyp() {
    use std::io::Cursor;
    use std::io::Write;
    let mut test: Vec<u8> = vec![0, 0, 0, 24]; // size
    write!(&mut test, "ftyp").unwrap(); // type
    write!(&mut test, "mp42").unwrap(); // major brand
    test.extend(vec![0, 0, 0, 0]);      // minor version
    write!(&mut test, "isom").unwrap(); // compatible brands...
    write!(&mut test, "mp42").unwrap();
    assert_eq!(test.len(), 24);

    let mut stream = Cursor::new(test);
    let header = read_box_header(&mut stream).unwrap();
    let parsed = read_ftyp(&mut stream, &header).unwrap();
    assert_eq!(parsed.name, 1718909296);
    assert_eq!(parsed.size, 24);
    assert_eq!(parsed.major_brand, 1836069938);
    assert_eq!(parsed.minor_version, 0);
    assert_eq!(parsed.compatible_brands.len(), 2);
    println!("box {}", parsed);
}
<commit_msg>Convert read_tkhd unwrap to try.<commit_after>// Module for parsing ISO Base Media Format aka video/mp4 streams.

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

/// Basic ISO box structure.
pub struct BoxHeader {
    /// Four character box type
pub name: u32,
    /// Size of the box in bytes
pub size: u64,
    /// Offset to the start of the contained data (or header size).
pub offset: u64,
}

/// File type box 'ftyp'.
pub struct FileTypeBox {
    name: u32,
    size: u64,
    major_brand: u32,
    minor_version: u32,
    compatible_brands: Vec<u32>,
}

/// Movie header box 'mvhd'.
pub struct MovieHeaderBox {
    pub name: u32,
    pub size: u64,
    pub timescale: u32,
    pub duration: u64,
    // Ignore other fields.
}

pub struct TrackHeaderBox {
    pub name: u32,
    pub size: u64,
    pub track_id: u32,
    pub duration: u64,
    pub width: u32,
    pub height: u32,
}

extern crate byteorder;
use byteorder::{BigEndian, ReadBytesExt};
use std::io::{Error, ErrorKind, Read, Result, Seek, SeekFrom, Take};
use std::io::Cursor;

/// Parse a box out of a data buffer.
pub fn read_box_header<T: ReadBytesExt>(src: &mut T) -> Result<BoxHeader> {
    let tmp_size = try!(src.read_u32::<BigEndian>());
    let name = try!(src.read_u32::<BigEndian>());
    let size = match tmp_size {
        1 => try!(src.read_u64::<BigEndian>()),
        _ => tmp_size as u64,
    };
    assert!(size >= 8);
    if tmp_size == 1 {
        assert!(size >= 16);
    }
    let offset = match tmp_size {
        1 => 4 + 4 + 8,
        _ => 4 + 4,
    };
    assert!(offset <= size);
    Ok(BoxHeader{
      name: name,
      size: size,
      offset: offset,
    })
}

/// Parse the extra header fields for a full box.
fn read_fullbox_extra<T: ReadBytesExt>(src: &mut T) -> (u8, u32) {
    let version = src.read_u8().unwrap();
    let flags_a = src.read_u8().unwrap();
    let flags_b = src.read_u8().unwrap();
    let flags_c = src.read_u8().unwrap();
    (version, (flags_a as u32) << 16 |
              (flags_b as u32) <<  8 |
              (flags_c as u32))
}

/// Skip over the contents of a box.
pub fn skip_box_content<T: ReadBytesExt + Seek>
  (src: &mut T, header: &BoxHeader)
  -> std::io::Result<u64>
{
    src.seek(SeekFrom::Current((header.size - header.offset) as i64))
}

/// Helper to construct a Take over the contents of a box.
fn limit<'a, T: Read>(f: &'a mut T, h: &BoxHeader) -> Take<&'a mut T> {
    f.take(h.size - h.offset)
}

/// Helper to construct a Cursor over the contents of a box.
fn recurse<T: Read>(f: &mut T, h: &BoxHeader) -> Result<()> {
    println!("{} -- recursing", h);
    // FIXME: I couldn't figure out how to do this without copying.
    // We use Seek on the Read we return in skip_box_content, but
    // that trait isn't implemented for a Take like our limit()
    // returns. Slurping the buffer and wrapping it in a Cursor
    // functions as a work around.
    let buf: Vec<u8> = limit(f, &h)
        .bytes()
        .map(|u| u.unwrap())
        .collect();
    let mut content = Cursor::new(buf);
    loop {
        match read_box(&mut content) {
            Ok(_) => {},
            Err(e) => {
                println!("Error '{:?}' reading box", e.kind());
                return Err(e);
            },
        }
    }
    println!("{} -- end", h);
    Ok(())
}

/// Read the contents of a box, including sub boxes.
/// Right now it just prints the box value rather than
/// returning anything.
pub fn read_box<T: Read + Seek>(f: &mut T) -> Result<()> {
    read_box_header(f).and_then(|h| {
        match &(fourcc_to_string(h.name))[..] {
            "ftyp" => {
                let mut content = limit(f, &h);
                let ftyp = read_ftyp(&mut content, &h).unwrap();
                println!("{}", ftyp);
            },
            "moov" => try!(recurse(f, &h)),
            "mvhd" => {
                let mut content = limit(f, &h);
                let mvhd = read_mvhd(&mut content, &h).unwrap();
                println!("  {}", mvhd);
            },
            "trak" => try!(recurse(f, &h)),
            "tkhd" => {
                let mut content = limit(f, &h);
                let tkhd = read_tkhd(&mut content, &h).unwrap();
                println!("  {}", tkhd);
            },
            _ => {
                // Skip the contents of unknown chunks.
                println!("{} (skipped)", h);
                try!(skip_box_content(f, &h).and(Ok(())));
            },
        };
        Ok(()) // and_then needs a Result to return.
    })
}

/// Entry point for C language callers.
/// Take a buffer and call read_box() on it.
#[no_mangle]
pub unsafe extern fn read_box_from_buffer(buffer: *const u8, size: usize)
  -> bool {
    use std::slice;
    use std::thread;

    // Validate arguments from C.
    if buffer.is_null() || size < 8 {
        return false;
    }

    // Wrap the buffer we've been give in a slice.
    let b = slice::from_raw_parts(buffer, size);
    let mut c = Cursor::new(b);

    // Parse in a subthread.
    let task = thread::spawn(move || {
        read_box(&mut c).unwrap();
    });
    // Catch any panics.
    task.join().is_ok()
}

/// Parse an ftype box.
pub fn read_ftyp<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<FileTypeBox> {
    let major = src.read_u32::<BigEndian>().unwrap();
    let minor = src.read_u32::<BigEndian>().unwrap();
    let brand_count = (head.size - 8 - 8) /4;
    let mut brands = Vec::new();
    for _ in 0..brand_count {
        brands.push(src.read_u32::<BigEndian>().unwrap());
    }
    Some(FileTypeBox{
        name: head.name,
        size: head.size,
        major_brand: major,
        minor_version: minor,
        compatible_brands: brands,
    })
}

/// Parse an mvhd box.
pub fn read_mvhd<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<MovieHeaderBox> {
    let (version, _) = read_fullbox_extra(src);
    match version {
        1 => {
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 16];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        0 => {
            // 32 bit creation and modification times.
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 8];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        _ => panic!("invalid mhdr version"),
    }
    let timescale = src.read_u32::<BigEndian>().unwrap();
    let duration = match version {
        1 => src.read_u64::<BigEndian>().unwrap(),
        0 => src.read_u32::<BigEndian>().unwrap() as u64,
        _ => panic!("invalid mhdr version"),
    };
    // Skip remaining fields.
    let mut skip: Vec<u8> = vec![0; 80];
    let r = src.read(&mut skip).unwrap();
    assert!(r == skip.len());
    Some(MovieHeaderBox {
        name: head.name,
        size: head.size,
        timescale: timescale,
        duration: duration,
    })
}

/// Parse a tkhd box.
pub fn read_tkhd<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Result<TrackHeaderBox> {
    let (version, flags) = read_fullbox_extra(src);
    if flags & 0x1u32 == 0 || flags & 0x2u32 == 0 {
        return Err(Error::new(ErrorKind::Other, "Track is disabled"));
    }
    match version {
        1 => {
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 16];
            let r = try!(src.read(&mut skip));
            assert!(r == skip.len());
        },
        0 => {
            // 32 bit creation and modification times.
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 8];
            let r = try!(src.read(&mut skip));
            assert!(r == skip.len());
        },
        _ => panic!("invalid tkhd version"),
    }
    let track_id = try!(src.read_u32::<BigEndian>());
    let _reserved = try!(src.read_u32::<BigEndian>());
    assert!(_reserved == 0);
    let duration = match version {
        1 => {
            try!(src.read_u64::<BigEndian>())
        },
        0 => try!(src.read_u32::<BigEndian>()) as u64,
        _ => panic!("invalid tkhd version"),
    };
    let _reserved = try!(src.read_u32::<BigEndian>());
    let _reserved = try!(src.read_u32::<BigEndian>());
    // Skip uninterested fields.
    let mut skip: Vec<u8> = vec![0; 44];
    let r = try!(src.read(&mut skip));
    assert!(r == skip.len());
    let width = try!(src.read_u32::<BigEndian>());
    let height = try!(src.read_u32::<BigEndian>());
    Ok(TrackHeaderBox {
       name: head.name,
       size: head.size,
       track_id: track_id,
       duration: duration,
       width: width,
       height: height,
    })
}

/// Convert the iso box type or other 4-character value to a string.
fn fourcc_to_string(name: u32) -> String {
    let u32_to_vec = |u| {
        vec!((u >> 24 & 0xffu32) as u8,
             (u >> 16 & 0xffu32) as u8,
             (u >>  8 & 0xffu32) as u8,
             (u & 0xffu32) as u8)
    };
    let name_bytes = u32_to_vec(name);
    String::from_utf8_lossy(&name_bytes).into_owned()
}

use std::fmt;
impl fmt::Display for BoxHeader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "'{}' {} bytes", fourcc_to_string(self.name), self.size)
    }
}

impl fmt::Display for FileTypeBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        let brand = fourcc_to_string(self.major_brand);
        write!(f, "'{}' {} bytes '{}' v{}", name, self.size,
            brand, self.minor_version)
    }
}

impl fmt::Display for MovieHeaderBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        write!(f, "'{}' {} bytes duration {}s", name, self.size,
            (self.duration as f64)/(self.timescale as f64))
    }
}

use std::u16;
impl fmt::Display for TrackHeaderBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        // Dimensions are 16.16 fixed-point.
        let base = u16::MAX as f64 + 1.0;
        let width = (self.width as f64) / base;
        let height = (self.height as f64) / base;
        write!(f, "'{}' {} bytes duration {} id {} {}x{}",
            name, self.size, self.duration, self.track_id,
            width, height)
    }
}

#[test]
fn test_read_box_header() {
    use std::io::Cursor;
    use std::io::Write;
    let mut test: Vec<u8> = vec![0, 0, 0, 8];  // minimal box length
    write!(&mut test, "test").unwrap(); // box type
    let mut stream = Cursor::new(test);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, 1952805748);
    assert_eq!(parsed.size, 8);
    println!("box {}", parsed);
}

#[test]
fn test_read_box_header_long() {
    use std::io::Cursor;
    let mut test: Vec<u8> = vec![0, 0, 0, 1]; // long box extension code
    test.extend("long".to_string().into_bytes()); // box type
    test.extend(vec![0, 0, 0, 0, 0, 0, 16, 0]); // 64 bit size
    // Skip generating box content.
    let mut stream = Cursor::new(test);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, 1819242087);
    assert_eq!(parsed.size, 4096);
    println!("box {}", parsed);
}

#[test]
fn test_read_ftyp() {
    use std::io::Cursor;
    use std::io::Write;
    let mut test: Vec<u8> = vec![0, 0, 0, 24]; // size
    write!(&mut test, "ftyp").unwrap(); // type
    write!(&mut test, "mp42").unwrap(); // major brand
    test.extend(vec![0, 0, 0, 0]);      // minor version
    write!(&mut test, "isom").unwrap(); // compatible brands...
    write!(&mut test, "mp42").unwrap();
    assert_eq!(test.len(), 24);

    let mut stream = Cursor::new(test);
    let header = read_box_header(&mut stream).unwrap();
    let parsed = read_ftyp(&mut stream, &header).unwrap();
    assert_eq!(parsed.name, 1718909296);
    assert_eq!(parsed.size, 24);
    assert_eq!(parsed.major_brand, 1836069938);
    assert_eq!(parsed.minor_version, 0);
    assert_eq!(parsed.compatible_brands.len(), 2);
    println!("box {}", parsed);
}
<|endoftext|>


<commit_before>//! #Example
//!
//! ```rust
//! extern crate conshash;
//!
//! use std::collections::hash_map::DefaultHasher;
//!
//! #[derive(Clone, Debug)]
//! struct TestNode {
//!     host_name: &'static str,
//!     ip_address: &'static str,
//!     port: u32,
//! }
//!
//! impl ToString for TestNode {
//!     fn to_string(&self) -> String {
//!         format!("{}{}", self.ip_address.to_string(), self.port.to_string())
//!     }
//! }
//!
//! let mut hash_ring = conshash::Ring::new(5);
//!
//! let test_node = TestNode{host_name: "Skynet", ip_address: "192.168.1.1", port: 42};
//! hash_ring.add_node(&test_node);
//! hash_ring.remove_node(&test_node);
//! hash_ring.add_node(&test_node);
//! let x = hash_ring.get_node(conshash::hash(&format!("{}{}", test_node.to_string(), 0.to_string())));
//! // x is the node in the form of an Option<T> where T: Clone + ToString + Debug
//! ```

use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::clone::Clone;
use std::fmt::Debug;
use std::string::ToString;
use std::collections::BTreeMap;

pub fn hash<T: Hash>(value: &T) -> u64 {
    let mut h = DefaultHasher::new();
    value.hash(&mut h);
    h.finish()
}

pub struct Ring <T: Clone + ToString + Debug> {
    num_replicas: usize,
    ring: BTreeMap<u64, T>,
}

impl <T> Ring<T> where T: Clone + ToString + Debug {
    pub fn new(num_replicas: usize) -> Ring<T> {
        Ring {
            num_replicas: num_replicas,
            ring: BTreeMap::new(),
        }
    }

    pub fn add_nodes(&mut self, nodes: &[T]) {
        if !nodes.is_empty() {
            for node in nodes.iter() { self.add_node(node); }
        }
    }

    pub fn remove_nodes(&mut self, nodes: &[T]) {
        if !nodes.is_empty() {
            for node in nodes.iter() { self.remove_node(node); }
        }
    }

    pub fn add_node(&mut self, node: &T) {
        for i in 0..self.num_replicas {
            let key = hash(&format!("{}{}", node.to_string(), i.to_string()));
            self.ring.insert(key, node.clone());
        }
    }

    pub fn remove_node(&mut self, node: &T) {
        assert!(!self.ring.is_empty());

        for i in 0..self.num_replicas {
            let key = hash(&format!("{}{}", node.to_string(), i.to_string()));
            self.ring.remove(&key);
        }
    }

    pub fn get_node(&self, key: u64) -> Option<&T> {
        assert!(!self.ring.is_empty());
        let mut keys = self.ring.keys();
        keys.find(|k| *k >= &key)
            .and_then(|k| self.ring.get(k))
            .or(keys.nth(0).and_then(|x| self.ring.get(x)))
    }
}

#[cfg (test)]
mod tests {

    use super::*;
    use std::string::ToString;

    #[derive(Clone, Debug)]
    struct TestNode {
        host_name: &'static str,
        ip_address: &'static str,
        port: u32,
    }

    impl ToString for TestNode {
        fn to_string(&self) -> String {
            format!("{}{}", self.ip_address.to_string(), self.port.to_string())
        }
    }

    #[test]
    fn test_add_node(){
        let mut hash_ring = Ring::new(3);
        assert_eq!(hash_ring.num_replicas, 3);

        let test_node = TestNode{host_name: "Skynet", ip_address: "192.168.1.1", port: 42};
        hash_ring.add_node(&test_node);

    }

    #[test]
    fn test_remove_node(){
        let mut hash_ring = Ring::new(3);
        assert_eq!(hash_ring.num_replicas, 3);

        let test_node = TestNode{host_name: "Skynet", ip_address: "192.168.1.1", port: 42};
        hash_ring.add_node(&test_node);
        hash_ring.remove_node(&test_node);
    }

    #[test]
    fn test_get_node(){
        let mut hash_ring = Ring::new(3);
        assert_eq!(hash_ring.num_replicas, 3);

        let test_node = TestNode{host_name: "Skynet", ip_address: "192.168.1.1", port: 42};
        hash_ring.add_node(&test_node);
        let my_node = hash_ring.get_node(hash(&test_node.to_string()));

        assert_eq!(my_node.unwrap().host_name, test_node.host_name);
        assert_eq!(my_node.unwrap().ip_address, test_node.ip_address);
        assert_eq!(my_node.unwrap().port, test_node.port);
    }

    #[test]
    fn test_add_nodes(){
        let mut hash_ring = Ring::new(3);
        assert_eq!(hash_ring.num_replicas, 3);

        let test_node1 = TestNode{host_name: "Skynet", ip_address: "192.168.1.1", port: 42};
        let test_node2 = TestNode{host_name: "Inferno", ip_address: "10.0.1.1", port: 666};
        let test_node3 = TestNode{host_name: "Klimt", ip_address: "127.0.0.1", port: 1};

        let v = vec![test_node1.clone(), test_node2.clone(), test_node3.clone()];
        hash_ring.add_nodes(&v);

        let node1 = hash_ring.get_node(hash(&format!("{}{}", test_node1.to_string(), 0.to_string())));
        let node2 = hash_ring.get_node(hash(&format!("{}{}", test_node2.to_string(), 0.to_string())));
        let node3 = hash_ring.get_node(hash(&format!("{}{}", test_node3.to_string(), 0.to_string())));

        assert_eq!(node1.unwrap().host_name, test_node1.host_name);
        assert_eq!(node1.unwrap().ip_address, test_node1.ip_address);
        assert_eq!(node1.unwrap().port, test_node1.port);

        assert_eq!(node2.unwrap().host_name, test_node2.host_name);
        assert_eq!(node2.unwrap().ip_address, test_node2.ip_address);
        assert_eq!(node2.unwrap().port, test_node2.port);

        assert_eq!(node3.unwrap().host_name, test_node3.host_name);
        assert_eq!(node3.unwrap().ip_address, test_node3.ip_address);
        assert_eq!(node3.unwrap().port, test_node3.port);
    }

    #[test]
    fn test_remove_nodes(){
        let mut hash_ring = Ring::new(3);
        assert_eq!(hash_ring.num_replicas, 3);

        let test_node1 = TestNode{host_name: "Skynet", ip_address: "192.168.1.1", port: 42};
        let test_node2 = TestNode{host_name: "Inferno", ip_address: "10.0.1.1", port: 666};
        let test_node3 = TestNode{host_name: "Klimt", ip_address: "127.0.0.1", port: 1};

        let v = vec![test_node1.clone(), test_node2.clone(), test_node3.clone()];
        hash_ring.add_nodes(&v);
        hash_ring.remove_nodes(&v);

        assert!(hash_ring.ring.is_empty());
    }
}
<commit_msg>Run rustfmt and fix clippy issues.<commit_after>//! #Example
//!
//! ```rust
//! extern crate conshash;
//!
//! use std::collections::hash_map::DefaultHasher;
//!
//! #[derive(Clone, Debug)]
//! struct TestNode {
//!     host_name: &'static str,
//!     ip_address: &'static str,
//!     port: u32,
//! }
//!
//! impl ToString for TestNode {
//!     fn to_string(&self) -> String {
//!         format!("{}{}", self.ip_address.to_string(), self.port.to_string())
//!     }
//! }
//!
//! let mut hash_ring = conshash::Ring::new(5);
//!
//! let test_node = TestNode{host_name: "Skynet", ip_address: "192.168.1.1", port: 42};
//! hash_ring.add_node(&test_node);
//! hash_ring.remove_node(&test_node);
//! hash_ring.add_node(&test_node);
//! let x = hash_ring.get_node(conshash::hash(&format!("{}{}", test_node.to_string(), 0.to_string())));
//! // x is the node in the form of an Option<T> where T: Clone + ToString + Debug
//! ```

use std::clone::Clone;
use std::collections::hash_map::DefaultHasher;
use std::collections::BTreeMap;
use std::fmt::Debug;
use std::hash::{Hash, Hasher};
use std::string::ToString;

pub fn hash<T: Hash>(value: &T) -> u64 {
    let mut h = DefaultHasher::new();
    value.hash(&mut h);
    h.finish()
}

pub struct Ring<T: Clone + ToString + Debug> {
    num_replicas: usize,
    ring: BTreeMap<u64, T>,
}

impl<T> Ring<T>
where
    T: Clone + ToString + Debug,
{
    pub fn new(num_replicas: usize) -> Ring<T> {
        Ring {
            num_replicas,
            ring: BTreeMap::new(),
        }
    }

    pub fn add_nodes(&mut self, nodes: &[T]) {
        if !nodes.is_empty() {
            for node in nodes.iter() {
                self.add_node(node);
            }
        }
    }

    pub fn remove_nodes(&mut self, nodes: &[T]) {
        if !nodes.is_empty() {
            for node in nodes.iter() {
                self.remove_node(node);
            }
        }
    }

    pub fn add_node(&mut self, node: &T) {
        for i in 0..self.num_replicas {
            let key = hash(&format!("{}{}", node.to_string(), i.to_string()));
            self.ring.insert(key, node.clone());
        }
    }

    pub fn remove_node(&mut self, node: &T) {
        assert!(!self.ring.is_empty());

        for i in 0..self.num_replicas {
            let key = hash(&format!("{}{}", node.to_string(), i.to_string()));
            self.ring.remove(&key);
        }
    }

    pub fn get_node(&self, key: u64) -> Option<&T> {
        assert!(!self.ring.is_empty());
        let mut keys = self.ring.keys();
        keys.find(|k| *k >= &key)
            .and_then(|k| self.ring.get(k))
            .or_else(|| keys.next().and_then(|x| self.ring.get(x)))
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use std::string::ToString;

    #[derive(Clone, Debug)]
    struct TestNode {
        host_name: &'static str,
        ip_address: &'static str,
        port: u32,
    }

    impl ToString for TestNode {
        fn to_string(&self) -> String {
            format!("{}{}", self.ip_address.to_string(), self.port.to_string())
        }
    }

    #[test]
    fn test_add_node() {
        let mut hash_ring = Ring::new(3);
        assert_eq!(hash_ring.num_replicas, 3);

        let test_node = TestNode {
            host_name: "Skynet",
            ip_address: "192.168.1.1",
            port: 42,
        };
        hash_ring.add_node(&test_node);
    }

    #[test]
    fn test_remove_node() {
        let mut hash_ring = Ring::new(3);
        assert_eq!(hash_ring.num_replicas, 3);

        let test_node = TestNode {
            host_name: "Skynet",
            ip_address: "192.168.1.1",
            port: 42,
        };
        hash_ring.add_node(&test_node);
        hash_ring.remove_node(&test_node);
    }

    #[test]
    fn test_get_node() {
        let mut hash_ring = Ring::new(3);
        assert_eq!(hash_ring.num_replicas, 3);

        let test_node = TestNode {
            host_name: "Skynet",
            ip_address: "192.168.1.1",
            port: 42,
        };
        hash_ring.add_node(&test_node);
        let my_node = hash_ring.get_node(hash(&test_node.to_string()));

        assert_eq!(my_node.unwrap().host_name, test_node.host_name);
        assert_eq!(my_node.unwrap().ip_address, test_node.ip_address);
        assert_eq!(my_node.unwrap().port, test_node.port);
    }

    #[test]
    fn test_add_nodes() {
        let mut hash_ring = Ring::new(3);
        assert_eq!(hash_ring.num_replicas, 3);

        let test_node1 = TestNode {
            host_name: "Skynet",
            ip_address: "192.168.1.1",
            port: 42,
        };
        let test_node2 = TestNode {
            host_name: "Inferno",
            ip_address: "10.0.1.1",
            port: 666,
        };
        let test_node3 = TestNode {
            host_name: "Klimt",
            ip_address: "127.0.0.1",
            port: 1,
        };

        let v = vec![test_node1.clone(), test_node2.clone(), test_node3.clone()];
        hash_ring.add_nodes(&v);

        let node1 = hash_ring.get_node(hash(&format!(
            "{}{}",
            test_node1.to_string(),
            0.to_string()
        )));
        let node2 = hash_ring.get_node(hash(&format!(
            "{}{}",
            test_node2.to_string(),
            0.to_string()
        )));
        let node3 = hash_ring.get_node(hash(&format!(
            "{}{}",
            test_node3.to_string(),
            0.to_string()
        )));

        assert_eq!(node1.unwrap().host_name, test_node1.host_name);
        assert_eq!(node1.unwrap().ip_address, test_node1.ip_address);
        assert_eq!(node1.unwrap().port, test_node1.port);

        assert_eq!(node2.unwrap().host_name, test_node2.host_name);
        assert_eq!(node2.unwrap().ip_address, test_node2.ip_address);
        assert_eq!(node2.unwrap().port, test_node2.port);

        assert_eq!(node3.unwrap().host_name, test_node3.host_name);
        assert_eq!(node3.unwrap().ip_address, test_node3.ip_address);
        assert_eq!(node3.unwrap().port, test_node3.port);
    }

    #[test]
    fn test_remove_nodes() {
        let mut hash_ring = Ring::new(3);
        assert_eq!(hash_ring.num_replicas, 3);

        let test_node1 = TestNode {
            host_name: "Skynet",
            ip_address: "192.168.1.1",
            port: 42,
        };
        let test_node2 = TestNode {
            host_name: "Inferno",
            ip_address: "10.0.1.1",
            port: 666,
        };
        let test_node3 = TestNode {
            host_name: "Klimt",
            ip_address: "127.0.0.1",
            port: 1,
        };

        let v = vec![test_node1.clone(), test_node2.clone(), test_node3.clone()];
        hash_ring.add_nodes(&v);
        hash_ring.remove_nodes(&v);

        assert!(hash_ring.ring.is_empty());
    }
}
<|endoftext|>