<commit_before>//! Implementation of Xen's ring mechanism

use core::{mem, cmp};

use core::marker::PhantomData;

use rlibc::memset;

use alloc_uni::__rust_allocate;

use hal::mmu::{Vaddr, Mfn};

use hal::arch::utils::{mb, wmb};

use hal::arch::defs::PAGE_SIZE;

use hal::xen::grant::Ref as GrantRef;
use hal::xen::grant::Table as GrantTable;

// Round a 32-bit unsigned constant down to the nearest power of two.
// Taken from xen/interface/io/ring.h
macro_rules! RD2 {
    ($x:expr) => (
        if $x & 0x2 != 0 {
            0x2
        } else {
            $x & 0x1
        }
    );
}

macro_rules! RD4 {
    ($x:expr) => (
        if $x & 0xC != 0 {
            RD2!($x >> 2) << 2
        } else {
            RD2!($x)
        }
    );
}

macro_rules! RD8 {
    ($x:expr) => (
        if $x & 0xF0 != 0 {
            RD4!($x >> 4) << 4
        } else {
            RD4!($x)
        }
    );
}

macro_rules! RD16 {
    ($x:expr) => (
        if $x & 0xFF00 != 0 {
            RD8!($x >> 8) << 8
        } else {
            RD8!($x)
        }
    );
}

macro_rules! RD32 {
    ($x:expr) => (
        if $x & 0xFFFF0000 != 0 {
            RD16!($x >> 16) << 16
        } else {
            RD16!($x)
        }
    );
}

/// Ring index type
pub type Idx = u32;

#[repr(C)]
/// Shared ring page
///
/// Equivalent of struct __name##_sring defined by DEFINE_RING_TYPES
struct XenSharedRing {
    req_prod: Idx,
    req_event: Idx,
    rsp_prod: Idx,
    rsp_event: Idx,
    pad: [u8; 48],
    ring: u8, // Variable length
}

// XXX: deallocate/end grant access on Drop ?
/// Xen's shared ring
///
/// `Req` is the type of the request
///
/// `Resp` is the type of the response
pub struct SharedRing<Req, Resp> {
    sring: *mut XenSharedRing,
    size: usize,
    elt_size: usize,
    grant_ref: Option<GrantRef>,
    _req: PhantomData<Req>,
    _resp: PhantomData<Resp>,
}

impl<Req, Resp> SharedRing<Req, Resp> {
    /// Allocate a new shared ring
    ///
    /// May return None for two reasons:
    /// - `size` is not aligned on PAGE_SIZE
    /// - The system is out of memory
    pub fn new(size: usize) -> Option<SharedRing<Req, Resp>> {
        if size < PAGE_SIZE || size % PAGE_SIZE != 0 {
            return None;
        }

        // Allocate memory for the ring
        let sring = __rust_allocate(size, PAGE_SIZE) as *mut XenSharedRing;

        if sring.is_null() {
            return None;
        }

        // Initialize the ring
        unsafe {
            memset(sring as *mut u8, 0, PAGE_SIZE);

            (*sring).req_prod = 0;
            (*sring).rsp_prod = 0;
            (*sring).req_event = 1;
            (*sring).rsp_event = 1;
        }

        Some(SharedRing {
            sring: sring,
            size: size,
            elt_size: cmp::max(mem::size_of::<Req>(), mem::size_of::<Resp>()),
            grant_ref: None,
            _req: PhantomData,
            _resp: PhantomData,
        })
    }

    /// Create a shared ring from an existing memory area
    pub unsafe fn from_ptr(ptr: *mut u8, size: usize) -> SharedRing<Req, Resp> {
        SharedRing {
            sring: ptr as *mut _,
            size: size,
            elt_size: cmp::max(mem::size_of::<Req>(), mem::size_of::<Resp>()),
            grant_ref: None,
            _req: PhantomData,
            _resp: PhantomData,
        }
    }

    #[inline]
    /// Returns the request production index
    pub fn req_prod(&self) -> Idx {
        unsafe {
            (*self.sring).req_prod
        }
    }

    #[inline]
    /// Returns the response production index
    pub fn rsp_prod(&self) -> Idx {
        unsafe {
            (*self.sring).rsp_prod
        }
    }

    #[inline]
    /// Returns the request event index
    pub fn req_event(&self) -> Idx {
        unsafe {
            (*self.sring).req_event
        }
    }

    #[inline]
    /// Returns the response event index
    pub fn rsp_event(&self) -> Idx {
        unsafe {
            (*self.sring).rsp_event
        }
    }

    /// Set the request production index
    ///
    /// This is unsafe as an invalid index can corrupt the ring
    pub unsafe fn req_prod_set(&mut self, req: Idx) {
        (*self.sring).req_prod = req;
        mb();
    }

    /// Set the response production index
    ///
    /// This is unsafe as an invalid index can corrupt the ring
    pub unsafe fn rsp_prod_set(&mut self, rsp: Idx) {
        (*self.sring).rsp_prod = rsp;
        mb();
    }

    #[inline]
    /// Set the request event index
    pub unsafe fn req_event_set(&mut self, req: Idx) {
        (*self.sring).req_event = req;
    }

    #[inline]
    /// Set the response event index
    pub unsafe fn rsp_event_set(&mut self, rsp: Idx) {
        (*self.sring).rsp_event = rsp;
    }

    /// Returns the size of the ring
    ///
    /// Equivalent to __RING_SIZE
    pub fn size(&self) -> usize {
        let ring_addr = unsafe { &(*self.sring).ring } as *const _ as isize;
        let sring_addr = self.sring as isize;

        RD32!((self.size as isize - ring_addr + sring_addr) as usize /
              self.elt_size)
    }

    /// Grant access to the ring to a foreign domain with id `domid`
    pub fn grant_access(&mut self, domid: u16) -> Option<GrantRef> {
        if let Some(ref r) = self.grant_ref {
            return Some(r.clone());
        }

        match GrantTable::alloc_ref() {
            None => None,
            Some(r) => {
                r.grant_access(domid, Mfn::from(Vaddr::from_ptr(self.sring)),
                               false);

                self.grant_ref = Some(r.clone());

                Some(r)
            }
        }
    }

    #[inline]
    fn ptr_from_index(&mut self, mut idx: usize) -> *mut u8 {
        idx = idx & (self.size() - 1);

        unsafe {
            let ring_ptr: *mut u8 = &mut (*self.sring).ring;

            ring_ptr.offset((idx * self.elt_size) as isize) as *mut u8
        }
    }

    /// Get a reference over the request at position `idx`
    pub fn request_from_index(&mut self, idx: usize) -> &mut Req {
        unsafe {
            &mut *(self.ptr_from_index(idx) as *mut Req)
        }
    }

    /// Get a reference over the response at position `idx`
    pub fn response_from_index(&mut self, idx: usize) -> &mut Resp {
        unsafe {
            &mut *(self.ptr_from_index(idx) as *mut Resp)
        }
    }
}

/// Xen's front end ring
///
/// `Req` is the type of the request
///
/// `Resp` is the type of the response
pub struct FrontRing<Req, Resp> {
    req_prod_pvt: Idx,
    rsp_cons_pvt: Idx,
    nr_ents: u32,
    sring: SharedRing<Req, Resp>,
}

impl<Req, Resp> FrontRing<Req, Resp> {
    /// Create a new frond end ring that operates on the shared ring `sring`
    pub fn new(sring: SharedRing<Req, Resp>) -> Self {
        FrontRing {
            req_prod_pvt: 0,
            rsp_cons_pvt: 0,
            nr_ents: sring.size() as u32,
            sring: sring,
        }
    }

    #[inline]
    /// Returns the size of the ring (i.e. the number of entries it contains)
    pub fn size(&self) -> usize {
        self.nr_ents as usize
    }

    #[inline]
    /// Returns the request production index
    pub fn req_prod(&self) -> Idx {
        self.req_prod_pvt
    }

    #[inline]
    /// Returns the response consumer index
    pub fn rsp_cons(&self) -> Idx {
        self.rsp_cons_pvt
    }

    #[inline]
    /// Returns the request production index as mutable
    ///
    /// This is unsafe as an invalid index can corrupt the ring
    pub unsafe fn req_prod_mut(&mut self) -> &mut Idx {
        &mut self.req_prod_pvt
    }

    #[inline]
    /// Returns the response production index as mutable
    ///
    /// This is unsafe as an invalid index can corrupt the ring
    pub unsafe fn rsp_cons_mut(&mut self) -> &mut Idx {
        &mut self.rsp_cons_pvt
    }

    #[inline]
    /// Returns a mutable reference over the underlying shared ring
    pub unsafe fn sring_mut(&mut self) -> &mut SharedRing<Req, Resp> {
        &mut self.sring
    }

    #[inline]
    /// Returns true if the ring has unconsumed responses
    ///
    /// Equivalent to RING_HAS_UNCONSUMED_RESPONSES
    pub fn has_unconsumed_responses(&self) -> bool {
        self.sring.rsp_prod() - self.rsp_cons_pvt > 0
    }

    /// Push requests down to the shared ring.
    ///
    /// This function will return true if the other end of the ring should be
    /// notified
    ///
    /// Equivalent to RING_PUSH_REQUESTS_AND_CHECK_NOTIFY
    pub fn push_requests(&mut self) -> bool {
        let old = self.sring.req_prod();
        let new = self.req_prod_pvt;

        wmb();

        unsafe {
            self.sring.req_prod_set(new);
        }
        mb();

        (new - self.sring.req_event()) < (new - old)
    }

    /// Returns true if there is still work to do
    ///
    /// Equivalent to RING_FINAL_CHECK_FOR_RESPONSES
    pub fn final_check_for_responses(&mut self) -> bool {
        if self.has_unconsumed_responses() {
            return true;
        }

        let rsp_event = self.rsp_cons() + 1;

        unsafe {
            self.sring.rsp_event_set(rsp_event);
        }

        self.has_unconsumed_responses()
    }
}
<commit_msg>xen/ring: Fix some issues detected by clippy<commit_after>//! Implementation of Xen's ring mechanism

use core::{mem, cmp};

use core::marker::PhantomData;

use rlibc::memset;

use alloc_uni::__rust_allocate;

use hal::mmu::{Vaddr, Mfn};

use hal::arch::utils::{mb, wmb};

use hal::arch::defs::PAGE_SIZE;

use hal::xen::grant::Ref as GrantRef;
use hal::xen::grant::Table as GrantTable;

// Round a 32-bit unsigned constant down to the nearest power of two.
// Taken from xen/interface/io/ring.h
macro_rules! RD2 {
    ($x:expr) => (
        if ($x) & 0x2 != 0 {
            0x2
        } else {
            ($x) & 0x1
        }
    );
}

macro_rules! RD4 {
    ($x:expr) => (
        if ($x) & 0xC != 0 {
            RD2!(($x) >> 2) << 2
        } else {
            RD2!($x)
        }
    );
}

macro_rules! RD8 {
    ($x:expr) => (
        if ($x) & 0xF0 != 0 {
            RD4!(($x) >> 4) << 4
        } else {
            RD4!($x)
        }
    );
}

macro_rules! RD16 {
    ($x:expr) => (
        if ($x) & 0xFF00 != 0 {
            RD8!(($x) >> 8) << 8
        } else {
            RD8!($x)
        }
    );
}

macro_rules! RD32 {
    ($x:expr) => (
        if ($x) & 0xFFFF0000 != 0 {
            RD16!(($x) >> 16) << 16
        } else {
            RD16!($x)
        }
    );
}

/// Ring index type
pub type Idx = u32;

#[repr(C)]
/// Shared ring page
///
/// Equivalent of struct __name##_sring defined by DEFINE_RING_TYPES
struct XenSharedRing {
    req_prod: Idx,
    req_event: Idx,
    rsp_prod: Idx,
    rsp_event: Idx,
    pad: [u8; 48],
    ring: u8, // Variable length
}

// XXX: deallocate/end grant access on Drop ?
/// Xen's shared ring
///
/// `Req` is the type of the request
///
/// `Resp` is the type of the response
pub struct SharedRing<Req, Resp> {
    sring: *mut XenSharedRing,
    size: usize,
    elt_size: usize,
    grant_ref: Option<GrantRef>,
    _req: PhantomData<Req>,
    _resp: PhantomData<Resp>,
}

impl<Req, Resp> SharedRing<Req, Resp> {
    /// Allocate a new shared ring
    ///
    /// May return None for two reasons:
    /// - `size` is not aligned on PAGE_SIZE
    /// - The system is out of memory
    pub fn new(size: usize) -> Option<SharedRing<Req, Resp>> {
        if size < PAGE_SIZE || size % PAGE_SIZE != 0 {
            return None;
        }

        // Allocate memory for the ring
        let sring = __rust_allocate(size, PAGE_SIZE) as *mut XenSharedRing;

        if sring.is_null() {
            return None;
        }

        // Initialize the ring
        unsafe {
            memset(sring as *mut u8, 0, PAGE_SIZE);

            (*sring).req_prod = 0;
            (*sring).rsp_prod = 0;
            (*sring).req_event = 1;
            (*sring).rsp_event = 1;
        }

        Some(SharedRing {
            sring: sring,
            size: size,
            elt_size: cmp::max(mem::size_of::<Req>(), mem::size_of::<Resp>()),
            grant_ref: None,
            _req: PhantomData,
            _resp: PhantomData,
        })
    }

    /// Create a shared ring from an existing memory area
    pub unsafe fn from_ptr(ptr: *mut u8, size: usize) -> SharedRing<Req, Resp> {
        SharedRing {
            sring: ptr as *mut _,
            size: size,
            elt_size: cmp::max(mem::size_of::<Req>(), mem::size_of::<Resp>()),
            grant_ref: None,
            _req: PhantomData,
            _resp: PhantomData,
        }
    }

    #[inline]
    /// Returns the request production index
    pub fn req_prod(&self) -> Idx {
        unsafe {
            (*self.sring).req_prod
        }
    }

    #[inline]
    /// Returns the response production index
    pub fn rsp_prod(&self) -> Idx {
        unsafe {
            (*self.sring).rsp_prod
        }
    }

    #[inline]
    /// Returns the request event index
    pub fn req_event(&self) -> Idx {
        unsafe {
            (*self.sring).req_event
        }
    }

    #[inline]
    /// Returns the response event index
    pub fn rsp_event(&self) -> Idx {
        unsafe {
            (*self.sring).rsp_event
        }
    }

    /// Set the request production index
    ///
    /// This is unsafe as an invalid index can corrupt the ring
    pub unsafe fn req_prod_set(&mut self, req: Idx) {
        (*self.sring).req_prod = req;
        mb();
    }

    /// Set the response production index
    ///
    /// This is unsafe as an invalid index can corrupt the ring
    pub unsafe fn rsp_prod_set(&mut self, rsp: Idx) {
        (*self.sring).rsp_prod = rsp;
        mb();
    }

    #[inline]
    /// Set the request event index
    pub unsafe fn req_event_set(&mut self, req: Idx) {
        (*self.sring).req_event = req;
    }

    #[inline]
    /// Set the response event index
    pub unsafe fn rsp_event_set(&mut self, rsp: Idx) {
        (*self.sring).rsp_event = rsp;
    }

    /// Returns the size of the ring
    ///
    /// Equivalent to __RING_SIZE
    pub fn size(&self) -> usize {
        let ring_addr = unsafe { &(*self.sring).ring } as *const _ as isize;
        let sring_addr = self.sring as isize;

        RD32!((self.size as isize - ring_addr + sring_addr) as usize /
              self.elt_size)
    }

    /// Grant access to the ring to a foreign domain with id `domid`
    pub fn grant_access(&mut self, domid: u16) -> Option<GrantRef> {
        if let Some(ref r) = self.grant_ref {
            return Some(r.clone());
        }

        match GrantTable::alloc_ref() {
            None => None,
            Some(r) => {
                r.grant_access(domid, Mfn::from(Vaddr::from_ptr(self.sring)),
                               false);

                self.grant_ref = Some(r.clone());

                Some(r)
            }
        }
    }

    #[inline]
    fn ptr_from_index(&mut self, mut idx: usize) -> *mut u8 {
        idx = idx & (self.size() - 1);

        unsafe {
            let ring_ptr: *mut u8 = &mut (*self.sring).ring;

            ring_ptr.offset((idx * self.elt_size) as isize) as *mut u8
        }
    }

    /// Get a reference over the request at position `idx`
    pub fn request_from_index(&mut self, idx: usize) -> &mut Req {
        unsafe {
            &mut *(self.ptr_from_index(idx) as *mut Req)
        }
    }

    /// Get a reference over the response at position `idx`
    pub fn response_from_index(&mut self, idx: usize) -> &mut Resp {
        unsafe {
            &mut *(self.ptr_from_index(idx) as *mut Resp)
        }
    }
}

/// Xen's front end ring
///
/// `Req` is the type of the request
///
/// `Resp` is the type of the response
pub struct FrontRing<Req, Resp> {
    req_prod_pvt: Idx,
    rsp_cons_pvt: Idx,
    nr_ents: u32,
    sring: SharedRing<Req, Resp>,
}

impl<Req, Resp> FrontRing<Req, Resp> {
    /// Create a new frond end ring that operates on the shared ring `sring`
    pub fn new(sring: SharedRing<Req, Resp>) -> Self {
        FrontRing {
            req_prod_pvt: 0,
            rsp_cons_pvt: 0,
            nr_ents: sring.size() as u32,
            sring: sring,
        }
    }

    #[inline]
    #[cfg_attr(feature = "clippy", allow(cyclomatic_complexity))]
    /// Returns the size of the ring (i.e. the number of entries it contains)
    pub fn size(&self) -> usize {
        self.nr_ents as usize
    }

    #[inline]
    /// Returns the request production index
    pub fn req_prod(&self) -> Idx {
        self.req_prod_pvt
    }

    #[inline]
    /// Returns the response consumer index
    pub fn rsp_cons(&self) -> Idx {
        self.rsp_cons_pvt
    }

    #[inline]
    /// Returns the request production index as mutable
    ///
    /// This is unsafe as an invalid index can corrupt the ring
    pub unsafe fn req_prod_mut(&mut self) -> &mut Idx {
        &mut self.req_prod_pvt
    }

    #[inline]
    /// Returns the response production index as mutable
    ///
    /// This is unsafe as an invalid index can corrupt the ring
    pub unsafe fn rsp_cons_mut(&mut self) -> &mut Idx {
        &mut self.rsp_cons_pvt
    }

    #[inline]
    /// Returns a mutable reference over the underlying shared ring
    pub unsafe fn sring_mut(&mut self) -> &mut SharedRing<Req, Resp> {
        &mut self.sring
    }

    #[inline]
    /// Returns true if the ring has unconsumed responses
    ///
    /// Equivalent to RING_HAS_UNCONSUMED_RESPONSES
    pub fn has_unconsumed_responses(&self) -> bool {
        self.sring.rsp_prod() - self.rsp_cons_pvt > 0
    }

    /// Push requests down to the shared ring.
    ///
    /// This function will return true if the other end of the ring should be
    /// notified
    ///
    /// Equivalent to RING_PUSH_REQUESTS_AND_CHECK_NOTIFY
    pub fn push_requests(&mut self) -> bool {
        let old = self.sring.req_prod();
        let new = self.req_prod_pvt;

        wmb();

        unsafe {
            self.sring.req_prod_set(new);
        }
        mb();

        (new - self.sring.req_event()) < (new - old)
    }

    /// Returns true if there is still work to do
    ///
    /// Equivalent to RING_FINAL_CHECK_FOR_RESPONSES
    pub fn final_check_for_responses(&mut self) -> bool {
        if self.has_unconsumed_responses() {
            return true;
        }

        let rsp_event = self.rsp_cons() + 1;

        unsafe {
            self.sring.rsp_event_set(rsp_event);
        }

        self.has_unconsumed_responses()
    }
}
<|endoftext|>


<commit_before>//! CSS funciri values.

use cssparser::Parser;

use crate::error::*;
use crate::parsers::Parse;
use crate::url_resolver::{Fragment, Href};

/// Used where style properties take a funciri or "none"
///
/// This is not to be used for values which don't come from properties.
/// For example, the `xlink:href` attribute in the `<image>` element
/// does not take a funciri value (which looks like `url(...)`), but rather
/// it takes a plain URL.  Use the `Href` type in that case.
#[derive(Debug, Clone, PartialEq)]
pub enum IRI {
    None,
    Resource(Fragment),
}

impl Default for IRI {
    fn default() -> IRI {
        IRI::None
    }
}

impl IRI {
    /// Returns the contents of an `IRI::Resource`, or `None`
    pub fn get(&self) -> Option<&Fragment> {
        match *self {
            IRI::None => None,
            IRI::Resource(ref f) => Some(f),
        }
    }
}

impl Parse for IRI {
    fn parse<'i>(parser: &mut Parser<'i, '_>) -> Result<IRI, ParseError<'i>> {
        if parser
            .try_parse(|i| i.expect_ident_matching("none"))
            .is_ok()
        {
            Ok(IRI::None)
        } else {
            let loc = parser.current_source_location();

            let url = parser.expect_url()?;

            let href =
                Href::parse(&url).map_err(|e| loc.new_custom_error(ValueErrorKind::from(e)))?;

            match href {
                Href::PlainUrl(_) => Err(loc.new_custom_error(ValueErrorKind::parse_error(
                    "href requires a fragment identifier",
                ))),
                Href::WithFragment(f) => Ok(IRI::Resource(f)),
            }
        }
    }
}

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

    #[test]
    fn parses_none() {
        assert_eq!(IRI::parse_str("none").unwrap(), IRI::None);
    }

    #[test]
    fn parses_url() {
        assert_eq!(
            IRI::parse_str("url(#bar)").unwrap(),
            IRI::Resource(Fragment::new(None, "bar".to_string()))
        );

        assert_eq!(
            IRI::parse_str("url(foo#bar)").unwrap(),
            IRI::Resource(Fragment::new(Some("foo".to_string()), "bar".to_string()))
        );

        // be permissive if the closing ) is missing
        assert_eq!(
            IRI::parse_str("url(#bar").unwrap(),
            IRI::Resource(Fragment::new(None, "bar".to_string()))
        );
        assert_eq!(
            IRI::parse_str("url(foo#bar").unwrap(),
            IRI::Resource(Fragment::new(Some("foo".to_string()), "bar".to_string()))
        );

        assert!(IRI::parse_str("").is_err());
        assert!(IRI::parse_str("foo").is_err());
        assert!(IRI::parse_str("url(foo)bar").is_err());
    }
}
<commit_msg>iri: use Fragment directly<commit_after>//! CSS funciri values.

use cssparser::Parser;

use crate::error::*;
use crate::parsers::Parse;
use crate::url_resolver::Fragment;

/// Used where style properties take a funciri or "none"
///
/// This is not to be used for values which don't come from properties.
/// For example, the `xlink:href` attribute in the `<image>` element
/// does not take a funciri value (which looks like `url(...)`), but rather
/// it takes a plain URL.  Use the `Href` type in that case.
#[derive(Debug, Clone, PartialEq)]
pub enum IRI {
    None,
    Resource(Fragment),
}

impl Default for IRI {
    fn default() -> IRI {
        IRI::None
    }
}

impl IRI {
    /// Returns the contents of an `IRI::Resource`, or `None`
    pub fn get(&self) -> Option<&Fragment> {
        match *self {
            IRI::None => None,
            IRI::Resource(ref f) => Some(f),
        }
    }
}

impl Parse for IRI {
    fn parse<'i>(parser: &mut Parser<'i, '_>) -> Result<IRI, ParseError<'i>> {
        if parser
            .try_parse(|i| i.expect_ident_matching("none"))
            .is_ok()
        {
            Ok(IRI::None)
        } else {
            let loc = parser.current_source_location();
            let url = parser.expect_url()?;
            let fragment =
                Fragment::parse(&url).map_err(|e| loc.new_custom_error(ValueErrorKind::from(e)))?;

            Ok(IRI::Resource(fragment))
        }
    }
}

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

    #[test]
    fn parses_none() {
        assert_eq!(IRI::parse_str("none").unwrap(), IRI::None);
    }

    #[test]
    fn parses_url() {
        assert_eq!(
            IRI::parse_str("url(#bar)").unwrap(),
            IRI::Resource(Fragment::new(None, "bar".to_string()))
        );

        assert_eq!(
            IRI::parse_str("url(foo#bar)").unwrap(),
            IRI::Resource(Fragment::new(Some("foo".to_string()), "bar".to_string()))
        );

        // be permissive if the closing ) is missing
        assert_eq!(
            IRI::parse_str("url(#bar").unwrap(),
            IRI::Resource(Fragment::new(None, "bar".to_string()))
        );
        assert_eq!(
            IRI::parse_str("url(foo#bar").unwrap(),
            IRI::Resource(Fragment::new(Some("foo".to_string()), "bar".to_string()))
        );

        assert!(IRI::parse_str("").is_err());
        assert!(IRI::parse_str("foo").is_err());
        assert!(IRI::parse_str("url(foo)bar").is_err());
    }
}
<|endoftext|>


<commit_before>/// Enable Interrupts.
pub unsafe fn enable()
{
    asm!("sti");
}

/// Disable Interrupts.
pub unsafe fn disable()
{
    asm!("cli");
}
<commit_msg>Added lidt and int! macro.<commit_after>/// Enable Interrupts.
pub unsafe fn enable()
{
    asm!("sti");
}

/// Disable Interrupts.
pub unsafe fn disable()
{
    asm!("cli");
}

/// Generate a software interrupt.
/// This is a macro argument needs to be an immediate.
#[macro_export]
macro_rules! int {
    ( $x:expr ) => {
        {
            asm!("int $0" :: "N" ($x));
        }
    };
}

/// Load IDT table.
pub unsafe fn lidt(idt: u64)
{
    asm!("lidt (%rax)" :: "{rax}" (idt));
}

<|endoftext|>


<commit_before>//! atty is a simple utility that answers one question
//! > is this a tty?
//!
//! usage is just as simple
//!
//! ```
//! if atty::is(atty::Stream::Stdout) {
//!   println!("i'm a tty")
//! }
//! ```
//!
//! ```
//! if atty::isnt(atty::Stream::Stdout) {
//!   println!("i'm not a tty")
//! }
//! ```

extern crate libc;

/// possible stream sources
pub enum Stream {
    Stdout,
    Stderr,
    Stdin
}

/// returns true if this is a tty
#[cfg(unix)]
pub fn is(stream: Stream) -> bool {
    let fd = match stream {
        Stream::Stdout => libc::STDOUT_FILENO,
        Stream::Stderr => libc::STDERR_FILENO,
        Stream::Stdin => libc::STDIN_FILENO
    };
    unsafe { libc::isatty(fd) != 0 }
}

/// returns true if this is a tty
#[cfg(windows)]
pub fn is(stream: Stream) -> bool {
    extern crate kernel32;
    extern crate winapi;

    let handle = match stream {
        Stream::Stdout => winapi::STD_OUTPUT_HANDLE,
        Stream::Stderr => winapi::STD_ERROR_HANDLE,
        Stream::Stdin => winapi::STD_INPUT_HANDLE
    };

    unsafe {
        let handle = kernel32::GetStdHandle(handle);
        let mut out = 0;
        panic!("is invalid? {:#?}  result {:#?} last err {:#?}",
        winapi::INVALID_HANDLE_VALUE == handle,
        kernel32::GetConsoleMode(handle, &mut out),
        kernel32::GetLastError());
        kernel32::GetConsoleMode(handle, &mut out) != 0
    }
}

/// returns true if this is _not_ a tty
pub fn isnt(stream: Stream) -> bool {
    !is(stream)
}

#[cfg(test)]
mod tests {
    use super::{is, Stream};

    #[test]
    fn is_err() {
        assert!(is(Stream::Stderr))
    }

    #[test]
    fn is_out() {
        assert!(is(Stream::Stdout))
    }

    #[test]
    fn is_in() {
        assert!(is(Stream::Stdin))
    }
}
<commit_msg>try GetConsoleScreenBufferInfo for output<commit_after>//! atty is a simple utility that answers one question
//! > is this a tty?
//!
//! usage is just as simple
//!
//! ```
//! if atty::is(atty::Stream::Stdout) {
//!   println!("i'm a tty")
//! }
//! ```
//!
//! ```
//! if atty::isnt(atty::Stream::Stdout) {
//!   println!("i'm not a tty")
//! }
//! ```

extern crate libc;

/// possible stream sources
pub enum Stream {
    Stdout,
    Stderr,
    Stdin
}

/// returns true if this is a tty
#[cfg(unix)]
pub fn is(stream: Stream) -> bool {
    let fd = match stream {
        Stream::Stdout => libc::STDOUT_FILENO,
        Stream::Stderr => libc::STDERR_FILENO,
        Stream::Stdin => libc::STDIN_FILENO
    };
    unsafe { libc::isatty(fd) != 0 }
}

/// returns true if this is a tty
#[cfg(windows)]
pub fn is(stream: Stream) -> bool {
    extern crate kernel32;
    extern crate winapi;

    let handle = match stream {
        Stream::Stdout => winapi::STD_OUTPUT_HANDLE,
        Stream::Stderr => winapi::STD_ERROR_HANDLE,
        Stream::Stdin => winapi::STD_INPUT_HANDLE
    };

    unsafe {
        let std_handle = kernel32::GetStdHandle(handle);
        match stream {
            Stream::Stdin => {
                let mut out = 0;
                panic!("is invalid? {:#?}  result {:#?} last err {:#?}",
                winapi::INVALID_HANDLE_VALUE == std_handle,
                kernel32::GetConsoleMode(std_handle, &mut out),
                kernel32::GetLastError());
                kernel32::GetConsoleMode(std_handle, &mut out) != 0
            },
            _ => {
                let mut buffer_info = unsafe { ::std::mem::uninitialized() };
                panic!("is invalid? {:#?}  result {:#?} last err {:#?}",
                winapi::INVALID_HANDLE_VALUE == std_handle,
                kernel32::GetConsoleMode(std_handle, &mut buffer_info),
                kernel32::GetLastError());
                kernel32::GetConsoleMode(std_handle, &mut buffer_info) != 0
            }
        }

    }
}

/// returns true if this is _not_ a tty
pub fn isnt(stream: Stream) -> bool {
    !is(stream)
}

#[cfg(test)]
mod tests {
    use super::{is, Stream};

    #[test]
    fn is_err() {
        assert!(is(Stream::Stderr))
    }

    #[test]
    fn is_out() {
        assert!(is(Stream::Stdout))
    }

    #[test]
    fn is_in() {
        assert!(is(Stream::Stdin))
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! update-ssh-keys library
//!
//! this library is meant to replace github.com/coreos/update-ssh-keys
//! it will provide the functionality for:
//! * listing authorized keys
//! * adding authorized keys
//! * removing an authorized key by name
//! * disabling an authorized key by name
//!
//! the library will take care of the file locking that is expected from users
//! of the authorized_keys.d directory.

#[macro_use]
extern crate error_chain;
extern crate fs2;
extern crate openssh_keys;
extern crate users;

#[allow(unused_doc_comment)]
pub mod errors {
    error_chain!{
        foreign_links {
            Io(::std::io::Error);
        }
        errors {
            KeysDisabled(name: String) {
                description("keys are disabled")
                display("keys with name '{}' are disabled", name)
            }
            KeysExist(name: String) {
                description("keys already exist")
                display("keys with name '{}' already exist", name)
            }
        }
    }
}

use errors::*;
use users::{User, switch};
use users::os::unix::UserExt;
use openssh_keys::PublicKey;
use std::fs::{File, self};
use fs2::FileExt;
use std::path::{Path, PathBuf};
use std::io::{BufReader, BufRead, Read, Write};
use std::collections::HashMap;

const SSH_DIR: &'static str = ".ssh";
const AUTHORIZED_KEYS_DIR: &'static str = "authorized_keys.d";
const AUTHORIZED_KEYS_FILE: &'static str = "authorized_keys";
const PRESERVED_KEYS_FILE: &'static str = "orig_authorized_keys";
const LOCK_FILE: &'static str = ".authorized_keys.d.lock";
const STAGE_FILE: &'static str = ".authorized_keys.d.stage_file";
const STAGE_DIR: &'static str = ".authorized_keys.d.stage_dir";

fn lock_file(user: &User) -> PathBuf {
    user.home_dir().join(LOCK_FILE)
}

fn ssh_dir(user: &User) -> PathBuf {
    user.home_dir().join(SSH_DIR)
}

fn authorized_keys_dir(user: &User) -> PathBuf {
    ssh_dir(user).join(AUTHORIZED_KEYS_DIR)
}

fn authorized_keys_file(user: &User) -> PathBuf {
    ssh_dir(user).join(AUTHORIZED_KEYS_FILE)
}

fn stage_dir(user: &User) -> PathBuf {
    ssh_dir(user).join(STAGE_DIR)
}

fn stage_file(user: &User) -> PathBuf {
    ssh_dir(user).join(STAGE_FILE)
}

fn switch_user(user: &User) -> Result<switch::SwitchUserGuard> {
    switch::switch_user_group(user.uid(), user.primary_group_id())
        .chain_err(|| format!("failed to switch to user '{}' with uid '{}'", user.name(), user.uid()))
}

#[derive(Debug)]
struct FileLock {
    pub lock: File,
}

impl Drop for FileLock {
    fn drop(&mut self) {
        self.unlock().unwrap();
    }
}

impl FileLock {
    fn new(path: &Path) -> Result<Self> {
        Ok(FileLock {
            lock: File::create(path)
                .chain_err(|| format!("failed to create lock file: {:?}", path))?,
        })
    }

    fn lock(&self) -> Result<()> {
        self.lock.lock_exclusive()
            .chain_err(|| "failed to lock file")
    }

    fn unlock(&self) -> Result<()> {
        self.lock.unlock()
            .chain_err(|| "failed to unlock file")
    }
}

#[derive(Debug)]
pub struct AuthorizedKeys {
    pub file: PathBuf,
    pub folder: PathBuf,
    pub keys: HashMap<String, AuthorizedKeySet>,
    pub user: User,
    lock: FileLock,
}

impl Drop for AuthorizedKeys {
    fn drop(&mut self) {}
}

#[derive(Clone, Debug)]
pub struct AuthorizedKeySet {
    pub filename: String,
    pub disabled: bool,
    pub keys: Vec<PublicKey>,
}

impl AuthorizedKeys {
    /// write writes all authorized_keys.d changes onto disk. it writes the
    /// current state to a staging directory and then moves that staging
    /// directory to the authorized_keys.d path.
    pub fn write(&self) -> Result<()> {
        // switch users
        let _guard = switch_user(&self.user)?;

        // get our staging directory
        let stage_dir = stage_dir(&self.user);
        if stage_dir.exists() {
            fs::remove_dir_all(&stage_dir)
                .chain_err(|| format!("failed to remove staging directory '{:?}'", stage_dir))?;
        }
        fs::create_dir(&stage_dir)
            .chain_err(|| format!("failed to create staging directory '{:?}'", stage_dir))?;

        // write all the keys to the staging directory
        for (_, keyset) in &self.keys {
            let keyfilename = stage_dir.join(&keyset.filename);
            let mut keyfile = File::create(&keyfilename)
                .chain_err(|| format!("failed to create file '{:?}'", keyfilename))?;
            // if the keyset is disabled, skip it. we still want to have a
            // zero-sized file with it's name though to signal that it's
            // disabled.
            if keyset.disabled {
                continue
            }
            for key in &keyset.keys {
                write!(keyfile, "{}\n", key)
                    .chain_err(|| format!("failed to write to file '{:?}'", keyfilename))?;
            }
        }

        // destroy the old authorized keys directory and move the staging one to
        // that location. rename expects an existing, empty directory
        fs::remove_dir_all(&self.folder)
            .chain_err(|| format!("failed to remove directory '{:?}'", stage_dir))?;
        fs::create_dir(&self.folder)
            .chain_err(|| format!("failed to create directory '{:?}'", stage_dir))?;
        fs::rename(&stage_dir, &self.folder)
            .chain_err(|| format!("failed to move '{:?}' to '{:?}'", stage_dir, self.folder))
    }

    /// sync writes all the keys we have to authorized_keys. it writes the
    /// current state to a staging file and then moves that staging file to the
    /// authorized_keys path
    pub fn sync(&self) -> Result<()> {
        // switch users
        let _guard = switch_user(&self.user)?;

        // get our staging directory
        let stage_filename = stage_file(&self.user);
        let mut stage_file = File::create(&stage_filename)
            .chain_err(|| format!("failed to create or truncate staging file '{:?}'", stage_filename))?;

        // write all the keys to the staging file
        for (_, keyset) in &self.keys {
            // if the keyset is disabled, skip it
            if keyset.disabled {
                continue
            }
            for key in &keyset.keys {
                write!(stage_file, "{}\n", key)
                    .chain_err(|| format!("failed to write to file '{:?}'", stage_filename))?;
            }
        }

        // destroy the old authorized keys file and move the staging one to that
        // location
        fs::rename(&stage_filename, &self.file)
            .chain_err(|| format!("failed to move '{:?}' to '{:?}'", stage_filename, self.file))
    }

    /// read_all_keys reads all of the authorized keys files in a given
    /// directory. it returns an error if there is a nested directory, if any
    /// file operations fail, or if it can't parse any of the authorized_keys
    /// files
    fn read_all_keys(dir: &Path) -> Result<HashMap<String, AuthorizedKeySet>> {
        let dir_contents = fs::read_dir(&dir)
            .chain_err(|| format!("failed to read from directory {:?}", dir))?;
        let mut keys = HashMap::new();
        for entry in dir_contents {
            let entry = entry.chain_err(|| format!("failed to read entry in directory {:?}", dir))?;
            let path = entry.path();
            if path.is_dir() {
                // if it's a directory, we don't know what to do
                return Err(format!("'{:?}' is a directory", path).into());
            } else {
                let name = path.file_name()
                    .ok_or_else(|| format!("failed to get filename for '{:?}'", path))?
                    .to_str()
                    .ok_or_else(|| format!("failed to convert filename '{:?}' to string", path))?;
                let from = File::open(&path)
                    .chain_err(|| format!("failed to open file {:?}", path))?;
                let keyset = AuthorizedKeys::read_keys(from)?;
                keys.insert(name.to_string(), AuthorizedKeySet {
                    filename: name.to_string(),
                    disabled: keyset.is_empty(),
                    keys: keyset,
                });
            }
        }
        Ok(keys)
    }

    /// read_keys reads a list of public keys from a reader. it returns an error
    /// of it can't read or parse any of the public keys in the list.
    pub fn read_keys<R>(r: R) -> Result<Vec<PublicKey>>
        where R: Read
    {
        let keybuf = BufReader::new(r);
        // authorized_keys files are newline-separated lists of public keys
        let mut keys = vec![];
        for key in keybuf.lines() {
            let key = key.chain_err(|| "failed to read public key")?;
            // skip any empty lines and any comment lines (prefixed with '#')
            if !key.is_empty() && !(key.trim().starts_with('#')) {
                keys.push(PublicKey::parse(&key)
                          .chain_err(|| "failed to parse public key")?);
            }
        }
        Ok(keys)
    }

    /// open creates a new authorized_keys object. if there is an existing
    /// authorized_keys directory on disk it reads all the keys from that. if
    /// there is no directory already and we are told to create it, we add the
    /// existing authorized keys file as an entry, if it exists.
    /// /// before open actually does any of that, it switches it's uid for the span
    /// of the function and then switched back. it also opens a file lock on the
    /// directory that other instances of `update-ssh-keys` will respect. the
    /// file lock will automatically close when this structure goes out of
    /// scope. you can make sure it is unlocked by calling `drop` yourself in
    /// cases where you think the memory may leak (like if you are tossing boxes
    /// around etc).
    ///
    /// open blocks until it can grab the file lock.
    ///
    /// open returns an error if any file operations fail, if it failes to parse
    /// any of the public keys in the existing files, if it failes to change
    /// users, if it failes to grab the lock, or if create is false but the
    /// directory doesn't exist.
    pub fn open(user: User, create: bool) -> Result<Self> {
        // switch users
        let _guard = switch_user(&user)?;
        // make a new file lock and lock it
        let lock = FileLock::new(&lock_file(&user))?;
        lock.lock()?;

        let akd = authorized_keys_dir(&user);

        let keys = if akd.is_dir() {
            // read the existing keysets from the dir
            AuthorizedKeys::read_all_keys(&akd)?
        } else if !akd.exists() && create {
            // read the existing keyset from the file
            let filename = authorized_keys_file(&user);
            if filename.exists() {
                let file = File::open(&filename)
                    .chain_err(|| format!("failed to open authorized keys file: '{:?}'", filename))?;
                let mut keys = HashMap::new();
                keys.insert(PRESERVED_KEYS_FILE.to_string(), AuthorizedKeySet {
                    filename: PRESERVED_KEYS_FILE.to_string(),
                    disabled: false,
                    keys: AuthorizedKeys::read_keys(file)?,
                });
                keys
            } else {
                // if the authorized_keys file doesn't exist, we don't start
                // with any keys
                HashMap::new()
            }
        } else {
            // either the akd doesn't exist and create is false, or it exists
            // and is not a directory
            return Err(format!("'{:?}' doesn't exist or is not a directory", akd).into())
        };

        Ok(AuthorizedKeys {
            file: authorized_keys_file(&user),
            folder: akd,
            user: user,
            keys: keys,
            lock: lock,
        })
    }

    /// get_keys gets the authorized keyset with the provided name
    pub fn get_keys(&self, name: &str) -> Option<&AuthorizedKeySet> {
        self.keys.get(name)
    }

    /// get_all_keys returns the hashmap from name to keyset containing all the
    /// keys we know about
    pub fn get_all_keys(&self) -> &HashMap<String, AuthorizedKeySet> {
        &self.keys
    }

    /// add_keys adds a list of public keys with the provide name. if replace is
    /// true, it will replace existing keys. if force is true, it will replace
    /// disabled keys.
    ///
    /// add_keys returns an error if the key already exists and replace is
    /// false, or if the key is disabled and force is false
    pub fn add_keys(&mut self, name: &str, keys: Vec<PublicKey>, replace: bool, force: bool) -> Result<()> {
        if let Some(keyset) = self.keys.get(name) {
            if keyset.disabled && !force {
                return Err(ErrorKind::KeysDisabled(name.to_string()).into());
            } else if !replace {
                return Err(ErrorKind::KeysExist(name.to_string()).into());
            }
        }
        self.keys.insert(name.to_string(), AuthorizedKeySet {
            filename: name.to_string(),
            disabled: false,
            keys: keys,
        });
        Ok(())
    }

    /// remove_keys removes the keyset with the given name.
    pub fn remove_keys(&mut self, name: &str) {
        self.keys.remove(name);
    }

    /// disable_keys disables keys with the given name. they can't be added
    /// again unless force is set to true when adding the set. disable_keys will
    /// succeed in disabling the key even if the key doesn't currently exist.
    pub fn disable_keys(&mut self, name: &str) {
        if let Some(keyset) = self.keys.get_mut(name) {
            keyset.disabled = true;
            return;
        }
        self.keys.insert(name.to_string(), AuthorizedKeySet {
            filename: name.to_string(),
            disabled: true,
            keys: vec![],
        });
    }
}
<commit_msg>lib: clippy cleanup<commit_after>// Copyright 2017 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! update-ssh-keys library
//!
//! this library is meant to replace github.com/coreos/update-ssh-keys
//! it will provide the functionality for:
//! * listing authorized keys
//! * adding authorized keys
//! * removing an authorized key by name
//! * disabling an authorized key by name
//!
//! the library will take care of the file locking that is expected from users
//! of the `authorized_keys.d` directory.

#[macro_use]
extern crate error_chain;
extern crate fs2;
extern crate openssh_keys;
extern crate users;

#[allow(unused_doc_comment)]
pub mod errors {
    error_chain!{
        foreign_links {
            Io(::std::io::Error);
        }
        errors {
            KeysDisabled(name: String) {
                description("keys are disabled")
                display("keys with name '{}' are disabled", name)
            }
            KeysExist(name: String) {
                description("keys already exist")
                display("keys with name '{}' already exist", name)
            }
        }
    }
}

use errors::*;
use users::{User, switch};
use users::os::unix::UserExt;
use openssh_keys::PublicKey;
use std::fs::{File, self};
use fs2::FileExt;
use std::path::{Path, PathBuf};
use std::io::{BufReader, BufRead, Read, Write};
use std::collections::HashMap;

const SSH_DIR: &'static str = ".ssh";
const AUTHORIZED_KEYS_DIR: &'static str = "authorized_keys.d";
const AUTHORIZED_KEYS_FILE: &'static str = "authorized_keys";
const PRESERVED_KEYS_FILE: &'static str = "orig_authorized_keys";
const LOCK_FILE: &'static str = ".authorized_keys.d.lock";
const STAGE_FILE: &'static str = ".authorized_keys.d.stage_file";
const STAGE_DIR: &'static str = ".authorized_keys.d.stage_dir";

fn lock_file(user: &User) -> PathBuf {
    user.home_dir().join(LOCK_FILE)
}

fn ssh_dir(user: &User) -> PathBuf {
    user.home_dir().join(SSH_DIR)
}

fn authorized_keys_dir(user: &User) -> PathBuf {
    ssh_dir(user).join(AUTHORIZED_KEYS_DIR)
}

fn authorized_keys_file(user: &User) -> PathBuf {
    ssh_dir(user).join(AUTHORIZED_KEYS_FILE)
}

fn stage_dir(user: &User) -> PathBuf {
    ssh_dir(user).join(STAGE_DIR)
}

fn stage_file(user: &User) -> PathBuf {
    ssh_dir(user).join(STAGE_FILE)
}

fn switch_user(user: &User) -> Result<switch::SwitchUserGuard> {
    switch::switch_user_group(user.uid(), user.primary_group_id())
        .chain_err(|| format!("failed to switch to user '{}' with uid '{}'", user.name(), user.uid()))
}

#[derive(Debug)]
struct FileLock {
    pub lock: File,
}

impl Drop for FileLock {
    fn drop(&mut self) {
        self.unlock().unwrap();
    }
}

impl FileLock {
    fn new(path: &Path) -> Result<Self> {
        Ok(FileLock {
            lock: File::create(path)
                .chain_err(|| format!("failed to create lock file: {:?}", path))?,
        })
    }

    fn lock(&self) -> Result<()> {
        self.lock.lock_exclusive()
            .chain_err(|| "failed to lock file")
    }

    fn unlock(&self) -> Result<()> {
        self.lock.unlock()
            .chain_err(|| "failed to unlock file")
    }
}

#[derive(Debug)]
pub struct AuthorizedKeys {
    pub file: PathBuf,
    pub folder: PathBuf,
    pub keys: HashMap<String, AuthorizedKeySet>,
    pub user: User,
    lock: FileLock,
}

impl Drop for AuthorizedKeys {
    fn drop(&mut self) {}
}

#[derive(Clone, Debug)]
pub struct AuthorizedKeySet {
    pub filename: String,
    pub disabled: bool,
    pub keys: Vec<PublicKey>,
}

impl AuthorizedKeys {
    /// write writes all authorized_keys.d changes onto disk. it writes the
    /// current state to a staging directory and then moves that staging
    /// directory to the authorized_keys.d path.
    pub fn write(&self) -> Result<()> {
        // switch users
        let _guard = switch_user(&self.user)?;

        // get our staging directory
        let stage_dir = stage_dir(&self.user);
        if stage_dir.exists() {
            fs::remove_dir_all(&stage_dir)
                .chain_err(|| format!("failed to remove staging directory '{:?}'", stage_dir))?;
        }
        fs::create_dir(&stage_dir)
            .chain_err(|| format!("failed to create staging directory '{:?}'", stage_dir))?;

        // write all the keys to the staging directory
        for keyset in self.keys.values() {
            let keyfilename = stage_dir.join(&keyset.filename);
            let mut keyfile = File::create(&keyfilename)
                .chain_err(|| format!("failed to create file '{:?}'", keyfilename))?;
            // if the keyset is disabled, skip it. we still want to have a
            // zero-sized file with it's name though to signal that it's
            // disabled.
            if keyset.disabled {
                continue
            }
            for key in &keyset.keys {
                write!(keyfile, "{}\n", key)
                    .chain_err(|| format!("failed to write to file '{:?}'", keyfilename))?;
            }
        }

        // destroy the old authorized keys directory and move the staging one to
        // that location. rename expects an existing, empty directory
        fs::remove_dir_all(&self.folder)
            .chain_err(|| format!("failed to remove directory '{:?}'", stage_dir))?;
        fs::create_dir(&self.folder)
            .chain_err(|| format!("failed to create directory '{:?}'", stage_dir))?;
        fs::rename(&stage_dir, &self.folder)
            .chain_err(|| format!("failed to move '{:?}' to '{:?}'", stage_dir, self.folder))
    }

    /// sync writes all the keys we have to authorized_keys. it writes the
    /// current state to a staging file and then moves that staging file to the
    /// authorized_keys path
    pub fn sync(&self) -> Result<()> {
        // switch users
        let _guard = switch_user(&self.user)?;

        // get our staging directory
        let stage_filename = stage_file(&self.user);
        let mut stage_file = File::create(&stage_filename)
            .chain_err(|| format!("failed to create or truncate staging file '{:?}'", stage_filename))?;

        // write all the keys to the staging file
        for keyset in self.keys.values() {
            // if the keyset is disabled, skip it
            if keyset.disabled {
                continue
            }
            for key in &keyset.keys {
                write!(stage_file, "{}\n", key)
                    .chain_err(|| format!("failed to write to file '{:?}'", stage_filename))?;
            }
        }

        // destroy the old authorized keys file and move the staging one to that
        // location
        fs::rename(&stage_filename, &self.file)
            .chain_err(|| format!("failed to move '{:?}' to '{:?}'", stage_filename, self.file))
    }

    /// read_all_keys reads all of the authorized keys files in a given
    /// directory. it returns an error if there is a nested directory, if any
    /// file operations fail, or if it can't parse any of the authorized_keys
    /// files
    fn read_all_keys(dir: &Path) -> Result<HashMap<String, AuthorizedKeySet>> {
        let dir_contents = fs::read_dir(&dir)
            .chain_err(|| format!("failed to read from directory {:?}", dir))?;
        let mut keys = HashMap::new();
        for entry in dir_contents {
            let entry = entry.chain_err(|| format!("failed to read entry in directory {:?}", dir))?;
            let path = entry.path();
            if path.is_dir() {
                // if it's a directory, we don't know what to do
                return Err(format!("'{:?}' is a directory", path).into());
            } else {
                let name = path.file_name()
                    .ok_or_else(|| format!("failed to get filename for '{:?}'", path))?
                    .to_str()
                    .ok_or_else(|| format!("failed to convert filename '{:?}' to string", path))?;
                let from = File::open(&path)
                    .chain_err(|| format!("failed to open file {:?}", path))?;
                let keyset = AuthorizedKeys::read_keys(from)?;
                keys.insert(name.to_string(), AuthorizedKeySet {
                    filename: name.to_string(),
                    disabled: keyset.is_empty(),
                    keys: keyset,
                });
            }
        }
        Ok(keys)
    }

    /// read_keys reads a list of public keys from a reader. it returns an error
    /// of it can't read or parse any of the public keys in the list.
    pub fn read_keys<R>(r: R) -> Result<Vec<PublicKey>>
        where R: Read
    {
        let keybuf = BufReader::new(r);
        // authorized_keys files are newline-separated lists of public keys
        let mut keys = vec![];
        for key in keybuf.lines() {
            let key = key.chain_err(|| "failed to read public key")?;
            // skip any empty lines and any comment lines (prefixed with '#')
            if !key.is_empty() && !(key.trim().starts_with('#')) {
                keys.push(PublicKey::parse(&key)
                          .chain_err(|| "failed to parse public key")?);
            }
        }
        Ok(keys)
    }

    /// open creates a new authorized_keys object. if there is an existing
    /// authorized_keys directory on disk it reads all the keys from that. if
    /// there is no directory already and we are told to create it, we add the
    /// existing authorized keys file as an entry, if it exists.
    /// /// before open actually does any of that, it switches it's uid for the span
    /// of the function and then switched back. it also opens a file lock on the
    /// directory that other instances of `update-ssh-keys` will respect. the
    /// file lock will automatically close when this structure goes out of
    /// scope. you can make sure it is unlocked by calling `drop` yourself in
    /// cases where you think the memory may leak (like if you are tossing boxes
    /// around etc).
    ///
    /// open blocks until it can grab the file lock.
    ///
    /// open returns an error if any file operations fail, if it failes to parse
    /// any of the public keys in the existing files, if it failes to change
    /// users, if it failes to grab the lock, or if create is false but the
    /// directory doesn't exist.
    pub fn open(user: User, create: bool) -> Result<Self> {
        // switch users
        let _guard = switch_user(&user)?;
        // make a new file lock and lock it
        let lock = FileLock::new(&lock_file(&user))?;
        lock.lock()?;

        let akd = authorized_keys_dir(&user);

        let keys = if akd.is_dir() {
            // read the existing keysets from the dir
            AuthorizedKeys::read_all_keys(&akd)?
        } else if !akd.exists() && create {
            // read the existing keyset from the file
            let filename = authorized_keys_file(&user);
            if filename.exists() {
                let file = File::open(&filename)
                    .chain_err(|| format!("failed to open authorized keys file: '{:?}'", filename))?;
                let mut keys = HashMap::new();
                keys.insert(PRESERVED_KEYS_FILE.to_string(), AuthorizedKeySet {
                    filename: PRESERVED_KEYS_FILE.to_string(),
                    disabled: false,
                    keys: AuthorizedKeys::read_keys(file)?,
                });
                keys
            } else {
                // if the authorized_keys file doesn't exist, we don't start
                // with any keys
                HashMap::new()
            }
        } else {
            // either the akd doesn't exist and create is false, or it exists
            // and is not a directory
            return Err(format!("'{:?}' doesn't exist or is not a directory", akd).into())
        };

        Ok(AuthorizedKeys {
            file: authorized_keys_file(&user),
            folder: akd,
            user: user,
            keys: keys,
            lock: lock,
        })
    }

    /// get_keys gets the authorized keyset with the provided name
    pub fn get_keys(&self, name: &str) -> Option<&AuthorizedKeySet> {
        self.keys.get(name)
    }

    /// get_all_keys returns the hashmap from name to keyset containing all the
    /// keys we know about
    pub fn get_all_keys(&self) -> &HashMap<String, AuthorizedKeySet> {
        &self.keys
    }

    /// add_keys adds a list of public keys with the provide name. if replace is
    /// true, it will replace existing keys. if force is true, it will replace
    /// disabled keys.
    ///
    /// add_keys returns an error if the key already exists and replace is
    /// false, or if the key is disabled and force is false
    pub fn add_keys(&mut self, name: &str, keys: Vec<PublicKey>, replace: bool, force: bool) -> Result<()> {
        if let Some(keyset) = self.keys.get(name) {
            if keyset.disabled && !force {
                return Err(ErrorKind::KeysDisabled(name.to_string()).into());
            } else if !replace {
                return Err(ErrorKind::KeysExist(name.to_string()).into());
            }
        }
        self.keys.insert(name.to_string(), AuthorizedKeySet {
            filename: name.to_string(),
            disabled: false,
            keys: keys,
        });
        Ok(())
    }

    /// remove_keys removes the keyset with the given name.
    pub fn remove_keys(&mut self, name: &str) {
        self.keys.remove(name);
    }

    /// disable_keys disables keys with the given name. they can't be added
    /// again unless force is set to true when adding the set. disable_keys will
    /// succeed in disabling the key even if the key doesn't currently exist.
    pub fn disable_keys(&mut self, name: &str) {
        if let Some(keyset) = self.keys.get_mut(name) {
            keyset.disabled = true;
            return;
        }
        self.keys.insert(name.to_string(), AuthorizedKeySet {
            filename: name.to_string(),
            disabled: true,
            keys: vec![],
        });
    }
}
<|endoftext|>


<commit_before>extern crate time;
extern crate byteorder;
extern crate uuid;

use std::option::Option;
use std::cmp::Ordering;
use std::io::Write;
use std::io::Read;
use std::net::TcpStream;
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::thread;
use std::thread::{JoinHandle, sleep_ms, Thread};
use std::sync::{Mutex, Condvar, Arc};
use std::sync::mpsc::channel;
use std::sync::mpsc::{Receiver, Sender};

use byteorder::{BigEndian, WriteBytesExt};
use byteorder::ReadBytesExt;
use uuid::Uuid;

pub mod constants;
use constants::*;

#[test]
fn constructor() {
    let c = Connection::new("localhost", 4730);

    assert!(c.host == "localhost");
    assert!(c.port == 4730);
    assert!(c.state == ConnectionState::INIT);

    let now = time::now_utc();
    match c.state_time {
        Some(c_time) => {
            let result = now.cmp(&c_time);  // Could fail if no nanoseconds
            assert_eq!(Ordering::Greater, result);
        },
        None => assert!(false)
    }
}

#[test]
fn change_state() {
    let mut c = Connection::new("localhost", 4730);

    c.change_state(ConnectionState::INIT);
    assert!(c.state == ConnectionState::INIT);
}

#[test]
fn packet_constructor() {
    let data: Box<[u8]> = Box::new([0x00u8]);
    let p = Packet::new(PacketCode::REQ, ECHO_REQ, data);
    assert!(p.code == PacketCode::REQ);
    assert!(p.ptype == ECHO_REQ);
    assert!(p.data.len() == 1);
}

#[test]
#[should_panic]
fn send_packet() {
    let mut conn = Connection::new("localost", 4730);
    let data: Box<[u8]> = Box::new([0x00u8]);
    let p = Packet::new(PacketCode::REQ, ECHO_REQ, data);
    conn.sendPacket(p)
}

#[test]
#[should_panic]
fn echo() {
    let mut test: Vec<u8> = Vec::new();
    test.extend("abc123".to_string().bytes());
    let mut conn = Connection::new("localhost", 4730);
    let result = String::from_utf8(conn.echo(test, 0)).unwrap();
    assert_eq!("abc123", result);
    let mut test2: Vec<u8> = Vec::new();
    let mut output = String::from_utf8(conn.echo(test2, 0)).unwrap();
    println!("UUID = {}", output);
}

#[test]
fn bcs_constructor() {
    let bcs = BaseClientServer::new("clientid".to_string().into_bytes());
}

#[derive(PartialEq)]
enum ConnectionState {
    INIT,
}

struct Connection {
    host: String,
    port: u16,
    conn: Option<TcpStream>,
    connect_time: Option<time::Tm>,
    state: ConnectionState,
    state_time: Option<time::Tm>,
    echo_conditions: HashMap<Vec<u8>, Arc<(Mutex<bool>, Condvar)>>,
}

struct Packet {
    code: PacketCode,
    ptype: u32,
    data: Box<[u8]>,
}

impl Connection {
    fn new(host: &str, port: u16) -> Connection {
        let mut c = Connection {
            host: host.to_string(),
            port: port,
            conn: None,
            connect_time: None,
            state: ConnectionState::INIT,
            state_time: None,
            echo_conditions: HashMap::new(),
        };
        c.change_state(ConnectionState::INIT);
        return c
    }

    fn change_state(&mut self, state: ConnectionState) {
        self.state_time = Some(time::now_utc());
        self.state = state;
    }

    fn connect(&mut self) {
        match self.conn {
            None => {
                self.conn = Some(TcpStream::connect((&*self.host, self.port)).unwrap());
                self.connect_time = Some(time::now_utc());
            },
            Some(_) => { /* log debug "already connected" */ },
        }
    }

    fn disconnect(&mut self) {
        match self.conn {
            Some(_) => self.conn = None,
            None => { /* log debug "already disconnected" */ },
        }
    }

    fn reconnect(&mut self) {
        self.disconnect();
        self.connect()
    }

    fn sendPacket(&mut self, packet: Packet) {
        match self.conn {
            Some(ref mut c) => {
                match packet.code {
                    PacketCode::REQ => { c.write(&REQ).unwrap(); }
                    PacketCode::RES => { c.write(&RES).unwrap(); }
                }
                c.write_u32::<BigEndian>(packet.ptype);
                c.write_u32::<BigEndian>(packet.data.len() as u32).unwrap();
                c.write_all(&packet.data).unwrap();
            },
            None => panic!("Attempted to send packet on disconnected socket")
        }
    }

    fn _code_from_vec(code: Vec<u8>) -> Result<PacketCode, &'static str> {
        if code[..] == REQ {
            return Ok(PacketCode::REQ);
        } else if code[..] == RES {
            return Ok(PacketCode::RES);
        }
        return Err("Invalid packet code");
    }

    fn readPacket(&mut self) -> Packet {
        match self.conn {
            Some(ref mut c) => {
                let mut code: Vec<u8> = Vec::new();
                {
                    let mut codestream = c.take(4);
                    codestream.read_to_end(&mut code).unwrap();
                }
                let code: PacketCode = Connection::_code_from_vec(code).unwrap();
                let ptype   = c.read_u32::<BigEndian>().unwrap();
                let datalen = c.read_u32::<BigEndian>().unwrap();
                let mut datastream = c.take(datalen as u64);
                let mut buf: Vec<u8> = Vec::new();
                datastream.read_to_end(&mut buf).unwrap();
                return Packet {
                    code: code,
                    ptype: ptype,
                    data: buf.into_boxed_slice(),
                }
            },
            None => panic!("Attempted to read packet on disconnected socket")
        }
    }

    fn echo(&mut self, mut data: Vec<u8>, timeout: u32) -> Vec<u8> {
        if data.len() == 0 {
            data.extend(Uuid::new_v4().to_simple_string().bytes());
        }
        let lock_cond = Arc::new((Mutex::new(false), Condvar::new()));
        let &(ref lock, ref cond) = &*lock_cond;
        let mut handled_echo = lock.lock().unwrap();
        self.sendEchoReq(data.clone());
        self.echo_conditions.insert(data.clone(), lock_cond.clone());
        if !*cond.wait_timeout_ms(handled_echo, timeout * 1000).unwrap().0 {
            panic!("Timed out waiting for response from echo.");
        };
        self.echo_conditions.remove(&data);
        data
    }

    fn sendEchoReq(&mut self, mut data: Vec<u8>) {
        let p = Packet::new(PacketCode::REQ, ECHO_REQ, data.into_boxed_slice());
        self.sendPacket(p)
    }

    fn handleEchoRes(&mut self, mut data: Vec<u8>) {
        if let Some(mutexlock) = self.echo_conditions.get_mut(&data) {
            let lock = &mutexlock.0;
            let cond = &mutexlock.1;
            let mut handled_echo = lock.lock().unwrap();
            *handled_echo = true;

            cond.notify_one();
        }
    }
}

struct BaseClientServer<'a> {
    client_id: Vec<u8>,
    running: bool,
    active_connections: HashMap<&'a String, Arc<Mutex<Box<Connection>>>>,
    host_io: (Sender<Arc<Mutex<Box<&'a String>>>>, Receiver<Arc<Mutex<Box<&'a String>>>>),
    conn_io: (Sender<Arc<Mutex<Box<Connection>>>>, Receiver<Arc<Mutex<Box<Connection>>>>),
}

impl <'a>BaseClientServer<'a> {

    fn new(client_id: Vec<u8>) -> BaseClientServer<'a> {
        BaseClientServer {
            client_id: client_id,
            running: true,
            active_connections: HashMap::new(),
            host_io: channel(),
            conn_io: channel(),
        }
    }

    /*
    fn start(mut self) {
        let mut bcs_poll_handle = Arc::new(Mutex::new(&mut self));
        let mut bcs_conn_handle = bcs_poll_handle.clone();
        let poll_join_handle = Some(thread::scoped(move || BaseClientServer::_doPollLoop(bcs_poll_handle)));
        let connect_join_handle = Some(thread::scoped(move || BaseClientServer::_doConnectLoop(bcs_conn_handle)));
        // implicit join
    }
    */

    fn connectionManager(&'a mut self) {
        /* Run as a dedicated thread to manage the list of active/inactive connections */
        let (_, ref mut rx) = self.host_io;
        let ac = &mut self.active_connections;
        while self.running {
            let container = rx.recv().unwrap();
            let host = **container.lock().unwrap();
            let mut insert = false;
            let conn = match ac.get(host) {
                Some(conn) => {
                    let mut real_conn = conn.lock().unwrap();
                    real_conn.reconnect();
                    conn.clone()
                },
                None => {
                    let conn: Arc<Mutex<Box<Connection>>> = Arc::new(Mutex::new(Box::new(Connection::new(&host, 4730)))).clone();
                    {
                        let mut real_conn = conn.lock().unwrap();
                        real_conn.connect();
                    }
                    insert = true;
                    conn
                }
            };
            if insert {
                ac.insert(host, conn.clone());
            }
        }
        for (_, conn) in ac.iter() {
            let mut conn = conn.lock().unwrap();
            conn.disconnect();
        }
    }

    fn addServer(&'a self, host: &'a String) {
        let (ref tx, _) = self.host_io;
        tx.send(Arc::new(Mutex::new(Box::new(host))));
    }

    fn pollingManager(&'a self) {
        let (_, ref rx) = self.conn_io;
        let (ref tx, _) = self.host_io;
        let mut threads: Vec<JoinHandle<()>> = Vec::new();
        while self.running {
            let conn = rx.recv().unwrap();
            // poll this conn
            threads.push(thread::spawn(move || {
                let mut conn = conn.lock().unwrap();
                let p = conn.readPacket();
                // handle packet
            }));
        };
        // scoped threads all joined here
        for mut thread in threads {
            thread.join();
        }
    }
}

impl PartialEq for Packet {
    fn eq(&self, other: &Packet) -> bool {
        self.code == other.code &&
        self.ptype == other.ptype &&
        self.data == other.data
    }
}

impl Packet {
    fn new(code: PacketCode, ptype: u32, data: Box<[u8]>) -> Packet {
        Packet {
            code: code,
            ptype: ptype,
            data: data,
        }
    }
}
<commit_msg>Less obtuse way of coercion<commit_after>extern crate time;
extern crate byteorder;
extern crate uuid;

use std::option::Option;
use std::cmp::Ordering;
use std::io::Write;
use std::io::Read;
use std::net::TcpStream;
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::thread;
use std::thread::{JoinHandle, sleep_ms, Thread};
use std::sync::{Mutex, Condvar, Arc};
use std::sync::mpsc::channel;
use std::sync::mpsc::{Receiver, Sender};

use byteorder::{BigEndian, WriteBytesExt};
use byteorder::ReadBytesExt;
use uuid::Uuid;

pub mod constants;
use constants::*;

#[test]
fn constructor() {
    let c = Connection::new("localhost", 4730);

    assert!(c.host == "localhost");
    assert!(c.port == 4730);
    assert!(c.state == ConnectionState::INIT);

    let now = time::now_utc();
    match c.state_time {
        Some(c_time) => {
            let result = now.cmp(&c_time);  // Could fail if no nanoseconds
            assert_eq!(Ordering::Greater, result);
        },
        None => assert!(false)
    }
}

#[test]
fn change_state() {
    let mut c = Connection::new("localhost", 4730);

    c.change_state(ConnectionState::INIT);
    assert!(c.state == ConnectionState::INIT);
}

#[test]
fn packet_constructor() {
    let data: Box<[u8]> = Box::new([0x00u8]);
    let p = Packet::new(PacketCode::REQ, ECHO_REQ, data);
    assert!(p.code == PacketCode::REQ);
    assert!(p.ptype == ECHO_REQ);
    assert!(p.data.len() == 1);
}

#[test]
#[should_panic]
fn send_packet() {
    let mut conn = Connection::new("localost", 4730);
    let data: Box<[u8]> = Box::new([0x00u8]);
    let p = Packet::new(PacketCode::REQ, ECHO_REQ, data);
    conn.sendPacket(p)
}

#[test]
#[should_panic]
fn echo() {
    let mut test: Vec<u8> = Vec::new();
    test.extend("abc123".to_string().bytes());
    let mut conn = Connection::new("localhost", 4730);
    let result = String::from_utf8(conn.echo(test, 0)).unwrap();
    assert_eq!("abc123", result);
    let mut test2: Vec<u8> = Vec::new();
    let mut output = String::from_utf8(conn.echo(test2, 0)).unwrap();
    println!("UUID = {}", output);
}

#[test]
fn bcs_constructor() {
    let bcs = BaseClientServer::new("clientid".to_string().into_bytes());
}

#[derive(PartialEq)]
enum ConnectionState {
    INIT,
}

struct Connection {
    host: String,
    port: u16,
    conn: Option<TcpStream>,
    connect_time: Option<time::Tm>,
    state: ConnectionState,
    state_time: Option<time::Tm>,
    echo_conditions: HashMap<Vec<u8>, Arc<(Mutex<bool>, Condvar)>>,
}

struct Packet {
    code: PacketCode,
    ptype: u32,
    data: Box<[u8]>,
}

impl Connection {
    fn new(host: &str, port: u16) -> Connection {
        let mut c = Connection {
            host: host.to_string(),
            port: port,
            conn: None,
            connect_time: None,
            state: ConnectionState::INIT,
            state_time: None,
            echo_conditions: HashMap::new(),
        };
        c.change_state(ConnectionState::INIT);
        return c
    }

    fn change_state(&mut self, state: ConnectionState) {
        self.state_time = Some(time::now_utc());
        self.state = state;
    }

    fn connect(&mut self) {
        match self.conn {
            None => {
                self.conn = Some(TcpStream::connect((&self.host as &str, self.port)).unwrap());
                self.connect_time = Some(time::now_utc());
            },
            Some(_) => { /* log debug "already connected" */ },
        }
    }

    fn disconnect(&mut self) {
        match self.conn {
            Some(_) => self.conn = None,
            None => { /* log debug "already disconnected" */ },
        }
    }

    fn reconnect(&mut self) {
        self.disconnect();
        self.connect()
    }

    fn sendPacket(&mut self, packet: Packet) {
        match self.conn {
            Some(ref mut c) => {
                match packet.code {
                    PacketCode::REQ => { c.write(&REQ).unwrap(); }
                    PacketCode::RES => { c.write(&RES).unwrap(); }
                }
                c.write_u32::<BigEndian>(packet.ptype);
                c.write_u32::<BigEndian>(packet.data.len() as u32).unwrap();
                c.write_all(&packet.data).unwrap();
            },
            None => panic!("Attempted to send packet on disconnected socket")
        }
    }

    fn _code_from_vec(code: Vec<u8>) -> Result<PacketCode, &'static str> {
        if code[..] == REQ {
            return Ok(PacketCode::REQ);
        } else if code[..] == RES {
            return Ok(PacketCode::RES);
        }
        return Err("Invalid packet code");
    }

    fn readPacket(&mut self) -> Packet {
        match self.conn {
            Some(ref mut c) => {
                let mut code: Vec<u8> = Vec::new();
                {
                    let mut codestream = c.take(4);
                    codestream.read_to_end(&mut code).unwrap();
                }
                let code: PacketCode = Connection::_code_from_vec(code).unwrap();
                let ptype   = c.read_u32::<BigEndian>().unwrap();
                let datalen = c.read_u32::<BigEndian>().unwrap();
                let mut datastream = c.take(datalen as u64);
                let mut buf: Vec<u8> = Vec::new();
                datastream.read_to_end(&mut buf).unwrap();
                return Packet {
                    code: code,
                    ptype: ptype,
                    data: buf.into_boxed_slice(),
                }
            },
            None => panic!("Attempted to read packet on disconnected socket")
        }
    }

    fn echo(&mut self, mut data: Vec<u8>, timeout: u32) -> Vec<u8> {
        if data.len() == 0 {
            data.extend(Uuid::new_v4().to_simple_string().bytes());
        }
        let lock_cond = Arc::new((Mutex::new(false), Condvar::new()));
        let &(ref lock, ref cond) = &*lock_cond;
        let mut handled_echo = lock.lock().unwrap();
        self.sendEchoReq(data.clone());
        self.echo_conditions.insert(data.clone(), lock_cond.clone());
        if !*cond.wait_timeout_ms(handled_echo, timeout * 1000).unwrap().0 {
            panic!("Timed out waiting for response from echo.");
        };
        self.echo_conditions.remove(&data);
        data
    }

    fn sendEchoReq(&mut self, mut data: Vec<u8>) {
        let p = Packet::new(PacketCode::REQ, ECHO_REQ, data.into_boxed_slice());
        self.sendPacket(p)
    }

    fn handleEchoRes(&mut self, mut data: Vec<u8>) {
        if let Some(mutexlock) = self.echo_conditions.get_mut(&data) {
            let lock = &mutexlock.0;
            let cond = &mutexlock.1;
            let mut handled_echo = lock.lock().unwrap();
            *handled_echo = true;

            cond.notify_one();
        }
    }
}

struct BaseClientServer<'a> {
    client_id: Vec<u8>,
    running: bool,
    active_connections: HashMap<&'a String, Arc<Mutex<Box<Connection>>>>,
    host_io: (Sender<Arc<Mutex<Box<&'a String>>>>, Receiver<Arc<Mutex<Box<&'a String>>>>),
    conn_io: (Sender<Arc<Mutex<Box<Connection>>>>, Receiver<Arc<Mutex<Box<Connection>>>>),
}

impl <'a>BaseClientServer<'a> {

    fn new(client_id: Vec<u8>) -> BaseClientServer<'a> {
        BaseClientServer {
            client_id: client_id,
            running: true,
            active_connections: HashMap::new(),
            host_io: channel(),
            conn_io: channel(),
        }
    }

    /*
    fn start(mut self) {
        let mut bcs_poll_handle = Arc::new(Mutex::new(&mut self));
        let mut bcs_conn_handle = bcs_poll_handle.clone();
        let poll_join_handle = Some(thread::scoped(move || BaseClientServer::_doPollLoop(bcs_poll_handle)));
        let connect_join_handle = Some(thread::scoped(move || BaseClientServer::_doConnectLoop(bcs_conn_handle)));
        // implicit join
    }
    */

    fn connectionManager(&'a mut self) {
        /* Run as a dedicated thread to manage the list of active/inactive connections */
        let (_, ref mut rx) = self.host_io;
        let ac = &mut self.active_connections;
        while self.running {
            let container = rx.recv().unwrap();
            let host = **container.lock().unwrap();
            let mut insert = false;
            let conn = match ac.get(host) {
                Some(conn) => {
                    let mut real_conn = conn.lock().unwrap();
                    real_conn.reconnect();
                    conn.clone()
                },
                None => {
                    let conn: Arc<Mutex<Box<Connection>>> = Arc::new(Mutex::new(Box::new(Connection::new(&host, 4730)))).clone();
                    {
                        let mut real_conn = conn.lock().unwrap();
                        real_conn.connect();
                    }
                    insert = true;
                    conn
                }
            };
            if insert {
                ac.insert(host, conn.clone());
            }
        }
        for (_, conn) in ac.iter() {
            let mut conn = conn.lock().unwrap();
            conn.disconnect();
        }
    }

    fn addServer(&'a self, host: &'a String) {
        let (ref tx, _) = self.host_io;
        tx.send(Arc::new(Mutex::new(Box::new(host))));
    }

    fn pollingManager(&'a self) {
        let (_, ref rx) = self.conn_io;
        let (ref tx, _) = self.host_io;
        let mut threads: Vec<JoinHandle<()>> = Vec::new();
        while self.running {
            let conn = rx.recv().unwrap();
            // poll this conn
            threads.push(thread::spawn(move || {
                let mut conn = conn.lock().unwrap();
                let p = conn.readPacket();
                // handle packet
            }));
        };
        // scoped threads all joined here
        for mut thread in threads {
            thread.join();
        }
    }
}

impl PartialEq for Packet {
    fn eq(&self, other: &Packet) -> bool {
        self.code == other.code &&
        self.ptype == other.ptype &&
        self.data == other.data
    }
}

impl Packet {
    fn new(code: PacketCode, ptype: u32, data: Box<[u8]>) -> Packet {
        Packet {
            code: code,
            ptype: ptype,
            data: data,
        }
    }
}
<|endoftext|>


<commit_before>//! A library providing a generic connection pool.
#![warn(missing_docs)]
#![doc(html_root_url="https://sfackler.github.io/r2d2/doc/master")]

#[macro_use]
extern crate log;
extern crate time;
extern crate debug_builders;

use debug_builders::DebugStruct;
use std::collections::VecDeque;
use std::error::Error;
use std::fmt;
use std::ops::{Deref, DerefMut};
use std::sync::{Arc, Mutex, MutexGuard, Condvar};
use time::{Duration, SteadyTime};

#[doc(inline)]
pub use config::Config;

use task::ScheduledThreadPool;

pub mod config;
mod task;
mod thunk;

/// A trait which provides connection-specific functionality.
pub trait ConnectionManager: Send + Sync + 'static {
    /// The connection type this manager deals with.
    type Connection: Send + 'static;

    /// The error type returned by `Connection`s.
    type Error: 'static;

    /// Attempts to create a new connection.
    fn connect(&self) -> Result<Self::Connection, Self::Error>;

    /// Determines if the connection is still connected to the database.
    ///
    /// A standard implementation would check if a simple query like `SELECT 1`
    /// succeeds.
    fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error>;

    /// *Quickly* determines if the connection is no longer usable.
    ///
    /// This will be called synchronously every time a connection is returned
    /// to the pool, so it should *not* block. If it returns `true`, the
    /// connection will be discarded.
    ///
    /// For example, an implementation might check if the underlying TCP socket
    /// has disconnected. Implementations that do not support this kind of
    /// fast health check may simply return `false`.
    fn has_broken(&self, conn: &mut Self::Connection) -> bool;
}

/// A trait which handles errors reported by the `ConnectionManager`.
pub trait ErrorHandler<E>: Send+Sync+'static {
    /// Handles an error.
    fn handle_error(&self, error: E);
}

/// An `ErrorHandler` which does nothing.
#[derive(Copy, Clone, Debug)]
pub struct NoopErrorHandler;

impl<E> ErrorHandler<E> for NoopErrorHandler {
    fn handle_error(&self, _: E) {}
}

/// An `ErrorHandler` which logs at the error level.
#[derive(Copy, Clone, Debug)]
pub struct LoggingErrorHandler;

impl<E> ErrorHandler<E> for LoggingErrorHandler where E: fmt::Debug {
    fn handle_error(&self, error: E) {
        error!("{:?}", error);
    }
}

struct PoolInternals<C> {
    conns: VecDeque<C>,
    num_conns: u32,
}

struct SharedPool<M> where M: ConnectionManager {
    config: Config,
    manager: M,
    error_handler: Box<ErrorHandler<M::Error>>,
    internals: Mutex<PoolInternals<M::Connection>>,
    cond: Condvar,
    thread_pool: ScheduledThreadPool,
}

fn add_connection<M>(delay: Duration, shared: &Arc<SharedPool<M>>) where M: ConnectionManager {
    let new_shared = shared.clone();
    shared.thread_pool.run_after(delay, move || {
        let shared = new_shared;
        match shared.manager.connect() {
            Ok(conn) => {
                let mut internals = shared.internals.lock().unwrap();
                internals.conns.push_back(conn);
                internals.num_conns += 1;
                shared.cond.notify_one();
            }
            Err(err) => {
                shared.error_handler.handle_error(err);
                add_connection(Duration::seconds(1), &shared);
            },
        }
    });
}

/// A generic connection pool.
pub struct Pool<M> where M: ConnectionManager {
    shared: Arc<SharedPool<M>>,
}

impl<M> Drop for Pool<M> where M: ConnectionManager {
    fn drop(&mut self) {
        self.shared.thread_pool.clear();
    }
}

impl<M> fmt::Debug for Pool<M> where M: ConnectionManager + fmt::Debug {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        DebugStruct::new(fmt, "Pool")
            .field("idle_connections", &self.shared.internals.lock().unwrap().conns.len())
            .field("config", &self.shared.config)
            .field("manager", &self.shared.manager)
            .finish()
    }
}

/// An error returned by `Pool::new` if it fails to initialize connections.
#[derive(Debug)]
pub struct InitializationError;

impl fmt::Display for InitializationError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str(self.description())
    }
}

impl Error for InitializationError {
    fn description(&self) -> &str {
        "Unable to initialize connections"
    }
}

/// An error returned by `Pool::get` if it times out without retrieving a connection.
#[derive(Debug)]
pub struct GetTimeout;

impl fmt::Display for GetTimeout {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str(self.description())
    }
}

impl Error for GetTimeout {
    fn description(&self) -> &str {
        "Timed out while waiting for a connection"
    }
}

impl<M> Pool<M> where M: ConnectionManager {
    /// Creates a new connection pool.
    ///
    /// Returns an `Err` value if `initialization_fail_fast` is set to true in
    /// the configuration and the pool is unable to open all of its
    /// connections.
    pub fn new(config: Config,
               manager: M,
               error_handler: Box<ErrorHandler<M::Error>>)
               -> Result<Pool<M>, InitializationError> {
        let internals = PoolInternals {
            conns: VecDeque::new(),
            num_conns: 0,
        };

        let shared = Arc::new(SharedPool {
            config: config,
            manager: manager,
            error_handler: error_handler,
            internals: Mutex::new(internals),
            cond: Condvar::new(),
            thread_pool: ScheduledThreadPool::new(config.helper_threads() as usize),
        });

        for _ in 0..config.pool_size() {
            add_connection(Duration::zero(), &shared);
        }

        if shared.config.initialization_fail_fast() {
            let end = SteadyTime::now() + shared.config.connection_timeout();
            let mut internals = shared.internals.lock().unwrap();

            while internals.num_conns != shared.config.pool_size() {
                let wait = end - SteadyTime::now();
                if wait <= Duration::zero() {
                    return Err(InitializationError);
                }
                internals = shared.cond.wait_timeout_ms(internals,
                                                        wait.num_milliseconds() as u32)
                    .unwrap().0;
            }
        }

        Ok(Pool {
            shared: shared,
        })
    }

    fn get_inner(&self) -> Result<M::Connection, GetTimeout> {
        let end = SteadyTime::now() + self.shared.config.connection_timeout();
        let mut internals = self.shared.internals.lock().unwrap();

        loop {
            match internals.conns.pop_front() {
                Some(mut conn) => {
                    drop(internals);

                    if self.shared.config.test_on_check_out() {
                        if let Err(e) = self.shared.manager.is_valid(&mut conn) {
                            self.shared.error_handler.handle_error(e);
                            internals = self.shared.internals.lock().unwrap();
                            self.handle_broken(&mut internals);
                            continue
                        }
                    }

                    return Ok(conn);
                }
                None => {
                    let now = SteadyTime::now();
                    let mut timeout = (end - now).num_milliseconds();
                    if timeout < 0 {
                        timeout = 0
                    };
                    let (new_internals, no_timeout) =
                        self.shared.cond.wait_timeout_ms(internals, timeout as u32).unwrap();
                    internals = new_internals;

                    if !no_timeout {
                        return Err(GetTimeout);
                    }
                }
            }
        }
    }

    /// Retrieves a connection from the pool.
    ///
    /// Waits for at most `Config::connection_timeout` before returning an
    /// error.
    pub fn get<'a>(&'a self) -> Result<PooledConnection<'a, M>, GetTimeout> {
        Ok(PooledConnection {
            pool: MaybeArcPool::Ref(self),
            conn: Some(try!(self.get_inner())),
        })
    }

    /// Retrieves a connection from a reference counted pool.
    ///
    /// Unlike `Pool::get`, the resulting `PooledConnection` is not bound by
    /// any lifetime, as it will store the `Arc` internally.
    ///
    /// # Note
    ///
    /// This is a static function rather than a method for technical reasons.
    pub fn get_arc(pool: Arc<Self>) -> Result<PooledConnection<'static, M>, GetTimeout> {
        Ok(PooledConnection {
            conn: Some(try!(pool.get_inner())),
            pool: MaybeArcPool::Arc(pool),
        })
    }

    fn handle_broken(&self, internals: &mut MutexGuard<PoolInternals<M::Connection>>) {
        internals.num_conns -= 1;
        add_connection(Duration::zero(), &self.shared);
    }

    fn put_back(&self, mut conn: M::Connection) {
        // This is specified to be fast, but call it before locking anyways
        let broken = self.shared.manager.has_broken(&mut conn);

        let mut internals = self.shared.internals.lock().unwrap();
        if broken {
            self.handle_broken(&mut internals);
        } else {
            internals.conns.push_back(conn);
            self.shared.cond.notify_one();
        }
    }
}

enum MaybeArcPool<'a, M> where M: ConnectionManager {
    Ref(&'a Pool<M>),
    Arc(Arc<Pool<M>>),
}

impl<'a, M> Deref for MaybeArcPool<'a, M> where M: ConnectionManager {
    type Target = Pool<M>;

    fn deref(&self) -> &Pool<M> {
        match *self {
            MaybeArcPool::Ref(p) => p,
            MaybeArcPool::Arc(ref p) => &**p,
        }
    }
}

/// A smart pointer wrapping a connection.
pub struct PooledConnection<'a, M> where M: ConnectionManager {
    pool: MaybeArcPool<'a, M>,
    conn: Option<M::Connection>,
}

impl<'a, M> fmt::Debug for PooledConnection<'a, M>
        where M: ConnectionManager + fmt::Debug, M::Connection: fmt::Debug {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        DebugStruct::new(fmt, "PooledConnection")
            .field("pool", &*self.pool)
            .field("connection", self.conn.as_ref().unwrap())
            .finish()
    }
}

impl<'a, M> Drop for PooledConnection<'a, M> where M: ConnectionManager {
    fn drop(&mut self) {
        self.pool.put_back(self.conn.take().unwrap());
    }
}

impl<'a, M> Deref for PooledConnection<'a, M> where M: ConnectionManager {
    type Target = M::Connection;

    fn deref(&self) -> &M::Connection {
        self.conn.as_ref().unwrap()
    }
}

impl<'a, M> DerefMut for PooledConnection<'a, M> where M: ConnectionManager {
    fn deref_mut(&mut self) -> &mut M::Connection {
        self.conn.as_mut().unwrap()
    }
}
<commit_msg>Remove note<commit_after>//! A library providing a generic connection pool.
#![warn(missing_docs)]
#![doc(html_root_url="https://sfackler.github.io/r2d2/doc/master")]

#[macro_use]
extern crate log;
extern crate time;
extern crate debug_builders;

use debug_builders::DebugStruct;
use std::collections::VecDeque;
use std::error::Error;
use std::fmt;
use std::ops::{Deref, DerefMut};
use std::sync::{Arc, Mutex, MutexGuard, Condvar};
use time::{Duration, SteadyTime};

#[doc(inline)]
pub use config::Config;

use task::ScheduledThreadPool;

pub mod config;
mod task;
mod thunk;

/// A trait which provides connection-specific functionality.
pub trait ConnectionManager: Send + Sync + 'static {
    /// The connection type this manager deals with.
    type Connection: Send + 'static;

    /// The error type returned by `Connection`s.
    type Error: 'static;

    /// Attempts to create a new connection.
    fn connect(&self) -> Result<Self::Connection, Self::Error>;

    /// Determines if the connection is still connected to the database.
    ///
    /// A standard implementation would check if a simple query like `SELECT 1`
    /// succeeds.
    fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error>;

    /// *Quickly* determines if the connection is no longer usable.
    ///
    /// This will be called synchronously every time a connection is returned
    /// to the pool, so it should *not* block. If it returns `true`, the
    /// connection will be discarded.
    ///
    /// For example, an implementation might check if the underlying TCP socket
    /// has disconnected. Implementations that do not support this kind of
    /// fast health check may simply return `false`.
    fn has_broken(&self, conn: &mut Self::Connection) -> bool;
}

/// A trait which handles errors reported by the `ConnectionManager`.
pub trait ErrorHandler<E>: Send+Sync+'static {
    /// Handles an error.
    fn handle_error(&self, error: E);
}

/// An `ErrorHandler` which does nothing.
#[derive(Copy, Clone, Debug)]
pub struct NoopErrorHandler;

impl<E> ErrorHandler<E> for NoopErrorHandler {
    fn handle_error(&self, _: E) {}
}

/// An `ErrorHandler` which logs at the error level.
#[derive(Copy, Clone, Debug)]
pub struct LoggingErrorHandler;

impl<E> ErrorHandler<E> for LoggingErrorHandler where E: fmt::Debug {
    fn handle_error(&self, error: E) {
        error!("{:?}", error);
    }
}

struct PoolInternals<C> {
    conns: VecDeque<C>,
    num_conns: u32,
}

struct SharedPool<M> where M: ConnectionManager {
    config: Config,
    manager: M,
    error_handler: Box<ErrorHandler<M::Error>>,
    internals: Mutex<PoolInternals<M::Connection>>,
    cond: Condvar,
    thread_pool: ScheduledThreadPool,
}

fn add_connection<M>(delay: Duration, shared: &Arc<SharedPool<M>>) where M: ConnectionManager {
    let new_shared = shared.clone();
    shared.thread_pool.run_after(delay, move || {
        let shared = new_shared;
        match shared.manager.connect() {
            Ok(conn) => {
                let mut internals = shared.internals.lock().unwrap();
                internals.conns.push_back(conn);
                internals.num_conns += 1;
                shared.cond.notify_one();
            }
            Err(err) => {
                shared.error_handler.handle_error(err);
                add_connection(Duration::seconds(1), &shared);
            },
        }
    });
}

/// A generic connection pool.
pub struct Pool<M> where M: ConnectionManager {
    shared: Arc<SharedPool<M>>,
}

impl<M> Drop for Pool<M> where M: ConnectionManager {
    fn drop(&mut self) {
        self.shared.thread_pool.clear();
    }
}

impl<M> fmt::Debug for Pool<M> where M: ConnectionManager + fmt::Debug {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        DebugStruct::new(fmt, "Pool")
            .field("idle_connections", &self.shared.internals.lock().unwrap().conns.len())
            .field("config", &self.shared.config)
            .field("manager", &self.shared.manager)
            .finish()
    }
}

/// An error returned by `Pool::new` if it fails to initialize connections.
#[derive(Debug)]
pub struct InitializationError;

impl fmt::Display for InitializationError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str(self.description())
    }
}

impl Error for InitializationError {
    fn description(&self) -> &str {
        "Unable to initialize connections"
    }
}

/// An error returned by `Pool::get` if it times out without retrieving a connection.
#[derive(Debug)]
pub struct GetTimeout;

impl fmt::Display for GetTimeout {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str(self.description())
    }
}

impl Error for GetTimeout {
    fn description(&self) -> &str {
        "Timed out while waiting for a connection"
    }
}

impl<M> Pool<M> where M: ConnectionManager {
    /// Creates a new connection pool.
    ///
    /// Returns an `Err` value if `initialization_fail_fast` is set to true in
    /// the configuration and the pool is unable to open all of its
    /// connections.
    pub fn new(config: Config,
               manager: M,
               error_handler: Box<ErrorHandler<M::Error>>)
               -> Result<Pool<M>, InitializationError> {
        let internals = PoolInternals {
            conns: VecDeque::new(),
            num_conns: 0,
        };

        let shared = Arc::new(SharedPool {
            config: config,
            manager: manager,
            error_handler: error_handler,
            internals: Mutex::new(internals),
            cond: Condvar::new(),
            thread_pool: ScheduledThreadPool::new(config.helper_threads() as usize),
        });

        for _ in 0..config.pool_size() {
            add_connection(Duration::zero(), &shared);
        }

        if shared.config.initialization_fail_fast() {
            let end = SteadyTime::now() + shared.config.connection_timeout();
            let mut internals = shared.internals.lock().unwrap();

            while internals.num_conns != shared.config.pool_size() {
                let wait = end - SteadyTime::now();
                if wait <= Duration::zero() {
                    return Err(InitializationError);
                }
                internals = shared.cond.wait_timeout_ms(internals,
                                                        wait.num_milliseconds() as u32)
                    .unwrap().0;
            }
        }

        Ok(Pool {
            shared: shared,
        })
    }

    fn get_inner(&self) -> Result<M::Connection, GetTimeout> {
        let end = SteadyTime::now() + self.shared.config.connection_timeout();
        let mut internals = self.shared.internals.lock().unwrap();

        loop {
            match internals.conns.pop_front() {
                Some(mut conn) => {
                    drop(internals);

                    if self.shared.config.test_on_check_out() {
                        if let Err(e) = self.shared.manager.is_valid(&mut conn) {
                            self.shared.error_handler.handle_error(e);
                            internals = self.shared.internals.lock().unwrap();
                            self.handle_broken(&mut internals);
                            continue
                        }
                    }

                    return Ok(conn);
                }
                None => {
                    let now = SteadyTime::now();
                    let mut timeout = (end - now).num_milliseconds();
                    if timeout < 0 {
                        timeout = 0
                    };
                    let (new_internals, no_timeout) =
                        self.shared.cond.wait_timeout_ms(internals, timeout as u32).unwrap();
                    internals = new_internals;

                    if !no_timeout {
                        return Err(GetTimeout);
                    }
                }
            }
        }
    }

    /// Retrieves a connection from the pool.
    ///
    /// Waits for at most `Config::connection_timeout` before returning an
    /// error.
    pub fn get<'a>(&'a self) -> Result<PooledConnection<'a, M>, GetTimeout> {
        Ok(PooledConnection {
            pool: MaybeArcPool::Ref(self),
            conn: Some(try!(self.get_inner())),
        })
    }

    /// Retrieves a connection from a reference counted pool.
    ///
    /// Unlike `Pool::get`, the resulting `PooledConnection` is not bound by
    /// any lifetime, as it will store the `Arc` internally.
    pub fn get_arc(pool: Arc<Self>) -> Result<PooledConnection<'static, M>, GetTimeout> {
        Ok(PooledConnection {
            conn: Some(try!(pool.get_inner())),
            pool: MaybeArcPool::Arc(pool),
        })
    }

    fn handle_broken(&self, internals: &mut MutexGuard<PoolInternals<M::Connection>>) {
        internals.num_conns -= 1;
        add_connection(Duration::zero(), &self.shared);
    }

    fn put_back(&self, mut conn: M::Connection) {
        // This is specified to be fast, but call it before locking anyways
        let broken = self.shared.manager.has_broken(&mut conn);

        let mut internals = self.shared.internals.lock().unwrap();
        if broken {
            self.handle_broken(&mut internals);
        } else {
            internals.conns.push_back(conn);
            self.shared.cond.notify_one();
        }
    }
}

enum MaybeArcPool<'a, M> where M: ConnectionManager {
    Ref(&'a Pool<M>),
    Arc(Arc<Pool<M>>),
}

impl<'a, M> Deref for MaybeArcPool<'a, M> where M: ConnectionManager {
    type Target = Pool<M>;

    fn deref(&self) -> &Pool<M> {
        match *self {
            MaybeArcPool::Ref(p) => p,
            MaybeArcPool::Arc(ref p) => &**p,
        }
    }
}

/// A smart pointer wrapping a connection.
pub struct PooledConnection<'a, M> where M: ConnectionManager {
    pool: MaybeArcPool<'a, M>,
    conn: Option<M::Connection>,
}

impl<'a, M> fmt::Debug for PooledConnection<'a, M>
        where M: ConnectionManager + fmt::Debug, M::Connection: fmt::Debug {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        DebugStruct::new(fmt, "PooledConnection")
            .field("pool", &*self.pool)
            .field("connection", self.conn.as_ref().unwrap())
            .finish()
    }
}

impl<'a, M> Drop for PooledConnection<'a, M> where M: ConnectionManager {
    fn drop(&mut self) {
        self.pool.put_back(self.conn.take().unwrap());
    }
}

impl<'a, M> Deref for PooledConnection<'a, M> where M: ConnectionManager {
    type Target = M::Connection;

    fn deref(&self) -> &M::Connection {
        self.conn.as_ref().unwrap()
    }
}

impl<'a, M> DerefMut for PooledConnection<'a, M> where M: ConnectionManager {
    fn deref_mut(&mut self) -> &mut M::Connection {
        self.conn.as_mut().unwrap()
    }
}
<|endoftext|>