<commit_before>//! Afterparty is a github webhook handler library for building custom integrations

#[macro_use]
extern crate log;
#[macro_use]
extern crate hyper;
extern crate crypto;
extern crate serde;
extern crate serde_json;

mod hook;
mod events;

pub use events::Event;
pub use hook::{AuthenticateHook, Hook};
use hyper::server::{Handler, Request, Response};
use std::collections::HashMap;
use std::io::Read;

/// signature for request
/// see [this document](https://developer.github.com/webhooks/securing/) for more information
header! {(XHubSignature, "X-Hub-Signature") => [String]}

/// name of Github event
/// see [this document](https://developer.github.com/webhooks/#events) for available types
header! {(XGithubEvent, "X-Github-Event") => [String]}

/// unique id for each delivery
header! {(XGithubDelivery, "X-Github-Delivery") => [String]}

/// A delivery encodes all information about web hook request
#[derive(Debug)]
pub struct Delivery<'a> {
    pub id: &'a str,
    pub event: &'a str,
    pub payload: Event,
    pub unparsed_payload: &'a str,
    pub signature: Option<&'a str>,
}

impl<'a> Delivery<'a> {
    pub fn new(id: &'a str,
               event: &'a str,
               payload: &'a str,
               signature: Option<&'a str>)
               -> Option<Delivery<'a>> {
        match serde_json::from_str::<Event>(&payload) {
            Ok(parsed) => {
                Some(Delivery {
                    id: id,
                    event: event,
                    payload: parsed,
                    unparsed_payload: payload,
                    signature: signature,
                })
            }
            _ => None,
        }
    }
}

/// A hub is a registry of hooks
#[derive(Default)]
pub struct Hub {
    hooks: HashMap<String, Vec<Box<Hook>>>,
}

impl Hub {
    /// construct a new hub instance
    pub fn new() -> Hub {
        Hub { ..Default::default() }
    }

    /// adds a new web hook which will only be applied
    /// when a delivery is revcieved with a valid
    /// request signature
    pub fn handle_authenticated<H, S>(&mut self, event: &str, secret: S, hook: H)
        where H: Hook + 'static,
              S: Into<String>
    {
        self.handle(event, AuthenticateHook::new(secret, hook))
    }

    /// add a need hook to list of hooks
    /// interested in a given event
    pub fn handle<H>(&mut self, event: &str, hook: H)
        where H: Hook + 'static
    {
        self.hooks
            .entry(event.to_owned())
            .or_insert(vec![])
            .push(Box::new(hook));
    }

    /// get all interested hooks for a given event
    pub fn hooks(&self, event: &str) -> Option<Vec<&Box<Hook>>> {
        let explicit = self.hooks.get(event);
        let implicit = self.hooks.get("*");
        let combined = match (explicit, implicit) {
            (Some(ex), Some(im)) => {
                Some(ex.iter().chain(im.iter()).into_iter().collect::<Vec<_>>())
            }
            (Some(ex), _) => Some(ex.into_iter().collect::<Vec<_>>()),
            (_, Some(im)) => Some(im.into_iter().collect::<Vec<_>>()),
            _ => None,
        };
        combined

    }

    /// serially handles delivery by applying each
    /// interested hook
    pub fn deliver(&self, delivery: &Delivery) {
        if let Some(hooks) = self.hooks(delivery.event) {
            for hook in hooks {
                hook.handle(&delivery)
            }
        }
    }
}

impl Handler for Hub {
    fn handle(&self, mut req: Request, res: Response) {
        let mut payload = String::new();
        let headers = req.headers.clone();
        if let (Some(&XGithubEvent(ref event)),
                Some(&XGithubDelivery(ref delivery))) = (headers.get::<XGithubEvent>(),
                                                         headers.get::<XGithubDelivery>()) {
            if let Ok(_) = req.read_to_string(&mut payload) {
                let signature = headers.get::<XHubSignature>();
                info!("recv '{}' event with signature '{:?}'", event, signature);
                match Delivery::new(delivery,
                                    event,
                                    payload.as_ref(),
                                    signature.map(|s| s.as_ref())) {
                    Some(delivery) => self.deliver(&delivery),
                    _ => {
                        error!("failed to parse event {:?} for delivery {:?}",
                               event,
                               delivery)
                    }
                }
            }
        }
        let _ = res.send(b"ok");
        ()
    }
}

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

    #[test]
    fn hub_hooks() {
        let mut hub = Hub::new();
        hub.handle("push", |_: &Delivery| {});
        hub.handle("*", |_: &Delivery| {});
        assert_eq!(Some(2),
                   hub.hooks("push").map(|hooks| hooks.into_iter().count()))
    }
}
<commit_msg>spacing<commit_after>//! Afterparty is a github webhook handler library for building custom integrations

#[macro_use]
extern crate log;
#[macro_use]
extern crate hyper;
extern crate crypto;
extern crate serde;
extern crate serde_json;

mod hook;
mod events;

pub use events::Event;
pub use hook::{AuthenticateHook, Hook};
use hyper::server::{Handler, Request, Response};
use std::collections::HashMap;
use std::io::Read;

/// signature for request
/// see [this document](https://developer.github.com/webhooks/securing/) for more information
header! {(XHubSignature, "X-Hub-Signature") => [String]}

/// name of Github event
/// see [this document](https://developer.github.com/webhooks/#events) for available types
header! {(XGithubEvent, "X-Github-Event") => [String]}

/// unique id for each delivery
header! {(XGithubDelivery, "X-Github-Delivery") => [String]}

/// A delivery encodes all information about web hook request
#[derive(Debug)]
pub struct Delivery<'a> {
    pub id: &'a str,
    pub event: &'a str,
    pub payload: Event,
    pub unparsed_payload: &'a str,
    pub signature: Option<&'a str>,
}

impl<'a> Delivery<'a> {
    pub fn new(id: &'a str,
               event: &'a str,
               payload: &'a str,
               signature: Option<&'a str>)
               -> Option<Delivery<'a>> {
        match serde_json::from_str::<Event>(&payload) {
            Ok(parsed) => {
                Some(Delivery {
                    id: id,
                    event: event,
                    payload: parsed,
                    unparsed_payload: payload,
                    signature: signature,
                })
            }
            _ => None,
        }
    }
}

/// A hub is a registry of hooks
#[derive(Default)]
pub struct Hub {
    hooks: HashMap<String, Vec<Box<Hook>>>,
}

impl Hub {
    /// construct a new hub instance
    pub fn new() -> Hub {
        Hub { ..Default::default() }
    }

    /// adds a new web hook which will only be applied
    /// when a delivery is revcieved with a valid
    /// request signature
    pub fn handle_authenticated<H, S>(&mut self, event: &str, secret: S, hook: H)
        where H: Hook + 'static,
              S: Into<String>
    {
        self.handle(event, AuthenticateHook::new(secret, hook))
    }

    /// add a need hook to list of hooks
    /// interested in a given event
    pub fn handle<H>(&mut self, event: &str, hook: H)
        where H: Hook + 'static
    {
        self.hooks
            .entry(event.to_owned())
            .or_insert(vec![])
            .push(Box::new(hook));
    }

    /// get all interested hooks for a given event
    pub fn hooks(&self, event: &str) -> Option<Vec<&Box<Hook>>> {
        let explicit = self.hooks.get(event);
        let implicit = self.hooks.get("*");
        let combined = match (explicit, implicit) {
            (Some(ex), Some(im)) => {
                Some(ex.iter().chain(im.iter()).into_iter().collect::<Vec<_>>())
            }
            (Some(ex), _) => Some(ex.into_iter().collect::<Vec<_>>()),
            (_, Some(im)) => Some(im.into_iter().collect::<Vec<_>>()),
            _ => None,
        };
        combined
    }

    /// serially handles delivery by applying each
    /// interested hook
    pub fn deliver(&self, delivery: &Delivery) {
        if let Some(hooks) = self.hooks(delivery.event) {
            for hook in hooks {
                hook.handle(&delivery)
            }
        }
    }
}

impl Handler for Hub {
    fn handle(&self, mut req: Request, res: Response) {
        let mut payload = String::new();
        let headers = req.headers.clone();
        if let (Some(&XGithubEvent(ref event)),
                Some(&XGithubDelivery(ref delivery))) = (headers.get::<XGithubEvent>(),
                                                         headers.get::<XGithubDelivery>()) {
            if let Ok(_) = req.read_to_string(&mut payload) {
                let signature = headers.get::<XHubSignature>();
                info!("recv '{}' event with signature '{:?}'", event, signature);
                match Delivery::new(delivery,
                                    event,
                                    payload.as_ref(),
                                    signature.map(|s| s.as_ref())) {
                    Some(delivery) => self.deliver(&delivery),
                    _ => {
                        error!("failed to parse event {:?} for delivery {:?}",
                               event,
                               delivery)
                    }
                }
            }
        }
        let _ = res.send(b"ok");
        ()
    }
}

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

    #[test]
    fn hub_hooks() {
        let mut hub = Hub::new();
        hub.handle("push", |_: &Delivery| {});
        hub.handle("*", |_: &Delivery| {});
        assert_eq!(Some(2),
                   hub.hooks("push").map(|hooks| hooks.into_iter().count()))
    }
}
<|endoftext|>


<commit_before>#![warn(missing_docs)]

//! input-map
//! =========
//!
//! A library for binding input keys to actions, and manipulating mouse input

extern crate input;
extern crate window;
extern crate rustc_serialize;
extern crate viewport;

mod builder;

use input::{Input, Button, Motion};
use window::Size;
use std::collections::HashMap;
use std::default::Default;
use std::hash::Hash;
use viewport::Viewport;

pub use builder::InputTranslatorBuilder;

/// Represents a logical action to be bound to a particular button press, e.g.
/// jump, attack, or move forward. Needs to be hashable, as it is used as a
/// lookup key when rebinding an action to a different button.
pub trait Action: Copy + PartialEq + Eq + Hash { }

/// A translated action.
#[derive(Debug, Copy, Clone, PartialEq)]
pub enum Translated<A: Action> {
    /// A keypress event which was bound to an action
    Press(A),

    /// A key release event which was bound to an action
    Release(A),

    /// A translated mouse motion. The logical origin of a translated MouseCursor event
    /// is in the top left corner of the window, and the logical scroll is non-natural.
    /// Relative events are unchanged for now.
    Move(Motion)
}

/// A three-element tuple of Option<Button>. For simplicity, a maximum number of 3
/// buttons can be bound to each action, and this is exposed through the `InputRebind`
/// struct.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct ButtonTuple(pub Option<Button>, pub Option<Button>, pub Option<Button>);

impl ButtonTuple {
    /// Creates a new tuple with no buttons in it (equivalent to `Default::default()`).
    pub fn new() -> Self { Default::default() }

    /// Check if the button is in the tuple.
    pub fn contains(&self, button: Button) -> bool {
        let sbtn = Some(button);
        self.0 == sbtn || self.1 == sbtn || self.2 == sbtn
    }

    /// Insert a button into the tuple if there is room, searching from left to right.
    /// If the button is inserted, returns true. Otherwise, if the button is not inserted,
    /// this function returns false.
    pub fn insert_inplace(&mut self, button: Button) -> bool {
        match self {
            &mut ButtonTuple(a, b, c) if a.is_none() => {*self = ButtonTuple(Some(button), b, c); true},
            &mut ButtonTuple(a, b, c) if b.is_none() => {*self = ButtonTuple(a, Some(button), c); true},
            &mut ButtonTuple(a, b, c) if c.is_none() => {*self = ButtonTuple(a, b, Some(button)); true}
            _ => false
        }
    }
}

impl Default for ButtonTuple {
    /// Creates a new tuple with no buttons in it.
    fn default() -> Self { ButtonTuple(None, None, None) }
}

/// An object which translates piston::input::Input events into input_map::Translated<A> events
#[derive(Clone)]
pub struct InputTranslator<A: Action> {
    keymap: HashMap<Button, A>,
    mouse_translator: MouseTranslator
}

impl<A: Action> InputTranslator<A> {
    /// Creates an empty InputTranslator.
    pub fn new(size: Size) -> Self {
        InputTranslator {
            keymap: HashMap::new(),
            mouse_translator: MouseTranslator::new(size)
        }
    }

    /// Translate an Input into a Translated<A> event. Returns `None` if there is no
    /// action associated with the `Input` variant.
    pub fn translate(&self, input: &Input) -> Option<Translated<A>> {
        macro_rules! translate_button(($but_state:ident, $but_var:ident) => (
            match self.keymap.get(&$but_var).map(|x| *x) {
                Some(act) => Some(Translated::$but_state(act)),
                None => None
            });
        );

        match input {
            &Input::Press(button) => translate_button!(Press, button),
            &Input::Release(button) => translate_button!(Release, button),
            &Input::Move(motion) =>
                Some(Translated::Move(self.mouse_translator.translate(motion))),
            _ => None
        }
    }

    /// Re-set the mouse bounds size used for calculating mouse events
    pub fn set_size(&mut self, size: Size) {
        self.mouse_translator.data.viewport_size = size
    }

    /// Re-set the mouse bounds size from a viewport
    pub fn set_size_from_viewport(&mut self, vp: Viewport) {
        self.set_size(Size::from(vp.draw_size));
    }
}

#[derive(Clone)]
struct MouseTranslationData {
    x_axis_motion_inverted: bool,
    y_axis_motion_inverted: bool,
    x_axis_scroll_inverted: bool,
    y_axis_scroll_inverted: bool,
    viewport_size: Size
}

impl MouseTranslationData {
    fn new(size: Size) -> Self {
        MouseTranslationData {
            x_axis_motion_inverted: false,
            y_axis_motion_inverted: false,
            x_axis_scroll_inverted: false,
            y_axis_scroll_inverted: false,
            viewport_size: size
        }
    }
}

#[derive(Clone)]
struct MouseTranslator {
    data: MouseTranslationData
}

impl MouseTranslator {
    fn new(size: Size) -> Self {
        MouseTranslator {
            data: MouseTranslationData::new(size)
        }
    }

    fn translate(&self, motion: Motion) -> Motion {
        match motion {
            Motion::MouseCursor(x, y) => {
                let (sw, sh) = {
                    let Size {width, height} = self.data.viewport_size;
                    (width as f64, height as f64)
                };

                let cx = if self.data.x_axis_motion_inverted { sw - x } else { x };
                let cy = if self.data.y_axis_motion_inverted { sh - y } else { y };

                Motion::MouseCursor(cx, cy)
            },
            Motion::MouseScroll(x, y) => {
                let mx = if self.data.x_axis_scroll_inverted { -1.0f64 } else { 1.0 };
                let my = if self.data.y_axis_scroll_inverted { -1.0f64 } else { 1.0 };
                Motion::MouseScroll(x * mx, y * my)
            },
            relative => relative
        }
    }
}

/// An interface for rebinding keys to actions. This is freely convertable to and
/// from an InputTranslator.
pub struct InputRebind<A: Action> {
    keymap: HashMap<A, ButtonTuple>,
    mouse_data: MouseTranslationData
}

impl<A: Action> InputRebind<A> {
    #[allow(missing_docs)]
    pub fn new(size: Size) -> Self {
        InputRebind {
            keymap: HashMap::new(),
            mouse_data: MouseTranslationData::new(size)
        }
    }

    #[allow(missing_docs)]
    pub fn insert_action(&mut self, action: A) -> Option<ButtonTuple> {
        self.keymap.insert(action, ButtonTuple::new())
    }

    #[allow(missing_docs)]
    pub fn insert_action_with_buttons(&mut self, action: A, buttons: ButtonTuple) -> Option<ButtonTuple> {
        self.keymap.insert(action, buttons)
    }

    #[allow(missing_docs)]
    pub fn get_bindings(&self, action: &A) -> Option<&ButtonTuple> {
        self.keymap.get(action)
    }

    #[allow(missing_docs)]
    pub fn get_bindings_mut(&mut self, action: &mut A) -> Option<&mut ButtonTuple> {
        self.keymap.get_mut(action)
    }

    #[allow(missing_docs)]
    pub fn get_x_scroll_inverted(&self) -> &bool {
        &self.mouse_data.x_axis_scroll_inverted
    }

    #[allow(missing_docs)]
    pub fn get_x_scroll_inverted_mut(&mut self) -> &mut bool {
        &mut self.mouse_data.x_axis_scroll_inverted
    }

    #[allow(missing_docs)]
    pub fn get_y_scroll_inverted(&self) -> &bool {
        &self.mouse_data.y_axis_scroll_inverted
    }

    #[allow(missing_docs)]
    pub fn get_y_scroll_inverted_mut(&mut self) -> &mut bool {
        &mut self.mouse_data.y_axis_scroll_inverted
    }

    #[allow(missing_docs)]
    pub fn get_x_motion_inverted(&self) -> &bool {
        &self.mouse_data.x_axis_motion_inverted
    }

    #[allow(missing_docs)]
    pub fn get_x_motion_inverted_mut(&mut self) -> &mut bool {
        &mut self.mouse_data.x_axis_motion_inverted
    }

    #[allow(missing_docs)]
    pub fn get_y_motion_inverted(&self) -> &bool {
        &self.mouse_data.y_axis_motion_inverted
    }

    #[allow(missing_docs)]
    pub fn get_y_motion_inverted_mut(&mut self) -> &mut bool {
        &mut self.mouse_data.y_axis_motion_inverted
    }

    #[allow(missing_docs)]
    pub fn get_viewport_size(&self) -> &Size {
        &self.mouse_data.viewport_size
    }

    #[allow(missing_docs)]
    pub fn get_viewport_size_mut(&mut self) -> &mut Size {
        &mut self.mouse_data.viewport_size
    }
}

impl<A: Action> Default for InputRebind<A> {
    #[allow(missing_docs)]
    fn default() -> Self {
        InputRebind::new((800, 600).into())
    }
}

impl<A: Action> Into<InputTranslator<A>> for InputRebind<A> {
    #[allow(missing_docs)]
    fn into(self) -> InputTranslator<A> {
        let mut input_translator = InputTranslator::new(self.mouse_data.viewport_size);
        input_translator.mouse_translator.data = self.mouse_data;
        input_translator.keymap = self.keymap.values()
                                             .flat_map(|&bt| bt.into_iter())
                                             .filter_map(|x| x)
                                             .zip(self.keymap.keys().map(|x| *x))
                                             .collect();

        input_translator
    }
}

impl<A: Action> Into<InputRebind<A>> for InputTranslator<A> {
    #[allow(missing_docs, unreachable_code)]
    fn into(self) -> InputRebind<A> {
        unimplemented!();

        let mut input_rebind = InputRebind::new(self.mouse_translator.data.viewport_size);
        input_rebind.mouse_data = self.mouse_translator.data;
        //input_rebind.keymap = self.keymap.btn_map.
        //input_rebind.keymap = self.keymap.btn_map.drain().map(|(k, v)| (v, k)).collect();
        input_rebind
    }
}

impl ButtonTuple {
    fn into_iter(self) -> ButtonTupleIterator {
        ButtonTupleIterator {
            button_tuple: self,
            i: 0
        }
    }
}

struct ButtonTupleIterator {
    button_tuple: ButtonTuple,
    i: usize
}

impl Iterator for ButtonTupleIterator {
    type Item = Option<Button>;

    fn next(&mut self) -> Option<Self::Item> {
        let i = self.i;
        self.i += 1;
        match i {
            0 => Some(self.button_tuple.0),
            1 => Some(self.button_tuple.1),
            2 => Some(self.button_tuple.2),
            _ => None
        }
    }
}

<commit_msg>Extract out Some(button) into a variable in insert_inplace() function<commit_after>#![warn(missing_docs)]

//! input-map
//! =========
//!
//! A library for binding input keys to actions, and manipulating mouse input

extern crate input;
extern crate window;
extern crate rustc_serialize;
extern crate viewport;

mod builder;

use input::{Input, Button, Motion};
use window::Size;
use std::collections::HashMap;
use std::default::Default;
use std::hash::Hash;
use viewport::Viewport;

pub use builder::InputTranslatorBuilder;

/// Represents a logical action to be bound to a particular button press, e.g.
/// jump, attack, or move forward. Needs to be hashable, as it is used as a
/// lookup key when rebinding an action to a different button.
pub trait Action: Copy + PartialEq + Eq + Hash { }

/// A translated action.
#[derive(Debug, Copy, Clone, PartialEq)]
pub enum Translated<A: Action> {
    /// A keypress event which was bound to an action
    Press(A),

    /// A key release event which was bound to an action
    Release(A),

    /// A translated mouse motion. The logical origin of a translated MouseCursor event
    /// is in the top left corner of the window, and the logical scroll is non-natural.
    /// Relative events are unchanged for now.
    Move(Motion)
}

/// A three-element tuple of Option<Button>. For simplicity, a maximum number of 3
/// buttons can be bound to each action, and this is exposed through the `InputRebind`
/// struct.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct ButtonTuple(pub Option<Button>, pub Option<Button>, pub Option<Button>);

impl ButtonTuple {
    /// Creates a new tuple with no buttons in it (equivalent to `Default::default()`).
    pub fn new() -> Self { Default::default() }

    /// Check if the button is in the tuple.
    pub fn contains(&self, button: Button) -> bool {
        let sbtn = Some(button);
        self.0 == sbtn || self.1 == sbtn || self.2 == sbtn
    }

    /// Insert a button into the tuple if there is room, searching from left to right.
    /// If the button is inserted, returns true. Otherwise, if the button is not inserted,
    /// this function returns false.
    pub fn insert_inplace(&mut self, button: Button) -> bool {
        let sbtn = Some(button);
        match self {
            &mut ButtonTuple(a, b, c) if a.is_none() => {*self = ButtonTuple(sbtn, b, c); true},
            &mut ButtonTuple(a, b, c) if b.is_none() => {*self = ButtonTuple(a, sbtn, c); true},
            &mut ButtonTuple(a, b, c) if c.is_none() => {*self = ButtonTuple(a, b, sbtn); true}
            _ => false
        }
    }
}

impl Default for ButtonTuple {
    /// Creates a new tuple with no buttons in it.
    fn default() -> Self { ButtonTuple(None, None, None) }
}

/// An object which translates piston::input::Input events into input_map::Translated<A> events
#[derive(Clone)]
pub struct InputTranslator<A: Action> {
    keymap: HashMap<Button, A>,
    mouse_translator: MouseTranslator
}

impl<A: Action> InputTranslator<A> {
    /// Creates an empty InputTranslator.
    pub fn new(size: Size) -> Self {
        InputTranslator {
            keymap: HashMap::new(),
            mouse_translator: MouseTranslator::new(size)
        }
    }

    /// Translate an Input into a Translated<A> event. Returns `None` if there is no
    /// action associated with the `Input` variant.
    pub fn translate(&self, input: &Input) -> Option<Translated<A>> {
        macro_rules! translate_button(($but_state:ident, $but_var:ident) => (
            match self.keymap.get(&$but_var).map(|x| *x) {
                Some(act) => Some(Translated::$but_state(act)),
                None => None
            });
        );

        match input {
            &Input::Press(button) => translate_button!(Press, button),
            &Input::Release(button) => translate_button!(Release, button),
            &Input::Move(motion) =>
                Some(Translated::Move(self.mouse_translator.translate(motion))),
            _ => None
        }
    }

    /// Re-set the mouse bounds size used for calculating mouse events
    pub fn set_size(&mut self, size: Size) {
        self.mouse_translator.data.viewport_size = size
    }

    /// Re-set the mouse bounds size from a viewport
    pub fn set_size_from_viewport(&mut self, vp: Viewport) {
        self.set_size(Size::from(vp.draw_size));
    }
}

#[derive(Clone)]
struct MouseTranslationData {
    x_axis_motion_inverted: bool,
    y_axis_motion_inverted: bool,
    x_axis_scroll_inverted: bool,
    y_axis_scroll_inverted: bool,
    viewport_size: Size
}

impl MouseTranslationData {
    fn new(size: Size) -> Self {
        MouseTranslationData {
            x_axis_motion_inverted: false,
            y_axis_motion_inverted: false,
            x_axis_scroll_inverted: false,
            y_axis_scroll_inverted: false,
            viewport_size: size
        }
    }
}

#[derive(Clone)]
struct MouseTranslator {
    data: MouseTranslationData
}

impl MouseTranslator {
    fn new(size: Size) -> Self {
        MouseTranslator {
            data: MouseTranslationData::new(size)
        }
    }

    fn translate(&self, motion: Motion) -> Motion {
        match motion {
            Motion::MouseCursor(x, y) => {
                let (sw, sh) = {
                    let Size {width, height} = self.data.viewport_size;
                    (width as f64, height as f64)
                };

                let cx = if self.data.x_axis_motion_inverted { sw - x } else { x };
                let cy = if self.data.y_axis_motion_inverted { sh - y } else { y };

                Motion::MouseCursor(cx, cy)
            },
            Motion::MouseScroll(x, y) => {
                let mx = if self.data.x_axis_scroll_inverted { -1.0f64 } else { 1.0 };
                let my = if self.data.y_axis_scroll_inverted { -1.0f64 } else { 1.0 };
                Motion::MouseScroll(x * mx, y * my)
            },
            relative => relative
        }
    }
}

/// An interface for rebinding keys to actions. This is freely convertable to and
/// from an InputTranslator.
pub struct InputRebind<A: Action> {
    keymap: HashMap<A, ButtonTuple>,
    mouse_data: MouseTranslationData
}

impl<A: Action> InputRebind<A> {
    #[allow(missing_docs)]
    pub fn new(size: Size) -> Self {
        InputRebind {
            keymap: HashMap::new(),
            mouse_data: MouseTranslationData::new(size)
        }
    }

    #[allow(missing_docs)]
    pub fn insert_action(&mut self, action: A) -> Option<ButtonTuple> {
        self.keymap.insert(action, ButtonTuple::new())
    }

    #[allow(missing_docs)]
    pub fn insert_action_with_buttons(&mut self, action: A, buttons: ButtonTuple) -> Option<ButtonTuple> {
        self.keymap.insert(action, buttons)
    }

    #[allow(missing_docs)]
    pub fn get_bindings(&self, action: &A) -> Option<&ButtonTuple> {
        self.keymap.get(action)
    }

    #[allow(missing_docs)]
    pub fn get_bindings_mut(&mut self, action: &mut A) -> Option<&mut ButtonTuple> {
        self.keymap.get_mut(action)
    }

    #[allow(missing_docs)]
    pub fn get_x_scroll_inverted(&self) -> &bool {
        &self.mouse_data.x_axis_scroll_inverted
    }

    #[allow(missing_docs)]
    pub fn get_x_scroll_inverted_mut(&mut self) -> &mut bool {
        &mut self.mouse_data.x_axis_scroll_inverted
    }

    #[allow(missing_docs)]
    pub fn get_y_scroll_inverted(&self) -> &bool {
        &self.mouse_data.y_axis_scroll_inverted
    }

    #[allow(missing_docs)]
    pub fn get_y_scroll_inverted_mut(&mut self) -> &mut bool {
        &mut self.mouse_data.y_axis_scroll_inverted
    }

    #[allow(missing_docs)]
    pub fn get_x_motion_inverted(&self) -> &bool {
        &self.mouse_data.x_axis_motion_inverted
    }

    #[allow(missing_docs)]
    pub fn get_x_motion_inverted_mut(&mut self) -> &mut bool {
        &mut self.mouse_data.x_axis_motion_inverted
    }

    #[allow(missing_docs)]
    pub fn get_y_motion_inverted(&self) -> &bool {
        &self.mouse_data.y_axis_motion_inverted
    }

    #[allow(missing_docs)]
    pub fn get_y_motion_inverted_mut(&mut self) -> &mut bool {
        &mut self.mouse_data.y_axis_motion_inverted
    }

    #[allow(missing_docs)]
    pub fn get_viewport_size(&self) -> &Size {
        &self.mouse_data.viewport_size
    }

    #[allow(missing_docs)]
    pub fn get_viewport_size_mut(&mut self) -> &mut Size {
        &mut self.mouse_data.viewport_size
    }
}

impl<A: Action> Default for InputRebind<A> {
    #[allow(missing_docs)]
    fn default() -> Self {
        InputRebind::new((800, 600).into())
    }
}

impl<A: Action> Into<InputTranslator<A>> for InputRebind<A> {
    #[allow(missing_docs)]
    fn into(self) -> InputTranslator<A> {
        let mut input_translator = InputTranslator::new(self.mouse_data.viewport_size);
        input_translator.mouse_translator.data = self.mouse_data;
        input_translator.keymap = self.keymap.values()
                                             .flat_map(|&bt| bt.into_iter())
                                             .filter_map(|x| x)
                                             .zip(self.keymap.keys().map(|x| *x))
                                             .collect();

        input_translator
    }
}

impl<A: Action> Into<InputRebind<A>> for InputTranslator<A> {
    #[allow(missing_docs, unreachable_code)]
    fn into(self) -> InputRebind<A> {
        unimplemented!();

        let mut input_rebind = InputRebind::new(self.mouse_translator.data.viewport_size);
        input_rebind.mouse_data = self.mouse_translator.data;
        //input_rebind.keymap = self.keymap.btn_map.
        //input_rebind.keymap = self.keymap.btn_map.drain().map(|(k, v)| (v, k)).collect();
        input_rebind
    }
}

impl ButtonTuple {
    fn into_iter(self) -> ButtonTupleIterator {
        ButtonTupleIterator {
            button_tuple: self,
            i: 0
        }
    }
}

struct ButtonTupleIterator {
    button_tuple: ButtonTuple,
    i: usize
}

impl Iterator for ButtonTupleIterator {
    type Item = Option<Button>;

    fn next(&mut self) -> Option<Self::Item> {
        let i = self.i;
        self.i += 1;
        match i {
            0 => Some(self.button_tuple.0),
            1 => Some(self.button_tuple.1),
            2 => Some(self.button_tuple.2),
            _ => None
        }
    }
}

<|endoftext|>


<commit_before>#![feature(catch_panic, fnbox)]

//! Work-stealing job-queue implementation.
//! Inspired by the Molecular Matters blog posts.

extern crate crossbeam;
extern crate rand;

use crossbeam::sync::TreiberStack;

use std::any::Any;
use std::boxed::FnBox;
use std::cell::RefCell;
use std::collections::VecDeque;
use std::sync::{Arc, Mutex, RwLock};
use std::sync::mpsc::{channel, Sender, Receiver, TryRecvError};
use std::thread;

use ::rand::thread_rng;
use ::rand::distributions::{IndependentSample, Range};

/// Messages that can be sent to workers.
enum WorkerMessage {
    Stop,
}

type Job = Box<FnBox(&Worker) + 'static + Send>;

/// A job's id in the pool.
pub struct JobId {
    idx: usize,
}

struct PoolEntry {
    job: Job,
    sender: Sender<JobExitStatus>,
}

enum JobExitStatus {
    Success,
    Panic(Box<Any>),
}

/// A mostly lock-free pool allocator of jobs.
/// It only locks when doubling the size of the vector.
/// Eventually, I may implement a lock-free vector based off of
/// Stroustroup-et-al's paper on the subject, but for now, I've
/// decided to "keep it simple, stupid".
struct JobPool {
    jobs: RwLock<Vec<RefCell<Option<PoolEntry>>>>,
    unused: TreiberStack<usize>,
}

impl JobPool {
    /// Initialize this pool with the given capacity.
    fn with_capacity(size: usize) -> Self {
        let mut v = Vec::with_capacity(size);
        let unused = TreiberStack::new();

        for i in (0..size).rev() {
            v.push(RefCell::new(None));

            unused.push(i);
        }

        JobPool {
            jobs: RwLock::new(v),
            unused: unused,
        }
    }

    /// Submit a job to be stored in the pool.
    fn submit(&self, job: Job) -> (JobId, Receiver<JobExitStatus>) {
        let next = if let Some(next) = self.unused.pop() {
            next
        } else {
            self.double();
            self.unused.pop().unwrap()
        };

        let jobs = match self.jobs.read() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        let (tx, rx) = channel();

        let entry = &jobs[next];
        *entry.borrow_mut() = Some(PoolEntry {
            job: job,
            sender: tx,
        });

        (JobId { idx: next, }, rx)
    }

    // double the vector size
    fn double(&self) {
        let mut jobs = match self.jobs.write() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        let old_len = jobs.len();
        let new_len = old_len * 2;

        jobs.reserve(new_len);

        for i in (old_len..new_len).rev() {
            jobs.push(RefCell::new(None));

            self.unused.push(i);
        }
    }

    /// Run the job given by id, passing the worker
    /// as a parameter.
    fn run(&self, id: JobId, worker: &Worker) {
        use std::mem;

        let jobs = match self.jobs.read() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        // take the job out before pushing the index back onto
        // the unused stack.
        let entry = jobs[id.idx].borrow_mut().take().unwrap();

        let PoolEntry { job, sender } = entry;

        // is there a better way to do this than to box again?
        let unbounded_fn: Box<FnBox() + Send> = Box::new(move || job.call_box((worker,)));
        let bounded_fn: Box<FnBox() + 'static + Send> = unsafe { mem::transmute(unbounded_fn) };

        let exit_status = match thread::catch_panic(|| bounded_fn()) {
            Ok(_) => { JobExitStatus::Success }
            Err(e) => { JobExitStatus::Panic(e) }
        };

        // We don't really care if they receive it or not.
        let _result = sender.send(exit_status);
        self.unused.push(id.idx);
    }
}

unsafe impl Send for JobPool {}
unsafe impl Sync for JobPool {}

/// A thread-safe double-ended queue for jobs.
/// Right now, it's just a dumb locking wrapper of VecDeque.
/// It stores indices of jobs in the global job pool.
struct JobQueue {
    inner: Mutex<VecDeque<JobId>>,
}

impl JobQueue {
    fn new() -> Self {
        JobQueue {
            inner: Mutex::new(VecDeque::new()),
        }
    }

    // pop a job from the private end of the queue.
    // only the owner of this queue should call this.
    fn pop(&self) -> Option<JobId> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.pop_back()
    }

    // steal a job from the public end of the queue.
    fn steal(&self) -> Option<JobId> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.pop_front()
    }

    // push another job to the queue.
    fn push(&self, job: JobId) {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.push_back(job);
    }
}

unsafe impl Send for JobQueue {}
unsafe impl Sync for JobQueue {}

/// A worker executes jobs.
pub struct Worker {
    pool: Arc<JobPool>,
    queue: Arc<JobQueue>,
    siblings: Box<[Arc<JobQueue>]>,
}

impl Worker {
    pub fn submit<F: FnOnce(&Worker) + 'static + Send>(&self, f: F) -> JoinHandle {
        let job = Box::new(f);
        let (id, notifier) = self.pool.submit(job);
        self.queue.push(id);

        JoinHandle {
            notifier: notifier,
            worker: Some(self),
        }
    }

    fn next_job(&self) -> Option<JobId> {
        if let Some(id) = self.queue.pop() {
            Some(id)
        } else {
            let mut rng = thread_rng();
            let idx = Range::new(0, self.siblings.len()).ind_sample(&mut rng);
            let other_queue: &Arc<JobQueue> = &self.siblings[idx];

            // make sure we don't steal from our own queue.
            if &*self.queue as *const JobQueue != &**other_queue as *const JobQueue {
                other_queue.steal()
            } else {
                None
            }
        }
    }

    fn run_job(&self) {
        if let Some(id) = self.next_job() {
            self.pool.run(id, self);
        }
    }

    // run this worker to completion.
    fn run_to_completion(&self) {
        while let Some(id) = self.queue.pop() {
            self.pool.run(id, self);
        }

        // help run all the sibling queues to completion as well.
        for sibling in self.siblings.iter() {
            while let Some(id) = sibling.steal() {
                self.pool.run(id, self);
            }
        }
    }
}

/// A JoinHandle can be used to manually wait for
/// a job to be completed.
pub struct JoinHandle<'a> {
    notifier: Receiver<JobExitStatus>,
    worker: Option<&'a Worker>,
}

impl<'a> JoinHandle<'a> {
    pub fn wait(self) -> Result<(), Box<Any>> {
        loop {
            match self.notifier.try_recv() {
                Ok(exit_status) => {
                    return match exit_status {
                        JobExitStatus::Success => Ok(()),
                        JobExitStatus::Panic(err) => Err(err),
                    }
                }

                // sender _never_ hangs up before the job is done.
                Err(TryRecvError::Disconnected) => unreachable!(),

                // nothing yet, just do some more work.
                _ => {
                    if let &Some(ref worker) = &self.worker {
                        worker.run_job()
                    }
                }
            }
        }
    }
}

pub struct JobSystem {
    handles: Vec<Sender<WorkerMessage>>,
    pool: Arc<JobPool>,
    queue: Arc<JobQueue>,
}

impl JobSystem {
    pub fn submit<F: FnOnce(&Worker) + 'static + Send>(&self, f: F) -> JoinHandle {
        let job = Box::new(f);
        let (id, notifier) = self.pool.submit(job);
        self.queue.push(id);

        JoinHandle {
            notifier: notifier,
            worker: None,
        }
    }
}

impl Drop for JobSystem {
    fn drop(&mut self) {
        for handle in &self.handles {
            match handle.send(WorkerMessage::Stop) {
                _ => {}
            }
        }
    }
}

/// Create a new `JobSystem` with `n_threads` worker threads.
/// This will fail if `n_threads` is less than 1.
pub fn make_pool(n_threads: usize) -> Option<JobSystem> {
    if n_threads < 1 {
        return None;
    }

    let job_pool = Arc::new(JobPool::with_capacity(4096));

    // make a queue for every thread, including this one.
    let queues: Vec<_> = (0..n_threads+1)
        .map(|_| Arc::new(JobQueue::new())).collect();

    let queues = queues.into_boxed_slice();

    let mut handles = Vec::new();

    // Spawn a worker thread for every queue but the last one,
    // since that belongs to the thread this is spawned on.
    for queue in queues.iter().cloned().take(n_threads) {
        let (tx, rx) = channel();
        let siblings = queues.clone();
        let pool = job_pool.clone();

        thread::spawn(move || {
            let worker = Worker {
                pool: pool,
                queue: queue,
                siblings: siblings,
            };

            worker_main(worker, rx);
        });
        handles.push(tx);
    }

    Some(JobSystem {
        handles: handles,
        pool: job_pool,
        queue: queues.last().unwrap().clone(),
    })
}

// the main thread for each worker
fn worker_main(worker: Worker, rx: Receiver<WorkerMessage>) {
    loop {
        // check for message
        match rx.try_recv() {
            // time to stop!
            Ok(WorkerMessage::Stop) => {
                worker.run_to_completion();
            }

            // somehow the JobSystem got dropped without disconnecting.
            Err(TryRecvError::Disconnected) => {
                worker.run_to_completion();
            }

            _ => {}
        }
        // run a job
        worker.run_job();
    }
}

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

    #[test]
    #[should_panic]
    fn no_zero_thread_pool() {
        make_pool(0).unwrap();
    }

    #[test]
    fn basic_spawning() {
        // four worker threads, including this one.
        let pool = make_pool(4).unwrap();

        for _ in 0..100 {
            // empty job
            pool.submit(|_| {});
        }
    }

    #[test]
    fn child_jobs() {
        let pool = make_pool(4).unwrap();

        pool.submit(|worker| {
            for _ in 0..100 {
                // empty job
                worker.submit(|_| {});
            };
        }).wait().unwrap();
    }

    #[test]
    fn panic_in_the_jobqueue() {
        use std::any::Any;

        let pool = make_pool(4).unwrap();

        let mut handles = Vec::new();
        for i in 0..100 {
            // 4 threads, 20 panics.
            let handle = pool.submit(move |_| match i % 5 {
                0 => panic!(i),
                _ => {}
            });

            handles.push(handle);
        }

        for handle in handles {
            match handle.wait() {
                Ok(_) => {},
                Err(b) => {
                    let panic_val = b.downcast_ref::<i32>().unwrap();
                    assert_eq!(panic_val % 5, 0);
                }
            }
        }

    }

/*    #[test]
    fn scoping() {
        let pool = make_pool(4).unwrap();
        let this_worker = pool.get_worker();
        this_worker.submit(|worker| {
            let mut v = vec![0; 256];

            worker.scope(|scope| {
                for chunk in v.chunks_mut(32) {
                    scope.submit(move |_| {
                        for i in chunk { *i += 1 }
                    });
                }
            });

            for i in v {
                assert_eq!(i, 1);
            }
        }).wait();
    }*/
}<commit_msg>Improved tests<commit_after>#![feature(catch_panic, fnbox)]

//! Work-stealing job-queue implementation.
//! Inspired by the Molecular Matters blog posts.

extern crate crossbeam;
extern crate rand;

use crossbeam::mem::CachePadded;
use crossbeam::sync::TreiberStack;

use std::any::Any;
use std::boxed::FnBox;
use std::cell::RefCell;
use std::collections::VecDeque;
use std::sync::{Arc, Mutex, RwLock};
use std::sync::mpsc::{channel, Sender, Receiver, TryRecvError};
use std::thread;

use ::rand::thread_rng;
use ::rand::distributions::{IndependentSample, Range};

/// Messages that can be sent to workers.
enum WorkerMessage {
    Stop,
}

type Job = Box<FnBox(&Worker) + 'static + Send>;

/// A job's id in the pool.
pub struct JobId {
    idx: usize,
}

struct PoolEntry {
    job: Job,
    sender: Sender<JobExitStatus>,
}

enum JobExitStatus {
    Success,
    Panic(Box<Any>),
}

/// A mostly lock-free pool allocator of jobs.
/// It only locks when doubling the size of the vector.
/// Eventually, I may implement a lock-free vector based off of
/// Stroustroup-et-al's paper on the subject, but for now, I've
/// decided to "keep it simple, stupid".
/// However, this should be mostly wait free unless *many* jobs
/// are being submitted.
struct JobPool {
    // we ensure with the lock-free stack of free indices that
    // an index is only used by one thread at a time.
    // mutable, disjoint borrows are therefore acceptable
    // as long as they occur within the umbrella of a read
    // lock and will not be interfered with by doubling of the
    // vector.
    jobs: RwLock<Vec<RefCell<Option<CachePadded<PoolEntry>>>>>,
    unused: TreiberStack<usize>,
}

impl JobPool {
    /// Initialize this pool with the given capacity.
    fn with_capacity(size: usize) -> Self {
        let mut v = Vec::with_capacity(size);
        let unused = TreiberStack::new();

        for i in (0..size).rev() {
            v.push(RefCell::new(None));

            unused.push(i);
        }

        JobPool {
            jobs: RwLock::new(v),
            unused: unused,
        }
    }

    /// Submit a job to be stored in the pool.
    fn submit(&self, job: Job) -> (JobId, Receiver<JobExitStatus>) {
        let next = if let Some(next) = self.unused.pop() {
            next
        } else {
            self.double();
            self.unused.pop().unwrap()
        };

        let jobs = match self.jobs.read() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        let (tx, rx) = channel();

        let entry = &jobs[next];
        *entry.borrow_mut() = Some(CachePadded::new(PoolEntry {
            job: job,
            sender: tx,
        }));

        (JobId { idx: next, }, rx)
    }

    // double the vector size
    fn double(&self) {
        let mut jobs = match self.jobs.write() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        let old_len = jobs.len();
        let new_len = old_len * 2;

        jobs.reserve(new_len);

        for i in (old_len..new_len).rev() {
            jobs.push(RefCell::new(None));

            self.unused.push(i);
        }
    }

    /// Run the job given by id, passing the worker
    /// as a parameter.
    fn run(&self, id: JobId, worker: &Worker) {
        use std::mem;
        use std::ptr;

        let jobs = match self.jobs.read() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        // take the job out before pushing the index back onto
        // the unused stack.
        let entry = jobs[id.idx].borrow_mut().take().unwrap();

        // CachePadded doesn't run destructors, so this is ok.
        let PoolEntry { job, sender } = unsafe { ptr::read(&*entry) } ;

        // is there a better way to do this than to box again?
        let unbounded_fn: Box<FnBox() + Send> = Box::new(move || job.call_box((worker,)));
        let bounded_fn: Box<FnBox() + 'static + Send> = unsafe { mem::transmute(unbounded_fn) };

        let exit_status = match thread::catch_panic(|| bounded_fn()) {
            Ok(_) => { JobExitStatus::Success }
            Err(e) => { JobExitStatus::Panic(e) }
        };

        // We don't really care if they receive it or not.
        let _result = sender.send(exit_status);
        self.unused.push(id.idx);
    }
}

unsafe impl Send for JobPool {}
unsafe impl Sync for JobPool {}

/// A thread-safe double-ended queue for jobs.
/// Right now, it's just a dumb locking wrapper of VecDeque.
/// It stores indices of jobs in the global job pool.
struct JobQueue {
    inner: Mutex<VecDeque<JobId>>,
}

impl JobQueue {
    fn new() -> Self {
        JobQueue {
            inner: Mutex::new(VecDeque::new()),
        }
    }

    // pop a job from the private end of the queue.
    // only the owner of this queue should call this.
    fn pop(&self) -> Option<JobId> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.pop_back()
    }

    // steal a job from the public end of the queue.
    fn steal(&self) -> Option<JobId> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.pop_front()
    }

    // push another job to the queue.
    fn push(&self, job: JobId) {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.push_back(job);
    }
}

unsafe impl Send for JobQueue {}
unsafe impl Sync for JobQueue {}

/// A worker executes jobs.
pub struct Worker {
    pool: Arc<JobPool>,
    queue: Arc<JobQueue>,
    siblings: Box<[Arc<JobQueue>]>,
}

impl Worker {
    pub fn submit<F: FnOnce(&Worker) + 'static + Send>(&self, f: F) -> JoinHandle {
        let job = Box::new(f);
        let (id, notifier) = self.pool.submit(job);
        self.queue.push(id);

        JoinHandle {
            notifier: notifier,
            worker: Some(self),
        }
    }

    fn next_job(&self) -> Option<JobId> {
        if let Some(id) = self.queue.pop() {
            Some(id)
        } else {
            let mut rng = thread_rng();
            let idx = Range::new(0, self.siblings.len()).ind_sample(&mut rng);
            let other_queue: &Arc<JobQueue> = &self.siblings[idx];

            // make sure we don't steal from our own queue.
            if &*self.queue as *const JobQueue != &**other_queue as *const JobQueue {
                other_queue.steal()
            } else {
                None
            }
        }
    }

    fn run_job(&self) {
        if let Some(id) = self.next_job() {
            self.pool.run(id, self);
        }
    }

    // run this worker to completion.
    fn run_to_completion(&self) {
        while let Some(id) = self.queue.pop() {
            self.pool.run(id, self);
        }

        // help run all the sibling queues to completion as well.
        for sibling in self.siblings.iter() {
            while let Some(id) = sibling.steal() {
                self.pool.run(id, self);
            }
        }
    }
}

/// A JoinHandle can be used to manually wait for
/// a job to be completed.
pub struct JoinHandle<'a> {
    notifier: Receiver<JobExitStatus>,
    worker: Option<&'a Worker>,
}

impl<'a> JoinHandle<'a> {
    pub fn wait(self) -> Result<(), Box<Any>> {
        loop {
            match self.notifier.try_recv() {
                Ok(exit_status) => {
                    return match exit_status {
                        JobExitStatus::Success => Ok(()),
                        JobExitStatus::Panic(err) => Err(err),
                    }
                }

                // sender _never_ hangs up before the job is done.
                Err(TryRecvError::Disconnected) => unreachable!(),

                // nothing yet, just do some more work.
                _ => {
                    if let &Some(ref worker) = &self.worker {
                        worker.run_job()
                    }
                }
            }
        }
    }
}

pub struct JobSystem {
    handles: Vec<Sender<WorkerMessage>>,
    pool: Arc<JobPool>,
    queue: Arc<JobQueue>,
}

impl JobSystem {
    pub fn submit<F: FnOnce(&Worker) + 'static + Send>(&self, f: F) -> JoinHandle {
        let job = Box::new(f);
        let (id, notifier) = self.pool.submit(job);
        self.queue.push(id);

        JoinHandle {
            notifier: notifier,
            worker: None,
        }
    }
}

impl Drop for JobSystem {
    fn drop(&mut self) {
        for handle in &self.handles {
            match handle.send(WorkerMessage::Stop) {
                _ => {}
            }
        }
    }
}

/// Create a new `JobSystem` with `n_threads` worker threads.
/// This will fail if `n_threads` is less than 1.
pub fn make_pool(n_threads: usize) -> Option<JobSystem> {
    if n_threads < 1 {
        return None;
    }

    let job_pool = Arc::new(JobPool::with_capacity(4096));

    // make a queue for every thread, including this one.
    let queues: Vec<_> = (0..n_threads+1)
        .map(|_| Arc::new(JobQueue::new())).collect();

    let queues = queues.into_boxed_slice();

    let mut handles = Vec::new();

    // Spawn a worker thread for every queue but the last one,
    // since that belongs to the thread this is spawned on.
    for queue in queues.iter().cloned().take(n_threads) {
        let (tx, rx) = channel();
        let siblings = queues.clone();
        let pool = job_pool.clone();

        thread::spawn(move || {
            let worker = Worker {
                pool: pool,
                queue: queue,
                siblings: siblings,
            };

            worker_main(worker, rx);
        });
        handles.push(tx);
    }

    Some(JobSystem {
        handles: handles,
        pool: job_pool,
        queue: queues.last().unwrap().clone(),
    })
}

// the main thread for each worker
fn worker_main(worker: Worker, rx: Receiver<WorkerMessage>) {
    loop {
        // check for message
        match rx.try_recv() {
            // time to stop!
            Ok(WorkerMessage::Stop) => {
                worker.run_to_completion();
            }

            // somehow the JobSystem got dropped without disconnecting.
            Err(TryRecvError::Disconnected) => {
                worker.run_to_completion();
            }

            _ => {}
        }
        // run a job
        worker.run_job();
    }
}

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

    #[test]
    #[should_panic]
    fn no_zero_thread_pool() {
        make_pool(0).unwrap();
    }

    #[test]
    fn basic_spawning() {
        // four worker threads, including this one.
        let pool = make_pool(4).unwrap();

        let mut handles = Vec::new();
        for _ in 0..100 {
            // empty job
            let handle = pool.submit(|_| {});
            handles.push(handle);
        }

        for handle in handles {
            // make sure every job runs and
            // no job panics.
            handle.wait().unwrap();
        }
    }

    #[test]
    fn child_jobs() {
        let pool = make_pool(4).unwrap();

        pool.submit(|worker| {
            for _ in 0..100 {
                // empty job
                worker.submit(|_| {});
            };
        }).wait().unwrap();
    }

    #[test]
    fn panic_in_the_jobqueue() {
        let pool = make_pool(4).unwrap();

        let mut handles = Vec::new();
        for i in 0..100 {
            // 4 threads, 20 panics.
            let handle = pool.submit(move |_| match i % 5 {
                0 => panic!(i),
                _ => {}
            });

            handles.push(handle);
        }

        for handle in handles {
            match handle.wait() {
                Ok(_) => {},
                Err(b) => {
                    let panic_val = b.downcast_ref::<i32>().unwrap();
                    assert_eq!(panic_val % 5, 0);
                }
            }
        }

    }

/*    #[test]
    fn scoping() {
        let pool = make_pool(4).unwrap();
        let this_worker = pool.get_worker();
        this_worker.submit(|worker| {
            let mut v = vec![0; 256];

            worker.scope(|scope| {
                for chunk in v.chunks_mut(32) {
                    scope.submit(move |_| {
                        for i in chunk { *i += 1 }
                    });
                }
            });

            for i in v {
                assert_eq!(i, 1);
            }
        }).wait();
    }*/
}<|endoftext|>


<commit_before>#![no_std]

const PRIME32_1: u32 = 2_654_435_761;
const PRIME32_2: u32 = 2_246_822_519;
const PRIME32_3: u32 = 3_266_489_917;
const PRIME32_4: u32 = 668_265_263;
const PRIME32_5: u32 = 374_761_393;

pub struct XXH32 {
    seed: u32,
    total_len: usize,
    v1: u32,
    v2: u32,
    v3: u32,
    v4: u32,
    memory: [u8; 16],
    memsize: usize,
}

fn calc_next_chunk(val: u32, bytes: &[u8]) -> u32 {
    u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
        .wrapping_mul(PRIME32_2)
        .wrapping_add(val)
        .rotate_left(13)
        .wrapping_mul(PRIME32_1)
}

impl Default for XXH32 {
    fn default() -> XXH32 {
        XXH32::new_with_seed(0)
    }
}

impl XXH32 {
    pub fn new_with_seed(seed: u32) -> XXH32 {
        XXH32 {
            seed,
            total_len: 0,
            v1: seed.wrapping_add(PRIME32_1).wrapping_add(PRIME32_2),
            v2: seed.wrapping_add(PRIME32_2),
            v3: seed,
            v4: seed.wrapping_sub(PRIME32_1),
            memory: [0; 16],
            memsize: 0,
        }
    }
    pub fn write(&mut self, bytes: &[u8]) {
        self.total_len += bytes.len();

        if self.memsize + bytes.len() < 16 {
            self.memory[self.memsize..self.memsize + bytes.len()].clone_from_slice(bytes);
            self.memsize += bytes.len();
            return;
        }

        let mut bytesview = &bytes[(16 - self.memsize) % 16..];

        if self.memsize > 0 {
            self.memory[self.memsize..].clone_from_slice(&bytes[..16 - self.memsize]);

            self.v1 = calc_next_chunk(self.v1, &self.memory[0..4]);
            self.v2 = calc_next_chunk(self.v2, &self.memory[4..8]);
            self.v3 = calc_next_chunk(self.v3, &self.memory[8..12]);
            self.v4 = calc_next_chunk(self.v4, &self.memory[12..16]);

            self.memsize = 0;
        }

        while bytesview.len() >= 16 {
            self.v1 = calc_next_chunk(self.v1, &bytesview[0..4]);
            self.v2 = calc_next_chunk(self.v2, &bytesview[4..8]);
            self.v3 = calc_next_chunk(self.v3, &bytesview[8..12]);
            self.v4 = calc_next_chunk(self.v4, &bytesview[12..16]);

            bytesview = &bytesview[16..];
        }

        self.memory[..bytesview.len()].clone_from_slice(bytesview);
        self.memsize += bytesview.len();
    }

    pub fn finish(&self) -> u32 {
        let mut h32 = if self.total_len >= 16 {
            self.v1
                .rotate_left(1)
                .wrapping_add(self.v2.rotate_left(7))
                .wrapping_add(self.v3.rotate_left(12))
                .wrapping_add(self.v4.rotate_left(18))
        } else {
            self.seed.wrapping_add(PRIME32_5)
        };

        h32 = h32.wrapping_add(self.total_len as u32);

        let mut memoryview = &self.memory[..self.memsize];

        while memoryview.len() >= 4 {
            h32 = h32.wrapping_add(
                u32::from_le_bytes([memoryview[0], memoryview[1], memoryview[2], memoryview[3]])
                    .wrapping_mul(PRIME32_3),
            );
            h32 = h32.rotate_left(17).wrapping_mul(PRIME32_4);

            memoryview = &memoryview[4..];
        }

        for byte in memoryview {
            let byte_u32 = u32::from(*byte);
            h32 = h32.wrapping_add(byte_u32.wrapping_mul(PRIME32_5));
            h32 = h32.rotate_left(11).wrapping_mul(PRIME32_1);
        }

        h32 = (h32 ^ (h32 >> 15)).wrapping_mul(PRIME32_2);
        h32 = (h32 ^ (h32 >> 13)).wrapping_mul(PRIME32_3);
        h32 ^ (h32 >> 16)
    }
}
<commit_msg>Use chunks_exact for iteration.<commit_after>#![no_std]

const PRIME32_1: u32 = 2_654_435_761;
const PRIME32_2: u32 = 2_246_822_519;
const PRIME32_3: u32 = 3_266_489_917;
const PRIME32_4: u32 = 668_265_263;
const PRIME32_5: u32 = 374_761_393;

pub struct XXH32 {
    seed: u32,
    total_len: usize,
    v1: u32,
    v2: u32,
    v3: u32,
    v4: u32,
    memory: [u8; 16],
    memsize: usize,
}

fn calc_next_chunk(val: u32, bytes: &[u8]) -> u32 {
    u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
        .wrapping_mul(PRIME32_2)
        .wrapping_add(val)
        .rotate_left(13)
        .wrapping_mul(PRIME32_1)
}

impl Default for XXH32 {
    fn default() -> XXH32 {
        XXH32::new_with_seed(0)
    }
}

impl XXH32 {
    pub fn new_with_seed(seed: u32) -> XXH32 {
        XXH32 {
            seed,
            total_len: 0,
            v1: seed.wrapping_add(PRIME32_1).wrapping_add(PRIME32_2),
            v2: seed.wrapping_add(PRIME32_2),
            v3: seed,
            v4: seed.wrapping_sub(PRIME32_1),
            memory: [0; 16],
            memsize: 0,
        }
    }
    pub fn write(&mut self, bytes: &[u8]) {
        self.total_len += bytes.len();

        if self.memsize + bytes.len() < 16 {
            self.memory[self.memsize..self.memsize + bytes.len()].copy_from_slice(bytes);
            self.memsize += bytes.len();
            return;
        }

        let bytesview = &bytes[(16 - self.memsize) % 16..];

        if self.memsize > 0 {
            self.memory[self.memsize..].copy_from_slice(&bytes[..16 - self.memsize]);

            self.v1 = calc_next_chunk(self.v1, &self.memory[0..4]);
            self.v2 = calc_next_chunk(self.v2, &self.memory[4..8]);
            self.v3 = calc_next_chunk(self.v3, &self.memory[8..12]);
            self.v4 = calc_next_chunk(self.v4, &self.memory[12..16]);

            self.memsize = 0;
        }

        let mut iter = bytesview.chunks_exact(16);
        for chunk in iter.by_ref() {
            self.v1 = calc_next_chunk(self.v1, &chunk[0..4]);
            self.v2 = calc_next_chunk(self.v2, &chunk[4..8]);
            self.v3 = calc_next_chunk(self.v3, &chunk[8..12]);
            self.v4 = calc_next_chunk(self.v4, &chunk[12..16]);
        }

        let bytesview = iter.remainder();

        self.memory[..bytesview.len()].copy_from_slice(bytesview);
        self.memsize += bytesview.len();
    }

    pub fn finish(&self) -> u32 {
        let mut h32 = if self.total_len >= 16 {
            self.v1
                .rotate_left(1)
                .wrapping_add(self.v2.rotate_left(7))
                .wrapping_add(self.v3.rotate_left(12))
                .wrapping_add(self.v4.rotate_left(18))
        } else {
            self.seed.wrapping_add(PRIME32_5)
        };

        h32 = h32.wrapping_add(self.total_len as u32);

        let mut iter = self.memory[..self.memsize].chunks_exact(4);
        for chunk in iter.by_ref() {
            h32 = h32.wrapping_add(
                u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])
                    .wrapping_mul(PRIME32_3),
            );
            h32 = h32.rotate_left(17).wrapping_mul(PRIME32_4);
        }

        for byte in iter.remainder() {
            h32 = h32.wrapping_add(u32::from(*byte).wrapping_mul(PRIME32_5));
            h32 = h32.rotate_left(11).wrapping_mul(PRIME32_1);
        }

        h32 = (h32 ^ (h32 >> 15)).wrapping_mul(PRIME32_2);
        h32 = (h32 ^ (h32 >> 13)).wrapping_mul(PRIME32_3);
        h32 ^ (h32 >> 16)
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;
extern crate env_logger;
extern crate getopts;
extern crate daemonize;
extern crate udt;
extern crate byteorder;
extern crate time;
extern crate sodiumoxide;
extern crate rustc_serialize;

use daemonize::{Daemonize};
use std::process::Command;
use std::net::{UdpSocket, SocketAddr, SocketAddrV4};
use std::str;
use std::env;
use std::fs::File;
use std::str::FromStr;
use std::path::Path;
use std::ffi::OsStr;
use std::io::{Cursor, Error, Seek, SeekFrom, ErrorKind, stderr, Read, Write};
use udt::*;
use byteorder::{ReadBytesExt, WriteBytesExt, LittleEndian};
use sodiumoxide::crypto::secretbox;
use sodiumoxide::crypto::secretbox::xsalsa20poly1305::{NONCEBYTES, Key, Nonce};
use rustc_serialize::hex::{FromHex, ToHex};

pub struct Server<'a> {
    sock: UdtSocket,
    key: Key,
    filename: &'a str
}

pub struct Client<'a> {
    remote_host: &'a str,
    remote_path: &'a str,
}

impl<'a> Server<'a> {
    pub fn new(filename: &str) -> Server {
        let sshconnstr = env::var("SSH_CONNECTION").unwrap_or(String::from("0.0.0.0 0 0.0.0.0 22")).trim().to_owned();
        let sshconn: Vec<&str> = sshconnstr.split(" ").collect();
        let ip = sshconn[2];
        let port = Server::get_open_port(55000, 55100).unwrap();
        let key = secretbox::gen_key();
        let Key(keybytes) = key;
        println!("shoop 0 {} {} {}", ip, port, keybytes.to_hex());
        let daemonize = Daemonize::new();
        match daemonize.start() {
            Ok(_) => { let _ = writeln!(&mut stderr(), "daemonized"); }
            Err(_) => { let _ = writeln!(&mut stderr(), "RWRWARWARARRR"); }
        }

        udt::init();
        let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Datagram).unwrap();
        sock.setsockopt(UdtOpts::UDP_RCVBUF, 1024000i32).unwrap();
        sock.setsockopt(UdtOpts::UDP_SNDBUF, 1024000i32).unwrap();
        sock.bind(SocketAddr::V4(SocketAddrV4::from_str(&format!("{}:{}", ip, port)[..]).unwrap())).unwrap();
        Server { sock: sock, key: key, filename: filename }
    }

    pub fn start(&self) {
        self.sock.listen(1).unwrap();

        let (stream, _) = self.sock.accept().unwrap();

        if let Ok(version) = stream.recvmsg(1) {
            if version[0] == 0x00 {
                self.send_file(stream).unwrap();
            } else {
                panic!("Unrecognized version.");
            }
        } else {
            panic!("Failed to receive version byte from client.");
        }
    }

    fn send_file(&self, stream: UdtSocket) -> Result<(), Error> {
        let mut f = File::open(self.filename).unwrap();

        let mut filesize = 0u64;
        let mut buf = vec![0; 1024 * 1024];
        loop {
            match try!(f.read(&mut buf)) {
                0    => { break },
                read => { filesize += read as u64 },
            }
        }

        let mut wtr = vec![];
        wtr.write_u64::<LittleEndian>(filesize).unwrap();
        match stream.sendmsg(&wtr[..]) {
            Ok(0) => {
                return Err(Error::new(ErrorKind::WriteZero, "failed to write filesize header before timeout"))
            },
            Err(e) => {
                return Err(Error::new(ErrorKind::Other, format!("{:?}", e)))
            }
            _ => {}
        }
        let mut payload = vec![0; 1300];
        f.seek(SeekFrom::Start(0)).unwrap();
        loop {
            match f.read(&mut payload) {
                Ok(0) => {
                    break;
                }
                Ok(read) => {
                    let nonce = secretbox::gen_nonce();
                    let Nonce(noncebytes) = nonce;
                    let mut hdr = vec![0u8; 1 + NONCEBYTES];
                    hdr[0] = NONCEBYTES as u8;
                    hdr[1..].clone_from_slice(&noncebytes);

                    let mut sealed = secretbox::seal(&payload[0..read], &nonce, &self.key);
                    let mut msg = Vec::with_capacity(hdr.len() + sealed.len());
                    msg.extend_from_slice(&hdr);
                    msg.append(&mut sealed);
                    match stream.sendmsg(&msg[..]) {
                        Ok(_) =>  { }
                        Err(e) => {
                            stream.close().expect("Error closing stream");
                            panic!("{:?}", e);
                        }
                    }
                },
                Err(e) => {
                    stream.close().expect("Error closing stream");
                    panic!("{:?}", e);
                }
            }
        }

        stream.close().expect("Error closing stream.");
        Ok(())
    }

    fn get_open_port(start: u16, end: u16) -> Result<u16, ()> {
        assert!(end >= start);
        let mut p = start;
        loop {
            match UdpSocket::bind(&format!("0.0.0.0:{}", p)[..]) {
                Ok(_) => {
                    return Ok(p);
                }
                Err(_) => {
                    p += 1;
                    if p > end {
                        return Err(());
                    }
                }
            }
        }
    }
}

impl<'a> Client<'a> {
    pub fn new(remote_ssh_host: &'a str, remote_path: &'a str) -> Client<'a> {
        Client{ remote_host: remote_ssh_host, remote_path: remote_path }
    }

    pub fn start(&self) {
        let cmd = format!("shoop -s '{}'", self.remote_path);
        // println!("addr: {}, path: {}, cmd: {}", addr, path, cmd);

        assert!(Client::command_exists("ssh"), "`ssh` is required!");
        let output = Command::new("ssh")
                             .arg(self.remote_host.to_owned())
                             .arg(cmd)
                             .output()
                             .unwrap_or_else(|e| {
                                 panic!("failed to execute process: {}", e);
                             });
        let infostring = String::from_utf8_lossy(&output.stdout).to_owned().trim().to_owned();
        let info: Vec<&str> = infostring.split(" ").collect();
        if info.len() != 5 {
            panic!("Unexpected response from server. Are you suuuuure shoop is setup on the server?");
        }

        let (magic, version, ip, port, keyhex) = (info[0], info[1], info[2], info[3], info[4]);
        println!("opening UDT connection to {}:{}", ip, port);
        if magic != "shoop" || version != "0" {
            panic!("Unexpected response from server. Are you suuuuure shoop is setup on the server?");
        }

        let mut keybytes = [0u8; 32];
        keybytes.copy_from_slice(&keyhex.from_hex().unwrap()[..]);
        let key = Key(keybytes);

        udt::init();
        let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Datagram).unwrap();
        sock.setsockopt(UdtOpts::UDP_RCVBUF, 1024000i32).unwrap();
        sock.setsockopt(UdtOpts::UDP_SNDBUF, 1024000i32).unwrap();
        let addr: SocketAddr = SocketAddr::V4(SocketAddrV4::from_str(&format!("{}:{}", ip, port)[..]).unwrap());
        match sock.connect(addr) {
           Ok(()) => {
               println!("connected!");
           },
           Err(e) => {
               panic!("errrrrrrr {:?}", e);
           }
        }

        sock.sendmsg(&[0u8; 1]).unwrap();
        // println!("checking if server is frand");

        match sock.recvmsg(8) {
           Ok(msg) => {
               if msg.len() == 0 {
                   panic!("failed to get filesize from server, probable timeout.");
               }
               let mut rdr = Cursor::new(msg);
               let filesize = rdr.read_u64::<LittleEndian>().unwrap();
               let filename = Path::new(&self.remote_path).file_name().unwrap_or(OsStr::new("outfile")).to_str().unwrap_or("outfile");
               println!("{}, {:.1}MB", filename, (filesize as f64)/(1024f64*1024f64));
               self.recv_file(sock, filesize, filename, key).unwrap();
           }
           Err(e) => {
               panic!("{:?}", e);
           }
        }
    }

    fn command_exists(command: &str) -> bool {
        match Command::new("which").arg(command).output() {
            Ok(output) => output.status.success(),
            Err(_)     => false
        }
    }

    fn recv_file(&self, sock: UdtSocket, filesize: u64, filename: &str, key: Key) -> Result<(), Error> {
        let mut f = File::create(filename).unwrap();
        let mut ts = time::precise_time_ns();
        let mut total = 0u64;
        loop {
            let buf = try!(sock.recvmsg(8192).map_err(|e| Error::new(ErrorKind::Other, format!("{:?}", e))));
            if buf.len() < 1 {
                return Err(Error::new(ErrorKind::InvalidInput, "empty message"));
            }
            let noncelen = buf[0] as usize;
            if noncelen != NONCEBYTES {
                return Err(Error::new(ErrorKind::InvalidInput, "nonce bytes unexpected len"));
            }
            if buf.len() < (1 + noncelen) {
                return Err(Error::new(ErrorKind::InvalidInput, "nonce != nonce_len"));
            }
            let mut noncebytes = [0u8; NONCEBYTES];
            noncebytes.copy_from_slice(&buf[1..1+noncelen]);
            let nonce = Nonce(noncebytes);

            let unsealed = try!(secretbox::open(&buf[1+noncelen..], &nonce, &key).map_err(|_| Error::new(ErrorKind::InvalidInput, "failed to decrypt")));

            total += unsealed.len() as u64;
            f.write_all(&unsealed[..]).unwrap();
            if time::precise_time_ns() > ts + 10000000 {
                print!("\x1b[2K\rreceived {:.1}M / {:.1}M ({:.1}%)", (total as f64)/(1024f64*1024f64), (filesize as f64)/(1024f64*1024f64), (total as f64) / (filesize as f64) * 100f64);
                ts = time::precise_time_ns();
            }
            if total >= filesize {
                println!("\nEOF");
                break;
            }
        }
        let _ = sock.close();
        Ok(())
    }
}
<commit_msg>stop using wildcards your hands will get hairy<commit_after>#[macro_use]
extern crate log;
extern crate env_logger;
extern crate getopts;
extern crate daemonize;
extern crate udt;
extern crate byteorder;
extern crate time;
extern crate sodiumoxide;
extern crate rustc_serialize;

use daemonize::{Daemonize};
use std::process::Command;
use std::net::{UdpSocket, SocketAddr, SocketAddrV4};
use std::str;
use std::env;
use std::fs::File;
use std::str::FromStr;
use std::path::Path;
use std::ffi::OsStr;
use std::io::{Cursor, Error, Seek, SeekFrom, ErrorKind, stderr, Read, Write};
use udt::{UdtSocket, UdtOpts, SocketType, SocketFamily};
use byteorder::{ReadBytesExt, WriteBytesExt, LittleEndian};
use sodiumoxide::crypto::secretbox;
use sodiumoxide::crypto::secretbox::xsalsa20poly1305::{NONCEBYTES, Key, Nonce};
use rustc_serialize::hex::{FromHex, ToHex};

pub struct Server<'a> {
    sock: UdtSocket,
    key: Key,
    filename: &'a str
}

pub struct Client<'a> {
    remote_host: &'a str,
    remote_path: &'a str,
}

impl<'a> Server<'a> {
    pub fn new(filename: &str) -> Server {
        let sshconnstr = env::var("SSH_CONNECTION").unwrap_or(String::from("0.0.0.0 0 0.0.0.0 22")).trim().to_owned();
        let sshconn: Vec<&str> = sshconnstr.split(" ").collect();
        let ip = sshconn[2];
        let port = Server::get_open_port(55000, 55100).unwrap();
        let key = secretbox::gen_key();
        let Key(keybytes) = key;
        println!("shoop 0 {} {} {}", ip, port, keybytes.to_hex());
        let daemonize = Daemonize::new();
        match daemonize.start() {
            Ok(_) => { let _ = writeln!(&mut stderr(), "daemonized"); }
            Err(_) => { let _ = writeln!(&mut stderr(), "RWRWARWARARRR"); }
        }

        udt::init();
        let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Datagram).unwrap();
        sock.setsockopt(UdtOpts::UDP_RCVBUF, 1024000i32).unwrap();
        sock.setsockopt(UdtOpts::UDP_SNDBUF, 1024000i32).unwrap();
        sock.bind(SocketAddr::V4(SocketAddrV4::from_str(&format!("{}:{}", ip, port)[..]).unwrap())).unwrap();
        Server { sock: sock, key: key, filename: filename }
    }

    pub fn start(&self) {
        self.sock.listen(1).unwrap();

        let (stream, _) = self.sock.accept().unwrap();

        if let Ok(version) = stream.recvmsg(1) {
            if version[0] == 0x00 {
                self.send_file(stream).unwrap();
            } else {
                panic!("Unrecognized version.");
            }
        } else {
            panic!("Failed to receive version byte from client.");
        }
    }

    fn send_file(&self, stream: UdtSocket) -> Result<(), Error> {
        let mut f = File::open(self.filename).unwrap();

        let mut filesize = 0u64;
        let mut buf = vec![0; 1024 * 1024];
        loop {
            match try!(f.read(&mut buf)) {
                0    => { break },
                read => { filesize += read as u64 },
            }
        }

        let mut wtr = vec![];
        wtr.write_u64::<LittleEndian>(filesize).unwrap();
        match stream.sendmsg(&wtr[..]) {
            Ok(0) => {
                return Err(Error::new(ErrorKind::WriteZero, "failed to write filesize header before timeout"))
            },
            Err(e) => {
                return Err(Error::new(ErrorKind::Other, format!("{:?}", e)))
            }
            _ => {}
        }
        let mut payload = vec![0; 1300];
        f.seek(SeekFrom::Start(0)).unwrap();
        loop {
            match f.read(&mut payload) {
                Ok(0) => {
                    break;
                }
                Ok(read) => {
                    let nonce = secretbox::gen_nonce();
                    let Nonce(noncebytes) = nonce;
                    let mut hdr = vec![0u8; 1 + NONCEBYTES];
                    hdr[0] = NONCEBYTES as u8;
                    hdr[1..].clone_from_slice(&noncebytes);

                    let mut sealed = secretbox::seal(&payload[0..read], &nonce, &self.key);
                    let mut msg = Vec::with_capacity(hdr.len() + sealed.len());
                    msg.extend_from_slice(&hdr);
                    msg.append(&mut sealed);
                    match stream.sendmsg(&msg[..]) {
                        Ok(_) =>  { }
                        Err(e) => {
                            stream.close().expect("Error closing stream");
                            panic!("{:?}", e);
                        }
                    }
                },
                Err(e) => {
                    stream.close().expect("Error closing stream");
                    panic!("{:?}", e);
                }
            }
        }

        stream.close().expect("Error closing stream.");
        Ok(())
    }

    fn get_open_port(start: u16, end: u16) -> Result<u16, ()> {
        assert!(end >= start);
        let mut p = start;
        loop {
            match UdpSocket::bind(&format!("0.0.0.0:{}", p)[..]) {
                Ok(_) => {
                    return Ok(p);
                }
                Err(_) => {
                    p += 1;
                    if p > end {
                        return Err(());
                    }
                }
            }
        }
    }
}

impl<'a> Client<'a> {
    pub fn new(remote_ssh_host: &'a str, remote_path: &'a str) -> Client<'a> {
        Client{ remote_host: remote_ssh_host, remote_path: remote_path }
    }

    pub fn start(&self) {
        let cmd = format!("shoop -s '{}'", self.remote_path);
        // println!("addr: {}, path: {}, cmd: {}", addr, path, cmd);

        assert!(Client::command_exists("ssh"), "`ssh` is required!");
        let output = Command::new("ssh")
                             .arg(self.remote_host.to_owned())
                             .arg(cmd)
                             .output()
                             .unwrap_or_else(|e| {
                                 panic!("failed to execute process: {}", e);
                             });
        let infostring = String::from_utf8_lossy(&output.stdout).to_owned().trim().to_owned();
        let info: Vec<&str> = infostring.split(" ").collect();
        if info.len() != 5 {
            panic!("Unexpected response from server. Are you suuuuure shoop is setup on the server?");
        }

        let (magic, version, ip, port, keyhex) = (info[0], info[1], info[2], info[3], info[4]);
        println!("opening UDT connection to {}:{}", ip, port);
        if magic != "shoop" || version != "0" {
            panic!("Unexpected response from server. Are you suuuuure shoop is setup on the server?");
        }

        let mut keybytes = [0u8; 32];
        keybytes.copy_from_slice(&keyhex.from_hex().unwrap()[..]);
        let key = Key(keybytes);

        udt::init();
        let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Datagram).unwrap();
        sock.setsockopt(UdtOpts::UDP_RCVBUF, 1024000i32).unwrap();
        sock.setsockopt(UdtOpts::UDP_SNDBUF, 1024000i32).unwrap();
        let addr: SocketAddr = SocketAddr::V4(SocketAddrV4::from_str(&format!("{}:{}", ip, port)[..]).unwrap());
        match sock.connect(addr) {
           Ok(()) => {
               println!("connected!");
           },
           Err(e) => {
               panic!("errrrrrrr {:?}", e);
           }
        }

        sock.sendmsg(&[0u8; 1]).unwrap();
        // println!("checking if server is frand");

        match sock.recvmsg(8) {
           Ok(msg) => {
               if msg.len() == 0 {
                   panic!("failed to get filesize from server, probable timeout.");
               }
               let mut rdr = Cursor::new(msg);
               let filesize = rdr.read_u64::<LittleEndian>().unwrap();
               let filename = Path::new(&self.remote_path).file_name().unwrap_or(OsStr::new("outfile")).to_str().unwrap_or("outfile");
               println!("{}, {:.1}MB", filename, (filesize as f64)/(1024f64*1024f64));
               self.recv_file(sock, filesize, filename, key).unwrap();
           }
           Err(e) => {
               panic!("{:?}", e);
           }
        }
    }

    fn command_exists(command: &str) -> bool {
        match Command::new("which").arg(command).output() {
            Ok(output) => output.status.success(),
            Err(_)     => false
        }
    }

    fn recv_file(&self, sock: UdtSocket, filesize: u64, filename: &str, key: Key) -> Result<(), Error> {
        let mut f = File::create(filename).unwrap();
        let mut ts = time::precise_time_ns();
        let mut total = 0u64;
        loop {
            let buf = try!(sock.recvmsg(8192).map_err(|e| Error::new(ErrorKind::Other, format!("{:?}", e))));
            if buf.len() < 1 {
                return Err(Error::new(ErrorKind::InvalidInput, "empty message"));
            }
            let noncelen = buf[0] as usize;
            if noncelen != NONCEBYTES {
                return Err(Error::new(ErrorKind::InvalidInput, "nonce bytes unexpected len"));
            }
            if buf.len() < (1 + noncelen) {
                return Err(Error::new(ErrorKind::InvalidInput, "nonce != nonce_len"));
            }
            let mut noncebytes = [0u8; NONCEBYTES];
            noncebytes.copy_from_slice(&buf[1..1+noncelen]);
            let nonce = Nonce(noncebytes);

            let unsealed = try!(secretbox::open(&buf[1+noncelen..], &nonce, &key).map_err(|_| Error::new(ErrorKind::InvalidInput, "failed to decrypt")));

            total += unsealed.len() as u64;
            f.write_all(&unsealed[..]).unwrap();
            if time::precise_time_ns() > ts + 10000000 {
                print!("\x1b[2K\rreceived {:.1}M / {:.1}M ({:.1}%)", (total as f64)/(1024f64*1024f64), (filesize as f64)/(1024f64*1024f64), (total as f64) / (filesize as f64) * 100f64);
                ts = time::precise_time_ns();
            }
            if total >= filesize {
                println!("\nEOF");
                break;
            }
        }
        let _ = sock.close();
        Ok(())
    }
}
<|endoftext|>


<commit_before>extern "C" {
    #![allow(improper_ctypes)] // we do not actually cross the FFI bound here

    fn rust_fuzzer_test_input(input: &[u8]);
}

#[export_name="LLVMFuzzerTestOneInput"]
pub fn test_input_wrap(data: *const u8, size: usize) -> i32 {
    ::std::panic::catch_unwind(|| unsafe {
        let data_slice = ::std::slice::from_raw_parts(data, size);
        rust_fuzzer_test_input(data_slice);
    })
    .err().map(|_|
        // hopefully the custom panic hook will be called before and abort the
        // process before the stack frames are unwinded.
        ::std::process::abort()
    );
    0
}

#[export_name="LLVMFuzzerInitialize"]
pub fn initialize(_argc: *const isize, _argv: *const *const *const u8) -> isize {
    // Registers a panic hook that aborts the process before unwinding.
    // It is useful to abort before unwinding so that the fuzzer will then be
    // able to analyse the process stack frames to tell different bugs appart.
    ::std::panic::set_hook(Box::new(|_| {
            ::std::process::abort();
    }));
    0
}

#[macro_export]
macro_rules! fuzz_target {
    (|$bytes:ident| $body:block) => {
        #[no_mangle]
        pub extern fn rust_fuzzer_test_input($bytes: &[u8]) {
            $body
        }
    };
    (|$data:ident: &[u8]| $body:block) => {
        fuzz_target!(|$data| $body);
    };
    (|$data:ident: $dty: ty| $body:block) => {
        extern crate arbitrary;

        #[no_mangle]
        pub extern fn rust_fuzzer_test_input(bytes: &[u8]) {
            use arbitrary::{Arbitrary, RingBuffer};

            let $data: $dty = if let Ok(d) = RingBuffer::new(bytes, bytes.len()).and_then(|mut b|{
                Arbitrary::arbitrary(&mut b).map_err(|_| "")
            }) {
                d
            } else {
                return
            };
            $body
        }
    };
}
<commit_msg>Add comment about panic::set_hook() and reference cargo bug report<commit_after>extern "C" {
    #![allow(improper_ctypes)] // we do not actually cross the FFI bound here

    fn rust_fuzzer_test_input(input: &[u8]);
}

#[export_name="LLVMFuzzerTestOneInput"]
pub fn test_input_wrap(data: *const u8, size: usize) -> i32 {
    ::std::panic::catch_unwind(|| unsafe {
        let data_slice = ::std::slice::from_raw_parts(data, size);
        rust_fuzzer_test_input(data_slice);
    })
    .err().map(|_|
        // hopefully the custom panic hook will be called before and abort the
        // process before the stack frames are unwinded.
        ::std::process::abort()
    );
    0
}

#[export_name="LLVMFuzzerInitialize"]
pub fn initialize(_argc: *const isize, _argv: *const *const *const u8) -> isize {
    // Registers a panic hook that aborts the process before unwinding.
    // It is useful to abort before unwinding so that the fuzzer will then be
    // able to analyse the process stack frames to tell different bugs appart.
    //
    // HACK / FIXME: it would be better to use `-C panic=abort` but it's currently
    // impossible to build code using compiler plugins with this flag.
    // We will be able to remove this code when
    // https://github.com/rust-lang/cargo/issues/5423 is fixed.
    ::std::panic::set_hook(Box::new(|_| {
            ::std::process::abort();
    }));
    0
}

#[macro_export]
macro_rules! fuzz_target {
    (|$bytes:ident| $body:block) => {
        #[no_mangle]
        pub extern fn rust_fuzzer_test_input($bytes: &[u8]) {
            $body
        }
    };
    (|$data:ident: &[u8]| $body:block) => {
        fuzz_target!(|$data| $body);
    };
    (|$data:ident: $dty: ty| $body:block) => {
        extern crate arbitrary;

        #[no_mangle]
        pub extern fn rust_fuzzer_test_input(bytes: &[u8]) {
            use arbitrary::{Arbitrary, RingBuffer};

            let $data: $dty = if let Ok(d) = RingBuffer::new(bytes, bytes.len()).and_then(|mut b|{
                Arbitrary::arbitrary(&mut b).map_err(|_| "")
            }) {
                d
            } else {
                return
            };
            $body
        }
    };
}
<|endoftext|>