<commit_before>//! Module: czmq-zcert

use {czmq_sys, Error, ErrorKind, RawInterface, Result, Sockish, zmq, ZList};
use std::ffi::{CStr, CString};
use std::os::raw::c_char;
use std::{error, fmt, ptr, result, slice, str};

const KEY_SIZE: usize = 32;

#[derive(Debug, Eq)]
pub struct ZCert {
    zcert: *mut czmq_sys::zcert_t,
    owned: bool,
}

unsafe impl Send for ZCert {}
unsafe impl Sync for ZCert {}

impl Drop for ZCert {
    fn drop(&mut self) {
        if self.owned {
            unsafe { czmq_sys::zcert_destroy(&mut self.zcert) };
        }
    }
}

impl PartialEq for ZCert {
    fn eq(&self, other: &ZCert) -> bool {
        ZCert::eq(self, other)
    }
}

impl ZCert {
    pub fn new() -> Result<ZCert> {
        let zcert = unsafe { czmq_sys::zcert_new() };

        if zcert == ptr::null_mut() {
            return Err(Error::new(ErrorKind::NullPtr, ZCertError::Instantiate));
        }

        Ok(ZCert {
            zcert: zcert,
            owned: true,
        })
    }

    pub fn from_keys(public_key: &[u8], secret_key: &[u8]) -> ZCert {
        ZCert {
            zcert: unsafe { czmq_sys::zcert_new_from(public_key.as_ptr(), secret_key.as_ptr()) },
            owned: true,
        }
    }

    pub fn from_txt(public_txt: &str, secret_txt: &str) -> ZCert {
        ZCert::from_keys(&zmq::z85_decode(public_txt), &zmq::z85_decode(secret_txt))
    }

    pub fn load(filename: &str) -> Result<ZCert> {
        let filename_c = try!(CString::new(filename));
        let zcert = unsafe { czmq_sys::zcert_load(filename_c.as_ptr()) };

        if zcert == ptr::null_mut() {
            return Err(Error::new(ErrorKind::NullPtr, ZCertError::InvalidCert(filename_c.into_string().unwrap())));
        }

        Ok(ZCert {
            zcert: zcert,
            owned: true,
        })
    }

    pub fn public_key<'a>(&'a self) -> &'a [u8] {
        unsafe {
            let ptr = czmq_sys::zcert_public_key(self.zcert);
            slice::from_raw_parts(ptr, KEY_SIZE)
        }
    }

    pub fn secret_key<'a>(&'a self) -> &'a [u8] {
        unsafe {
            let ptr = czmq_sys::zcert_secret_key(self.zcert);
            slice::from_raw_parts(ptr, KEY_SIZE)
        }
    }

    pub fn public_txt<'a>(&'a self) -> &'a str {
        unsafe {
            let ptr = czmq_sys::zcert_public_txt(self.zcert);
            CStr::from_ptr(ptr as *const c_char).to_str().unwrap_or("")
        }
    }

    pub fn secret_txt<'a>(&'a self) -> &'a str {
        unsafe {
            let ptr = czmq_sys::zcert_secret_txt(self.zcert);
            CStr::from_ptr(ptr as *const c_char).to_str().unwrap_or("")
        }
    }

    pub fn set_meta(&self, key: &str, value: &str) {
        let key_c = CString::new(key).unwrap_or(CString::new("").unwrap());
        let value_c = CString::new(value).unwrap_or(CString::new("").unwrap());

        unsafe { czmq_sys::zcert_set_meta(self.zcert, key_c.as_ptr(), "%s\0".as_ptr() as *const i8, value_c.as_ptr()) };
    }

    pub fn meta(&self, key: &str) -> Option<result::Result<String, Vec<u8>>> {
        let key_c = CString::new(key).unwrap_or(CString::new("").unwrap());

        let ptr = unsafe { czmq_sys::zcert_meta(self.zcert, key_c.as_ptr()) };

        if ptr == ptr::null_mut() {
            None
        } else {
            let c_string = unsafe { CStr::from_ptr(ptr).to_owned() };
            let bytes = c_string.as_bytes().to_vec();
            match c_string.into_string() {
                Ok(s) => Some(Ok(s)),
                Err(_) => Some(Err(bytes)),
            }
        }
    }

    pub fn meta_keys(&self) -> ZList<&'static str> {
        let ptr = unsafe { czmq_sys::zcert_meta_keys(self.zcert) };
        ZList::<&'static str>::from_raw(ptr)
    }

    /// Encode certificate metadata into ZMQ wire format.
    ///
    /// ```no_run
    /// # use czmq::{ZCert, ZFrame, ZSock};
    /// # let sock = ZSock::new_rep("...").unwrap();
    /// let cert = ZCert::new().unwrap();
    /// cert.set_meta("key", "value");
    ///
    /// let encoded = cert.encode_meta();
    /// let frame = ZFrame::new(&encoded).unwrap();
    /// frame.send(&sock, None).unwrap();
    /// ```
    pub fn encode_meta(&self) -> Vec<u8> {
        let mut encoded: Vec<u8> = Vec::new();

        for metakey in self.meta_keys() {
            if let Some(Ok(metaval)) = self.meta(metakey) {
                encoded.push(metakey.len() as u8);
                encoded.extend_from_slice(metakey.as_bytes());
                encoded.push(((metaval.len() >> 24) & 0xff) as u8);
                encoded.push(((metaval.len() >> 16) & 0xff) as u8);
                encoded.push(((metaval.len() >> 8) & 0xff) as u8);
                encoded.push((metaval.len() & 0xff) as u8);
                encoded.extend_from_slice(metaval.as_bytes());
            }
        }

        encoded
    }

    /// Decode and set certificate metadata from ZMQ wire format.
    ///
    /// ```no_run
    /// # use czmq::{ZCert, ZFrame, ZSock};
    /// # let sock = ZSock::new_rep("...").unwrap();
    /// let frame = ZFrame::recv(&sock).unwrap();
    /// let encoded = match frame.data().unwrap() {
    ///     Ok(str) => str.into_bytes(),
    ///     Err(bytes) => bytes,
    /// };
    ///
    /// let cert = ZCert::new().unwrap();
    /// cert.decode_meta(&encoded);
    /// ```
    pub fn decode_meta(&self, encoded: &[u8]) -> Result<()> {
        let mut bytes_left = encoded.len();
        let mut index = 0;

        while bytes_left > 1 {
            let name_length: usize = *try!(encoded.get(index).ok_or(Error::new(ErrorKind::InvalidArg, ZCertError::InvalidMetaEncoded))) as usize;
            index += 1;
            bytes_left -= 1;

            if bytes_left < name_length {
                return Err(Error::new(ErrorKind::InvalidArg, ZCertError::InvalidMetaEncoded));
            }

            let name = try!(str::from_utf8(&encoded[index..index + name_length]));
            index += name_length;
            bytes_left -= name_length;

            if bytes_left < 4 {
                return Err(Error::new(ErrorKind::InvalidArg, ZCertError::InvalidMetaEncoded));
            }

            let value_length = (((encoded[index] as u32) << 24) |
                                ((encoded[index + 1] as u32) << 16) |
                                ((encoded[index + 2] as u32) << 8) |
                                  encoded[index + 3] as u32) as usize;
            index += 4;
            bytes_left -= 4;

            if bytes_left < value_length {
                return Err(Error::new(ErrorKind::InvalidArg, ZCertError::InvalidMetaEncoded));
            }

            let value = try!(str::from_utf8(&encoded[index..index + value_length]));
            index += value_length;
            bytes_left -= value_length;

            self.set_meta(name, value);
        }

        Ok(())
    }

    pub fn save(&self, filename: &str) -> Result<()> {
        let filename_c = try!(CString::new(filename));

        unsafe {
            let rc = czmq_sys::zcert_save(self.zcert, filename_c.as_ptr());
            if rc == -1 {
                Err(Error::new(ErrorKind::NonZero, ZCertError::SavePath(filename_c.into_string().unwrap())))
            } else {
                Ok(())
            }
        }
    }

    pub fn save_public(&self, filename: &str) -> Result<()> {
        let filename_c = try!(CString::new(filename));

        unsafe {
            let rc = czmq_sys::zcert_save_public(self.zcert, filename_c.as_ptr());
            if rc == -1 {
                Err(Error::new(ErrorKind::NonZero, ZCertError::SavePath(filename_c.into_string().unwrap())))
            } else {
                Ok(())
            }
        }
    }

    pub fn save_secret(&self, filename: &str) -> Result<()> {
        let filename_c = try!(CString::new(filename));

        unsafe {
            let rc = czmq_sys::zcert_save_secret(self.zcert, filename_c.as_ptr());
            if rc == -1 {
                Err(Error::new(ErrorKind::NonZero, ZCertError::SavePath(filename_c.into_string().unwrap())))
            } else {
                Ok(())
            }
        }
    }

    pub fn apply<S: Sockish>(&self, sock: &S) {
        unsafe { czmq_sys::zcert_apply(self.zcert, sock.borrow_raw()) };
    }

    pub fn dup(&self) -> ZCert {
        let ptr = unsafe { czmq_sys::zcert_dup(self.zcert) };

        ZCert {
            zcert: ptr,
            owned: true,
        }
    }

    pub fn eq(&self, cert: &ZCert) -> bool {
        let result = unsafe { czmq_sys::zcert_eq(self.zcert, cert.zcert) };
        result == 1
    }

    pub fn print(&self) {
        unsafe { czmq_sys::zcert_print(self.zcert) };
    }
}

impl RawInterface<czmq_sys::zcert_t> for ZCert {
    fn from_raw(ptr: *mut czmq_sys::zcert_t, owned: bool) -> ZCert {
        ZCert {
            zcert: ptr,
            owned: owned,
        }
    }

    fn into_raw(mut self) -> *mut czmq_sys::zcert_t {
        self.owned = false;
        self.zcert
    }

    fn borrow_raw(&self) -> *mut czmq_sys::zcert_t {
        self.zcert
    }
}

#[derive(Debug)]
pub enum ZCertError {
    Instantiate,
    InvalidCert(String),
    InvalidMetaEncoded,
    SavePath(String),
}

impl fmt::Display for ZCertError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ZCertError::Instantiate => write!(f, "Could not instantiate new ZCert struct"),
            ZCertError::InvalidCert(ref e) => write!(f, "Could not open certificate at path: {}", e),
            ZCertError::InvalidMetaEncoded => write!(f, "Encoded metadata is invalid"),
            ZCertError::SavePath(ref e) => write!(f, "Could not save certificate file to path: {}", e),
        }
    }
}

impl error::Error for ZCertError {
    fn description(&self) -> &str {
        match *self {
            ZCertError::Instantiate => "Could not instantiate new ZCert struct",
            ZCertError::InvalidCert(_) => "Certificate was invalid or non-existent",
            ZCertError::InvalidMetaEncoded => "Encoded metadata is invalid",
            ZCertError::SavePath(_) => "Could not save certificate file to given path",
        }
    }
}

#[cfg(test)]
mod tests {
    use {ZSock, zsys_init};
    use super::*;
    use zmq;

    const PUBLIC_TXT: &'static str = "Ko9/&3Uw)$U]Zyp>+4$-i/yaDea2QqlDPGl-&V1s";
    const SECRET_TXT: &'static str = "[MfOo!1^1N}zZY/x{[A6^9>VRC.+O6vX&]zYvDC-";

    #[test]
    fn test_public_key() {
        let cert = create_cert();
        let key = cert.public_key();
        let test_key = zmq::z85_decode(PUBLIC_TXT);

        let mut iter = 0;
        for _ in key.iter() {
            assert_eq!(key[iter], test_key[iter]);
            iter += 1;
        }
    }

    #[test]
    fn test_secret_key() {
        let cert = create_cert();
        let key = cert.secret_key();
        let test_key = zmq::z85_decode(SECRET_TXT);

        let mut iter = 0;
        for _ in key.iter() {
            assert_eq!(key[iter], test_key[iter]);
            iter += 1;
        }
    }

    #[test]
    fn test_public_txt() {
        let cert = create_cert();
        assert_eq!(cert.public_txt(), PUBLIC_TXT);
    }

    #[test]
    fn test_secret_txt() {
        let cert = create_cert();
        assert_eq!(cert.secret_txt(), SECRET_TXT);
    }

    #[test]
    fn test_getset_meta() {
        let cert = create_cert();
        cert.set_meta("moo", "cow");
        assert_eq!(cert.meta("moo").unwrap().unwrap(), "cow");
    }

    #[test]
    fn test_meta_keys() {
        let cert = create_cert();
        cert.set_meta("moo", "cow");

        let mut keys = cert.meta_keys();
        assert_eq!(keys.next().unwrap(), "moo");
    }

    #[test]
    fn test_encode_decode_meta() {
        let cert = create_cert();
        cert.set_meta("moo", "cow");
        cert.set_meta("foobar", "baz");
        cert.set_meta("ka", "BOOM!!");
        let encoded = cert.encode_meta();

        let cert = create_cert();
        cert.decode_meta(&encoded).unwrap();

        assert_eq!(cert.meta("moo").unwrap().unwrap(), "cow");
        assert_eq!(cert.meta("foobar").unwrap().unwrap(), "baz");
        assert_eq!(cert.meta("ka").unwrap().unwrap(), "BOOM!!");
    }

    #[test]
    fn test_apply_zmq() {
        let cert = create_cert();
        let mut ctx = zmq::Context::new();
        let sock = ctx.socket(zmq::REQ).unwrap();
        cert.apply(&sock);
        assert_eq!(sock.get_curve_publickey().unwrap().unwrap(), PUBLIC_TXT);
        assert_eq!(sock.get_curve_secretkey().unwrap().unwrap(), SECRET_TXT);
    }

    #[test]
    fn test_apply_zsock() {
        zsys_init();

        let cert = create_cert();
        let sock = ZSock::new_rep("inproc://zcert_test_apply_zsock").unwrap();
        cert.apply(&sock);
        assert_eq!(sock.curve_publickey().unwrap().unwrap(), PUBLIC_TXT);
        assert_eq!(sock.curve_secretkey().unwrap().unwrap(), SECRET_TXT);
    }

    #[test]
    fn test_dup() {
        let cert = create_cert();
        let dup = cert.dup();
        assert_eq!(cert.secret_txt(), dup.secret_txt());
    }

    #[test]
    fn test_eq() {
        let c1 = create_cert();
        let c2 = create_cert();
        assert_eq!(c1, c2);
    }

    fn create_cert() -> ZCert {
        ZCert::from_txt(PUBLIC_TXT, SECRET_TXT)
    }
}
<commit_msg>Squash warning on decode_meta doc example<commit_after>//! Module: czmq-zcert

use {czmq_sys, Error, ErrorKind, RawInterface, Result, Sockish, zmq, ZList};
use std::ffi::{CStr, CString};
use std::os::raw::c_char;
use std::{error, fmt, ptr, result, slice, str};

const KEY_SIZE: usize = 32;

#[derive(Debug, Eq)]
pub struct ZCert {
    zcert: *mut czmq_sys::zcert_t,
    owned: bool,
}

unsafe impl Send for ZCert {}
unsafe impl Sync for ZCert {}

impl Drop for ZCert {
    fn drop(&mut self) {
        if self.owned {
            unsafe { czmq_sys::zcert_destroy(&mut self.zcert) };
        }
    }
}

impl PartialEq for ZCert {
    fn eq(&self, other: &ZCert) -> bool {
        ZCert::eq(self, other)
    }
}

impl ZCert {
    pub fn new() -> Result<ZCert> {
        let zcert = unsafe { czmq_sys::zcert_new() };

        if zcert == ptr::null_mut() {
            return Err(Error::new(ErrorKind::NullPtr, ZCertError::Instantiate));
        }

        Ok(ZCert {
            zcert: zcert,
            owned: true,
        })
    }

    pub fn from_keys(public_key: &[u8], secret_key: &[u8]) -> ZCert {
        ZCert {
            zcert: unsafe { czmq_sys::zcert_new_from(public_key.as_ptr(), secret_key.as_ptr()) },
            owned: true,
        }
    }

    pub fn from_txt(public_txt: &str, secret_txt: &str) -> ZCert {
        ZCert::from_keys(&zmq::z85_decode(public_txt), &zmq::z85_decode(secret_txt))
    }

    pub fn load(filename: &str) -> Result<ZCert> {
        let filename_c = try!(CString::new(filename));
        let zcert = unsafe { czmq_sys::zcert_load(filename_c.as_ptr()) };

        if zcert == ptr::null_mut() {
            return Err(Error::new(ErrorKind::NullPtr, ZCertError::InvalidCert(filename_c.into_string().unwrap())));
        }

        Ok(ZCert {
            zcert: zcert,
            owned: true,
        })
    }

    pub fn public_key<'a>(&'a self) -> &'a [u8] {
        unsafe {
            let ptr = czmq_sys::zcert_public_key(self.zcert);
            slice::from_raw_parts(ptr, KEY_SIZE)
        }
    }

    pub fn secret_key<'a>(&'a self) -> &'a [u8] {
        unsafe {
            let ptr = czmq_sys::zcert_secret_key(self.zcert);
            slice::from_raw_parts(ptr, KEY_SIZE)
        }
    }

    pub fn public_txt<'a>(&'a self) -> &'a str {
        unsafe {
            let ptr = czmq_sys::zcert_public_txt(self.zcert);
            CStr::from_ptr(ptr as *const c_char).to_str().unwrap_or("")
        }
    }

    pub fn secret_txt<'a>(&'a self) -> &'a str {
        unsafe {
            let ptr = czmq_sys::zcert_secret_txt(self.zcert);
            CStr::from_ptr(ptr as *const c_char).to_str().unwrap_or("")
        }
    }

    pub fn set_meta(&self, key: &str, value: &str) {
        let key_c = CString::new(key).unwrap_or(CString::new("").unwrap());
        let value_c = CString::new(value).unwrap_or(CString::new("").unwrap());

        unsafe { czmq_sys::zcert_set_meta(self.zcert, key_c.as_ptr(), "%s\0".as_ptr() as *const i8, value_c.as_ptr()) };
    }

    pub fn meta(&self, key: &str) -> Option<result::Result<String, Vec<u8>>> {
        let key_c = CString::new(key).unwrap_or(CString::new("").unwrap());

        let ptr = unsafe { czmq_sys::zcert_meta(self.zcert, key_c.as_ptr()) };

        if ptr == ptr::null_mut() {
            None
        } else {
            let c_string = unsafe { CStr::from_ptr(ptr).to_owned() };
            let bytes = c_string.as_bytes().to_vec();
            match c_string.into_string() {
                Ok(s) => Some(Ok(s)),
                Err(_) => Some(Err(bytes)),
            }
        }
    }

    pub fn meta_keys(&self) -> ZList<&'static str> {
        let ptr = unsafe { czmq_sys::zcert_meta_keys(self.zcert) };
        ZList::<&'static str>::from_raw(ptr)
    }

    /// Encode certificate metadata into ZMQ wire format.
    ///
    /// ```no_run
    /// # use czmq::{ZCert, ZFrame, ZSock};
    /// # let sock = ZSock::new_rep("...").unwrap();
    /// let cert = ZCert::new().unwrap();
    /// cert.set_meta("key", "value");
    ///
    /// let encoded = cert.encode_meta();
    /// let frame = ZFrame::new(&encoded).unwrap();
    /// frame.send(&sock, None).unwrap();
    /// ```
    pub fn encode_meta(&self) -> Vec<u8> {
        let mut encoded: Vec<u8> = Vec::new();

        for metakey in self.meta_keys() {
            if let Some(Ok(metaval)) = self.meta(metakey) {
                encoded.push(metakey.len() as u8);
                encoded.extend_from_slice(metakey.as_bytes());
                encoded.push(((metaval.len() >> 24) & 0xff) as u8);
                encoded.push(((metaval.len() >> 16) & 0xff) as u8);
                encoded.push(((metaval.len() >> 8) & 0xff) as u8);
                encoded.push((metaval.len() & 0xff) as u8);
                encoded.extend_from_slice(metaval.as_bytes());
            }
        }

        encoded
    }

    /// Decode and set certificate metadata from ZMQ wire format.
    ///
    /// ```no_run
    /// # use czmq::{ZCert, ZFrame, ZSock};
    /// # let sock = ZSock::new_rep("...").unwrap();
    /// let frame = ZFrame::recv(&sock).unwrap();
    /// let encoded = match frame.data().unwrap() {
    ///     Ok(str) => str.into_bytes(),
    ///     Err(bytes) => bytes,
    /// };
    ///
    /// let cert = ZCert::new().unwrap();
    /// cert.decode_meta(&encoded).unwrap();
    /// ```
    pub fn decode_meta(&self, encoded: &[u8]) -> Result<()> {
        let mut bytes_left = encoded.len();
        let mut index = 0;

        while bytes_left > 1 {
            let name_length: usize = *try!(encoded.get(index).ok_or(Error::new(ErrorKind::InvalidArg, ZCertError::InvalidMetaEncoded))) as usize;
            index += 1;
            bytes_left -= 1;

            if bytes_left < name_length {
                return Err(Error::new(ErrorKind::InvalidArg, ZCertError::InvalidMetaEncoded));
            }

            let name = try!(str::from_utf8(&encoded[index..index + name_length]));
            index += name_length;
            bytes_left -= name_length;

            if bytes_left < 4 {
                return Err(Error::new(ErrorKind::InvalidArg, ZCertError::InvalidMetaEncoded));
            }

            let value_length = (((encoded[index] as u32) << 24) |
                                ((encoded[index + 1] as u32) << 16) |
                                ((encoded[index + 2] as u32) << 8) |
                                  encoded[index + 3] as u32) as usize;
            index += 4;
            bytes_left -= 4;

            if bytes_left < value_length {
                return Err(Error::new(ErrorKind::InvalidArg, ZCertError::InvalidMetaEncoded));
            }

            let value = try!(str::from_utf8(&encoded[index..index + value_length]));
            index += value_length;
            bytes_left -= value_length;

            self.set_meta(name, value);
        }

        Ok(())
    }

    pub fn save(&self, filename: &str) -> Result<()> {
        let filename_c = try!(CString::new(filename));

        unsafe {
            let rc = czmq_sys::zcert_save(self.zcert, filename_c.as_ptr());
            if rc == -1 {
                Err(Error::new(ErrorKind::NonZero, ZCertError::SavePath(filename_c.into_string().unwrap())))
            } else {
                Ok(())
            }
        }
    }

    pub fn save_public(&self, filename: &str) -> Result<()> {
        let filename_c = try!(CString::new(filename));

        unsafe {
            let rc = czmq_sys::zcert_save_public(self.zcert, filename_c.as_ptr());
            if rc == -1 {
                Err(Error::new(ErrorKind::NonZero, ZCertError::SavePath(filename_c.into_string().unwrap())))
            } else {
                Ok(())
            }
        }
    }

    pub fn save_secret(&self, filename: &str) -> Result<()> {
        let filename_c = try!(CString::new(filename));

        unsafe {
            let rc = czmq_sys::zcert_save_secret(self.zcert, filename_c.as_ptr());
            if rc == -1 {
                Err(Error::new(ErrorKind::NonZero, ZCertError::SavePath(filename_c.into_string().unwrap())))
            } else {
                Ok(())
            }
        }
    }

    pub fn apply<S: Sockish>(&self, sock: &S) {
        unsafe { czmq_sys::zcert_apply(self.zcert, sock.borrow_raw()) };
    }

    pub fn dup(&self) -> ZCert {
        let ptr = unsafe { czmq_sys::zcert_dup(self.zcert) };

        ZCert {
            zcert: ptr,
            owned: true,
        }
    }

    pub fn eq(&self, cert: &ZCert) -> bool {
        let result = unsafe { czmq_sys::zcert_eq(self.zcert, cert.zcert) };
        result == 1
    }

    pub fn print(&self) {
        unsafe { czmq_sys::zcert_print(self.zcert) };
    }
}

impl RawInterface<czmq_sys::zcert_t> for ZCert {
    fn from_raw(ptr: *mut czmq_sys::zcert_t, owned: bool) -> ZCert {
        ZCert {
            zcert: ptr,
            owned: owned,
        }
    }

    fn into_raw(mut self) -> *mut czmq_sys::zcert_t {
        self.owned = false;
        self.zcert
    }

    fn borrow_raw(&self) -> *mut czmq_sys::zcert_t {
        self.zcert
    }
}

#[derive(Debug)]
pub enum ZCertError {
    Instantiate,
    InvalidCert(String),
    InvalidMetaEncoded,
    SavePath(String),
}

impl fmt::Display for ZCertError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ZCertError::Instantiate => write!(f, "Could not instantiate new ZCert struct"),
            ZCertError::InvalidCert(ref e) => write!(f, "Could not open certificate at path: {}", e),
            ZCertError::InvalidMetaEncoded => write!(f, "Encoded metadata is invalid"),
            ZCertError::SavePath(ref e) => write!(f, "Could not save certificate file to path: {}", e),
        }
    }
}

impl error::Error for ZCertError {
    fn description(&self) -> &str {
        match *self {
            ZCertError::Instantiate => "Could not instantiate new ZCert struct",
            ZCertError::InvalidCert(_) => "Certificate was invalid or non-existent",
            ZCertError::InvalidMetaEncoded => "Encoded metadata is invalid",
            ZCertError::SavePath(_) => "Could not save certificate file to given path",
        }
    }
}

#[cfg(test)]
mod tests {
    use {ZSock, zsys_init};
    use super::*;
    use zmq;

    const PUBLIC_TXT: &'static str = "Ko9/&3Uw)$U]Zyp>+4$-i/yaDea2QqlDPGl-&V1s";
    const SECRET_TXT: &'static str = "[MfOo!1^1N}zZY/x{[A6^9>VRC.+O6vX&]zYvDC-";

    #[test]
    fn test_public_key() {
        let cert = create_cert();
        let key = cert.public_key();
        let test_key = zmq::z85_decode(PUBLIC_TXT);

        let mut iter = 0;
        for _ in key.iter() {
            assert_eq!(key[iter], test_key[iter]);
            iter += 1;
        }
    }

    #[test]
    fn test_secret_key() {
        let cert = create_cert();
        let key = cert.secret_key();
        let test_key = zmq::z85_decode(SECRET_TXT);

        let mut iter = 0;
        for _ in key.iter() {
            assert_eq!(key[iter], test_key[iter]);
            iter += 1;
        }
    }

    #[test]
    fn test_public_txt() {
        let cert = create_cert();
        assert_eq!(cert.public_txt(), PUBLIC_TXT);
    }

    #[test]
    fn test_secret_txt() {
        let cert = create_cert();
        assert_eq!(cert.secret_txt(), SECRET_TXT);
    }

    #[test]
    fn test_getset_meta() {
        let cert = create_cert();
        cert.set_meta("moo", "cow");
        assert_eq!(cert.meta("moo").unwrap().unwrap(), "cow");
    }

    #[test]
    fn test_meta_keys() {
        let cert = create_cert();
        cert.set_meta("moo", "cow");

        let mut keys = cert.meta_keys();
        assert_eq!(keys.next().unwrap(), "moo");
    }

    #[test]
    fn test_encode_decode_meta() {
        let cert = create_cert();
        cert.set_meta("moo", "cow");
        cert.set_meta("foobar", "baz");
        cert.set_meta("ka", "BOOM!!");
        let encoded = cert.encode_meta();

        let cert = create_cert();
        cert.decode_meta(&encoded).unwrap();

        assert_eq!(cert.meta("moo").unwrap().unwrap(), "cow");
        assert_eq!(cert.meta("foobar").unwrap().unwrap(), "baz");
        assert_eq!(cert.meta("ka").unwrap().unwrap(), "BOOM!!");
    }

    #[test]
    fn test_apply_zmq() {
        let cert = create_cert();
        let mut ctx = zmq::Context::new();
        let sock = ctx.socket(zmq::REQ).unwrap();
        cert.apply(&sock);
        assert_eq!(sock.get_curve_publickey().unwrap().unwrap(), PUBLIC_TXT);
        assert_eq!(sock.get_curve_secretkey().unwrap().unwrap(), SECRET_TXT);
    }

    #[test]
    fn test_apply_zsock() {
        zsys_init();

        let cert = create_cert();
        let sock = ZSock::new_rep("inproc://zcert_test_apply_zsock").unwrap();
        cert.apply(&sock);
        assert_eq!(sock.curve_publickey().unwrap().unwrap(), PUBLIC_TXT);
        assert_eq!(sock.curve_secretkey().unwrap().unwrap(), SECRET_TXT);
    }

    #[test]
    fn test_dup() {
        let cert = create_cert();
        let dup = cert.dup();
        assert_eq!(cert.secret_txt(), dup.secret_txt());
    }

    #[test]
    fn test_eq() {
        let c1 = create_cert();
        let c2 = create_cert();
        assert_eq!(c1, c2);
    }

    fn create_cert() -> ZCert {
        ZCert::from_txt(PUBLIC_TXT, SECRET_TXT)
    }
}
<|endoftext|>


<commit_before>//! Parse and decode COSE signatures.

use cbor::CborType;
use cbor::decoder::decode;
use {CoseError, SignatureAlgorithm};
use util::get_sig_struct_bytes;
use std::collections::BTreeMap;

const COSE_SIGN_TAG: u64 = 98;

/// The result of `decode_signature` holding a decoded COSE signature.
#[derive(Debug)]
pub struct CoseSignature {
    pub signature_type: SignatureAlgorithm,
    pub signature: Vec<u8>,
    pub signer_cert: Vec<u8>,
    pub certs: Vec<Vec<u8>>,
    pub to_verify: Vec<u8>,
}

pub const COSE_TYPE_ES256: i64 = -7;
pub const COSE_TYPE_ES384: i64 = -35;
pub const COSE_TYPE_ES512: i64 = -36;
pub const COSE_TYPE_PS256: i64 = -37;

macro_rules! unpack {
   ($to:tt, $var:ident) => (
        match *$var {
            CborType::$to(ref cbor_object) => {
                cbor_object
            }
            _ => return Err(CoseError::UnexpectedType),
        };
    )
}

fn get_map_value(
    map: &BTreeMap<CborType, CborType>,
    key: &CborType,
) -> Result<CborType, CoseError> {
    match map.get(key) {
        Some(x) => Ok(x.clone()),
        _ => Err(CoseError::MissingHeader),
    }
}

fn validate_empty_map(map: &CborType) -> Result<(), CoseError> {
    let unpacked = unpack!(Map, map);
    if unpacked.len() != 0 {
        return Err(CoseError::MalformedInput);
    }
    Ok(())
}

// This syntax is a little unintuitive. Taken together, the two previous definitions essentially
// mean:
//
// COSE_Sign = [
//     protected : empty_or_serialized_map,
//     unprotected : header_map
//     payload : bstr / nil,
//     signatures : [+ COSE_Signature]
// ]
//
// (COSE_Sign is an array. The first element is an empty or serialized map (in our case, it is
// never expected to be empty). The second element is a map (it is expected to be empty. The third
// element is a bstr or nil (it is expected to be nil). The fourth element is an array of
// COSE_Signature.)
//
// COSE_Signature =  [
//     Headers,
//     signature : bstr
// ]
//
// but again, unpacking this:
//
// COSE_Signature =  [
//     protected : empty_or_serialized_map,
//     unprotected : header_map
//     signature : bstr
// ]
fn decode_signature_struct(
    cose_signature: &CborType,
    payload: &[u8],
    protected_body_head: &CborType,
) -> Result<CoseSignature, CoseError> {
    let cose_signature = unpack!(Array, cose_signature);
    if cose_signature.len() != 3 {
        return Err(CoseError::MalformedInput);
    }
    let protected_signature_header_serialized = &cose_signature[0];
    let protected_signature_header_bytes = unpack!(Bytes, protected_signature_header_serialized);

    // Parse the protected signature header.
    let protected_signature_header = &match decode(&protected_signature_header_bytes) {
        Err(_) => return Err(CoseError::DecodingFailure),
        Ok(value) => value,
    };
    let protected_signature_header = unpack!(Map, protected_signature_header);
    if protected_signature_header.len() != 2 {
        return Err(CoseError::MalformedInput);
    }
    let signature_algorithm = get_map_value(&protected_signature_header, &CborType::Integer(1))?;
    let signature_algorithm = match signature_algorithm {
        CborType::SignedInteger(val) => {
            match val {
                COSE_TYPE_ES256 => SignatureAlgorithm::ES256,
                COSE_TYPE_ES384 => SignatureAlgorithm::ES384,
                COSE_TYPE_ES512 => SignatureAlgorithm::ES512,
                COSE_TYPE_PS256 => SignatureAlgorithm::PS256,
                _ => return Err(CoseError::UnexpectedHeaderValue),
            }
        }
        _ => return Err(CoseError::UnexpectedType),
    };

    let ee_cert = &get_map_value(&protected_signature_header, &CborType::Integer(4))?;
    let ee_cert = unpack!(Bytes, ee_cert).clone();

    // The unprotected header section is expected to be an empty map.
    validate_empty_map(&cose_signature[1])?;

    // Build signature structure to verify.
    let signature_bytes = &cose_signature[2];
    let signature_bytes = unpack!(Bytes, signature_bytes).clone();
    let sig_structure_bytes = get_sig_struct_bytes(
        protected_body_head.clone(),
        protected_signature_header_serialized.clone(),
        payload,
    );

    // Read intermediate certificates from protected_body_head.
    // Any tampering of the protected header during transport will be detected
    // because it is input to the signature verification.
    // Note that a protected header has to be present and hold a kid with an
    // empty list of intermediate certificates.
    let protected_body_head_bytes = unpack!(Bytes, protected_body_head);
    let protected_body_head_map = &match decode(protected_body_head_bytes) {
        Ok(value) => value,
        Err(_) => return Err(CoseError::DecodingFailure),
    };
    let protected_body_head_map = unpack!(Map, protected_body_head_map);
    if protected_body_head_map.len() != 1 {
        return Err(CoseError::MalformedInput);
    }
    let intermediate_certs_array = &get_map_value(&protected_body_head_map, &CborType::Integer(4))?;
    let intermediate_certs = unpack!(Array, intermediate_certs_array);
    let mut certs: Vec<Vec<u8>> = Vec::new();
    for cert in intermediate_certs {
        let cert = unpack!(Bytes, cert);
        certs.push(cert.clone());
    }

    Ok(CoseSignature {
        signature_type: signature_algorithm,
        signature: signature_bytes,
        signer_cert: ee_cert,
        certs: certs,
        to_verify: sig_structure_bytes,
    })
}

/// Decode COSE signature bytes and return a vector of `CoseSignature`.
///
///```rust,ignore
/// COSE_Sign = [
///     Headers,
///     payload : bstr / nil,
///     signatures : [+ COSE_Signature]
/// ]
///
/// Headers = (
///     protected : empty_or_serialized_map,
///     unprotected : header_map
/// )
///```
pub fn decode_signature(bytes: &[u8], payload: &[u8]) -> Result<Vec<CoseSignature>, CoseError> {
    // This has to be a COSE_Sign object, which is a tagged array.
    let tagged_cose_sign = match decode(&bytes) {
        Err(_) => return Err(CoseError::DecodingFailure),
        Ok(value) => value,
    };
    let cose_sign_array = match tagged_cose_sign {
        CborType::Tag(tag, cose_sign) => {
            if tag != COSE_SIGN_TAG {
                return Err(CoseError::UnexpectedTag);
            }
            match *cose_sign {
                CborType::Array(values) => values,
                _ => return Err(CoseError::UnexpectedType),
            }
        }
        _ => return Err(CoseError::UnexpectedType),
    };
    if cose_sign_array.len() != 4 {
        return Err(CoseError::MalformedInput);
    }

    // The unprotected header section is expected to be an empty map.
    validate_empty_map(&cose_sign_array[1])?;

    // The payload is expected to be Null (i.e. this is a detached signature).
    match &cose_sign_array[2] {
        &CborType::Null => {}
        _ => return Err(CoseError::UnexpectedType),
    };

    let signatures = &cose_sign_array[3];
    let signatures = unpack!(Array, signatures);

    // Decode COSE_Signatures.
    // There has to be at least one signature to make this a valid COSE signature.
    if signatures.len() < 1 {
        return Err(CoseError::MalformedInput);
    }
    let mut result = Vec::new();
    for cose_signature in signatures {
        // cose_sign_array[0] holds the protected body header.
        let signature = decode_signature_struct(cose_signature, payload, &cose_sign_array[0])?;
        result.push(signature);
    }

    Ok(result)
}
<commit_msg>address clippy comments<commit_after>//! Parse and decode COSE signatures.

use cbor::CborType;
use cbor::decoder::decode;
use {CoseError, SignatureAlgorithm};
use util::get_sig_struct_bytes;
use std::collections::BTreeMap;

const COSE_SIGN_TAG: u64 = 98;

/// The result of `decode_signature` holding a decoded COSE signature.
#[derive(Debug)]
pub struct CoseSignature {
    pub signature_type: SignatureAlgorithm,
    pub signature: Vec<u8>,
    pub signer_cert: Vec<u8>,
    pub certs: Vec<Vec<u8>>,
    pub to_verify: Vec<u8>,
}

pub const COSE_TYPE_ES256: i64 = -7;
pub const COSE_TYPE_ES384: i64 = -35;
pub const COSE_TYPE_ES512: i64 = -36;
pub const COSE_TYPE_PS256: i64 = -37;

macro_rules! unpack {
   ($to:tt, $var:ident) => (
        match *$var {
            CborType::$to(ref cbor_object) => {
                cbor_object
            }
            _ => return Err(CoseError::UnexpectedType),
        };
    )
}

fn get_map_value(
    map: &BTreeMap<CborType, CborType>,
    key: &CborType,
) -> Result<CborType, CoseError> {
    match map.get(key) {
        Some(x) => Ok(x.clone()),
        _ => Err(CoseError::MissingHeader),
    }
}

fn validate_empty_map(map: &CborType) -> Result<(), CoseError> {
    let unpacked = unpack!(Map, map);
    if !unpacked.is_empty() {
        return Err(CoseError::MalformedInput);
    }
    Ok(())
}

// This syntax is a little unintuitive. Taken together, the two previous definitions essentially
// mean:
//
// COSE_Sign = [
//     protected : empty_or_serialized_map,
//     unprotected : header_map
//     payload : bstr / nil,
//     signatures : [+ COSE_Signature]
// ]
//
// (COSE_Sign is an array. The first element is an empty or serialized map (in our case, it is
// never expected to be empty). The second element is a map (it is expected to be empty. The third
// element is a bstr or nil (it is expected to be nil). The fourth element is an array of
// COSE_Signature.)
//
// COSE_Signature =  [
//     Headers,
//     signature : bstr
// ]
//
// but again, unpacking this:
//
// COSE_Signature =  [
//     protected : empty_or_serialized_map,
//     unprotected : header_map
//     signature : bstr
// ]
fn decode_signature_struct(
    cose_signature: &CborType,
    payload: &[u8],
    protected_body_head: &CborType,
) -> Result<CoseSignature, CoseError> {
    let cose_signature = unpack!(Array, cose_signature);
    if cose_signature.len() != 3 {
        return Err(CoseError::MalformedInput);
    }
    let protected_signature_header_serialized = &cose_signature[0];
    let protected_signature_header_bytes = unpack!(Bytes, protected_signature_header_serialized);

    // Parse the protected signature header.
    let protected_signature_header = &match decode(protected_signature_header_bytes) {
        Err(_) => return Err(CoseError::DecodingFailure),
        Ok(value) => value,
    };
    let protected_signature_header = unpack!(Map, protected_signature_header);
    if protected_signature_header.len() != 2 {
        return Err(CoseError::MalformedInput);
    }
    let signature_algorithm = get_map_value(protected_signature_header, &CborType::Integer(1))?;
    let signature_algorithm = match signature_algorithm {
        CborType::SignedInteger(val) => {
            match val {
                COSE_TYPE_ES256 => SignatureAlgorithm::ES256,
                COSE_TYPE_ES384 => SignatureAlgorithm::ES384,
                COSE_TYPE_ES512 => SignatureAlgorithm::ES512,
                COSE_TYPE_PS256 => SignatureAlgorithm::PS256,
                _ => return Err(CoseError::UnexpectedHeaderValue),
            }
        }
        _ => return Err(CoseError::UnexpectedType),
    };

    let ee_cert = &get_map_value(protected_signature_header, &CborType::Integer(4))?;
    let ee_cert = unpack!(Bytes, ee_cert).clone();

    // The unprotected header section is expected to be an empty map.
    validate_empty_map(&cose_signature[1])?;

    // Build signature structure to verify.
    let signature_bytes = &cose_signature[2];
    let signature_bytes = unpack!(Bytes, signature_bytes).clone();
    let sig_structure_bytes = get_sig_struct_bytes(
        protected_body_head.clone(),
        protected_signature_header_serialized.clone(),
        payload,
    );

    // Read intermediate certificates from protected_body_head.
    // Any tampering of the protected header during transport will be detected
    // because it is input to the signature verification.
    // Note that a protected header has to be present and hold a kid with an
    // empty list of intermediate certificates.
    let protected_body_head_bytes = unpack!(Bytes, protected_body_head);
    let protected_body_head_map = &match decode(protected_body_head_bytes) {
        Ok(value) => value,
        Err(_) => return Err(CoseError::DecodingFailure),
    };
    let protected_body_head_map = unpack!(Map, protected_body_head_map);
    if protected_body_head_map.len() != 1 {
        return Err(CoseError::MalformedInput);
    }
    let intermediate_certs_array = &get_map_value(protected_body_head_map, &CborType::Integer(4))?;
    let intermediate_certs = unpack!(Array, intermediate_certs_array);
    let mut certs: Vec<Vec<u8>> = Vec::new();
    for cert in intermediate_certs {
        let cert = unpack!(Bytes, cert);
        certs.push(cert.clone());
    }

    Ok(CoseSignature {
        signature_type: signature_algorithm,
        signature: signature_bytes,
        signer_cert: ee_cert,
        certs: certs,
        to_verify: sig_structure_bytes,
    })
}

/// Decode COSE signature bytes and return a vector of `CoseSignature`.
///
///```rust,ignore
/// COSE_Sign = [
///     Headers,
///     payload : bstr / nil,
///     signatures : [+ COSE_Signature]
/// ]
///
/// Headers = (
///     protected : empty_or_serialized_map,
///     unprotected : header_map
/// )
///```
pub fn decode_signature(bytes: &[u8], payload: &[u8]) -> Result<Vec<CoseSignature>, CoseError> {
    // This has to be a COSE_Sign object, which is a tagged array.
    let tagged_cose_sign = match decode(bytes) {
        Err(_) => return Err(CoseError::DecodingFailure),
        Ok(value) => value,
    };
    let cose_sign_array = match tagged_cose_sign {
        CborType::Tag(tag, cose_sign) => {
            if tag != COSE_SIGN_TAG {
                return Err(CoseError::UnexpectedTag);
            }
            match *cose_sign {
                CborType::Array(values) => values,
                _ => return Err(CoseError::UnexpectedType),
            }
        }
        _ => return Err(CoseError::UnexpectedType),
    };
    if cose_sign_array.len() != 4 {
        return Err(CoseError::MalformedInput);
    }

    // The unprotected header section is expected to be an empty map.
    validate_empty_map(&cose_sign_array[1])?;

    // The payload is expected to be Null (i.e. this is a detached signature).
    match cose_sign_array[2] {
        CborType::Null => {}
        _ => return Err(CoseError::UnexpectedType),
    };

    let signatures = &cose_sign_array[3];
    let signatures = unpack!(Array, signatures);

    // Decode COSE_Signatures.
    // There has to be at least one signature to make this a valid COSE signature.
    if signatures.len() < 1 {
        return Err(CoseError::MalformedInput);
    }
    let mut result = Vec::new();
    for cose_signature in signatures {
        // cose_sign_array[0] holds the protected body header.
        let signature = decode_signature_struct(cose_signature, payload, &cose_sign_array[0])?;
        result.push(signature);
    }

    Ok(result)
}
<|endoftext|>


<commit_before>use types::{OscPacket, OscType, Result, OscMessage, OscMidiMessage, OscColor, OscBundle};
use errors::OscError;
use utils;

use std::{io, char};
use std::io::{Read, BufRead};

use byteorder::{BigEndian, ReadBytesExt};

/// Common MTP size for ethernet
pub const MTP: usize = 1536;

pub fn decode(msg: &[u8]) -> Result<OscPacket> {
    match msg[0] as char {
        '/' => {
            decode_message(msg)
        }
        '#' => {
            decode_bundle(msg)
        }
        _ => Err(OscError::BadPacket("Unknown message format.")),
    }
}

fn decode_message(msg: &[u8]) -> Result<OscPacket> {
    let mut cursor: io::Cursor<&[u8]> = io::Cursor::new(msg);

    let addr: String = try!(read_osc_string(&mut cursor));
    let type_tags: String = try!(read_osc_string(&mut cursor));

    if type_tags.len() > 1 {
        let args: Vec<OscType> = try!(read_osc_args(&mut cursor, type_tags));

        Ok(OscPacket::Message(OscMessage {
            addr: addr,
            args: Some(args),
        }))
    } else {
        Ok(OscPacket::Message(OscMessage {
            addr: addr,
            args: None,
        }))
    }
}

fn decode_bundle(msg: &[u8]) -> Result<OscPacket> {
    let mut cursor: io::Cursor<&[u8]> = io::Cursor::new(msg);

    let b = try!(read_osc_string(&mut cursor));
    if b != "bundle" {
        return Err(OscError::BadBundle(format!("Wrong bundle specifier: {}", bundle_tag)));
    }

    let time_tag = try!(read_time_tag(&mut cursor));
    let mut bundle: Vec<OscPacket> = Vec::new();
    let size: usize = try!(cursor.read_u32::<BigEndian>()
                                 .map_err(OscError::ByteOrderError)) as usize;
    let mut buf: Vec<u8> = Vec::with_capacity(size);
    let cnt: usize = try!(cursor.take(size as u64)
                                .read_to_end(&mut buf)
                                .map_err(OscError::ReadError));
    if cnt == size {
        try!(decode(&mut buf).map(|p| bundle.push(p)));
    } else {
        return Err(OscError::BadBundle("Bundle shorter than expected!".to_string()));
    }

    Ok(OscPacket::Bundle(OscBundle {
        timetag: time_tag,
        content: bundle,
    }))
}

fn read_osc_string(cursor: &mut io::Cursor<&[u8]>) -> Result<String> {
    let mut str_buf: Vec<u8> = Vec::new();
    try!(cursor.read_until(0, &mut str_buf)
               .map_err(OscError::ReadError)); // ignore returned byte count
    pad_cursor(cursor);
    // convert to String and remove nul bytes
    String::from_utf8(str_buf)
        .map_err(|e| OscError::StringError(e))
        .map(|s| {
            s.trim_matches(0u8 as char)
             .to_string()
        })
}

fn read_osc_args(cursor: &mut io::Cursor<&[u8]>, raw_type_tags: String) -> Result<Vec<OscType>> {
    let type_tags: Vec<char> = raw_type_tags.chars()
                                            .skip(1)
                                            .map(|c| c as char)
                                            .collect();
    let mut args: Vec<OscType> = Vec::with_capacity(type_tags.len());
    for tag in type_tags {
        let arg: OscType = try!(read_osc_arg(cursor, tag));
        args.push(arg);
    }
    Ok(args)
}

fn read_osc_arg(cursor: &mut io::Cursor<&[u8]>, tag: char) -> Result<OscType> {
    match tag {
        'f' => {
            cursor.read_f32::<BigEndian>()
                  .map(OscType::Float)
                  .map_err(OscError::ByteOrderError)
        }
        'd' => {
            cursor.read_f64::<BigEndian>()
                  .map(OscType::Double)
                  .map_err(OscError::ByteOrderError)
        }
        'i' => {
            cursor.read_i32::<BigEndian>()
                  .map(OscType::Int)
                  .map_err(OscError::ByteOrderError)
        }
        'h' => {
            cursor.read_i64::<BigEndian>()
                  .map(OscType::Long)
                  .map_err(OscError::ByteOrderError)
        }
        's' => {
            read_osc_string(cursor).map(|s| OscType::String(s))
        }
        't' => {
            read_time_tag(cursor)
        }
        'b' => {
            read_blob(cursor)
        }
        'r' => {
            read_osc_color(cursor)
        }
        'T' => {
            Ok(OscType::Bool(true))
        }
        'F' => {
            Ok(OscType::Bool(false))
        }
        'N' => {
            Ok(OscType::Nil)
        }
        'I' => {
            Ok(OscType::Inf)
        }
        'c' => {
            read_char(cursor)
        }
        'm' => {
            read_midi_message(cursor)
        }
        _ => Err(OscError::BadArg(format!("Type tag \"{}\" is not implemented!", tag))),
    }
}

fn read_char(cursor: &mut io::Cursor<&[u8]>) -> Result<OscType> {
    let opt_char = try!(cursor.read_u32::<BigEndian>()
                              .map(char::from_u32)
                              .map_err(OscError::ByteOrderError));
    match opt_char {
        Some(c) => Ok(OscType::Char(c)),
        None => Err(OscError::BadArg("Argument is not a char!".to_string())),
    }
}

fn read_blob(cursor: &mut io::Cursor<&[u8]>) -> Result<OscType> {
    let size: usize = try!(cursor.read_u32::<BigEndian>()
                                 .map_err(OscError::ByteOrderError)) as usize;
    let mut byte_buf: Vec<u8> = Vec::with_capacity(size);

    try!(cursor.take(size as u64)
               .read_to_end(&mut byte_buf)
               .map_err(OscError::ReadError));

    pad_cursor(cursor);

    Ok(OscType::Blob(byte_buf))
}

fn read_time_tag(cursor: &mut io::Cursor<&[u8]>) -> Result<OscType> {
    let date = try!(cursor.read_u32::<BigEndian>()
                          .map_err(OscError::ByteOrderError));
    let frac = try!(cursor.read_u32::<BigEndian>()
                          .map_err(OscError::ByteOrderError));

    Ok(OscType::Time(date, frac))
}

fn read_midi_message(cursor: &mut io::Cursor<&[u8]>) -> Result<OscType> {
    let mut buf: Vec<u8> = Vec::with_capacity(4);
    try!(cursor.take(4).read_to_end(&mut buf).map_err(OscError::ReadError));

    Ok(OscType::Midi(OscMidiMessage {
        port: buf[0],
        status: buf[1],
        data1: buf[2],
        data2: buf[3],
    }))

}

fn read_osc_color(cursor: &mut io::Cursor<&[u8]>) -> Result<OscType> {
    let mut buf: Vec<u8> = Vec::with_capacity(4);
    try!(cursor.take(4).read_to_end(&mut buf).map_err(OscError::ReadError));

    Ok(OscType::Color(OscColor {
        red: buf[0],
        green: buf[1],
        blue: buf[2],
        alpha: buf[3],
    }))
}

fn pad_cursor(cursor: &mut io::Cursor<&[u8]>) {
    let pos = cursor.position();
    cursor.set_position(utils::pad(pos));
}
<commit_msg>Fix index exception by checking for an empty message.<commit_after>use types::{OscPacket, OscType, Result, OscMessage, OscMidiMessage, OscColor, OscBundle};
use errors::OscError;
use utils;

use std::{io, char};
use std::io::{Read, BufRead};

use byteorder::{BigEndian, ReadBytesExt};

/// Common MTP size for ethernet
pub const MTP: usize = 1536;

pub fn decode(msg: &[u8]) -> Result<OscPacket> {
    if msg.len() == 0 {
        return Err(OscError::BadPacket("Empty packet."));
    }

    match msg[0] as char {
        '/' => {
            decode_message(msg)
        }
        '#' => {
            decode_bundle(msg)
        }
        _ => Err(OscError::BadPacket("Unknown message format.")),
    }
}

fn decode_message(msg: &[u8]) -> Result<OscPacket> {
    let mut cursor: io::Cursor<&[u8]> = io::Cursor::new(msg);

    let addr: String = try!(read_osc_string(&mut cursor));
    let type_tags: String = try!(read_osc_string(&mut cursor));

    if type_tags.len() > 1 {
        let args: Vec<OscType> = try!(read_osc_args(&mut cursor, type_tags));

        Ok(OscPacket::Message(OscMessage {
            addr: addr,
            args: Some(args),
        }))
    } else {
        Ok(OscPacket::Message(OscMessage {
            addr: addr,
            args: None,
        }))
    }
}

fn decode_bundle(msg: &[u8]) -> Result<OscPacket> {
    let mut cursor: io::Cursor<&[u8]> = io::Cursor::new(msg);

    let b = try!(read_osc_string(&mut cursor));
    if b != "bundle" {
        return Err(OscError::BadBundle(format!("Wrong bundle specifier: {}", bundle_tag)));
    }

    let time_tag = try!(read_time_tag(&mut cursor));
    let mut bundle: Vec<OscPacket> = Vec::new();
    let size: usize = try!(cursor.read_u32::<BigEndian>()
                                 .map_err(OscError::ByteOrderError)) as usize;
    let mut buf: Vec<u8> = Vec::with_capacity(size);
    let cnt: usize = try!(cursor.take(size as u64)
                                .read_to_end(&mut buf)
                                .map_err(OscError::ReadError));
    if cnt == size {
        try!(decode(&mut buf).map(|p| bundle.push(p)));
    } else {
        return Err(OscError::BadBundle("Bundle shorter than expected!".to_string()));
    }

    Ok(OscPacket::Bundle(OscBundle {
        timetag: time_tag,
        content: bundle,
    }))
}

fn read_osc_string(cursor: &mut io::Cursor<&[u8]>) -> Result<String> {
    let mut str_buf: Vec<u8> = Vec::new();
    try!(cursor.read_until(0, &mut str_buf)
               .map_err(OscError::ReadError)); // ignore returned byte count
    pad_cursor(cursor);
    // convert to String and remove nul bytes
    String::from_utf8(str_buf)
        .map_err(|e| OscError::StringError(e))
        .map(|s| {
            s.trim_matches(0u8 as char)
             .to_string()
        })
}

fn read_osc_args(cursor: &mut io::Cursor<&[u8]>, raw_type_tags: String) -> Result<Vec<OscType>> {
    let type_tags: Vec<char> = raw_type_tags.chars()
                                            .skip(1)
                                            .map(|c| c as char)
                                            .collect();
    let mut args: Vec<OscType> = Vec::with_capacity(type_tags.len());
    for tag in type_tags {
        let arg: OscType = try!(read_osc_arg(cursor, tag));
        args.push(arg);
    }
    Ok(args)
}

fn read_osc_arg(cursor: &mut io::Cursor<&[u8]>, tag: char) -> Result<OscType> {
    match tag {
        'f' => {
            cursor.read_f32::<BigEndian>()
                  .map(OscType::Float)
                  .map_err(OscError::ByteOrderError)
        }
        'd' => {
            cursor.read_f64::<BigEndian>()
                  .map(OscType::Double)
                  .map_err(OscError::ByteOrderError)
        }
        'i' => {
            cursor.read_i32::<BigEndian>()
                  .map(OscType::Int)
                  .map_err(OscError::ByteOrderError)
        }
        'h' => {
            cursor.read_i64::<BigEndian>()
                  .map(OscType::Long)
                  .map_err(OscError::ByteOrderError)
        }
        's' => {
            read_osc_string(cursor).map(|s| OscType::String(s))
        }
        't' => {
            read_time_tag(cursor)
        }
        'b' => {
            read_blob(cursor)
        }
        'r' => {
            read_osc_color(cursor)
        }
        'T' => {
            Ok(OscType::Bool(true))
        }
        'F' => {
            Ok(OscType::Bool(false))
        }
        'N' => {
            Ok(OscType::Nil)
        }
        'I' => {
            Ok(OscType::Inf)
        }
        'c' => {
            read_char(cursor)
        }
        'm' => {
            read_midi_message(cursor)
        }
        _ => Err(OscError::BadArg(format!("Type tag \"{}\" is not implemented!", tag))),
    }
}

fn read_char(cursor: &mut io::Cursor<&[u8]>) -> Result<OscType> {
    let opt_char = try!(cursor.read_u32::<BigEndian>()
                              .map(char::from_u32)
                              .map_err(OscError::ByteOrderError));
    match opt_char {
        Some(c) => Ok(OscType::Char(c)),
        None => Err(OscError::BadArg("Argument is not a char!".to_string())),
    }
}

fn read_blob(cursor: &mut io::Cursor<&[u8]>) -> Result<OscType> {
    let size: usize = try!(cursor.read_u32::<BigEndian>()
                                 .map_err(OscError::ByteOrderError)) as usize;
    let mut byte_buf: Vec<u8> = Vec::with_capacity(size);

    try!(cursor.take(size as u64)
               .read_to_end(&mut byte_buf)
               .map_err(OscError::ReadError));

    pad_cursor(cursor);

    Ok(OscType::Blob(byte_buf))
}

fn read_time_tag(cursor: &mut io::Cursor<&[u8]>) -> Result<OscType> {
    let date = try!(cursor.read_u32::<BigEndian>()
                          .map_err(OscError::ByteOrderError));
    let frac = try!(cursor.read_u32::<BigEndian>()
                          .map_err(OscError::ByteOrderError));

    Ok(OscType::Time(date, frac))
}

fn read_midi_message(cursor: &mut io::Cursor<&[u8]>) -> Result<OscType> {
    let mut buf: Vec<u8> = Vec::with_capacity(4);
    try!(cursor.take(4).read_to_end(&mut buf).map_err(OscError::ReadError));

    Ok(OscType::Midi(OscMidiMessage {
        port: buf[0],
        status: buf[1],
        data1: buf[2],
        data2: buf[3],
    }))

}

fn read_osc_color(cursor: &mut io::Cursor<&[u8]>) -> Result<OscType> {
    let mut buf: Vec<u8> = Vec::with_capacity(4);
    try!(cursor.take(4).read_to_end(&mut buf).map_err(OscError::ReadError));

    Ok(OscType::Color(OscColor {
        red: buf[0],
        green: buf[1],
        blue: buf[2],
        alpha: buf[3],
    }))
}

fn pad_cursor(cursor: &mut io::Cursor<&[u8]>) {
    let pos = cursor.position();
    cursor.set_position(utils::pad(pos));
}
<|endoftext|>


<commit_before>use cargo_registry::views::krate_publish as u;
use std::{collections::HashMap, io::Read};

use flate2::{write::GzEncoder, Compression};

use super::DependencyBuilder;

lazy_static! {
    // The bytes of an empty tarball is not an empty vector of bytes because of tarball headers.
    // Unless files are added to a PublishBuilder, the `.crate` tarball that gets uploaded
    // will be empty, so precompute the empty tarball bytes to use as a default.
    static ref EMPTY_TARBALL_BYTES: Vec<u8> = {
        let mut empty_tarball = vec![];
        {
            let mut ar =
                tar::Builder::new(GzEncoder::new(&mut empty_tarball, Compression::default()));
            assert_ok!(ar.finish());
        }
        empty_tarball
    };
}

/// A builder for constructing a crate for the purposes of testing publishing. If you only need
/// a crate to exist and don't need to test behavior caused by the publish request, inserting
/// a crate into the database directly by using CrateBuilder will be faster.
pub struct PublishBuilder {
    authors: Vec<String>,
    badges: HashMap<String, HashMap<String, String>>,
    categories: Vec<String>,
    deps: Vec<u::EncodableCrateDependency>,
    desc: Option<String>,
    doc_url: Option<String>,
    keywords: Vec<String>,
    pub krate_name: String,
    license: Option<String>,
    license_file: Option<String>,
    readme: Option<String>,
    tarball: Vec<u8>,
    version: semver::Version,
}

impl PublishBuilder {
    /// Create a request to publish a crate with the given name, version 1.0.0, and no files
    /// in its tarball.
    pub fn new(krate_name: &str) -> Self {
        PublishBuilder {
            authors: vec!["foo".to_string()],
            badges: HashMap::new(),
            categories: vec![],
            deps: vec![],
            desc: Some("description".to_string()),
            doc_url: None,
            keywords: vec![],
            krate_name: krate_name.into(),
            license: Some("MIT".to_string()),
            license_file: None,
            readme: None,
            tarball: EMPTY_TARBALL_BYTES.to_vec(),
            version: semver::Version::parse("1.0.0").unwrap(),
        }
    }

    /// Set the version of the crate being published to something other than the default of 1.0.0.
    pub fn version(mut self, version: &str) -> Self {
        self.version = semver::Version::parse(version).unwrap();
        self
    }

    /// Set the files in the crate's tarball.
    pub fn files(self, files: &[(&str, &[u8])]) -> Self {
        let mut slices = files.iter().map(|p| p.1).collect::<Vec<_>>();
        let mut files = files
            .iter()
            .zip(&mut slices)
            .map(|(&(name, _), data)| {
                let len = data.len() as u64;
                (name, data as &mut dyn Read, len)
            })
            .collect::<Vec<_>>();

        self.files_with_io(&mut files)
    }

    /// Set the tarball from a Read trait object
    pub fn files_with_io(mut self, files: &mut [(&str, &mut dyn Read, u64)]) -> Self {
        let mut tarball = Vec::new();
        {
            let mut ar = tar::Builder::new(GzEncoder::new(&mut tarball, Compression::default()));
            for &mut (name, ref mut data, size) in files {
                let mut header = tar::Header::new_gnu();
                assert_ok!(header.set_path(name));
                header.set_size(size);
                header.set_cksum();
                assert_ok!(ar.append(&header, data));
            }
            assert_ok!(ar.finish());
        }

        self.tarball = tarball;
        self
    }

    /// Set the tarball directly to the given Vec of bytes
    pub fn tarball(mut self, tarball: Vec<u8>) -> Self {
        self.tarball = tarball;
        self
    }

    /// Add a dependency to this crate. Make sure the dependency already exists in the
    /// database or publish will fail.
    pub fn dependency(mut self, dep: DependencyBuilder) -> Self {
        self.deps.push(dep.build());
        self
    }

    /// Set the description of this crate
    pub fn description(mut self, description: &str) -> Self {
        self.desc = Some(description.to_string());
        self
    }

    /// Unset the description of this crate. Publish will fail unless description is reset.
    pub fn unset_description(mut self) -> Self {
        self.desc = None;
        self
    }

    /// Set the readme of this crate
    pub fn readme(mut self, readme: &str) -> Self {
        self.readme = Some(readme.to_string());
        self
    }

    /// Set the documentation URL of this crate
    pub fn documentation(mut self, documentation: &str) -> Self {
        self.doc_url = Some(documentation.to_string());
        self
    }

    /// Add a keyword to this crate.
    pub fn keyword(mut self, keyword: &str) -> Self {
        self.keywords.push(keyword.into());
        self
    }

    /// Add a category to this crate. Make sure the category already exists in the
    /// database or it will be ignored.
    pub fn category(mut self, slug: &str) -> Self {
        self.categories.push(slug.into());
        self
    }

    /// Add badges to this crate.
    pub fn badges(mut self, badges: HashMap<String, HashMap<String, String>>) -> Self {
        self.badges = badges;
        self
    }

    /// Remove the license from this crate. Publish will fail unless license or license file is set.
    pub fn unset_license(mut self) -> Self {
        self.license = None;
        self
    }

    /// Set the license file for this crate
    pub fn license_file(mut self, license_file: &str) -> Self {
        self.license_file = Some(license_file.into());
        self
    }

    /// Add an author to this crate
    pub fn author(mut self, author: &str) -> Self {
        self.authors.push(author.into());
        self
    }

    /// Remove the authors from this crate. Publish will fail unless authors are reset.
    pub fn unset_authors(mut self) -> Self {
        self.authors = vec![];
        self
    }

    /// Consume this builder to make the Put request body
    pub fn body(self) -> Vec<u8> {
        let new_crate = u::EncodableCrateUpload {
            name: u::EncodableCrateName(self.krate_name.clone()),
            vers: u::EncodableCrateVersion(self.version),
            features: HashMap::new(),
            deps: self.deps,
            authors: self.authors,
            description: self.desc,
            homepage: None,
            documentation: self.doc_url,
            readme: self.readme,
            readme_file: None,
            keywords: u::EncodableKeywordList(
                self.keywords.into_iter().map(u::EncodableKeyword).collect(),
            ),
            categories: u::EncodableCategoryList(
                self.categories
                    .into_iter()
                    .map(u::EncodableCategory)
                    .collect(),
            ),
            license: self.license,
            license_file: self.license_file,
            repository: None,
            badges: Some(self.badges),
            links: None,
        };

        let json = serde_json::to_string(&new_crate).unwrap();
        let mut body = Vec::new();
        body.extend(
            [
                json.len() as u8,
                (json.len() >> 8) as u8,
                (json.len() >> 16) as u8,
                (json.len() >> 24) as u8,
            ]
            .iter()
            .cloned(),
        );
        body.extend(json.as_bytes().iter().cloned());

        let tarball = &self.tarball;
        body.extend(&[
            tarball.len() as u8,
            (tarball.len() >> 8) as u8,
            (tarball.len() >> 16) as u8,
            (tarball.len() >> 24) as u8,
        ]);
        body.extend(tarball);
        body
    }
}
<commit_msg>tests/builders/publish: Extract `build()` and `create_publish_body()` functions<commit_after>use cargo_registry::views::krate_publish as u;
use std::{collections::HashMap, io::Read};

use flate2::{write::GzEncoder, Compression};

use super::DependencyBuilder;

lazy_static! {
    // The bytes of an empty tarball is not an empty vector of bytes because of tarball headers.
    // Unless files are added to a PublishBuilder, the `.crate` tarball that gets uploaded
    // will be empty, so precompute the empty tarball bytes to use as a default.
    static ref EMPTY_TARBALL_BYTES: Vec<u8> = {
        let mut empty_tarball = vec![];
        {
            let mut ar =
                tar::Builder::new(GzEncoder::new(&mut empty_tarball, Compression::default()));
            assert_ok!(ar.finish());
        }
        empty_tarball
    };
}

/// A builder for constructing a crate for the purposes of testing publishing. If you only need
/// a crate to exist and don't need to test behavior caused by the publish request, inserting
/// a crate into the database directly by using CrateBuilder will be faster.
pub struct PublishBuilder {
    authors: Vec<String>,
    badges: HashMap<String, HashMap<String, String>>,
    categories: Vec<String>,
    deps: Vec<u::EncodableCrateDependency>,
    desc: Option<String>,
    doc_url: Option<String>,
    keywords: Vec<String>,
    pub krate_name: String,
    license: Option<String>,
    license_file: Option<String>,
    readme: Option<String>,
    tarball: Vec<u8>,
    version: semver::Version,
}

impl PublishBuilder {
    /// Create a request to publish a crate with the given name, version 1.0.0, and no files
    /// in its tarball.
    pub fn new(krate_name: &str) -> Self {
        PublishBuilder {
            authors: vec!["foo".to_string()],
            badges: HashMap::new(),
            categories: vec![],
            deps: vec![],
            desc: Some("description".to_string()),
            doc_url: None,
            keywords: vec![],
            krate_name: krate_name.into(),
            license: Some("MIT".to_string()),
            license_file: None,
            readme: None,
            tarball: EMPTY_TARBALL_BYTES.to_vec(),
            version: semver::Version::parse("1.0.0").unwrap(),
        }
    }

    /// Set the version of the crate being published to something other than the default of 1.0.0.
    pub fn version(mut self, version: &str) -> Self {
        self.version = semver::Version::parse(version).unwrap();
        self
    }

    /// Set the files in the crate's tarball.
    pub fn files(self, files: &[(&str, &[u8])]) -> Self {
        let mut slices = files.iter().map(|p| p.1).collect::<Vec<_>>();
        let mut files = files
            .iter()
            .zip(&mut slices)
            .map(|(&(name, _), data)| {
                let len = data.len() as u64;
                (name, data as &mut dyn Read, len)
            })
            .collect::<Vec<_>>();

        self.files_with_io(&mut files)
    }

    /// Set the tarball from a Read trait object
    pub fn files_with_io(mut self, files: &mut [(&str, &mut dyn Read, u64)]) -> Self {
        let mut tarball = Vec::new();
        {
            let mut ar = tar::Builder::new(GzEncoder::new(&mut tarball, Compression::default()));
            for &mut (name, ref mut data, size) in files {
                let mut header = tar::Header::new_gnu();
                assert_ok!(header.set_path(name));
                header.set_size(size);
                header.set_cksum();
                assert_ok!(ar.append(&header, data));
            }
            assert_ok!(ar.finish());
        }

        self.tarball = tarball;
        self
    }

    /// Set the tarball directly to the given Vec of bytes
    pub fn tarball(mut self, tarball: Vec<u8>) -> Self {
        self.tarball = tarball;
        self
    }

    /// Add a dependency to this crate. Make sure the dependency already exists in the
    /// database or publish will fail.
    pub fn dependency(mut self, dep: DependencyBuilder) -> Self {
        self.deps.push(dep.build());
        self
    }

    /// Set the description of this crate
    pub fn description(mut self, description: &str) -> Self {
        self.desc = Some(description.to_string());
        self
    }

    /// Unset the description of this crate. Publish will fail unless description is reset.
    pub fn unset_description(mut self) -> Self {
        self.desc = None;
        self
    }

    /// Set the readme of this crate
    pub fn readme(mut self, readme: &str) -> Self {
        self.readme = Some(readme.to_string());
        self
    }

    /// Set the documentation URL of this crate
    pub fn documentation(mut self, documentation: &str) -> Self {
        self.doc_url = Some(documentation.to_string());
        self
    }

    /// Add a keyword to this crate.
    pub fn keyword(mut self, keyword: &str) -> Self {
        self.keywords.push(keyword.into());
        self
    }

    /// Add a category to this crate. Make sure the category already exists in the
    /// database or it will be ignored.
    pub fn category(mut self, slug: &str) -> Self {
        self.categories.push(slug.into());
        self
    }

    /// Add badges to this crate.
    pub fn badges(mut self, badges: HashMap<String, HashMap<String, String>>) -> Self {
        self.badges = badges;
        self
    }

    /// Remove the license from this crate. Publish will fail unless license or license file is set.
    pub fn unset_license(mut self) -> Self {
        self.license = None;
        self
    }

    /// Set the license file for this crate
    pub fn license_file(mut self, license_file: &str) -> Self {
        self.license_file = Some(license_file.into());
        self
    }

    /// Add an author to this crate
    pub fn author(mut self, author: &str) -> Self {
        self.authors.push(author.into());
        self
    }

    /// Remove the authors from this crate. Publish will fail unless authors are reset.
    pub fn unset_authors(mut self) -> Self {
        self.authors = vec![];
        self
    }

    pub fn build(self) -> (String, Vec<u8>) {
        let new_crate = u::EncodableCrateUpload {
            name: u::EncodableCrateName(self.krate_name.clone()),
            vers: u::EncodableCrateVersion(self.version),
            features: HashMap::new(),
            deps: self.deps,
            authors: self.authors,
            description: self.desc,
            homepage: None,
            documentation: self.doc_url,
            readme: self.readme,
            readme_file: None,
            keywords: u::EncodableKeywordList(
                self.keywords.into_iter().map(u::EncodableKeyword).collect(),
            ),
            categories: u::EncodableCategoryList(
                self.categories
                    .into_iter()
                    .map(u::EncodableCategory)
                    .collect(),
            ),
            license: self.license,
            license_file: self.license_file,
            repository: None,
            badges: Some(self.badges),
            links: None,
        };

        (serde_json::to_string(&new_crate).unwrap(), self.tarball)
    }

    /// Consume this builder to make the Put request body
    pub fn body(self) -> Vec<u8> {
        let (json, tarball) = self.build();
        PublishBuilder::create_publish_body(&json, &tarball)
    }

    pub fn create_publish_body(json: &str, tarball: &[u8]) -> Vec<u8> {
        let mut body = Vec::new();
        body.extend(
            [
                json.len() as u8,
                (json.len() >> 8) as u8,
                (json.len() >> 16) as u8,
                (json.len() >> 24) as u8,
            ]
            .iter()
            .cloned(),
        );
        body.extend(json.as_bytes().iter().cloned());

        body.extend(&[
            tarball.len() as u8,
            (tarball.len() >> 8) as u8,
            (tarball.len() >> 16) as u8,
            (tarball.len() >> 24) as u8,
        ]);
        body.extend(tarball);
        body
    }
}
<|endoftext|>


<commit_before>
use {ffi, DataSource, Return, Result, Raii, Handle, Connected};
use ffi::SQLRETURN::*;
use std::marker::PhantomData;
use std::ptr::null_mut;
use std;

/// `Statement` state used to represent a freshly allocated connection
pub enum Allocated {}
/// `Statement` state used to represent a statement with a result set cursor
///
/// A statement is most likely to enter this state after a `SELECT` query.
pub enum HasResult {}
/// `Statement` state used to represent a statement with no result set
///
/// A statement is likely to enter this state after executing e.g. a `CREATE TABLE` statement
type NoResult = Allocated; // pub enum NoResult {}

/// Holds a `Statement` after execution of a query.Allocated
///
/// A executed statement may be in one of two states. Either the statement has yielded a result set
/// or not. Keep in mind that some ODBC drivers just yield empty result sets on e.g. `INSERT`
/// Statements
pub enum Executed<'a> {
    Data(Statement<'a, HasResult>),
    NoData(Statement<'a, NoResult>),
}
pub use Executed::*;

/// RAII wrapper around ODBC statement
pub struct Statement<'a, S> {
    raii: Raii<ffi::Stmt>,
    // we use phantom data to tell the borrow checker that we need to keep the data source alive
    // for the lifetime of the statement
    parent: PhantomData<&'a DataSource<'a, Connected>>,
    state: PhantomData<S>,
}

/// Used to retrieve data from the fields of a query resul
pub struct Cursor<'a, 'b: 'a> {
    stmt: &'a mut Statement<'b, HasResult>,
    buffer: [u8; 512],
}

impl<'a, S> Handle for Statement<'a, S> {
    type To = ffi::Stmt;
    unsafe fn handle(&self) -> ffi::SQLHSTMT {
        self.raii.handle()
    }
}

impl<'a, S> Statement<'a, S> {
    fn with_raii(raii: Raii<ffi::Stmt>) -> Self {
        Statement {
            raii: raii,
            parent: PhantomData,
            state: PhantomData,
        }
    }
}

impl<'a> Statement<'a, Allocated> {
    pub fn with_parent(ds: &'a DataSource<Connected>) -> Result<Self> {
        let raii = Raii::with_parent(ds).into_result(ds)?;
        Ok(Self::with_raii(raii))
    }

    pub fn tables<'b>(mut self) -> Result<Statement<'a, HasResult>> {
        self.raii.tables().into_result(&self)?;
        Ok(Statement::with_raii(self.raii))
    }

    /// Executes a preparable statement, using the current values of the parameter marker variables
    /// if any parameters exist in the statement.
    ///
    /// `SQLExecDirect` is the fastest way to submit an SQL statement for one-time execution.
    pub fn exec_direct(mut self, statement_text: &str) -> Result<Executed<'a>> {
        if self.raii.exec_direct(statement_text).into_result(&self)? {
            Ok(Executed::Data(Statement::with_raii(self.raii)))
        } else {
            Ok(Executed::NoData(Statement::with_raii(self.raii)))
        }
    }
}

impl<'a> Statement<'a, HasResult> {
    /// The number of columns in a result set
    ///
    /// Can be called successfully only when the statement is in the prepared, executed, or
    /// positioned state. If the statement does not return columns the result will be 0.
    pub fn num_result_cols(&self) -> Result<i16> {
        self.raii.num_result_cols().into_result(self)
    }

    /// Fetches the next rowset of data from the result set and returns data for all bound columns.
    ///
    /// # Return
    /// Returns false on the last row
    pub fn fetch<'b>(&'b mut self) -> Result<Option<Cursor<'b, 'a>>> {
        if self.raii.fetch().into_result(self)? {
            Ok(Some(Cursor {
                        stmt: self,
                        buffer: [0u8; 512],
                    }))
        } else {
            Ok(None)
        }
    }

    /// Call this method to reuse the statement to execute another query.
    ///
    /// Only call this method if you have already read the result set returned by the previous
    /// query, or if you do no not intend to read it.
    ///
    /// # Example
    ///
    /// ```
    /// # use odbc::*;
    /// # fn reuse () -> Result<()> {
    /// let env = Environment::new().unwrap().set_odbc_version_3()?;
    /// let conn = DataSource::with_parent(&env)?.connect("TestDataSource", "", "")?;
    /// let stmt = Statement::with_parent(&conn)?;
    /// let stmt = match stmt.exec_direct("CREATE TABLE STAGE (A TEXT, B TEXT);")?{
    ///     Data(stmt) => stmt.close_cursor()?,
    ///     NoData(stmt) => stmt,
    /// };
    /// let stmt = stmt.exec_direct("INSERT INTO STAGE (A, B) VALUES ('Hello', 'World');")?;
    /// //...
    /// # Ok(())
    /// # };
    /// ```
    pub fn close_cursor(mut self) -> Result<Statement<'a, NoResult>> {
        self.raii.close_cursor().into_result(&self)?;
        Ok(Statement::with_raii(self.raii))
    }
}

impl<'a, 'b> Cursor<'a, 'b> {
    /// Retrieves data for a single column in the result set
    pub fn get_data(&mut self, col_or_param_num: u16) -> Result<Option<String>> {
        self.stmt.raii.get_data(col_or_param_num, &mut self.buffer).into_result(self.stmt)
    }
}

impl Raii<ffi::Stmt> {
    fn num_result_cols(&self) -> Return<i16> {
        let mut num_cols: ffi::SQLSMALLINT = 0;
        unsafe {
            match ffi::SQLNumResultCols(self.handle(), &mut num_cols as *mut ffi::SQLSMALLINT) {
                SQL_SUCCESS => Return::Success(num_cols),
                SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(num_cols),
                SQL_ERROR => Return::Error,
                r => panic!("SQLNumResultCols returned unexpected result: {:?}", r),
            }
        }
    }

    fn exec_direct(&mut self, statement_text: &str) -> Return<bool> {
        let length = statement_text.len();
        if length > ffi::SQLINTEGER::max_value() as usize {
            panic!("Statement text too long");
        }
        match unsafe {
                  ffi::SQLExecDirect(self.handle(),
                                     statement_text.as_ptr(),
                                     length as ffi::SQLINTEGER)
              } {
            ffi::SQL_SUCCESS => Return::Success(true),
            ffi::SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(true),
            ffi::SQL_ERROR => Return::Error,
            ffi::SQL_NEED_DATA => panic!("SQLExecDirec returned SQL_NEED_DATA"),
            ffi::SQL_NO_DATA => Return::Success(false),
            r => panic!("SQLExecDirect returned unexpected result: {:?}", r),
        }
    }

    /// Fetches the next rowset of data from the result set and returns data for all bound columns.
    fn fetch(&mut self) -> Return<bool> {
        match unsafe { ffi::SQLFetch(self.handle()) } {
            ffi::SQL_SUCCESS => Return::Success(true),
            ffi::SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(true),
            ffi::SQL_ERROR => Return::Error,
            ffi::SQL_NO_DATA => Return::Success(false),
            r => panic!("SQLFetch returned unexpected result: {:?}", r),
        }
    }

    fn tables(&mut self) -> Return<()> {
        let catalog_name = "";
        let schema_name = "";
        let table_name = "";
        let table_type = "TABLE";
        unsafe {
            match ffi::SQLTables(self.handle(),
                                 catalog_name.as_ptr(),
                                 catalog_name.as_bytes().len() as ffi::SQLSMALLINT,
                                 schema_name.as_ptr(),
                                 schema_name.as_bytes().len() as ffi::SQLSMALLINT,
                                 table_name.as_ptr(),
                                 table_name.as_bytes().len() as ffi::SQLSMALLINT,
                                 table_type.as_ptr(),
                                 table_type.as_bytes().len() as ffi::SQLSMALLINT) {
                SQL_SUCCESS => Return::Success(()),
                SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(()),
                SQL_ERROR => Return::Error,
                r => panic!("SQLTables returned: {:?}", r),
            }
        }
    }

    fn get_data(&mut self, col_or_param_num: u16, buffer: &mut [u8]) -> Return<Option<String>> {
        if buffer.len() == 0 {
            panic!("buffer length may not be null");
        }
        if buffer.len() > ffi::SQLLEN::max_value() as usize {
            panic!("buffer is larger than {} bytes", ffi::SQLLEN::max_value());
        }

        let mut indicator: ffi::SQLLEN = 0;
        unsafe {
            // Get buffer length...
            let result = ffi::SQLGetData(self.handle(),
                                         col_or_param_num,
                                         ffi::SQL_C_CHAR,
                                         buffer.as_mut_ptr() as ffi::SQLPOINTER,
                                         buffer.len() as ffi::SQLLEN,
                                         &mut indicator as *mut ffi::SQLLEN);
            match result {
                ffi::SQL_SUCCESS => {
                    if indicator == ffi::SQL_NULL_DATA {
                        Return::Success(None)
                    } else {
                        Return::Success(Some(std::str::from_utf8(&buffer[..(indicator as usize)])
                                                 .unwrap()
                                                 .to_owned()))
                    }
                }
                ffi::SQL_SUCCESS_WITH_INFO => {
                    if indicator == ffi::SQL_NO_TOTAL {
                        Return::SuccessWithInfo(None)
                    } else {
                        // Check if string has been truncated. String is also truncated if
                        // indicator is equal to BUF_LENGTH because of terminating nul
                        if indicator >= buffer.len() as ffi::SQLLEN {
                            let extra_space = (indicator as usize + 1) - (buffer.len() - 1);
                            let mut heap_buf = Vec::with_capacity((indicator as usize) + 1);
                            // Copy everything but the terminating zero into the new buffer
                            heap_buf.extend_from_slice(&buffer[..(buffer.len() - 1)]);
                            // increase length
                            heap_buf.extend(std::iter::repeat(0).take(extra_space));
                            // Get remainder of string
                            let ret = ffi::SQLGetData(self.handle(),
                                                      col_or_param_num,
                                                      ffi::SQL_C_CHAR,
                                                      heap_buf.as_mut_slice()[buffer.len() - 1..]
                                                          .as_mut_ptr() as
                                                      ffi::SQLPOINTER,
                                                      extra_space as ffi::SQLLEN,
                                                      null_mut());
                            heap_buf.pop();
                            let value = String::from_utf8(heap_buf).unwrap();
                            match ret {
                                ffi::SQL_SUCCESS => Return::Success(Some(value)),
                                ffi::SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(Some(value)),
                                ffi::SQL_ERROR => Return::Error,
                                r => panic!("SQLGetData returned {:?}", r),
                            }
                        } else {
                            // No truncation. Warning may be due to some other issue.
                            Return::SuccessWithInfo(Some(std::str::from_utf8(&buffer[..(indicator as
                                                                                usize)])
                                                                 .unwrap()
                                                                 .to_owned()))
                        }
                    }
                }
                ffi::SQL_ERROR => Return::Error,
                ffi::SQL_NO_DATA => panic!("SQLGetData has already returned the colmun data"),
                r => panic!("unexpected return value from SQLGetData: {:?}", r),
            }
        }
    }

    fn close_cursor(&mut self) -> Return<()> {
        unsafe {
            match ffi::SQLCloseCursor(self.handle()) {
                ffi::SQL_SUCCESS => Return::Success(()),
                ffi::SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(()),
                ffi::SQL_ERROR => Return::Error,
                r => panic!("unexpected return value from SQLCloseCursor: {:?}", r),
            }
        }
    }
}

<commit_msg>improve example code<commit_after>
use {ffi, DataSource, Return, Result, Raii, Handle, Connected};
use ffi::SQLRETURN::*;
use std::marker::PhantomData;
use std::ptr::null_mut;
use std;

/// `Statement` state used to represent a freshly allocated connection
pub enum Allocated {}
/// `Statement` state used to represent a statement with a result set cursor
///
/// A statement is most likely to enter this state after a `SELECT` query.
pub enum HasResult {}
/// `Statement` state used to represent a statement with no result set
///
/// A statement is likely to enter this state after executing e.g. a `CREATE TABLE` statement
type NoResult = Allocated; // pub enum NoResult {}

/// Holds a `Statement` after execution of a query.Allocated
///
/// A executed statement may be in one of two states. Either the statement has yielded a result set
/// or not. Keep in mind that some ODBC drivers just yield empty result sets on e.g. `INSERT`
/// Statements
pub enum Executed<'a> {
    Data(Statement<'a, HasResult>),
    NoData(Statement<'a, NoResult>),
}
pub use Executed::*;

/// RAII wrapper around ODBC statement
pub struct Statement<'a, S> {
    raii: Raii<ffi::Stmt>,
    // we use phantom data to tell the borrow checker that we need to keep the data source alive
    // for the lifetime of the statement
    parent: PhantomData<&'a DataSource<'a, Connected>>,
    state: PhantomData<S>,
}

/// Used to retrieve data from the fields of a query resul
pub struct Cursor<'a, 'b: 'a> {
    stmt: &'a mut Statement<'b, HasResult>,
    buffer: [u8; 512],
}

impl<'a, S> Handle for Statement<'a, S> {
    type To = ffi::Stmt;
    unsafe fn handle(&self) -> ffi::SQLHSTMT {
        self.raii.handle()
    }
}

impl<'a, S> Statement<'a, S> {
    fn with_raii(raii: Raii<ffi::Stmt>) -> Self {
        Statement {
            raii: raii,
            parent: PhantomData,
            state: PhantomData,
        }
    }
}

impl<'a> Statement<'a, Allocated> {
    pub fn with_parent(ds: &'a DataSource<Connected>) -> Result<Self> {
        let raii = Raii::with_parent(ds).into_result(ds)?;
        Ok(Self::with_raii(raii))
    }

    pub fn tables<'b>(mut self) -> Result<Statement<'a, HasResult>> {
        self.raii.tables().into_result(&self)?;
        Ok(Statement::with_raii(self.raii))
    }

    /// Executes a preparable statement, using the current values of the parameter marker variables
    /// if any parameters exist in the statement.
    ///
    /// `SQLExecDirect` is the fastest way to submit an SQL statement for one-time execution.
    pub fn exec_direct(mut self, statement_text: &str) -> Result<Executed<'a>> {
        if self.raii.exec_direct(statement_text).into_result(&self)? {
            Ok(Executed::Data(Statement::with_raii(self.raii)))
        } else {
            Ok(Executed::NoData(Statement::with_raii(self.raii)))
        }
    }
}

impl<'a> Statement<'a, HasResult> {
    /// The number of columns in a result set
    ///
    /// Can be called successfully only when the statement is in the prepared, executed, or
    /// positioned state. If the statement does not return columns the result will be 0.
    pub fn num_result_cols(&self) -> Result<i16> {
        self.raii.num_result_cols().into_result(self)
    }

    /// Fetches the next rowset of data from the result set and returns data for all bound columns.
    ///
    /// # Return
    /// Returns false on the last row
    pub fn fetch<'b>(&'b mut self) -> Result<Option<Cursor<'b, 'a>>> {
        if self.raii.fetch().into_result(self)? {
            Ok(Some(Cursor {
                        stmt: self,
                        buffer: [0u8; 512],
                    }))
        } else {
            Ok(None)
        }
    }

    /// Call this method to reuse the statement to execute another query.
    ///
    /// Only call this method if you have already read the result set returned by the previous
    /// query, or if you do no not intend to read it.
    ///
    /// # Example
    ///
    /// ```
    /// # use odbc::*;
    /// # fn reuse () -> Result<()> {
    /// let env = Environment::new().unwrap().set_odbc_version_3()?;
    /// let conn = DataSource::with_parent(&env)?.connect("TestDataSource", "", "")?;
    /// let stmt = Statement::with_parent(&conn)?;
    /// let stmt = match stmt.exec_direct("CREATE TABLE STAGE (A TEXT, B TEXT);")?{
    ///     // Some drivers will return an empty result set. We need to close it before we can use
    ///     // statement again.
    ///     Data(stmt) => stmt.close_cursor()?,
    ///     NoData(stmt) => stmt,
    /// };
    /// let stmt = stmt.exec_direct("INSERT INTO STAGE (A, B) VALUES ('Hello', 'World');")?;
    /// //...
    /// # Ok(())
    /// # };
    /// ```
    pub fn close_cursor(mut self) -> Result<Statement<'a, NoResult>> {
        self.raii.close_cursor().into_result(&self)?;
        Ok(Statement::with_raii(self.raii))
    }
}

impl<'a, 'b> Cursor<'a, 'b> {
    /// Retrieves data for a single column in the result set
    pub fn get_data(&mut self, col_or_param_num: u16) -> Result<Option<String>> {
        self.stmt.raii.get_data(col_or_param_num, &mut self.buffer).into_result(self.stmt)
    }
}

impl Raii<ffi::Stmt> {
    fn num_result_cols(&self) -> Return<i16> {
        let mut num_cols: ffi::SQLSMALLINT = 0;
        unsafe {
            match ffi::SQLNumResultCols(self.handle(), &mut num_cols as *mut ffi::SQLSMALLINT) {
                SQL_SUCCESS => Return::Success(num_cols),
                SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(num_cols),
                SQL_ERROR => Return::Error,
                r => panic!("SQLNumResultCols returned unexpected result: {:?}", r),
            }
        }
    }

    fn exec_direct(&mut self, statement_text: &str) -> Return<bool> {
        let length = statement_text.len();
        if length > ffi::SQLINTEGER::max_value() as usize {
            panic!("Statement text too long");
        }
        match unsafe {
                  ffi::SQLExecDirect(self.handle(),
                                     statement_text.as_ptr(),
                                     length as ffi::SQLINTEGER)
              } {
            ffi::SQL_SUCCESS => Return::Success(true),
            ffi::SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(true),
            ffi::SQL_ERROR => Return::Error,
            ffi::SQL_NEED_DATA => panic!("SQLExecDirec returned SQL_NEED_DATA"),
            ffi::SQL_NO_DATA => Return::Success(false),
            r => panic!("SQLExecDirect returned unexpected result: {:?}", r),
        }
    }

    /// Fetches the next rowset of data from the result set and returns data for all bound columns.
    fn fetch(&mut self) -> Return<bool> {
        match unsafe { ffi::SQLFetch(self.handle()) } {
            ffi::SQL_SUCCESS => Return::Success(true),
            ffi::SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(true),
            ffi::SQL_ERROR => Return::Error,
            ffi::SQL_NO_DATA => Return::Success(false),
            r => panic!("SQLFetch returned unexpected result: {:?}", r),
        }
    }

    fn tables(&mut self) -> Return<()> {
        let catalog_name = "";
        let schema_name = "";
        let table_name = "";
        let table_type = "TABLE";
        unsafe {
            match ffi::SQLTables(self.handle(),
                                 catalog_name.as_ptr(),
                                 catalog_name.as_bytes().len() as ffi::SQLSMALLINT,
                                 schema_name.as_ptr(),
                                 schema_name.as_bytes().len() as ffi::SQLSMALLINT,
                                 table_name.as_ptr(),
                                 table_name.as_bytes().len() as ffi::SQLSMALLINT,
                                 table_type.as_ptr(),
                                 table_type.as_bytes().len() as ffi::SQLSMALLINT) {
                SQL_SUCCESS => Return::Success(()),
                SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(()),
                SQL_ERROR => Return::Error,
                r => panic!("SQLTables returned: {:?}", r),
            }
        }
    }

    fn get_data(&mut self, col_or_param_num: u16, buffer: &mut [u8]) -> Return<Option<String>> {
        if buffer.len() == 0 {
            panic!("buffer length may not be null");
        }
        if buffer.len() > ffi::SQLLEN::max_value() as usize {
            panic!("buffer is larger than {} bytes", ffi::SQLLEN::max_value());
        }

        let mut indicator: ffi::SQLLEN = 0;
        unsafe {
            // Get buffer length...
            let result = ffi::SQLGetData(self.handle(),
                                         col_or_param_num,
                                         ffi::SQL_C_CHAR,
                                         buffer.as_mut_ptr() as ffi::SQLPOINTER,
                                         buffer.len() as ffi::SQLLEN,
                                         &mut indicator as *mut ffi::SQLLEN);
            match result {
                ffi::SQL_SUCCESS => {
                    if indicator == ffi::SQL_NULL_DATA {
                        Return::Success(None)
                    } else {
                        Return::Success(Some(std::str::from_utf8(&buffer[..(indicator as usize)])
                                                 .unwrap()
                                                 .to_owned()))
                    }
                }
                ffi::SQL_SUCCESS_WITH_INFO => {
                    if indicator == ffi::SQL_NO_TOTAL {
                        Return::SuccessWithInfo(None)
                    } else {
                        // Check if string has been truncated. String is also truncated if
                        // indicator is equal to BUF_LENGTH because of terminating nul
                        if indicator >= buffer.len() as ffi::SQLLEN {
                            let extra_space = (indicator as usize + 1) - (buffer.len() - 1);
                            let mut heap_buf = Vec::with_capacity((indicator as usize) + 1);
                            // Copy everything but the terminating zero into the new buffer
                            heap_buf.extend_from_slice(&buffer[..(buffer.len() - 1)]);
                            // increase length
                            heap_buf.extend(std::iter::repeat(0).take(extra_space));
                            // Get remainder of string
                            let ret = ffi::SQLGetData(self.handle(),
                                                      col_or_param_num,
                                                      ffi::SQL_C_CHAR,
                                                      heap_buf.as_mut_slice()[buffer.len() - 1..]
                                                          .as_mut_ptr() as
                                                      ffi::SQLPOINTER,
                                                      extra_space as ffi::SQLLEN,
                                                      null_mut());
                            heap_buf.pop();
                            let value = String::from_utf8(heap_buf).unwrap();
                            match ret {
                                ffi::SQL_SUCCESS => Return::Success(Some(value)),
                                ffi::SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(Some(value)),
                                ffi::SQL_ERROR => Return::Error,
                                r => panic!("SQLGetData returned {:?}", r),
                            }
                        } else {
                            // No truncation. Warning may be due to some other issue.
                            Return::SuccessWithInfo(Some(std::str::from_utf8(&buffer[..(indicator as
                                                                                usize)])
                                                                 .unwrap()
                                                                 .to_owned()))
                        }
                    }
                }
                ffi::SQL_ERROR => Return::Error,
                ffi::SQL_NO_DATA => panic!("SQLGetData has already returned the colmun data"),
                r => panic!("unexpected return value from SQLGetData: {:?}", r),
            }
        }
    }

    fn close_cursor(&mut self) -> Return<()> {
        unsafe {
            match ffi::SQLCloseCursor(self.handle()) {
                ffi::SQL_SUCCESS => Return::Success(()),
                ffi::SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(()),
                ffi::SQL_ERROR => Return::Error,
                r => panic!("unexpected return value from SQLCloseCursor: {:?}", r),
            }
        }
    }
}

<|endoftext|>