<commit_before>// Copyright 2012 The Rust Project Developers.
// Copyright 2015 Guillaume Gomez
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Library to interface with chunks of memory allocated in C.
//!
//! It is often desirable to safely interface with memory allocated from C,
//! encapsulating the unsafety into allocation and destruction time.  Indeed,
//! allocating memory externally is currently the only way to give Rust shared
//! mut state with C programs that keep their own references; vectors are
//! unsuitable because they could be reallocated or moved at any time, and
//! importing C memory into a vector takes a one-time snapshot of the memory.
//!
//! This module simplifies the usage of such external blocks of memory.  Memory
//! is encapsulated into an opaque object after creation; the lifecycle of the
//! memory can be optionally managed by Rust, if an appropriate destructor
//! closure is provided.  Safety is ensured by bounds-checking accesses, which
//! are marshalled through get and set functions.
//!
//! There are three unsafe functions: the two constructors, and the
//! unwrapping method. The constructors are unsafe for the
//! obvious reason (they act on a pointer that cannot be checked inside the
//! method), but `into_inner()` is somewhat more subtle in its unsafety.
//! It returns the contained pointer, but at the same time destroys the CVec
//! without running its destructor. This can be used to pass memory back to
//! C, but care must be taken that the ownership of underlying resources are
//! handled correctly, i.e. that allocated memory is eventually freed
//! if necessary.

#![feature(unsafe_destructor, unique, core)]
#![cfg_attr(test, feature(alloc))]

use std::boxed::FnBox;
use std::ptr;
use std::slice;
use std::ptr::Unique;
use std::ops::{Index, IndexMut};

/// The type representing a foreign chunk of memory
pub struct CVec<T> {
    base: Unique<T>,
    len: usize,
    dtor: Option<Box<FnBox(*mut T)>>
}

#[unsafe_destructor]
impl<T> Drop for CVec<T> {
    fn drop(&mut self) {
        match self.dtor.take() {
            None => (),
            Some(f) => f(*self.base),
        }
    }
}

impl<T> CVec<T> {
    /// Create a `CVec` from a raw pointer to a buffer with a given length.
    ///
    /// Panics if the given pointer is null. The returned vector will not attempt
    /// to deallocate the vector when dropped.
    ///
    /// # Arguments
    ///
    /// * base - A unique pointer to a buffer
    /// * len - The number of elements in the buffer
    pub unsafe fn new(base: Unique<T>, len: usize) -> CVec<T> {
        assert!(*base != ptr::null_mut());
        CVec {
            base: base,
            len: len,
            dtor: None,
        }
    }

    /// Create a `CVec` from a foreign buffer, with a given length,
    /// and a function to run upon destruction.
    ///
    /// Panics if the given pointer is null.
    ///
    /// # Arguments
    ///
    /// * base - A unique pointer to a buffer
    /// * len - The number of elements in the buffer
    /// * dtor - A fn to run when the value is destructed, useful
    ///          for freeing the buffer, etc. `base` will be passed
    ///          to it as an argument.
    pub unsafe fn new_with_dtor<F>(base: Unique<T>,
                                   len: usize,
                                   dtor: F)
                                   -> CVec<T>
        where F: FnBox(*mut T) + 'static
    {
        assert!(*base != ptr::null_mut());
        let dtor = Box::new(dtor);
        CVec {
            base: base,
            len: len,
            dtor: Some(dtor)
        }
    }

    /// Retrieves an element at a given index, returning `None` if the requested
    /// index is greater than the length of the vector.
    pub fn get<'a>(&'a self, ofs: usize) -> Option<&'a T> {
        if ofs < self.len {
            Some(unsafe { &*self.base.offset(ofs as isize) })
        } else {
            None
        }
    }

    /// Retrieves a mutable element at a given index, returning `None` if the
    /// requested index is greater than the length of the vector.
    pub fn get_mut<'a>(&'a mut self, ofs: usize) -> Option<&'a mut T> {
        if ofs < self.len {
            Some(unsafe { &mut *self.base.offset(ofs as isize) })
        } else {
            None
        }
    }

    /// Unwrap the pointer without running the destructor
    ///
    /// This method retrieves the underlying pointer, and in the process
    /// destroys the CVec but without running the destructor. A use case
    /// would be transferring ownership of the buffer to a C function, as
    /// in this case you would not want to run the destructor.
    ///
    /// Note that if you want to access the underlying pointer without
    /// cancelling the destructor, you can simply call `transmute` on the return
    /// value of `get(0)`.
    pub unsafe fn into_inner(mut self) -> *mut T {
        self.dtor = None;
        *self.base
    }

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

    /// Returns whether this vector is empty.
    pub fn is_empty(&self) -> bool { self.len() == 0 }

    /// Convert to CSlice
    pub fn as_cslice(&self) -> CSlice<T> {
        CSlice {
            base: *self.base,
            len: self.len
        }
    }
}

impl<T> AsRef<[T]> for CVec<T> {
    /// View the stored data as a slice.
    fn as_ref(&self) -> &[T] {
        unsafe {
            slice::from_raw_parts(*self.base as *const T, self.len)
        }
    }
}

impl<T> AsMut<[T]> for CVec<T> {
    /// View the stored data as a slice.
    fn as_mut(&mut self) -> &mut [T] {
        unsafe {
            slice::from_raw_parts_mut(*self.base, self.len)
        }
    }
}

/// The type representing an 'unsafe' foreign chunk of memory
pub struct CSlice<T> {
    base: *mut T,
    len: usize
}

impl<T> CSlice<T> {
    /// Create a `CSlice` from a raw pointer to a buffer with a given length.
    ///
    /// Panics if the given pointer is null. The returned vector will not attempt
    /// to deallocate the vector when dropped.
    ///
    /// # Arguments
    ///
    /// * base - A raw pointer to a buffer
    /// * len - The number of elements in the buffer
    pub unsafe fn new(base: *mut T, len: usize) -> CSlice<T> {
        assert!(base != ptr::null_mut());
        CSlice {
            base: base,
            len: len
        }
    }

    /// Retrieves an element at a given index, returning `None` if the requested
    /// index is greater than the length of the vector.
    pub fn get<'a>(&'a self, ofs: usize) -> Option<&'a T> {
        if ofs < self.len {
            Some(unsafe { &*self.base.offset(ofs as isize) })
        } else {
            None
        }
    }

    /// Retrieves a mutable element at a given index, returning `None` if the
    /// requested index is greater than the length of the vector.
    pub fn get_mut<'a>(&'a mut self, ofs: usize) -> Option<&'a mut T> {
        if ofs < self.len {
            Some(unsafe { &mut *self.base.offset(ofs as isize) })
        } else {
            None
        }
    }

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

    /// Returns whether this vector is empty.
    pub fn is_empty(&self) -> bool { self.len() == 0 }
}

impl<T> AsRef<[T]> for CSlice<T> {
    /// View the stored data as a slice.
    fn as_ref(&self) -> &[T] {
        unsafe {
            slice::from_raw_parts(self.base as *const T, self.len)
        }
    }
}

impl<T> AsMut<[T]> for CSlice<T> {
    /// View the stored data as a slice.
    fn as_mut(&mut self) -> &mut [T] {
        unsafe {
            slice::from_raw_parts_mut(self.base, self.len)
        }
    }
}

impl<T> Index<usize> for CSlice<T> {
    type Output = T;

    fn index<'a>(&'a self, _index: usize) -> &'a T {
        assert!(_index < self.len);
        unsafe { &*self.base.offset(_index as isize) }
    }
}

impl<T> IndexMut<usize> for CSlice<T> {
    fn index_mut<'a>(&'a mut self, _index: usize) -> &'a mut T {
        assert!(_index < self.len);
        unsafe { &mut *self.base.offset(_index as isize) }
    }
}

#[cfg(test)]
mod tests {
    extern crate alloc;
    extern crate libc;

    use super::{CVec, CSlice};
    use std::ptr::{self, Unique};

    // allocation of CVec
    fn v_malloc(n: usize) -> CVec<u8> {
        unsafe {
            let mem: Unique<u8> = Unique::new(
                libc::malloc(n as libc::size_t) as *mut _);
            if (*mem).is_null() { alloc::oom() }

            CVec::new_with_dtor(mem, n,
                                |mem| { libc::free((mem) as *mut _); })
        }
    }

    // allocation of CSlice
    fn s_malloc(n: usize) -> CSlice<u8> {
        unsafe {
            let mem: *mut u8 = libc::malloc(n as libc::size_t) as *mut _;
            if mem.is_null() { alloc::oom() }

            CSlice::new(mem, n)
        }
    }

    #[test]
    fn vec_test_basic() {
        let mut cv = v_malloc(16);

        *cv.get_mut(3).unwrap() = 8;
        *cv.get_mut(4).unwrap() = 9;
        assert_eq!(*cv.get(3).unwrap(), 8);
        assert_eq!(*cv.get(4).unwrap(), 9);
        assert_eq!(cv.len(), 16);
    }

    #[test]
    fn slice_test_basic() {
        let mut cs = s_malloc(16);

        cs[3] = 8;
        cs[4] = 9;
        assert_eq!(cs[3], 8);
        assert_eq!(cs[4], 9);
        assert_eq!(cs.len(), 16);
    }

    #[test]
    #[should_panic]
    fn vec_test_panic_at_null() {
        unsafe {
            CVec::new(Unique::new(ptr::null_mut::<u8>()), 9);
        }
    }

    #[test]
    #[should_panic]
    fn slice_test_panic_at_null() {
        unsafe {
            CSlice::new(ptr::null_mut::<u8>(), 9);
        }
    }

    #[test]
    fn vec_test_overrun_get() {
        let cv = v_malloc(16);

        assert!(cv.get(17).is_none());
    }

    #[test]
    #[should_panic]
    fn slice_test_overrun_get() {
        let cs = s_malloc(16);

        assert!(cs[17] == 18);
    }

    #[test]
    fn vec_test_overrun_set() {
        let mut cv = v_malloc(16);

        assert!(cv.get_mut(17).is_none());
    }

    #[test]
    fn vec_test_unwrap() {
        unsafe {
            let cv = CVec::new_with_dtor(Unique::new(1 as *mut isize),
                                         0,
                                         |_| panic!("Don't run this destructor!"));
            let p = cv.into_inner();
            assert_eq!(p, 1 as *mut isize);
        }
    }

    #[test]
    fn vec_to_slice_test() {
        let mut cv = v_malloc(2);

        *cv.get_mut(0).unwrap() = 10;
        *cv.get_mut(1).unwrap() = 12;
        let cs = cv.as_cslice();

        assert_eq!(cs[0], 10);
        assert_eq!(cs[1], 12);
    }

    #[test]
    fn slice_to_vec_test() {
        let cv = v_malloc(2);
        let mut cs = cv.as_cslice();

        cs[0] = 13;
        cs[1] = 26;
        assert_eq!(*cv.get(0).unwrap(), 13);
        assert_eq!(*cv.get(1).unwrap(), 26);
    }
}
<commit_msg>Avoid any unstable in Rust 1.0 features<commit_after>// Copyright 2012 The Rust Project Developers.
// Copyright 2015 Guillaume Gomez
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Library to interface with chunks of memory allocated in C.
//!
//! It is often desirable to safely interface with memory allocated from C,
//! encapsulating the unsafety into allocation and destruction time.  Indeed,
//! allocating memory externally is currently the only way to give Rust shared
//! mut state with C programs that keep their own references; vectors are
//! unsuitable because they could be reallocated or moved at any time, and
//! importing C memory into a vector takes a one-time snapshot of the memory.
//!
//! This module simplifies the usage of such external blocks of memory.  Memory
//! is encapsulated into an opaque object after creation; the lifecycle of the
//! memory can be optionally managed by Rust, if an appropriate destructor
//! closure is provided.  Safety is ensured by bounds-checking accesses, which
//! are marshalled through get and set functions.
//!
//! There are three unsafe functions: the two constructors, and the
//! unwrapping method. The constructors are unsafe for the
//! obvious reason (they act on a pointer that cannot be checked inside the
//! method), but `into_inner()` is somewhat more subtle in its unsafety.
//! It returns the contained pointer, but at the same time destroys the CVec
//! without running its destructor. This can be used to pass memory back to
//! C, but care must be taken that the ownership of underlying resources are
//! handled correctly, i.e. that allocated memory is eventually freed
//! if necessary.

use std::ptr;
use std::slice;
use std::ops::{Index, IndexMut};

/// The type representing a foreign chunk of memory
pub struct CVec<T> {
    base: *mut T,
    len: usize,
    dtor: Option<Box<FnMut(*mut T)>>
}

impl<T> Drop for CVec<T> {
    fn drop(&mut self) {
        if let Some(mut f) = self.dtor.take() {
            f(self.base);
        }
    }
}

impl<T> CVec<T> {
    /// Create a `CVec` from a raw pointer to a buffer with a given length.
    ///
    /// Panics if the given pointer is null. The returned vector will not attempt
    /// to deallocate the vector when dropped.
    ///
    /// # Arguments
    ///
    /// * base - A unique pointer to a buffer
    /// * len - The number of elements in the buffer
    pub unsafe fn new(base: *mut T, len: usize) -> CVec<T> {
        assert!(base != ptr::null_mut());
        CVec {
            base: base,
            len: len,
            dtor: None,
        }
    }

    /// Create a `CVec` from a foreign buffer, with a given length,
    /// and a function to run upon destruction.
    ///
    /// Panics if the given pointer is null.
    ///
    /// # Arguments
    ///
    /// * base - A unique pointer to a buffer
    /// * len - The number of elements in the buffer
    /// * dtor - A fn to run when the value is destructed, useful
    ///          for freeing the buffer, etc. `base` will be passed
    ///          to it as an argument.
    pub unsafe fn new_with_dtor<F>(base: *mut T,
                                   len: usize,
                                   dtor: F)
                                   -> CVec<T>
        where F: FnMut(*mut T) + 'static
    {
        assert!(base != ptr::null_mut());
        let dtor = Box::new(dtor);
        CVec {
            base: base,
            len: len,
            dtor: Some(dtor)
        }
    }

    /// Retrieves an element at a given index, returning `None` if the requested
    /// index is greater than the length of the vector.
    pub fn get<'a>(&'a self, ofs: usize) -> Option<&'a T> {
        if ofs < self.len {
            Some(unsafe { &*self.base.offset(ofs as isize) })
        } else {
            None
        }
    }

    /// Retrieves a mutable element at a given index, returning `None` if the
    /// requested index is greater than the length of the vector.
    pub fn get_mut<'a>(&'a mut self, ofs: usize) -> Option<&'a mut T> {
        if ofs < self.len {
            Some(unsafe { &mut *self.base.offset(ofs as isize) })
        } else {
            None
        }
    }

    /// Unwrap the pointer without running the destructor
    ///
    /// This method retrieves the underlying pointer, and in the process
    /// destroys the CVec but without running the destructor. A use case
    /// would be transferring ownership of the buffer to a C function, as
    /// in this case you would not want to run the destructor.
    ///
    /// Note that if you want to access the underlying pointer without
    /// cancelling the destructor, you can simply call `transmute` on the return
    /// value of `get(0)`.
    pub unsafe fn into_inner(mut self) -> *mut T {
        self.dtor = None;
        self.base
    }

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

    /// Returns whether this vector is empty.
    pub fn is_empty(&self) -> bool { self.len() == 0 }

    /// Convert to CSlice
    pub fn as_cslice(&self) -> CSlice<T> {
        CSlice {
            base: self.base,
            len: self.len
        }
    }
}

impl<T> AsRef<[T]> for CVec<T> {
    /// View the stored data as a slice.
    fn as_ref(&self) -> &[T] {
        unsafe {
            slice::from_raw_parts(self.base as *const T, self.len)
        }
    }
}

impl<T> AsMut<[T]> for CVec<T> {
    /// View the stored data as a slice.
    fn as_mut(&mut self) -> &mut [T] {
        unsafe {
            slice::from_raw_parts_mut(self.base, self.len)
        }
    }
}

/// The type representing an 'unsafe' foreign chunk of memory
pub struct CSlice<T> {
    base: *mut T,
    len: usize
}

impl<T> CSlice<T> {
    /// Create a `CSlice` from a raw pointer to a buffer with a given length.
    ///
    /// Panics if the given pointer is null. The returned vector will not attempt
    /// to deallocate the vector when dropped.
    ///
    /// # Arguments
    ///
    /// * base - A raw pointer to a buffer
    /// * len - The number of elements in the buffer
    pub unsafe fn new(base: *mut T, len: usize) -> CSlice<T> {
        assert!(base != ptr::null_mut());
        CSlice {
            base: base,
            len: len
        }
    }

    /// Retrieves an element at a given index, returning `None` if the requested
    /// index is greater than the length of the vector.
    pub fn get<'a>(&'a self, ofs: usize) -> Option<&'a T> {
        if ofs < self.len {
            Some(unsafe { &*self.base.offset(ofs as isize) })
        } else {
            None
        }
    }

    /// Retrieves a mutable element at a given index, returning `None` if the
    /// requested index is greater than the length of the vector.
    pub fn get_mut<'a>(&'a mut self, ofs: usize) -> Option<&'a mut T> {
        if ofs < self.len {
            Some(unsafe { &mut *self.base.offset(ofs as isize) })
        } else {
            None
        }
    }

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

    /// Returns whether this vector is empty.
    pub fn is_empty(&self) -> bool { self.len() == 0 }
}

impl<T> AsRef<[T]> for CSlice<T> {
    /// View the stored data as a slice.
    fn as_ref(&self) -> &[T] {
        unsafe {
            slice::from_raw_parts(self.base as *const T, self.len)
        }
    }
}

impl<T> AsMut<[T]> for CSlice<T> {
    /// View the stored data as a slice.
    fn as_mut(&mut self) -> &mut [T] {
        unsafe {
            slice::from_raw_parts_mut(self.base, self.len)
        }
    }
}

impl<T> Index<usize> for CSlice<T> {
    type Output = T;

    fn index<'a>(&'a self, _index: usize) -> &'a T {
        assert!(_index < self.len);
        unsafe { &*self.base.offset(_index as isize) }
    }
}

impl<T> IndexMut<usize> for CSlice<T> {
    fn index_mut<'a>(&'a mut self, _index: usize) -> &'a mut T {
        assert!(_index < self.len);
        unsafe { &mut *self.base.offset(_index as isize) }
    }
}

#[cfg(test)]
mod tests {
    extern crate libc;

    use super::{CVec, CSlice};
    use std::ptr;

    // allocation of CVec
    fn v_malloc(n: usize) -> CVec<u8> {
        unsafe {
            let mem = libc::malloc(n as libc::size_t) as *mut u8;
            CVec::new_with_dtor(mem, n, |mem| { libc::free((mem) as *mut _); })
        }
    }

    // allocation of CSlice
    fn s_malloc(n: usize) -> CSlice<u8> {
        unsafe {
            let mem: *mut u8 = libc::malloc(n as libc::size_t) as *mut _;
            CSlice::new(mem, n)
        }
    }

    #[test]
    fn vec_test_basic() {
        let mut cv = v_malloc(16);

        *cv.get_mut(3).unwrap() = 8;
        *cv.get_mut(4).unwrap() = 9;
        assert_eq!(*cv.get(3).unwrap(), 8);
        assert_eq!(*cv.get(4).unwrap(), 9);
        assert_eq!(cv.len(), 16);
    }

    #[test]
    fn slice_test_basic() {
        let mut cs = s_malloc(16);

        cs[3] = 8;
        cs[4] = 9;
        assert_eq!(cs[3], 8);
        assert_eq!(cs[4], 9);
        assert_eq!(cs.len(), 16);
    }

    #[test]
    #[should_panic]
    fn vec_test_panic_at_null() {
        unsafe {
            CVec::new(ptr::null_mut::<u8>(), 9);
        }
    }

    #[test]
    #[should_panic]
    fn slice_test_panic_at_null() {
        unsafe {
            CSlice::new(ptr::null_mut::<u8>(), 9);
        }
    }

    #[test]
    fn vec_test_overrun_get() {
        let cv = v_malloc(16);

        assert!(cv.get(17).is_none());
    }

    #[test]
    #[should_panic]
    fn slice_test_overrun_get() {
        let cs = s_malloc(16);

        assert!(cs[17] == 18);
    }

    #[test]
    fn vec_test_overrun_set() {
        let mut cv = v_malloc(16);

        assert!(cv.get_mut(17).is_none());
    }

    #[test]
    fn vec_test_unwrap() {
        unsafe {
            let cv = CVec::new_with_dtor(1 as *mut isize,
                                         0,
                                         |_| panic!("Don't run this destructor!"));
            let p = cv.into_inner();
            assert_eq!(p, 1 as *mut isize);
        }
    }

    #[test]
    fn vec_to_slice_test() {
        let mut cv = v_malloc(2);

        *cv.get_mut(0).unwrap() = 10;
        *cv.get_mut(1).unwrap() = 12;
        let cs = cv.as_cslice();

        assert_eq!(cs[0], 10);
        assert_eq!(cs[1], 12);
    }

    #[test]
    fn slice_to_vec_test() {
        let cv = v_malloc(2);
        let mut cs = cv.as_cslice();

        cs[0] = 13;
        cs[1] = 26;
        assert_eq!(*cv.get(0).unwrap(), 13);
        assert_eq!(*cv.get(1).unwrap(), 26);
    }
}
<|endoftext|>


<commit_before>
use syntax::ast;
use syntax::ast::{MetaItemKind,NestedMetaItemKind};
use rustc::hir::def_id::DefId;
use rustc::hir::def;
use syntax::codemap::Span;
use rustc::lint::{self, LintContext, LintPass, LateLintPass, LintArray};
use rustc::ty::{self, TyCtxt};
use rustc::hir;

declare_lint!(NOT_TAGGED_SAFE, Warn, "Warn about use of non-tagged methods within tagged function");

#[derive(Default)]
pub struct Pass
{
    visit_stack: Vec<ast::NodeId>,
}

impl LintPass for Pass {
    fn get_lints(&self) -> LintArray {
        lint_array!(NOT_TAGGED_SAFE)
    }
}

impl<'a,'b> LateLintPass<'a,'b> for Pass {
    fn check_fn(&mut self, cx: &lint::LateContext, _kind: hir::intravisit::FnKind, _decl: &hir::FnDecl, body: &hir::Body, _: Span, id: ast::NodeId) {
        let attrs = cx.tcx.hir.attrs(id);

        // If this function is tagged with a particular safety, store
        {
            let mut lh = ::database::CACHE.write().unwrap();
            for tag_name in get_tags(attrs, "is_safe")
            {
                let tag = lh.get_tag_or_add(&tag_name.as_str());
                lh.mark(id, tag,  true);
            }
            for tag_name in get_tags(attrs, "not_safe")
            {
                let tag = lh.get_tag_or_add(&tag_name.as_str());
                lh.mark(id, tag,  false);
            }
        }

        // For each required safety, check
        for tag_name in get_tags(attrs, "req_safe")
        {
            let ty_tag = {
                let mut lh = ::database::CACHE.write().unwrap();
                let tag = lh.get_tag_or_add(&tag_name.as_str());
                //let tag = if let Some(v) = lh.get_tag_opt(&tag_name.as_str()) { v } else { error!("Tag {} unknown", ty_name);  continue; };
                lh.mark(id, tag,  true);
                tag
                };

            // Search body for calls to non safe methods
            let mut v = Visitor {
                    pass: self, cx: cx, tag: ty_tag,
                    cb: |span| {
                            cx.span_lint(NOT_TAGGED_SAFE, *span,
                                &format!("Calling {0}-unsafe method from a #[req_safe({0})] method", tag_name)[..]
                                );
                        },
                    };
            debug!("Method {:?} is marked safe '{}'", id, tag_name);
            hir::intravisit::walk_body(&mut v, body);
        }

        // TODO: For all known safeties (that aren't already set) populate.
        // - Requires knowing all safeties (which we don't, ... yet)
    }
}

impl Pass
{
    fn fill_cache_for(&mut self, tcx: &TyCtxt, node_id: ast::NodeId)
    {
        debug!("Filling cache for node {:?}", node_id);
        let attrs = tcx.hir.attrs(node_id);
        let mut lh = ::database::CACHE.write().unwrap();
        for tag_name in Iterator::chain( get_tags(attrs, "is_safe"), get_tags(attrs, "req_safe") )
        {
            debug!("#[is_safe/req_safe] {} - {}", tag_name, node_id);
            let tag = lh.get_tag_or_add(&tag_name.as_str());
            lh.mark(node_id, tag,  true);
        }
        for tag_name in get_tags(attrs, "not_safe")
        {
            debug!("#[not_safe] {} - {}", tag_name, node_id);
            let tag = lh.get_tag_or_add(&tag_name.as_str());
            lh.mark(node_id, tag,  false);
        }
    }

    /// Recursively check that the provided function is either safe or unsafe.
    // Used to avoid excessive annotating
    fn recurse_fcn_body(&mut self, cx: &lint::LateContext, node_id: ast::NodeId, tag: ::database::Tag) -> bool
    {
        // and apply a visitor to all
        match cx.tcx.hir.get(node_id)
        {
        hir::map::NodeItem(i) =>
            match i.node {
            hir::ItemFn(_, _, _, _, _, ref body) => {
                // Enumerate this function's code, recursively checking for a call to an unsafe method
                let mut is_safe = true;
                {
                    let mut v = Visitor {
                        pass: self, cx: cx, tag: tag,
                        cb: |_| { is_safe = false; }
                        };
                    hir::intravisit::walk_body(&mut v, cx.tcx.hir.body(*body));
                }
                is_safe
                },
            ref v @ _ => {
                error!("Node ID {} points to a non-function item {:?}", node_id, v);
                true
                },
            },
        hir::map::NodeImplItem(i) =>
            match i.node {
            hir::ImplItemKind::Method(_, ref body) => {

                let mut is_safe = true;
                {
                    let mut v = Visitor {
                        pass: self, cx: cx, tag: tag,
                        cb: |_| { is_safe = false; }
                        };
                    hir::intravisit::walk_body(&mut v, cx.tcx.hir.body(*body));
                }
                is_safe
                },
            _ => true,
            },
        //hir::map::NodeForeignItem(i) =>
        //    if Self::check_for_marker(tcx, i.id, "tag_safe", name) {
        //        true
        //    }
        //    else if Self::check_for_marker(tcx, i.id, "tag_unsafe", name) {
        //        false
        //    }
        //    else {
        //        unknown_assume
        //    },
        ref v @ _ => {
            error!("Node ID {} points to non-item {:?}", node_id, v);
            true
            }
        }
    }

    /// Locate a #[tag_safe(<name>)] attribute on the passed item
    pub fn method_is_safe(&mut self, cx: &lint::LateContext, id: DefId, tag: ::database::Tag) -> bool
    {
        if ! id.is_local()
        {
            // TODO: Get the entry from the crate cache
            if let Some(v) = ::database::CACHE.read().unwrap().get_extern(&cx.tcx,id.krate, id.index, tag) {
                debug!("{:?} - {} (extern cached)", id, v);
                v
            }
            else {
                debug!("{:?} - {} (extern assumed)", id, true);
                true
            }
        }
        else
        {
            let node_id = cx.tcx.hir.as_local_node_id(id).expect("Unable to locate node");
            let mut local_opt = ::database::CACHE.read().unwrap().get_local(node_id, tag);
            // NOTE: This only fires once (ideally)
            if local_opt.is_none() {
                self.fill_cache_for(&cx.tcx, node_id);
                local_opt = ::database::CACHE.read().unwrap().get_local(node_id, tag);
            }
            if let Some(v) = local_opt {
                debug!("{} - {} (cached)", node_id, v);
                v
            }
            else {
                // If this node is currently being checked, assume it's valid.
                // TODO: This can lead to a false positive being stored.
                if self.visit_stack.iter().position(|x| *x == node_id).is_some() {
                    warn!("Recursion, assuming true");
                    true
                }
                else {
                    self.visit_stack.push(node_id);
                    let rv = self.recurse_fcn_body(cx, node_id, tag);
                    self.visit_stack.pop();

                    debug!("{} - {} (recursed)", node_id, rv);
                    ::database::CACHE.write().unwrap().mark(node_id, tag,  rv);
                    rv
                }
            }
        }
    }
}

struct Visitor<'a, 'tcx: 'a, F: FnMut(&Span) + 'a>
{
    pass: &'a mut Pass,
	cx: &'a lint::LateContext<'a,'tcx>,
    tag: ::database::Tag,
    cb: F,
}

impl<'a, 'tcx: 'a, F: FnMut(&Span)> hir::intravisit::Visitor<'a> for Visitor<'a, 'tcx, F>
{
	fn nested_visit_map<'this>(&'this mut self) -> hir::intravisit::NestedVisitorMap<'this, 'a> {
		hir::intravisit::NestedVisitorMap::None
	}

    // Locate function/method calls in a code block
    fn visit_expr(&mut self, ex: &'a hir::Expr) {
        debug!("visit node - {:?}", ex);
        match ex.node
        {
        // Call expressions - check that it's a path call
        hir::ExprCall(ref fcn, ..) =>
			match fcn.node
			{
			hir::ExprPath(ref qp, ..) =>
				match self.cx.tables.qpath_def(qp, fcn.id)
				{
				def::Def::Fn(did) | def::Def::Method(did) =>
					// Check for a safety tag
					if !self.pass.method_is_safe(self.cx, did, self.tag)
					{
						(self.cb)(&ex.span);
					}
					else {
						debug!("Safe call {:?}", ex);
					},
				_ => {
					info!("Call ExprPath with an unknown Def type");
					},
				},
			_ => {
				info!("Call without ExprPath");
				},
			},

        // Method call expressions - get the relevant method
        hir::ExprMethodCall(ref _id, ref _tys, ref _exprs) =>
			match self.cx.tables.type_dependent_defs.get(&ex.id)
			{
			Some(callee) => {
                let id = callee.def_id();

				// Check for a safety tag
				if !self.pass.method_is_safe(self.cx, id, self.tag) {
					(self.cb)(&ex.span);
				}
				},
			None => info!("ExprMethodCall with unknown callee"),
			},

        // Ignore any other type of node
        _ => {},
        }
        hir::intravisit::walk_expr(self, ex);
    }
}

fn get_tags<'a>(meta_items: &'a [ast::Attribute], attr_name: &'a str) -> impl Iterator<Item=::syntax::symbol::Symbol>+'a {
    meta_items.iter()
        .filter(move |attr| attr.name().map(move |n| n == attr_name).unwrap_or(false))
        .flat_map(|attr|
			if let Some(v) = attr.meta() {
				if let MetaItemKind::List(v) = v.node {
					v.into_iter()
				}
				else {
					panic!("");
				}
			}
			else {
				panic!("");
			}
            )
        .filter_map(|tag_meta|
            if let NestedMetaItemKind::MetaItem(ref ptr) = tag_meta.node {
                if let MetaItemKind::Word = ptr.node {
                    Some(ptr.name)
                }
                else {
                    warn!("");
                    None
                }
            }
            else {
                warn!("");
                None
            }
            )
}
<commit_msg>Nightly upgrade<commit_after>
use syntax::ast;
use syntax::ast::{MetaItemKind,NestedMetaItemKind};
use rustc::hir::def_id::DefId;
use rustc::hir::def;
use syntax::codemap::Span;
use rustc::lint::{self, LintContext, LintPass, LateLintPass, LintArray};
use rustc::ty::{self, TyCtxt};
use rustc::hir;

declare_lint!(NOT_TAGGED_SAFE, Warn, "Warn about use of non-tagged methods within tagged function");

#[derive(Default)]
pub struct Pass
{
    visit_stack: Vec<ast::NodeId>,
}

impl LintPass for Pass {
    fn get_lints(&self) -> LintArray {
        lint_array!(NOT_TAGGED_SAFE)
    }
}

impl<'a,'b> LateLintPass<'a,'b> for Pass {
    fn check_fn(&mut self, cx: &lint::LateContext, _kind: hir::intravisit::FnKind, _decl: &hir::FnDecl, body: &hir::Body, _: Span, id: ast::NodeId) {
        let attrs = cx.tcx.hir.attrs(id);

        // If this function is tagged with a particular safety, store
        {
            let mut lh = ::database::CACHE.write().unwrap();
            for tag_name in get_tags(attrs, "is_safe")
            {
                let tag = lh.get_tag_or_add(&tag_name.as_str());
                lh.mark(id, tag,  true);
            }
            for tag_name in get_tags(attrs, "not_safe")
            {
                let tag = lh.get_tag_or_add(&tag_name.as_str());
                lh.mark(id, tag,  false);
            }
        }

        // For each required safety, check
        for tag_name in get_tags(attrs, "req_safe")
        {
            let ty_tag = {
                let mut lh = ::database::CACHE.write().unwrap();
                let tag = lh.get_tag_or_add(&tag_name.as_str());
                //let tag = if let Some(v) = lh.get_tag_opt(&tag_name.as_str()) { v } else { error!("Tag {} unknown", ty_name);  continue; };
                lh.mark(id, tag,  true);
                tag
                };

            // Search body for calls to non safe methods
            let mut v = Visitor {
                    pass: self, cx: cx, tag: ty_tag,
                    cb: |span| {
                            cx.span_lint(NOT_TAGGED_SAFE, *span,
                                &format!("Calling {0}-unsafe method from a #[req_safe({0})] method", tag_name)[..]
                                );
                        },
                    };
            debug!("Method {:?} is marked safe '{}'", id, tag_name);
            hir::intravisit::walk_body(&mut v, body);
        }

        // TODO: For all known safeties (that aren't already set) populate.
        // - Requires knowing all safeties (which we don't, ... yet)
    }
}

impl Pass
{
    fn fill_cache_for(&mut self, tcx: &TyCtxt, node_id: ast::NodeId)
    {
        debug!("Filling cache for node {:?}", node_id);
        let attrs = tcx.hir.attrs(node_id);
        let mut lh = ::database::CACHE.write().unwrap();
        for tag_name in Iterator::chain( get_tags(attrs, "is_safe"), get_tags(attrs, "req_safe") )
        {
            debug!("#[is_safe/req_safe] {} - {}", tag_name, node_id);
            let tag = lh.get_tag_or_add(&tag_name.as_str());
            lh.mark(node_id, tag,  true);
        }
        for tag_name in get_tags(attrs, "not_safe")
        {
            debug!("#[not_safe] {} - {}", tag_name, node_id);
            let tag = lh.get_tag_or_add(&tag_name.as_str());
            lh.mark(node_id, tag,  false);
        }
    }

    /// Recursively check that the provided function is either safe or unsafe.
    // Used to avoid excessive annotating
    fn recurse_fcn_body(&mut self, cx: &lint::LateContext, node_id: ast::NodeId, tag: ::database::Tag) -> bool
    {
        // and apply a visitor to all
        match cx.tcx.hir.get(node_id)
        {
        hir::map::NodeItem(i) =>
            match i.node {
            hir::ItemFn(_, _, _, _, _, ref body) => {
                // Enumerate this function's code, recursively checking for a call to an unsafe method
                let mut is_safe = true;
                {
                    let mut v = Visitor {
                        pass: self, cx: cx, tag: tag,
                        cb: |_| { is_safe = false; }
                        };
                    hir::intravisit::walk_body(&mut v, cx.tcx.hir.body(*body));
                }
                is_safe
                },
            ref v @ _ => {
                error!("Node ID {} points to a non-function item {:?}", node_id, v);
                true
                },
            },
        hir::map::NodeImplItem(i) =>
            match i.node {
            hir::ImplItemKind::Method(_, ref body) => {

                let mut is_safe = true;
                {
                    let mut v = Visitor {
                        pass: self, cx: cx, tag: tag,
                        cb: |_| { is_safe = false; }
                        };
                    hir::intravisit::walk_body(&mut v, cx.tcx.hir.body(*body));
                }
                is_safe
                },
            _ => true,
            },
        //hir::map::NodeForeignItem(i) =>
        //    if Self::check_for_marker(tcx, i.id, "tag_safe", name) {
        //        true
        //    }
        //    else if Self::check_for_marker(tcx, i.id, "tag_unsafe", name) {
        //        false
        //    }
        //    else {
        //        unknown_assume
        //    },
        ref v @ _ => {
            error!("Node ID {} points to non-item {:?}", node_id, v);
            true
            }
        }
    }

    /// Locate a #[tag_safe(<name>)] attribute on the passed item
    pub fn method_is_safe(&mut self, cx: &lint::LateContext, id: DefId, tag: ::database::Tag) -> bool
    {
        if ! id.is_local()
        {
            // TODO: Get the entry from the crate cache
            if let Some(v) = ::database::CACHE.read().unwrap().get_extern(&cx.tcx,id.krate, id.index, tag) {
                debug!("{:?} - {} (extern cached)", id, v);
                v
            }
            else {
                debug!("{:?} - {} (extern assumed)", id, true);
                true
            }
        }
        else
        {
            let node_id = cx.tcx.hir.as_local_node_id(id).expect("Unable to locate node");
            let mut local_opt = ::database::CACHE.read().unwrap().get_local(node_id, tag);
            // NOTE: This only fires once (ideally)
            if local_opt.is_none() {
                self.fill_cache_for(&cx.tcx, node_id);
                local_opt = ::database::CACHE.read().unwrap().get_local(node_id, tag);
            }
            if let Some(v) = local_opt {
                debug!("{} - {} (cached)", node_id, v);
                v
            }
            else {
                // If this node is currently being checked, assume it's valid.
                // TODO: This can lead to a false positive being stored.
                if self.visit_stack.iter().position(|x| *x == node_id).is_some() {
                    warn!("Recursion, assuming true");
                    true
                }
                else {
                    self.visit_stack.push(node_id);
                    let rv = self.recurse_fcn_body(cx, node_id, tag);
                    self.visit_stack.pop();

                    debug!("{} - {} (recursed)", node_id, rv);
                    ::database::CACHE.write().unwrap().mark(node_id, tag,  rv);
                    rv
                }
            }
        }
    }
}

struct Visitor<'a, 'tcx: 'a, F: FnMut(&Span) + 'a>
{
    pass: &'a mut Pass,
	cx: &'a lint::LateContext<'a,'tcx>,
    tag: ::database::Tag,
    cb: F,
}

impl<'a, 'tcx: 'a, F: FnMut(&Span)> hir::intravisit::Visitor<'a> for Visitor<'a, 'tcx, F>
{
	fn nested_visit_map<'this>(&'this mut self) -> hir::intravisit::NestedVisitorMap<'this, 'a> {
		hir::intravisit::NestedVisitorMap::None
	}

    // Locate function/method calls in a code block
    fn visit_expr(&mut self, ex: &'a hir::Expr) {
        debug!("visit node - {:?}", ex);
        match ex.node
        {
        // Call expressions - check that it's a path call
        hir::ExprCall(ref fcn, ..) =>
			match fcn.node
			{
			hir::ExprPath(ref qp, ..) =>
				match self.cx.tables.qpath_def(qp, fcn.hir_id)
				{
				def::Def::Fn(did) | def::Def::Method(did) =>
					// Check for a safety tag
					if !self.pass.method_is_safe(self.cx, did, self.tag)
					{
						(self.cb)(&ex.span);
					}
					else {
						debug!("Safe call {:?}", ex);
					},
				_ => {
					info!("Call ExprPath with an unknown Def type");
					},
				},
			_ => {
				info!("Call without ExprPath");
				},
			},

        // Method call expressions - get the relevant method
        hir::ExprMethodCall(ref _id, ref _tys, ref _exprs) =>
			match self.cx.tables.type_dependent_defs().get(ex.hir_id)
			{
			Some(callee) => {
                let id = callee.def_id();

				// Check for a safety tag
				if !self.pass.method_is_safe(self.cx, id, self.tag) {
					(self.cb)(&ex.span);
				}
				},
			None => info!("ExprMethodCall with unknown callee"),
			},

        // Ignore any other type of node
        _ => {},
        }
        hir::intravisit::walk_expr(self, ex);
    }
}

fn get_tags<'a>(meta_items: &'a [ast::Attribute], attr_name: &'a str) -> impl Iterator<Item=::syntax::symbol::Symbol>+'a {
    meta_items.iter()
        .filter(move |attr| attr.name().map(move |n| n == attr_name).unwrap_or(false))
        .flat_map(|attr|
			if let Some(v) = attr.meta() {
				if let MetaItemKind::List(v) = v.node {
					v.into_iter()
				}
				else {
					panic!("");
				}
			}
			else {
				panic!("");
			}
            )
        .filter_map(|tag_meta|
            if let NestedMetaItemKind::MetaItem(ref ptr) = tag_meta.node {
                if let MetaItemKind::Word = ptr.node {
                    Some(ptr.name)
                }
                else {
                    warn!("");
                    None
                }
            }
            else {
                warn!("");
                None
            }
            )
}
<|endoftext|>


<commit_before>use crate::error::ZmqError;
use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fmt::Display;
use tokio_util::codec::{Decoder, Encoder};

use crate::SocketType;

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqMechanism {
    NULL,
    PLAIN,
    CURVE,
}

impl Display for ZmqMechanism {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return match self {
            ZmqMechanism::NULL => write!(f, "NULL"),
            ZmqMechanism::PLAIN => write!(f, "PLAIN"),
            ZmqMechanism::CURVE => write!(f, "CURVE"),
        };
    }
}

impl TryFrom<Vec<u8>> for ZmqMechanism {
    type Error = ZmqError;

    fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
        let mech = value.split(|x| *x == 0x0).next().unwrap_or(b"");
        // mechanism-char = "A"-"Z" | DIGIT
        //                  | "-" | "_" | "." | "+" | %x0
        // according to https://rfc.zeromq.org/spec:23/ZMTP/
        let mechanism = unsafe { String::from_utf8_unchecked(mech.to_vec()) };
        match mechanism.as_str() {
            "NULL" => Ok(ZmqMechanism::NULL),
            "PLAIN" => Ok(ZmqMechanism::PLAIN),
            "CURVE" => Ok(ZmqMechanism::CURVE),
            _ => Err(ZmqError::Other("Failed to parse ZmqMechanism")),
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub(crate) struct ZmqGreeting {
    pub version: (u8, u8),
    pub mechanism: ZmqMechanism,
    pub as_server: bool,
}

impl Default for ZmqGreeting {
    fn default() -> Self {
        Self {
            version: (3, 0),
            mechanism: ZmqMechanism::NULL,
            as_server: false,
        }
    }
}

impl TryFrom<Bytes> for ZmqGreeting {
    type Error = ZmqError;

    fn try_from(value: Bytes) -> Result<Self, Self::Error> {
        if !(value[0] == 0xff && value[9] == 0x7f) {
            return Err(ZmqError::Codec("Failed to parse greeting"));
        }
        Ok(ZmqGreeting {
            version: (value[10], value[11]),
            mechanism: ZmqMechanism::try_from(value[12..32].to_vec())?,
            as_server: value[32] == 0x01,
        })
    }
}

impl From<ZmqGreeting> for BytesMut {
    fn from(greet: ZmqGreeting) -> Self {
        let mut data: [u8; 64] = [0; 64];
        data[0] = 0xff;
        data[9] = 0x7f;
        data[10] = greet.version.0;
        data[11] = greet.version.1;
        let mech = format!("{}", greet.mechanism);
        data[12..12 + mech.len()].copy_from_slice(mech.as_bytes());
        data[32] = greet.as_server.into();
        let mut bytes = BytesMut::new();
        bytes.extend_from_slice(&data);
        bytes
    }
}

#[derive(Debug, Clone)]
pub struct ZmqMessage {
    pub data: Bytes,
}

#[derive(Debug, Clone)]
pub(crate) enum Message {
    Greeting(ZmqGreeting),
    Command(ZmqCommand),
    Message(ZmqMessage),
    MultipartMessage(Vec<ZmqMessage>),
}

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqCommandName {
    READY,
}

impl From<ZmqCommandName> for String {
    fn from(c_name: ZmqCommandName) -> Self {
        match c_name {
            ZmqCommandName::READY => "READY".into(),
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct ZmqCommand {
    pub name: ZmqCommandName,
    pub properties: HashMap<String, String>,
}

impl ZmqCommand {
    pub fn ready(socket: SocketType) -> Self {
        let mut properties = HashMap::new();
        properties.insert("Socket-Type".into(), format!("{}", socket));
        Self {
            name: ZmqCommandName::READY,
            properties,
        }
    }
}

impl TryFrom<BytesMut> for ZmqCommand {
    type Error = ZmqError;

    fn try_from(mut buf: BytesMut) -> Result<Self, Self::Error> {
        let command_len = buf.get_u8() as usize;
        // command-name-char = ALPHA according to https://rfc.zeromq.org/spec:23/ZMTP/
        let command_name =
            unsafe { String::from_utf8_unchecked(buf.split_to(command_len).to_vec()) };
        let command = match command_name.as_str() {
            "READY" => ZmqCommandName::READY,
            _ => return Err(ZmqError::Codec("Uknown command received")),
        };
        let mut properties = HashMap::new();

        while !buf.is_empty() {
            // Collect command properties
            let prop_len = buf.get_u8() as usize;
            let property = unsafe { String::from_utf8_unchecked(buf.split_to(prop_len).to_vec()) };

            let prop_val_len = buf.get_u32() as usize;
            let prop_value =
                unsafe { String::from_utf8_unchecked(buf.split_to(prop_val_len).to_vec()) };
            properties.insert(property, prop_value);
        }
        Ok(Self {
            name: command,
            properties,
        })
    }
}

impl From<ZmqCommand> for BytesMut {
    fn from(command: ZmqCommand) -> Self {
        let mut message_len = 0;

        let command_name: String = command.name.into();
        message_len += command_name.len() + 1;
        for (prop, val) in command.properties.iter() {
            message_len += prop.len() + 1;
            message_len += val.len() + 4;
        }

        let long_message = message_len > 255;

        let mut bytes = BytesMut::new();
        if long_message {
            bytes.reserve(message_len + 9);
            bytes.put_u8(0x06);
            bytes.put_u64(message_len as u64);
        } else {
            bytes.reserve(message_len + 2);
            bytes.put_u8(0x04);
            bytes.put_u8(message_len as u8);
        };
        bytes.put_u8(command_name.len() as u8);
        bytes.extend_from_slice(command_name.as_ref());
        for (prop, val) in command.properties.iter() {
            bytes.put_u8(prop.len() as u8);
            bytes.extend_from_slice(prop.as_ref());
            bytes.put_u32(val.len() as u32);
            bytes.extend_from_slice(val.as_ref());
        }
        bytes
    }
}

#[derive(Debug, Clone, Copy)]
struct Frame {
    command: bool,
    long: bool,
    more: bool,
}

#[derive(Debug)]
enum DecoderState {
    Greeting,
    FrameHeader,
    FrameLen(Frame),
    Frame(Frame),
}

pub(crate) struct ZmqCodec {
    state: DecoderState,
    waiting_for: usize, // Number of bytes needed to decode frame
    // Needed to store incoming multipart message
    // This allows to incapsulate it's processing inside codec and not expose
    // internal details to higher levels
    buffered_message: Option<Message>,
}

impl ZmqCodec {
    pub fn new() -> Self {
        Self {
            state: DecoderState::Greeting,
            waiting_for: 64, // len of the greeting frame,
            buffered_message: None,
        }
    }
}

impl Decoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        dbg!(&src);
        if src.len() < self.waiting_for {
            src.reserve(self.waiting_for - src.len());
            return Ok(None);
        }
        match self.state {
            DecoderState::Greeting => {
                if src[0] != 0xff {
                    return Err(ZmqError::Codec("Bad first byte of greeting"));
                }
                self.state = DecoderState::FrameHeader;
                self.waiting_for = 1;
                Ok(Some(Message::Greeting(ZmqGreeting::try_from(
                    src.split_to(64).freeze(),
                )?)))
            }
            DecoderState::FrameHeader => {
                let flags = src.get_u8();
                let command = (flags & 0b0000_0100) != 0;
                let long = (flags & 0b0000_0010) != 0;
                let more = (flags & 0b0000_0001) != 0;

                if more && self.buffered_message.is_none() {
                    self.buffered_message = Some(Message::MultipartMessage(Vec::new()));
                }
                let frame = Frame {
                    command,
                    long,
                    more,
                };
                self.state = DecoderState::FrameLen(frame);
                self.waiting_for = if frame.long { 8 } else { 1 };
                return self.decode(src);
            }
            DecoderState::FrameLen(frame) => {
                self.state = DecoderState::Frame(frame);
                self.waiting_for = if frame.long {
                    src.get_u64() as usize
                } else {
                    src.get_u8() as usize
                };
                return self.decode(src);
            }
            DecoderState::Frame(frame) => {
                let data = src.split_to(self.waiting_for);
                self.state = DecoderState::FrameHeader;
                self.waiting_for = 1;
                if frame.command {
                    Ok(Some(Message::Command(ZmqCommand::try_from(data)?)))
                } else if frame.more {
                    // cache incoming multipart message
                    match &mut self.buffered_message {
                        Some(Message::MultipartMessage(message)) => message.push(ZmqMessage {
                            data: data.freeze(),
                        }),
                        _ => panic!("Corrupted decoder state"),
                    }
                    return self.decode(src);
                } else {
                    let mess = ZmqMessage { data: data.freeze() };
                    if let Some(Message::MultipartMessage(mut message)) =
                        self.buffered_message.take()
                    {
                        message.push(mess);
                        Ok(Some(Message::MultipartMessage(message)))
                    } else {
                        Ok(Some(Message::Message(mess)))
                    }
                }
            }
        }
    }
}

impl ZmqCodec {
    fn _encode_message(&mut self, message: ZmqMessage, dst: &mut BytesMut, more: bool) {
        let mut flags: u8 = 0;
        if more {
            flags |= 0b0000_0001;
        }
        let len = message.data.len();
        if len > 255 {
            flags |= 0b0000_0010;
            dst.reserve(len + 9);
        } else {
            dst.reserve(len + 2);
        }
        dst.put_u8(flags);
        if len > 255 {
            dst.put_u64(len as u64);
        } else {
            dst.put_u8(len as u8);
        }
        dst.extend_from_slice(message.data.as_ref());
    }
}

impl Encoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn encode(&mut self, message: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
        match message {
            Message::Greeting(payload) => dst.unsplit(payload.into()),
            Message::Message(message) => self._encode_message(message, dst, false),
            Message::Command(command) => dst.unsplit(command.into()),
            Message::MultipartMessage(parts) => {
                let last_element = parts.len() - 1;
                for (idx, part) in parts.into_iter().enumerate() {
                    self._encode_message(part, dst, idx != last_element);
                }
            }
        }
        Ok(())
    }
}
<commit_msg>Added basic converters to simplify working with messages<commit_after>use crate::error::ZmqError;
use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fmt::Display;
use tokio_util::codec::{Decoder, Encoder};

use crate::SocketType;
use std::string::FromUtf8Error;

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqMechanism {
    NULL,
    PLAIN,
    CURVE,
}

impl Display for ZmqMechanism {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return match self {
            ZmqMechanism::NULL => write!(f, "NULL"),
            ZmqMechanism::PLAIN => write!(f, "PLAIN"),
            ZmqMechanism::CURVE => write!(f, "CURVE"),
        };
    }
}

impl TryFrom<Vec<u8>> for ZmqMechanism {
    type Error = ZmqError;

    fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
        let mech = value.split(|x| *x == 0x0).next().unwrap_or(b"");
        // mechanism-char = "A"-"Z" | DIGIT
        //                  | "-" | "_" | "." | "+" | %x0
        // according to https://rfc.zeromq.org/spec:23/ZMTP/
        let mechanism = unsafe { String::from_utf8_unchecked(mech.to_vec()) };
        match mechanism.as_str() {
            "NULL" => Ok(ZmqMechanism::NULL),
            "PLAIN" => Ok(ZmqMechanism::PLAIN),
            "CURVE" => Ok(ZmqMechanism::CURVE),
            _ => Err(ZmqError::Other("Failed to parse ZmqMechanism")),
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub(crate) struct ZmqGreeting {
    pub version: (u8, u8),
    pub mechanism: ZmqMechanism,
    pub as_server: bool,
}

impl Default for ZmqGreeting {
    fn default() -> Self {
        Self {
            version: (3, 0),
            mechanism: ZmqMechanism::NULL,
            as_server: false,
        }
    }
}

impl TryFrom<Bytes> for ZmqGreeting {
    type Error = ZmqError;

    fn try_from(value: Bytes) -> Result<Self, Self::Error> {
        if !(value[0] == 0xff && value[9] == 0x7f) {
            return Err(ZmqError::Codec("Failed to parse greeting"));
        }
        Ok(ZmqGreeting {
            version: (value[10], value[11]),
            mechanism: ZmqMechanism::try_from(value[12..32].to_vec())?,
            as_server: value[32] == 0x01,
        })
    }
}

impl From<ZmqGreeting> for BytesMut {
    fn from(greet: ZmqGreeting) -> Self {
        let mut data: [u8; 64] = [0; 64];
        data[0] = 0xff;
        data[9] = 0x7f;
        data[10] = greet.version.0;
        data[11] = greet.version.1;
        let mech = format!("{}", greet.mechanism);
        data[12..12 + mech.len()].copy_from_slice(mech.as_bytes());
        data[32] = greet.as_server.into();
        let mut bytes = BytesMut::new();
        bytes.extend_from_slice(&data);
        bytes
    }
}

#[derive(Debug, Clone)]
pub struct ZmqMessage {
    pub data: Bytes,
}

impl From<Bytes> for ZmqMessage {
    fn from(data: Bytes) -> Self {
        Self { data }
    }
}

impl From<BytesMut> for ZmqMessage {
    fn from(data: BytesMut) -> Self {
        data.freeze().into()
    }
}

impl From<Vec<u8>> for ZmqMessage {
    fn from(data: Vec<u8>) -> Self {
        Bytes::from(data).into()
    }
}

impl From<String> for ZmqMessage {
    fn from(data: String) -> Self {
        data.into_bytes().into()
    }
}

impl From<ZmqMessage> for Vec<u8> {
    fn from(m: ZmqMessage) -> Self {
        m.data.to_vec()
    }
}

impl TryFrom<ZmqMessage> for String {
    type Error = FromUtf8Error;

    fn try_from(m: ZmqMessage) -> Result<Self, Self::Error> {
        String::from_utf8(m.into())
    }
}

#[derive(Debug, Clone)]
pub(crate) enum Message {
    Greeting(ZmqGreeting),
    Command(ZmqCommand),
    Message(ZmqMessage),
    MultipartMessage(Vec<ZmqMessage>),
}

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqCommandName {
    READY,
}

impl From<ZmqCommandName> for String {
    fn from(c_name: ZmqCommandName) -> Self {
        match c_name {
            ZmqCommandName::READY => "READY".into(),
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct ZmqCommand {
    pub name: ZmqCommandName,
    pub properties: HashMap<String, String>,
}

impl ZmqCommand {
    pub fn ready(socket: SocketType) -> Self {
        let mut properties = HashMap::new();
        properties.insert("Socket-Type".into(), format!("{}", socket));
        Self {
            name: ZmqCommandName::READY,
            properties,
        }
    }
}

impl TryFrom<BytesMut> for ZmqCommand {
    type Error = ZmqError;

    fn try_from(mut buf: BytesMut) -> Result<Self, Self::Error> {
        let command_len = buf.get_u8() as usize;
        // command-name-char = ALPHA according to https://rfc.zeromq.org/spec:23/ZMTP/
        let command_name =
            unsafe { String::from_utf8_unchecked(buf.split_to(command_len).to_vec()) };
        let command = match command_name.as_str() {
            "READY" => ZmqCommandName::READY,
            _ => return Err(ZmqError::Codec("Uknown command received")),
        };
        let mut properties = HashMap::new();

        while !buf.is_empty() {
            // Collect command properties
            let prop_len = buf.get_u8() as usize;
            let property = unsafe { String::from_utf8_unchecked(buf.split_to(prop_len).to_vec()) };

            let prop_val_len = buf.get_u32() as usize;
            let prop_value =
                unsafe { String::from_utf8_unchecked(buf.split_to(prop_val_len).to_vec()) };
            properties.insert(property, prop_value);
        }
        Ok(Self {
            name: command,
            properties,
        })
    }
}

impl From<ZmqCommand> for BytesMut {
    fn from(command: ZmqCommand) -> Self {
        let mut message_len = 0;

        let command_name: String = command.name.into();
        message_len += command_name.len() + 1;
        for (prop, val) in command.properties.iter() {
            message_len += prop.len() + 1;
            message_len += val.len() + 4;
        }

        let long_message = message_len > 255;

        let mut bytes = BytesMut::new();
        if long_message {
            bytes.reserve(message_len + 9);
            bytes.put_u8(0x06);
            bytes.put_u64(message_len as u64);
        } else {
            bytes.reserve(message_len + 2);
            bytes.put_u8(0x04);
            bytes.put_u8(message_len as u8);
        };
        bytes.put_u8(command_name.len() as u8);
        bytes.extend_from_slice(command_name.as_ref());
        for (prop, val) in command.properties.iter() {
            bytes.put_u8(prop.len() as u8);
            bytes.extend_from_slice(prop.as_ref());
            bytes.put_u32(val.len() as u32);
            bytes.extend_from_slice(val.as_ref());
        }
        bytes
    }
}

#[derive(Debug, Clone, Copy)]
struct Frame {
    command: bool,
    long: bool,
    more: bool,
}

#[derive(Debug)]
enum DecoderState {
    Greeting,
    FrameHeader,
    FrameLen(Frame),
    Frame(Frame),
}

pub(crate) struct ZmqCodec {
    state: DecoderState,
    waiting_for: usize, // Number of bytes needed to decode frame
    // Needed to store incoming multipart message
    // This allows to incapsulate it's processing inside codec and not expose
    // internal details to higher levels
    buffered_message: Option<Message>,
}

impl ZmqCodec {
    pub fn new() -> Self {
        Self {
            state: DecoderState::Greeting,
            waiting_for: 64, // len of the greeting frame,
            buffered_message: None,
        }
    }
}

impl Decoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        dbg!(&src);
        if src.len() < self.waiting_for {
            src.reserve(self.waiting_for - src.len());
            return Ok(None);
        }
        match self.state {
            DecoderState::Greeting => {
                if src[0] != 0xff {
                    return Err(ZmqError::Codec("Bad first byte of greeting"));
                }
                self.state = DecoderState::FrameHeader;
                self.waiting_for = 1;
                Ok(Some(Message::Greeting(ZmqGreeting::try_from(
                    src.split_to(64).freeze(),
                )?)))
            }
            DecoderState::FrameHeader => {
                let flags = src.get_u8();
                let command = (flags & 0b0000_0100) != 0;
                let long = (flags & 0b0000_0010) != 0;
                let more = (flags & 0b0000_0001) != 0;

                if more && self.buffered_message.is_none() {
                    self.buffered_message = Some(Message::MultipartMessage(Vec::new()));
                }
                let frame = Frame {
                    command,
                    long,
                    more,
                };
                self.state = DecoderState::FrameLen(frame);
                self.waiting_for = if frame.long { 8 } else { 1 };
                return self.decode(src);
            }
            DecoderState::FrameLen(frame) => {
                self.state = DecoderState::Frame(frame);
                self.waiting_for = if frame.long {
                    src.get_u64() as usize
                } else {
                    src.get_u8() as usize
                };
                return self.decode(src);
            }
            DecoderState::Frame(frame) => {
                let data = src.split_to(self.waiting_for);
                self.state = DecoderState::FrameHeader;
                self.waiting_for = 1;
                if frame.command {
                    Ok(Some(Message::Command(ZmqCommand::try_from(data)?)))
                } else if frame.more {
                    // cache incoming multipart message
                    match &mut self.buffered_message {
                        Some(Message::MultipartMessage(message)) => message.push(ZmqMessage {
                            data: data.freeze(),
                        }),
                        _ => panic!("Corrupted decoder state"),
                    }
                    return self.decode(src);
                } else {
                    let mess = ZmqMessage { data: data.freeze() };
                    if let Some(Message::MultipartMessage(mut message)) =
                        self.buffered_message.take()
                    {
                        message.push(mess);
                        Ok(Some(Message::MultipartMessage(message)))
                    } else {
                        Ok(Some(Message::Message(mess)))
                    }
                }
            }
        }
    }
}

impl ZmqCodec {
    fn _encode_message(&mut self, message: ZmqMessage, dst: &mut BytesMut, more: bool) {
        let mut flags: u8 = 0;
        if more {
            flags |= 0b0000_0001;
        }
        let len = message.data.len();
        if len > 255 {
            flags |= 0b0000_0010;
            dst.reserve(len + 9);
        } else {
            dst.reserve(len + 2);
        }
        dst.put_u8(flags);
        if len > 255 {
            dst.put_u64(len as u64);
        } else {
            dst.put_u8(len as u8);
        }
        dst.extend_from_slice(message.data.as_ref());
    }
}

impl Encoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn encode(&mut self, message: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
        match message {
            Message::Greeting(payload) => dst.unsplit(payload.into()),
            Message::Message(message) => self._encode_message(message, dst, false),
            Message::Command(command) => dst.unsplit(command.into()),
            Message::MultipartMessage(parts) => {
                let last_element = parts.len() - 1;
                for (idx, part) in parts.into_iter().enumerate() {
                    self._encode_message(part, dst, idx != last_element);
                }
            }
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::mem;
use std::ops::Add;

use bytes::{alloc, MutByteBuf, ByteBuf, Buf, MutBuf};
use mio::{TryWrite, TryRead};

pub trait Codec<In: ?Sized, Out: ?Sized>
{
    fn decode(&mut self, buf: &mut In) -> Vec<Out>;
    fn encode(&self, a: Out) -> In;
}

/*
impl <In: ?Sized, Mid: ?Sized, Out: ?Sized> Add<Codec<Mid, Out>> for Codec<In, Mid> {
    type Output = Codec<In, Out>;

    fn add(self, right: Codec<Mid, Out>) -> Codec<In, Out> {
        CodecStack {
            left: Box::new(self),
            right: Box::new(right),
        }
    }
}

pub struct CodecStack<In, Mid, Out> {
    left: Box<Codec<In, Mid>>,
    right: Box<Codec<Mid, Out>>,
}

impl <In, Mid, Out> Codec<In, Out> for CodecStack<In, Mid, Out> {
    fn decode(&mut self, buf: &mut In) -> Vec<Out> {
        self.left.decode(buf).and_then(|mut d| self.right.decode(&mut d))
    }

    fn encode(&self, out: Out) -> In {
        self.left.encode(self.right.encode(out))
    }
}
*/

pub struct Framed {
    sz_buf: MutByteBuf,
    msg: Option<MutByteBuf>,
}

impl Framed {
    pub fn new() -> Framed {
        Framed {
            sz_buf: ByteBuf::mut_with_capacity(4),
            msg: None,
        }
    }
}

impl Codec<ByteBuf, ByteBuf> for Framed {

    fn decode(&mut self, buf: &mut ByteBuf) -> Vec<ByteBuf> {
        let mut res = vec![];
        loop {
            // read size if we don't have a message yet
            if self.msg.is_none() {
                // TODO(tyler) incrementally build size, or this will explode
                let sz_read = buf.try_read_buf(&mut self.sz_buf);
                debug!("read {} bytes into the sz buffer", sz_read.unwrap().unwrap());
                // if we've read 4 bytes for the size, create a msg
                if self.sz_buf.remaining() != 0 {
                    break;
                }

                let sz_buf = self.sz_buf.bytes();
                let size = array_to_usize([sz_buf[0], sz_buf[1], sz_buf[2], sz_buf[3]]);
                self.msg = unsafe {
                    // manually create bytebuf so we can have exact cap and lim
                    Some(ByteBuf::from_mem_ref(
                        alloc::heap(size.next_power_of_two()),
                        size as u32, // cap
                        0,           // pos
                        size as u32  // lim
                    ).flip())
                };
            }

            if self.msg.is_none() {
                break;
            }

            let mut msg = self.msg.take().unwrap();

            // read actual message
            match buf.try_read_buf(&mut msg) {
                Ok(Some(read)) => {
                    // if we're done, return our Item
                    if msg.remaining() == 0 {
                        // get ready to read a new size
                        self.sz_buf.clear();
                        // return the message
                        res.push(msg.flip())
                    } else {
                        self.msg = Some(msg);
                        break
                    }
                },
                _ => break,
            }
        }
        res
    }

    fn encode(&self, item: ByteBuf) -> ByteBuf {
        let b = item.bytes();
        println!("encoding: {:?}", b);
        let mut res = ByteBuf::mut_with_capacity(4 + b.len());
        assert!(res.write_slice(&usize_to_array(b.len())) == 4);
        assert!(res.write_slice(b) == b.len());
        res.flip()
    }
}

pub fn usize_to_array(u: usize) -> [u8;4] {
    [(u >> 24) as u8, (u >> 16) as u8, (u >> 8) as u8, u as u8]
}

pub fn array_to_usize(ip: [u8;4]) -> usize {
    ((ip[0] as usize) << 24) as usize +
        ((ip[1] as usize) << 16) as usize +
        ((ip[2] as usize) << 8) as usize +
        (ip[3] as usize)
}

#[cfg(test)]
mod tests {
    extern crate quickcheck;
    use rand::{Rng, thread_rng};

    use codec;
    use codec::Codec;
    use bytes::{MutByteBuf, ByteBuf, Buf};

    fn array_prop(u: usize) -> bool {
        codec::array_to_usize(codec::usize_to_array(u)) == u
    }

    #[test]
    fn test_usize_to_array_to_usize() {
        quickcheck::quickcheck(array_prop as fn(usize)->bool);
        let ip = [250,1,2,3];
        assert!(codec::usize_to_array(codec::array_to_usize(ip)) == ip);
    }

    fn framed_prop(sz: usize) -> bool {
        if sz == 0 {
            // TODO(tyler) currently, feeding an empty slice to
            // ByteBuf::from_slice causes a segfault...
            return true;
        }
        let mut rng = thread_rng();
        let mut v: Vec<u8> = rng.gen_iter::<u8>().take(sz).collect();
        let mut c = codec::Framed::new();
        let mut bytes = ByteBuf::from_slice(&*v);
        let mut encoded = codec::Framed::encode(bytes);
        let decoded = c.decode(&mut encoded).unwrap();
        decoded.bytes() == &*v
    }

    #[test]
    fn test_framed_codec() {
        quickcheck::quickcheck(framed_prop as fn(usize)->bool);
    }
}
<commit_msg>create CodecStack that can compose multiple codecs<commit_after>use std::io;
use std::mem;
use std::ops::Add;

use bytes::{alloc, MutByteBuf, ByteBuf, Buf, MutBuf};
use mio::{TryWrite, TryRead};

pub trait Codec<In: ?Sized, Out: ?Sized>
{
    fn decode(&mut self, buf: &mut In) -> Vec<Out>;
    fn encode(&self, a: Out) -> In;
}

/*
impl <In: ?Sized, Mid: ?Sized, Out: ?Sized> Add<Codec<Mid, Out>> for Codec<In, Mid> {
    type Output = CodecStack<In, Mid, Out>;

    fn add(self, right: Codec<Mid, Out>) -> CodecStack<In, Mid, Out> {
        CodecStack {
            left: Box::new(self),
            right: Box::new(right),
        }
    }
}
*/

pub struct CodecStack<In, Mid, Out> {
    left: Box<Codec<In, Mid>>,
    right: Box<Codec<Mid, Out>>,
}

impl <In, Mid, Out> Codec<In, Out> for CodecStack<In, Mid, Out> {
    fn decode(&mut self, buf: &mut In) -> Vec<Out> {
        self.left.decode(buf).iter_mut().flat_map( |mut d|
            self.right.decode(&mut d)).collect()
    }

    fn encode(&self, out: Out) -> In {
        self.left.encode(self.right.encode(out))
    }
}

pub struct Framed {
    sz_buf: MutByteBuf,
    msg: Option<MutByteBuf>,
}

impl Framed {
    pub fn new() -> Framed {
        Framed {
            sz_buf: ByteBuf::mut_with_capacity(4),
            msg: None,
        }
    }
}

impl Codec<ByteBuf, ByteBuf> for Framed {

    fn decode(&mut self, buf: &mut ByteBuf) -> Vec<ByteBuf> {
        let mut res = vec![];
        loop {
            // read size if we don't have a message yet
            if self.msg.is_none() {
                // TODO(tyler) incrementally build size, or this will explode
                let sz_read = buf.try_read_buf(&mut self.sz_buf);
                debug!("read {} bytes into the sz buffer", sz_read.unwrap().unwrap());
                // if we've read 4 bytes for the size, create a msg
                if self.sz_buf.remaining() != 0 {
                    break;
                }

                let sz_buf = self.sz_buf.bytes();
                let size = array_to_usize([sz_buf[0], sz_buf[1], sz_buf[2], sz_buf[3]]);
                self.msg = unsafe {
                    // manually create bytebuf so we can have exact cap and lim
                    Some(ByteBuf::from_mem_ref(
                        alloc::heap(size.next_power_of_two()),
                        size as u32, // cap
                        0,           // pos
                        size as u32  // lim
                    ).flip())
                };
            }

            if self.msg.is_none() {
                break;
            }

            let mut msg = self.msg.take().unwrap();

            // read actual message
            match buf.try_read_buf(&mut msg) {
                Ok(Some(read)) => {
                    // if we're done, return our Item
                    if msg.remaining() == 0 {
                        // get ready to read a new size
                        self.sz_buf.clear();
                        // return the message
                        res.push(msg.flip())
                    } else {
                        self.msg = Some(msg);
                        break
                    }
                },
                _ => break,
            }
        }
        res
    }

    fn encode(&self, item: ByteBuf) -> ByteBuf {
        let b = item.bytes();
        println!("encoding: {:?}", b);
        let mut res = ByteBuf::mut_with_capacity(4 + b.len());
        assert!(res.write_slice(&usize_to_array(b.len())) == 4);
        assert!(res.write_slice(b) == b.len());
        res.flip()
    }
}

pub fn usize_to_array(u: usize) -> [u8;4] {
    [(u >> 24) as u8, (u >> 16) as u8, (u >> 8) as u8, u as u8]
}

pub fn array_to_usize(ip: [u8;4]) -> usize {
    ((ip[0] as usize) << 24) as usize +
        ((ip[1] as usize) << 16) as usize +
        ((ip[2] as usize) << 8) as usize +
        (ip[3] as usize)
}

#[cfg(test)]
mod tests {
    extern crate quickcheck;
    use rand::{Rng, thread_rng};

    use codec;
    use codec::Codec;
    use bytes::{MutByteBuf, ByteBuf, Buf};

    fn array_prop(u: usize) -> bool {
        codec::array_to_usize(codec::usize_to_array(u)) == u
    }

    #[test]
    fn test_usize_to_array_to_usize() {
        quickcheck::quickcheck(array_prop as fn(usize)->bool);
        let ip = [250,1,2,3];
        assert!(codec::usize_to_array(codec::array_to_usize(ip)) == ip);
    }

    fn framed_prop(sz: usize) -> bool {
        if sz == 0 {
            // TODO(tyler) currently, feeding an empty slice to
            // ByteBuf::from_slice causes a segfault...
            return true;
        }
        let mut rng = thread_rng();
        let mut v: Vec<u8> = rng.gen_iter::<u8>().take(sz).collect();
        let mut c = codec::Framed::new();
        let mut bytes = ByteBuf::from_slice(&*v);
        let mut encoded = codec::Framed::encode(bytes);
        let decoded = c.decode(&mut encoded).unwrap();
        decoded.bytes() == &*v
    }

    #[test]
    fn test_framed_codec() {
        quickcheck::quickcheck(framed_prop as fn(usize)->bool);
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 tokio-jsonrpc Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! The codecs to encode and decode messages from a stream of bytes.
//!
//! You can choose to use either line separated one ([Line](struct.Line.html)) or
//! boundary separated one ([Boundary](struct.Boundary.html)). The first one needs the
//! messages to be separated by newlines and not to contain newlines in their representation. On
//! the other hand, it can recover from syntax error in a message and you can respond with an error
//! instead of terminating the connection.

use std::io::{Result as IoResult, Error, ErrorKind};

use tokio_core::io::{Codec, EasyBuf};
use serde_json::ser::to_vec;
use serde_json::error::Error as SerdeError;

use message::{Message, Parsed, from_slice};

/// A helper to wrap the error
fn err_map(e: SerdeError) -> Error {
    Error::new(ErrorKind::Other, e)
}

/// A codec working with JSONRPC 2.0 messages.
///
/// This produces or encodes [Message](../message/enum.Message.hmtl). It separates the records by
/// newlines, so it can recover from syntax error.s
///
/// Note that the produced items is a `Result`, to allow not terminating the stream on
/// protocol-level errors.
#[derive(Debug, Default)]
pub struct Line(usize);

impl Line {
    pub fn new() -> Self {
        Line(0)
    }
}

impl Codec for Line {
    type In = Parsed;
    type Out = Message;
    fn decode(&mut self, buf: &mut EasyBuf) -> IoResult<Option<Parsed>> {
        // Where did we stop scanning before? Scan only the new part
        let start_pos = self.0;
        if let Some(i) = buf.as_slice()[start_pos..].iter().position(|&b| b == b'\n') {
            let end_pos = start_pos + i;
            let line = buf.drain_to(end_pos);
            buf.drain_to(1);
            // We'll start from the beginning next time.
            self.0 = 0;
            Ok(Some(from_slice(line.as_slice())))
        } else {
            // Mark where we ended scanning.
            self.0 = buf.len();
            Ok(None)
        }
    }
    fn encode(&mut self, msg: Message, buf: &mut Vec<u8>) -> IoResult<()> {
        let mut encoded = to_vec(&msg).map_err(err_map)?;
        // As discovered the hard way, we must not overwrite buf, but append to it.
        buf.reserve(encoded.len() + 1);
        buf.append(&mut encoded);
        buf.push(b'\n');
        Ok(())
    }
}

/// A codec working with JSONRPC 2.0 messages.
///
/// This produces or encodes [Message](../message/enum.Message.hmtl). It takes the JSON object boundaries,
/// so it works with both newline-separated and object-separated encoding. It produces
/// newline-separated stream, which is more generic.
///
/// TODO: This is not implemented yet.
pub struct Boundary;

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

    #[test]
    fn encode() {
        let mut output = Vec::new();
        let mut codec = Line::new();
        codec.encode(Message::notification("notif".to_owned(), None), &mut output).unwrap();
        assert_eq!(Vec::from(&b"{\"jsonrpc\":\"2.0\",\"method\":\"notif\"}\n"[..]), output);
    }

    #[test]
    fn decode() {
        fn one(input: &[u8], rest: &[u8]) -> IoResult<Option<Parsed>> {
            let mut codec = Line::new();
            let mut buf = EasyBuf::new();
            buf.get_mut().extend_from_slice(input);
            let result = codec.decode(&mut buf);
            assert_eq!(rest, buf.as_slice());
            result
        }

        let notif = Message::notification("notif".to_owned(), None);
        let msgstring = Vec::from(&b"{\"jsonrpc\":\"2.0\",\"method\":\"notif\"}\n"[..]);
        // A single message, nothing is left
        assert_eq!(one(&msgstring, b"").unwrap(), Some(Ok(notif.clone())));
        // The first message is decoded, the second stays in the buffer
        let mut twomsgs = msgstring.clone();
        twomsgs.extend_from_slice(&msgstring);
        assert_eq!(one(&twomsgs, &msgstring).unwrap(), Some(Ok(notif.clone())));
        // The second message is incomplete, but stays there
        let incomplete = Vec::from(&br#"{"jsonrpc": "2.0", "method":""#[..]);
        let mut oneandhalf = msgstring.clone();
        oneandhalf.extend_from_slice(&incomplete);
        assert_eq!(one(&oneandhalf, &incomplete).unwrap(), Some(Ok(notif.clone())));
        // An incomplete message ‒ nothing gets out and everything stays
        assert_eq!(one(&incomplete, &incomplete).unwrap(), None);
        // A syntax error is reported as an error (and eaten, but that's no longer interesting)
        match one(b"{]\n", b"") {
            Ok(Some(Err(Broken::SyntaxError(_)))) => (),
            other => panic!("Something unexpected: {:?}", other),
        };
    }
}
<commit_msg>The DirtyLine codec<commit_after>// Copyright 2017 tokio-jsonrpc Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! The codecs to encode and decode messages from a stream of bytes.
//!
//! You can choose to use either line separated one ([Line](struct.Line.html)) or
//! boundary separated one ([Boundary](struct.Boundary.html)). The first one needs the
//! messages to be separated by newlines and not to contain newlines in their representation. On
//! the other hand, it can recover from syntax error in a message and you can respond with an error
//! instead of terminating the connection.

use std::io::{Result as IoResult, Error, ErrorKind};

use tokio_core::io::{Codec, EasyBuf};
use serde_json::ser::to_vec;
use serde_json::error::Error as SerdeError;

use message::{Message, Parsed, from_slice, from_str};

/// A helper to wrap the error
fn err_map(e: SerdeError) -> Error {
    Error::new(ErrorKind::Other, e)
}

/// A helper trait to unify `Line` and `DirtyLine`
trait PositionCache {
    fn position(&mut self) -> &mut usize;
}

/// An encoding function reused by `Line` and `DirtyLine`
fn encode_codec(msg: &Message, buf: &mut Vec<u8>) -> IoResult<()> {
    let mut encoded = to_vec(&msg).map_err(err_map)?;
    // As discovered the hard way, we must not overwrite buf, but append to it.
    buf.reserve(encoded.len() + 1);
    buf.append(&mut encoded);
    buf.push(b'\n');
    Ok(())
}

fn decode_codec<Cache, Convert>(cache: &mut Cache, buf: &mut EasyBuf, convert: Convert) -> IoResult<Option<Parsed>>
    where
        Cache: PositionCache,
        Convert: FnOnce(&[u8]) -> Parsed
{
    // Where did we stop scanning before? Scan only the new part
    let start_pos = cache.position();
    if let Some(i) = buf.as_slice()[*start_pos..].iter().position(|&b| b == b'\n') {
        let end_pos = *start_pos + i;
        let line = buf.drain_to(end_pos);
        buf.drain_to(1);
        // We'll start from the beginning next time.
        *start_pos = 0;
        Ok(Some(convert(line.as_slice())))
    } else {
        // Mark where we ended scanning.
        *start_pos = buf.len();
        Ok(None)
    }
}

/// A codec working with JSONRPC 2.0 messages.
///
/// This produces or encodes [Message](../message/enum.Message.hmtl). It separates the records by
/// newlines, so it can recover from syntax error.s
///
/// Note that the produced items is a `Result`, to allow not terminating the stream on
/// protocol-level errors.
#[derive(Debug, Default)]
pub struct Line(usize);

impl Line {
    /// A constructor
    pub fn new() -> Self {
        Line(0)
    }
}

impl PositionCache for Line {
    fn position(&mut self) -> &mut usize {
        &mut self.0
    }
}

impl Codec for Line {
    type In = Parsed;
    type Out = Message;
    fn decode(&mut self, buf: &mut EasyBuf) -> IoResult<Option<Parsed>> {
        decode_codec(self, buf, from_slice)
    }
    fn encode(&mut self, msg: Message, buf: &mut Vec<u8>) -> IoResult<()> {
        encode_codec(&msg, buf)
    }
}

/// A codec working with JSONRPC 2.0 messages on top of badly encoded utf-8.
///
/// This works like the [Line](struct.Line.html) codec. However, it can cope with the input not
/// being valid utf-8. That is arguably broken, nevertheless found in the wild and sometimes the
/// only thing left to be done is to cope with it. This copes with it by running the input through
/// the `String::from_utf8_lossy` conversion, effectively replacing anything that is not valid with
/// these special utf-8 WTF question marks (U+FFFD).
///
/// In contrast, Line errors on such invalid inputs. Encoding is the same for both codecs, however.
#[derive(Debug, Default)]
pub struct DirtyLine(usize);

impl DirtyLine {
    /// A constructor
    pub fn new() -> Self {
        DirtyLine(0)
    }
}

impl PositionCache for DirtyLine {
    fn position(&mut self) -> &mut usize {
        &mut self.0
    }
}

impl Codec for DirtyLine {
    type In = Parsed;
    type Out = Message;
    fn decode(&mut self, buf: &mut EasyBuf) -> IoResult<Option<Parsed>> {
        decode_codec(self, buf, |bytes| from_str(String::from_utf8_lossy(bytes).as_ref()))
    }
    fn encode(&mut self, msg: Message, buf: &mut Vec<u8>) -> IoResult<()> {
        encode_codec(&msg, buf)
    }
}

/// A codec working with JSONRPC 2.0 messages.
///
/// This produces or encodes [Message](../message/enum.Message.hmtl). It takes the JSON object boundaries,
/// so it works with both newline-separated and object-separated encoding. It produces
/// newline-separated stream, which is more generic.
///
/// TODO: This is not implemented yet.
pub struct Boundary;

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

    #[test]
    fn encode() {
        let mut output = Vec::new();
        let mut codec = Line::new();
        let msg = Message::notification("notif".to_owned(), None);
        let encoded = Vec::from(&b"{\"jsonrpc\":\"2.0\",\"method\":\"notif\"}\n"[..]);
        codec.encode(msg.clone(), &mut output).unwrap();
        assert_eq!(encoded, output);
        let mut dirty_codec = DirtyLine::new();
        output.clear();
        dirty_codec.encode(msg, &mut output).unwrap();
        assert_eq!(encoded, output);
    }

    fn get_buf(input: &[u8]) -> EasyBuf {
        let mut result = EasyBuf::new();
        result.get_mut().extend_from_slice(input);
        result
    }

    #[test]
    fn decode() {
        fn one(input: &[u8], rest: &[u8]) -> IoResult<Option<Parsed>> {
            let mut codec = Line::new();
            let mut buf = get_buf(input);
            let result = codec.decode(&mut buf);
            assert_eq!(rest, buf.as_slice());
            // On all the valid inputs, DirtyLine should act the same as Line
            let mut dirty_codec = DirtyLine::new();
            let mut buf = get_buf(input);
            let dirty = dirty_codec.decode(&mut buf);
            assert_eq!(rest, buf.as_slice());
            assert_eq!(result.as_ref().unwrap(), dirty.as_ref().unwrap());
            result
        }

        let notif = Message::notification("notif".to_owned(), None);
        let msgstring = Vec::from(&b"{\"jsonrpc\":\"2.0\",\"method\":\"notif\"}\n"[..]);
        // A single message, nothing is left
        assert_eq!(one(&msgstring, b"").unwrap(), Some(Ok(notif.clone())));
        // The first message is decoded, the second stays in the buffer
        let mut twomsgs = msgstring.clone();
        twomsgs.extend_from_slice(&msgstring);
        assert_eq!(one(&twomsgs, &msgstring).unwrap(), Some(Ok(notif.clone())));
        // The second message is incomplete, but stays there
        let incomplete = Vec::from(&br#"{"jsonrpc": "2.0", "method":""#[..]);
        let mut oneandhalf = msgstring.clone();
        oneandhalf.extend_from_slice(&incomplete);
        assert_eq!(one(&oneandhalf, &incomplete).unwrap(), Some(Ok(notif.clone())));
        // An incomplete message ‒ nothing gets out and everything stays
        assert_eq!(one(&incomplete, &incomplete).unwrap(), None);
        // A syntax error is reported as an error (and eaten, but that's no longer interesting)
        match one(b"{]\n", b"") {
            Ok(Some(Err(Broken::SyntaxError(_)))) => (),
            other => panic!("Something unexpected: {:?}", other),
        };
    }

    /// Test with invalid utf-8 in a string
    #[test]
    fn decode_nonunicode() {
        let broken_input = b"{\"jsonrpc\":\"2.0\",\"method\":\"Hello \xF0\x90\x80World\"}\n";
        let mut codec = Line::new();
        let mut buf = get_buf(broken_input);
        // The ordinary line codec gives up
        let result = codec.decode(&mut buf).unwrap();
        match result {
            Some(Err(Broken::SyntaxError(_))) => (),
            other => panic!("Something unexpected: {:?}", other),
        };
        buf = get_buf(broken_input);
        // But the dirty one just keeps going on
        let mut dirty = DirtyLine::new();
        let result = dirty.decode(&mut buf).unwrap();
        assert_eq!(result, Some(Ok(Message::notification("Hello �World".to_owned(), None))));
    }
}
<|endoftext|>