<commit_before>use super::super::commands::CommandResult;
use super::super::{LayoutTree, TreeError};
use super::super::core::Direction;
use super::super::core::container::{Container, ContainerType, Layout};

use petgraph::graph::NodeIndex;
use rustwlc::WlcView;
use uuid::Uuid;

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum FocusError {
    /// Reached a container where we can keep climbing the tree no longer.
    ///
    /// Usually this is a workspace, because it doesn't make sense to move a
    /// container out of a workspace
    ReachedLimit(NodeIndex),
    /// Tried to focus on a container that was not a view.
    NotAView(Uuid),
    /// Tried to focus on a container (first one),
    /// but that container was superseded by a fullscreen container (second one)
    BlockedByFullscreen(Uuid, Uuid)
}

impl LayoutTree {
    /// Focuses on the container by the uuid, if it points to a View.
    /// Otherwise, an error is returned.
    pub fn focus_on(&mut self, uuid: Uuid) -> CommandResult {
        if let Some(fullscreen_id) = try!(self.in_fullscreen_workspace(uuid)) {
            if fullscreen_id != uuid {
                return Err(TreeError::Focus(FocusError::BlockedByFullscreen(uuid, fullscreen_id)))
            }
        }
        let container = try!(self.lookup(uuid));
        match *container {
            Container::View { handle, .. } => {
                handle.focus();
                Ok(())
            },
            _ => Err(TreeError::Focus(FocusError::NotAView(uuid)))
        }
    }
    /// Focus on the container relative to the active container.
    ///
    /// If Horizontal, left and right will move within siblings.
    /// If Vertical, up and down will move within siblings.
    /// Other wise, it moves to the next sibling of the parent container.
    ///
    /// If the edge of the children is hit, it does not wrap around,
    /// but moves between ancestor siblings.
    pub fn move_focus(&mut self, direction: Direction) -> CommandResult {
        if let Some(prev_active_ix) = self.active_container {
            let active_id = self.tree[prev_active_ix].get_id();
            if let Some(fullscreen_id) = try!(self.in_fullscreen_workspace(active_id)) {
                return Err(TreeError::Focus(
                    FocusError::BlockedByFullscreen(active_id, fullscreen_id)))
            }
            let new_active_ix = self.move_focus_recurse(prev_active_ix, direction)
                .unwrap_or(prev_active_ix);
            try!(self.set_active_node(new_active_ix));
            match self.tree[self.active_container.unwrap()] {
                Container::View { ref handle, .. } => handle.focus(),
                _ => warn!("move_focus returned a non-view, cannot focus")
            }
        } else {
            return Err(TreeError::NoActiveContainer)
        }
        self.validate();
        Ok(())
    }

    fn move_focus_recurse(&mut self, node_ix: NodeIndex, direction: Direction)
                          -> Result<NodeIndex, TreeError> {
        match self.tree[node_ix].get_type() {
            ContainerType::View | ContainerType::Container => { /* continue */ },
            _ => return Err(TreeError::UuidWrongType(self.tree[node_ix].get_id(),
                                                     vec!(ContainerType::View, ContainerType::Container)))
        }
        let parent_ix = self.tree.parent_of(node_ix)
            .expect("Active ix had no parent");
        match self.tree[parent_ix] {
            Container::Container { layout, .. } => {
                match (layout, direction) {
                    (Layout::Horizontal, Direction::Left) |
                    (Layout::Horizontal, Direction::Right) |
                    (Layout::Vertical, Direction::Up) |
                    (Layout::Vertical, Direction::Down) => {
                        let siblings = self.tree.children_of(parent_ix);
                        let cur_index = siblings.iter().position(|node| {
                            *node == node_ix
                        }).expect("Could not find self in parent");
                        let maybe_new_index = match direction {
                            Direction::Right | Direction::Down => {
                                cur_index.checked_add(1)
                            }
                            Direction::Left  | Direction::Up => {
                                cur_index.checked_sub(1)
                            }
                        };
                        if maybe_new_index.is_some() &&
                            maybe_new_index.unwrap() < siblings.len() {
                                // There is a sibling to move to.
                                let new_index = maybe_new_index.unwrap();
                                let new_active_ix = siblings[new_index];
                                match self.tree[new_active_ix].get_type() {
                                    ContainerType::Container => {
                                        let path_ix = self.tree.follow_path(new_active_ix);
                                        // If the path wasn't complete, find the first view and focus on that
                                        let node_ix = try!(self.tree.descendant_of_type(path_ix, ContainerType::View)
                                                           .map_err(|err| TreeError::PetGraph(err)));
                                        let parent_ix = try!(self.tree.parent_of(node_ix)
                                                             .map_err(|err| TreeError::PetGraph(err)));
                                        match self.tree[node_ix].get_type() {
                                            ContainerType::View | ContainerType::Container => {},
                                            _ => panic!("Following path did not lead to a container or a view!")
                                        }
                                        trace!("Moving to different view {:?} in container {:?}",
                                                self.tree[node_ix], self.tree[parent_ix]);
                                        return Ok(node_ix);
                                    },
                                    ContainerType::View => {
                                        trace!("Moving to other view {:?}", self.tree[new_active_ix]);
                                        return Ok(new_active_ix)
                                    },
                                    _ => unreachable!()
                                };
                            }
                    },
                    _ => { /* We are moving out of siblings, recurse */ }
                }
            }
            Container::Workspace { .. } => {
                return Err(TreeError::Focus(FocusError::ReachedLimit(parent_ix)));
            }
            _ => unreachable!()
        }
        let parent_ix = self.tree.parent_of(node_ix)
            .expect("Node had no parent");
        return self.move_focus_recurse(parent_ix, direction);
    }

    /// Updates the current active container to be the next container or view
    /// to focus on after the previous view/container was moved/removed.
    ///
    /// A new view will tried to be set, starting with the children of the
    /// parent node. If a view cannot be found there, it starts climbing the
    /// tree until either a view is found or the workspace is (in which case
    /// it set the active container to the root container of the workspace)
    ///
    /// If there are siblings, the chosen node is the one with the lowest
    /// active count.
    ///
    /// If there is a fullscreen container in this workspace, that is focused on next,
    /// with the active path updated accordingly.
    pub fn focus_on_next_container(&mut self, mut parent_ix: NodeIndex) {
        let last_ix = self.tree.active_path().last()
            .expect("Active path did not lead anywhere").0;
        let id = self.tree[last_ix].get_id();
        match self.in_fullscreen_workspace(id) {
            Ok(Some(fullscreen_id)) => {
                self.set_active_container(fullscreen_id)
                    .unwrap_or_else(|err| {
                        error!("Could not set {:?} to be the active container, \
                                even though it's the fullscreen container! {:?}",
                               fullscreen_id, err);
                        error!("{:#?}", self);
                        panic!("Could not set the fullscreen container to be \
                                the active container!")
                    });
                return;
            },
            _ => {}
        }
        match self.tree[last_ix] {
            Container::View { handle, .. } => {
                handle.focus();
                self.active_container = Some(last_ix);
                return
            },
            Container::Container { .. } => {
                parent_ix = last_ix;
            },
            _ => {}
        }
        while self.tree.node_type(parent_ix)
            .expect("Node not part of the tree") != ContainerType::Workspace {
                if let Some(node_ix) = self.tree.lowest_active_view(parent_ix) {
                    match self.tree[node_ix] {
                        Container::View { .. } => {
                            trace!("Active container set to view at {:?}", node_ix);
                            let id = self.tree[node_ix].get_id();
                            self.set_active_container(id)
                                .expect("Could not set active container");
                            return;
                        },
                        _ => {}
                    }
                }
                parent_ix = self.tree.ancestor_of_type(parent_ix,
                                                       ContainerType::Container)
                    .unwrap_or_else(|_| {
                        self.tree.ancestor_of_type(parent_ix, ContainerType::Workspace)
                            .expect("Container was not part of a workspace")
                    });
        }
        // If this is reached, parent is workspace
        let container_ix = self.tree.children_of(parent_ix).get(0).cloned();
        if container_ix.is_none() {
            trace!("There were no other containers to focus on, \
                    focusing on nothing in particular!");
            return;
        }
        let container_ix = container_ix.unwrap();
        let root_c_children = self.tree.grounded_children(container_ix);
        if root_c_children.len() > 0 {
            // Only searches first child of root container, can't be floating view.
            let new_active_ix = self.tree.descendant_of_type(root_c_children[0],
                                                                ContainerType::View)
                .unwrap_or(root_c_children[0]);
            let id = self.tree[new_active_ix].get_id();
            self.set_active_container(id)
                .expect("Could not set active container");
            match self.tree[new_active_ix] {
                Container::View { ref handle, .. } => {
                    debug!("Focusing on {:?}", handle);
                    handle.focus();
                },
                Container::Container { .. } => {
                    debug!("No view found, focusing on nothing in workspace {:?}", parent_ix);
                    WlcView::root().focus();
                }
                _ => unreachable!()
            };
            return;
        } else {
            let floating_children = self.tree.floating_children(container_ix);
            for child_ix in floating_children {
                if let Ok(view_ix) = self.tree.descendant_of_type(child_ix,
                                                                    ContainerType::View) {
                    match self.tree[view_ix] {
                        Container::View { handle, id, .. } => {
                            debug!("Floating view found, focusing on {:#?}", handle);
                            handle.focus();
                            self.set_active_container(id)
                                .expect("Could not set active container");
                            return;
                        },
                        _ => unreachable!()
                    };
                }
            }
        }
        trace!("Active container set to container {:?}", container_ix);
        let id = self.tree[container_ix].get_id();
        self.set_active_container(id)
            .expect("Could not set active container");

        // Update focus to new container
        self.get_active_container().map(|con| match *con {
            Container::View { ref handle, .. } => handle.focus(),
            Container::Container { .. } => WlcView::root().focus(),
            _ => panic!("Active container not view or container!")
        });
    }

    /// If the currently focused view is floating, then the non-floating at the end of
    /// the path becomes the focused view. Otherwise, the first floating view becomes
    /// the focused view.
    ///
    /// If there is no currently focused view, does nothing.
    pub fn toggle_floating_focus(&mut self) -> CommandResult {
        let active_ix = try!(self.active_container.ok_or(TreeError::NoActiveContainer));
        let id; let floating;
        {
            let container = &self.tree[active_ix];
            id = container.get_id();
            floating = container.floating();
        }
        if let Some(fullscreen_id) = try!(self.in_fullscreen_workspace(id)) {
            return Err(TreeError::Focus(FocusError::BlockedByFullscreen(id, fullscreen_id)))
        }
        if floating {
            let root_ix = self.tree.root_ix();
            let new_ix = self.tree.follow_path(root_ix);
            match self.tree[new_ix].get_type() {
                ContainerType::View | ContainerType::Container  => {
                    try!(self.set_active_node(new_ix));
                    Ok(())
                },
                type_ => {
                    error!("Path lead to the wrong container, {:#?}\n{:#?}\n{:#?}",
                           active_ix, type_, self);
                    panic!("toggle_floating_focused: bad path");
                }
            }
        } else {
            // Current view is not floating, gotta focus on floating view
            let root_c_ix = self.root_container_ix()
                .expect("No root container ancestor of active container");
            let floating_children = self.tree.floating_children(root_c_ix);
            if floating_children.len() > 0 {
                try!(self.set_active_node(floating_children[0]));
            }
            Ok(())
        }
    }
}
<commit_msg>Made focus_on actually update active c and path<commit_after>use super::super::commands::CommandResult;
use super::super::{LayoutTree, TreeError};
use super::super::core::Direction;
use super::super::core::container::{Container, ContainerType, Layout};

use petgraph::graph::NodeIndex;
use rustwlc::WlcView;
use uuid::Uuid;

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum FocusError {
    /// Reached a container where we can keep climbing the tree no longer.
    ///
    /// Usually this is a workspace, because it doesn't make sense to move a
    /// container out of a workspace
    ReachedLimit(NodeIndex),
    /// Tried to focus on a container that was not a view.
    NotAView(Uuid),
    /// Tried to focus on a container (first one),
    /// but that container was superseded by a fullscreen container (second one)
    BlockedByFullscreen(Uuid, Uuid)
}

impl LayoutTree {
    /// Focuses on the container by the uuid, if it points to a View.
    /// Otherwise, an error is returned.
    pub fn focus_on(&mut self, uuid: Uuid) -> CommandResult {
        if let Some(fullscreen_id) = try!(self.in_fullscreen_workspace(uuid)) {
            if fullscreen_id != uuid {
                return Err(TreeError::Focus(FocusError::BlockedByFullscreen(uuid, fullscreen_id)))
            }
        }
        let node_ix = self.tree.lookup_id(uuid)
            .ok_or(TreeError::NodeNotFound(uuid))?;
        match self.tree[node_ix] {
            Container::View { handle, .. } => {
                handle.focus();
                self.active_container = Some(node_ix);
            },
            _ => return Err(TreeError::Focus(FocusError::NotAView(uuid)))
        }
        if !self.tree[node_ix].floating() {
            self.tree.set_ancestor_paths_active(node_ix);
        }
        Ok(())
    }
    /// Focus on the container relative to the active container.
    ///
    /// If Horizontal, left and right will move within siblings.
    /// If Vertical, up and down will move within siblings.
    /// Other wise, it moves to the next sibling of the parent container.
    ///
    /// If the edge of the children is hit, it does not wrap around,
    /// but moves between ancestor siblings.
    pub fn move_focus(&mut self, direction: Direction) -> CommandResult {
        if let Some(prev_active_ix) = self.active_container {
            let active_id = self.tree[prev_active_ix].get_id();
            if let Some(fullscreen_id) = try!(self.in_fullscreen_workspace(active_id)) {
                return Err(TreeError::Focus(
                    FocusError::BlockedByFullscreen(active_id, fullscreen_id)))
            }
            let new_active_ix = self.move_focus_recurse(prev_active_ix, direction)
                .unwrap_or(prev_active_ix);
            try!(self.set_active_node(new_active_ix));
            match self.tree[self.active_container.unwrap()] {
                Container::View { ref handle, .. } => handle.focus(),
                _ => warn!("move_focus returned a non-view, cannot focus")
            }
        } else {
            return Err(TreeError::NoActiveContainer)
        }
        self.validate();
        Ok(())
    }

    fn move_focus_recurse(&mut self, node_ix: NodeIndex, direction: Direction)
                          -> Result<NodeIndex, TreeError> {
        match self.tree[node_ix].get_type() {
            ContainerType::View | ContainerType::Container => { /* continue */ },
            _ => return Err(TreeError::UuidWrongType(self.tree[node_ix].get_id(),
                                                     vec!(ContainerType::View, ContainerType::Container)))
        }
        let parent_ix = self.tree.parent_of(node_ix)
            .expect("Active ix had no parent");
        match self.tree[parent_ix] {
            Container::Container { layout, .. } => {
                match (layout, direction) {
                    (Layout::Horizontal, Direction::Left) |
                    (Layout::Horizontal, Direction::Right) |
                    (Layout::Vertical, Direction::Up) |
                    (Layout::Vertical, Direction::Down) => {
                        let siblings = self.tree.children_of(parent_ix);
                        let cur_index = siblings.iter().position(|node| {
                            *node == node_ix
                        }).expect("Could not find self in parent");
                        let maybe_new_index = match direction {
                            Direction::Right | Direction::Down => {
                                cur_index.checked_add(1)
                            }
                            Direction::Left  | Direction::Up => {
                                cur_index.checked_sub(1)
                            }
                        };
                        if maybe_new_index.is_some() &&
                            maybe_new_index.unwrap() < siblings.len() {
                                // There is a sibling to move to.
                                let new_index = maybe_new_index.unwrap();
                                let new_active_ix = siblings[new_index];
                                match self.tree[new_active_ix].get_type() {
                                    ContainerType::Container => {
                                        let path_ix = self.tree.follow_path(new_active_ix);
                                        // If the path wasn't complete, find the first view and focus on that
                                        let node_ix = try!(self.tree.descendant_of_type(path_ix, ContainerType::View)
                                                           .map_err(|err| TreeError::PetGraph(err)));
                                        let parent_ix = try!(self.tree.parent_of(node_ix)
                                                             .map_err(|err| TreeError::PetGraph(err)));
                                        match self.tree[node_ix].get_type() {
                                            ContainerType::View | ContainerType::Container => {},
                                            _ => panic!("Following path did not lead to a container or a view!")
                                        }
                                        trace!("Moving to different view {:?} in container {:?}",
                                                self.tree[node_ix], self.tree[parent_ix]);
                                        return Ok(node_ix);
                                    },
                                    ContainerType::View => {
                                        trace!("Moving to other view {:?}", self.tree[new_active_ix]);
                                        return Ok(new_active_ix)
                                    },
                                    _ => unreachable!()
                                };
                            }
                    },
                    _ => { /* We are moving out of siblings, recurse */ }
                }
            }
            Container::Workspace { .. } => {
                return Err(TreeError::Focus(FocusError::ReachedLimit(parent_ix)));
            }
            _ => unreachable!()
        }
        let parent_ix = self.tree.parent_of(node_ix)
            .expect("Node had no parent");
        return self.move_focus_recurse(parent_ix, direction);
    }

    /// Updates the current active container to be the next container or view
    /// to focus on after the previous view/container was moved/removed.
    ///
    /// A new view will tried to be set, starting with the children of the
    /// parent node. If a view cannot be found there, it starts climbing the
    /// tree until either a view is found or the workspace is (in which case
    /// it set the active container to the root container of the workspace)
    ///
    /// If there are siblings, the chosen node is the one with the lowest
    /// active count.
    ///
    /// If there is a fullscreen container in this workspace, that is focused on next,
    /// with the active path updated accordingly.
    pub fn focus_on_next_container(&mut self, mut parent_ix: NodeIndex) {
        let last_ix = self.tree.active_path().last()
            .expect("Active path did not lead anywhere").0;
        let id = self.tree[last_ix].get_id();
        match self.in_fullscreen_workspace(id) {
            Ok(Some(fullscreen_id)) => {
                self.set_active_container(fullscreen_id)
                    .unwrap_or_else(|err| {
                        error!("Could not set {:?} to be the active container, \
                                even though it's the fullscreen container! {:?}",
                               fullscreen_id, err);
                        error!("{:#?}", self);
                        panic!("Could not set the fullscreen container to be \
                                the active container!")
                    });
                return;
            },
            _ => {}
        }
        match self.tree[last_ix] {
            Container::View { handle, .. } => {
                handle.focus();
                self.active_container = Some(last_ix);
                return
            },
            Container::Container { .. } => {
                parent_ix = last_ix;
            },
            _ => {}
        }
        while self.tree.node_type(parent_ix)
            .expect("Node not part of the tree") != ContainerType::Workspace {
                if let Some(node_ix) = self.tree.lowest_active_view(parent_ix) {
                    match self.tree[node_ix] {
                        Container::View { .. } => {
                            trace!("Active container set to view at {:?}", node_ix);
                            let id = self.tree[node_ix].get_id();
                            self.set_active_container(id)
                                .expect("Could not set active container");
                            return;
                        },
                        _ => {}
                    }
                }
                parent_ix = self.tree.ancestor_of_type(parent_ix,
                                                       ContainerType::Container)
                    .unwrap_or_else(|_| {
                        self.tree.ancestor_of_type(parent_ix, ContainerType::Workspace)
                            .expect("Container was not part of a workspace")
                    });
        }
        // If this is reached, parent is workspace
        let container_ix = self.tree.children_of(parent_ix).get(0).cloned();
        if container_ix.is_none() {
            trace!("There were no other containers to focus on, \
                    focusing on nothing in particular!");
            return;
        }
        let container_ix = container_ix.unwrap();
        let root_c_children = self.tree.grounded_children(container_ix);
        if root_c_children.len() > 0 {
            // Only searches first child of root container, can't be floating view.
            let new_active_ix = self.tree.descendant_of_type(root_c_children[0],
                                                                ContainerType::View)
                .unwrap_or(root_c_children[0]);
            let id = self.tree[new_active_ix].get_id();
            self.set_active_container(id)
                .expect("Could not set active container");
            match self.tree[new_active_ix] {
                Container::View { ref handle, .. } => {
                    debug!("Focusing on {:?}", handle);
                    handle.focus();
                },
                Container::Container { .. } => {
                    debug!("No view found, focusing on nothing in workspace {:?}", parent_ix);
                    WlcView::root().focus();
                }
                _ => unreachable!()
            };
            return;
        } else {
            let floating_children = self.tree.floating_children(container_ix);
            for child_ix in floating_children {
                if let Ok(view_ix) = self.tree.descendant_of_type(child_ix,
                                                                    ContainerType::View) {
                    match self.tree[view_ix] {
                        Container::View { handle, id, .. } => {
                            debug!("Floating view found, focusing on {:#?}", handle);
                            handle.focus();
                            self.set_active_container(id)
                                .expect("Could not set active container");
                            return;
                        },
                        _ => unreachable!()
                    };
                }
            }
        }
        trace!("Active container set to container {:?}", container_ix);
        let id = self.tree[container_ix].get_id();
        self.set_active_container(id)
            .expect("Could not set active container");

        // Update focus to new container
        self.get_active_container().map(|con| match *con {
            Container::View { ref handle, .. } => handle.focus(),
            Container::Container { .. } => WlcView::root().focus(),
            _ => panic!("Active container not view or container!")
        });
    }

    /// If the currently focused view is floating, then the non-floating at the end of
    /// the path becomes the focused view. Otherwise, the first floating view becomes
    /// the focused view.
    ///
    /// If there is no currently focused view, does nothing.
    pub fn toggle_floating_focus(&mut self) -> CommandResult {
        let active_ix = try!(self.active_container.ok_or(TreeError::NoActiveContainer));
        let id; let floating;
        {
            let container = &self.tree[active_ix];
            id = container.get_id();
            floating = container.floating();
        }
        if let Some(fullscreen_id) = try!(self.in_fullscreen_workspace(id)) {
            return Err(TreeError::Focus(FocusError::BlockedByFullscreen(id, fullscreen_id)))
        }
        if floating {
            let root_ix = self.tree.root_ix();
            let new_ix = self.tree.follow_path(root_ix);
            match self.tree[new_ix].get_type() {
                ContainerType::View | ContainerType::Container  => {
                    try!(self.set_active_node(new_ix));
                    Ok(())
                },
                type_ => {
                    error!("Path lead to the wrong container, {:#?}\n{:#?}\n{:#?}",
                           active_ix, type_, self);
                    panic!("toggle_floating_focused: bad path");
                }
            }
        } else {
            // Current view is not floating, gotta focus on floating view
            let root_c_ix = self.root_container_ix()
                .expect("No root container ancestor of active container");
            let floating_children = self.tree.floating_children(root_c_ix);
            if floating_children.len() > 0 {
                try!(self.set_active_node(floating_children[0]));
            }
            Ok(())
        }
    }
}
<|endoftext|>


<commit_before>/* Multi-producer/single-consumer queue
 * Copyright (c) 2010-2011 Dmitry Vyukov. All rights reserved.
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *    1. Redistributions of source code must retain the above copyright notice,
 *       this list of conditions and the following disclaimer.
 *
 *    2. Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY DMITRY VYUKOV "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
 * SHALL DMITRY VYUKOV OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * The views and conclusions contained in the software and documentation are
 * those of the authors and should not be interpreted as representing official
 * policies, either expressed or implied, of Dmitry Vyukov.
 */

//! A mostly lock-free multi-producer, single consumer queue.

use unstable::sync::UnsafeArc;
use unstable::atomics::{AtomicPtr,Relaxed,Release,Acquire};
use ptr::{mut_null, to_mut_unsafe_ptr};
use cast;
use option::*;
use clone::Clone;
use default::Default;
use kinds::Send;
use fmt;

struct Node<T> {
    next: AtomicPtr<Node<T>>,
    value: Option<T>,
}

impl<T> Node<T> {
    fn new(value: T) -> Node<T> {
        Node{next: AtomicPtr::new(mut_null()), value: Some(value)}
    }
}

impl<T> Default for Node<T> {
    fn default() -> Node<T> {
        Node{next: AtomicPtr::new(mut_null()), value: None}
    }
}

struct State<T> {
    pad0: [u8, ..64],
    head: AtomicPtr<Node<T>>,
    pad1: [u8, ..64],
    stub: Node<T>,
    pad2: [u8, ..64],
    tail: *mut Node<T>,
    pad3: [u8, ..64],
}

struct Queue<T> {
    priv state: UnsafeArc<State<T>>,
}

impl<T: Send> Clone for Queue<T> {
    fn clone(&self) -> Queue<T> {
        Queue {
            state: self.state.clone()
        }
    }
}

impl<T: Send> fmt::Default for Queue<T> {
    fn fmt(value: &Queue<T>, f: &mut fmt::Formatter) {
        write!(f.buf, "Queue({})", value.state.get());
    }
}

impl<T: Send> Queue<T> {
    pub fn new() -> Queue<T> {
        let mut q = Queue{state: UnsafeArc::new(State {
            pad0: [0, ..64],
            head: AtomicPtr::new(mut_null()),
            pad1: [0, ..64],
            stub: Default::default(),
            pad2: [0, ..64],
            tail: mut_null(),
            pad3: [0, ..64],
        })};
        let stub = q.get_stub_unsafe();
        q.get_head().store(stub, Relaxed);
        q.set_tail(stub);
        q
    }

    pub fn push(&mut self, value: T) {
        unsafe {
            let node = cast::transmute(~Node::new(value));
            self.push_node(node);
        }
    }

    fn push_node(&mut self, node: *mut Node<T>) {
        unsafe {
            (*node).next.store(mut_null(), Release);
            let prev = self.get_head().swap(node, Relaxed);
            (*prev).next.store(node, Release);
        }
    }

    fn get_stub_unsafe(&mut self) -> *mut Node<T> {
        unsafe { to_mut_unsafe_ptr(&mut (*self.state.get()).stub) }
    }

    fn get_head(&mut self) -> &mut AtomicPtr<Node<T>> {
        unsafe { &mut (*self.state.get()).head }
    }

    fn get_tail(&mut self) -> *mut Node<T> {
        unsafe { (*self.state.get()).tail }
    }

    fn set_tail(&mut self, tail: *mut Node<T>) {
        unsafe { (*self.state.get()).tail = tail }
    }

    pub fn casual_pop(&mut self) -> Option<T> {
        self.pop()
    }

    pub fn pop(&mut self) -> Option<T> {
        unsafe {
            let mut tail = self.get_tail();
            let mut next = (*tail).next.load(Acquire);
            let stub = self.get_stub_unsafe();
            if tail == stub {
                if mut_null() == next {
                    return None
                }
                self.set_tail(next);
                tail = next;
                next = (*next).next.load(Acquire);
            }
            if next != mut_null() {
                let tail: ~Node<T> = cast::transmute(tail);
                self.set_tail(next);
                return tail.value
            }
            let head = self.get_head().load(Relaxed);
            if tail != head {
                return None
            }
            self.push_node(stub);
            next = (*tail).next.load(Acquire);
            if next != mut_null() {
                let tail: ~Node<T> = cast::transmute(tail);
                self.set_tail(next);
                return tail.value
            }
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use prelude::*;
    use option::*;
    use task;
    use comm;
    use super::Queue;

    #[test]
    fn test() {
        let nthreads = 8u;
        let nmsgs = 1000u;
        let mut q = Queue::new();
        assert_eq!(None, q.pop());

        for _ in range(0, nthreads) {
            let (port, chan)  = comm::stream();
            chan.send(q.clone());
            do task::spawn_sched(task::SingleThreaded) {
                let mut q = port.recv();
                for i in range(0, nmsgs) {
                    q.push(i);
                }
            }
        }

        let mut i = 0u;
        loop {
            match q.pop() {
                None => {},
                Some(_) => {
                    i += 1;
                    if i == nthreads*nmsgs { break }
                }
            }
        }
    }
}

<commit_msg>clean up<commit_after>/* Multi-producer/single-consumer queue
 * Copyright (c) 2010-2011 Dmitry Vyukov. All rights reserved.
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *    1. Redistributions of source code must retain the above copyright notice,
 *       this list of conditions and the following disclaimer.
 *
 *    2. Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY DMITRY VYUKOV "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
 * SHALL DMITRY VYUKOV OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * The views and conclusions contained in the software and documentation are
 * those of the authors and should not be interpreted as representing official
 * policies, either expressed or implied, of Dmitry Vyukov.
 */

//! A mostly lock-free multi-producer, single consumer queue.

use unstable::sync::UnsafeArc;
use unstable::atomics::{AtomicPtr,Relaxed,Release,Acquire};
use ptr::{mut_null, to_mut_unsafe_ptr};
use cast;
use option::*;
use clone::Clone;
use default::Default;
use kinds::Send;
use fmt;

struct Node<T> {
    next: AtomicPtr<Node<T>>,
    value: Option<T>,
}

impl<T> Node<T> {
    fn new(value: T) -> Node<T> {
        Node{next: AtomicPtr::new(mut_null()), value: Some(value)}
    }
}

impl<T> Default for Node<T> {
    fn default() -> Node<T> {
        Node{next: AtomicPtr::new(mut_null()), value: None}
    }
}

struct State<T> {
    pad0: [u8, ..64],
    head: AtomicPtr<Node<T>>,
    pad1: [u8, ..64],
    stub: Node<T>,
    pad2: [u8, ..64],
    tail: *mut Node<T>,
    pad3: [u8, ..64],
}

struct Queue<T> {
    priv state: UnsafeArc<State<T>>,
}

impl<T: Send> Clone for Queue<T> {
    fn clone(&self) -> Queue<T> {
        Queue {
            state: self.state.clone()
        }
    }
}

impl<T: Send> fmt::Default for Queue<T> {
    fn fmt(value: &Queue<T>, f: &mut fmt::Formatter) {
        write!(f.buf, "Queue({})", value.state.get());
    }
}

impl<T: Send> State<T> {
    pub fn new() -> State<T> {
        let mut state = State {
            pad0: [0, ..64],
            head: AtomicPtr::new(mut_null()),
            pad1: [0, ..64],
            stub: Default::default(),
            pad2: [0, ..64],
            tail: mut_null(),
            pad3: [0, ..64],
        };
        let stub = state.get_stub_unsafe();
        state.head.store(stub, Relaxed);
        state.tail = stub;
        state
    }

    fn get_stub_unsafe(&mut self) -> *mut Node<T> {
        unsafe { to_mut_unsafe_ptr(&mut self.stub) }
    }

    fn push(&mut self, value: T) {
        unsafe {
            let node = cast::transmute(~Node::new(value));
            self.push_node(node);
        }
    }

    fn push_node(&mut self, node: *mut Node<T>) {
        unsafe {
            (*node).next.store(mut_null(), Release);
            let prev = self.head.swap(node, Relaxed);
            (*prev).next.store(node, Release);
        }
    }

    fn pop(&mut self) -> Option<T> {
        unsafe {
            let mut tail = self.tail;
            let mut next = (*tail).next.load(Acquire);
            let stub = self.get_stub_unsafe();
            if tail == stub {
                if mut_null() == next {
                    return None
                }
                self.tail = next;
                tail = next;
                next = (*next).next.load(Acquire);
            }
            if next != mut_null() {
                let tail: ~Node<T> = cast::transmute(tail);
                self.tail = next;
                return tail.value
            }
            let head = self.head.load(Relaxed);
            if tail != head {
                return None
            }
            self.push_node(stub);
            next = (*tail).next.load(Acquire);
            if next != mut_null() {
                let tail: ~Node<T> = cast::transmute(tail);
                self.tail = next;
                return tail.value
            }
        }
        None
    }
}

impl<T: Send> Queue<T> {
    pub fn new() -> Queue<T> {
        Queue{state: UnsafeArc::new(State::new())}
    }

    pub fn push(&mut self, value: T) {
        unsafe { (*self.state.get()).push(value) }
    }

    pub fn casual_pop(&mut self) -> Option<T> {
        unsafe { (*self.state.get()).pop() }
    }

    pub fn pop(&mut self) -> Option<T> {
        unsafe{ (*self.state.get()).pop() }
    }
}

#[cfg(test)]
mod tests {
    use prelude::*;
    use option::*;
    use task;
    use comm;
    use super::Queue;

    #[test]
    fn test() {
        let nthreads = 8u;
        let nmsgs = 1000u;
        let mut q = Queue::new();
        assert_eq!(None, q.pop());

        for _ in range(0, nthreads) {
            let (port, chan)  = comm::stream();
            chan.send(q.clone());
            do task::spawn_sched(task::SingleThreaded) {
                let mut q = port.recv();
                for i in range(0, nmsgs) {
                    q.push(i);
                }
            }
        }

        let mut i = 0u;
        loop {
            match q.pop() {
                None => {},
                Some(_) => {
                    i += 1;
                    if i == nthreads*nmsgs { break }
                }
            }
        }
    }
}

<|endoftext|>


<commit_before>use std::fmt;
use std::mem::transmute;
use std::io::{Cursor, Read, Write};
use std::default::Default;
use std::iter::FromIterator;
use std::string::{String, FromUtf8Error};
use std::result::Result as StdResult;
use byteorder::{ByteOrder, NetworkEndian};
use bytes::BufMut;

use rand;

use error::{Error, Result};
use super::coding::{OpCode, Control, Data, CloseCode};

fn apply_mask(buf: &mut [u8], mask: &[u8; 4]) {
    let iter = buf.iter_mut().zip(mask.iter().cycle());
    for (byte, &key) in iter {
        *byte ^= key
    }
}

#[inline]
fn generate_mask() -> [u8; 4] {
    unsafe { transmute(rand::random::<u32>()) }
}

/// A struct representing a WebSocket frame.
#[derive(Debug, Clone)]
pub struct Frame {
    finished: bool,
    rsv1: bool,
    rsv2: bool,
    rsv3: bool,
    opcode: OpCode,

    mask: Option<[u8; 4]>,

    payload: Vec<u8>,
}

impl Frame {

    /// Get the length of the frame.
    /// This is the length of the header + the length of the payload.
    #[inline]
    pub fn len(&self) -> usize {
        let mut header_length = 2;
        let payload_len = self.payload().len();
        if payload_len > 125 {
            if payload_len <= u16::max_value() as usize {
                header_length += 2;
            } else {
                header_length += 8;
            }
        }

        if self.is_masked() {
            header_length += 4;
        }

        header_length + payload_len
    }

    /// Test whether the frame is a final frame.
    #[inline]
    pub fn is_final(&self) -> bool {
        self.finished
    }

    /// Test whether the first reserved bit is set.
    #[inline]
    pub fn has_rsv1(&self) -> bool {
        self.rsv1
    }

    /// Test whether the second reserved bit is set.
    #[inline]
    pub fn has_rsv2(&self) -> bool {
        self.rsv2
    }

    /// Test whether the third reserved bit is set.
    #[inline]
    pub fn has_rsv3(&self) -> bool {
        self.rsv3
    }

    /// Get the OpCode of the frame.
    #[inline]
    pub fn opcode(&self) -> OpCode {
        self.opcode
    }

    /// Get a reference to the frame's payload.
    #[inline]
    pub fn payload(&self) -> &Vec<u8> {
        &self.payload
    }

    // Test whether the frame is masked.
    #[doc(hidden)]
    #[inline]
    pub fn is_masked(&self) -> bool {
        self.mask.is_some()
    }

    // Get an optional reference to the frame's mask.
    #[doc(hidden)]
    #[allow(dead_code)]
    #[inline]
    pub fn mask(&self) -> Option<&[u8; 4]> {
        self.mask.as_ref()
    }

    /// Make this frame a final frame.
    #[allow(dead_code)]
    #[inline]
    pub fn set_final(&mut self, is_final: bool) -> &mut Frame {
        self.finished = is_final;
        self
    }

    /// Set the first reserved bit.
    #[inline]
    pub fn set_rsv1(&mut self, has_rsv1: bool) -> &mut Frame {
        self.rsv1 = has_rsv1;
        self
    }

    /// Set the second reserved bit.
    #[inline]
    pub fn set_rsv2(&mut self, has_rsv2: bool) -> &mut Frame {
        self.rsv2 = has_rsv2;
        self
    }

    /// Set the third reserved bit.
    #[inline]
    pub fn set_rsv3(&mut self, has_rsv3: bool) -> &mut Frame {
        self.rsv3 = has_rsv3;
        self
    }

    /// Set the OpCode.
    #[allow(dead_code)]
    #[inline]
    pub fn set_opcode(&mut self, opcode: OpCode) -> &mut Frame {
        self.opcode = opcode;
        self
    }

    /// Edit the frame's payload.
    #[allow(dead_code)]
    #[inline]
    pub fn payload_mut(&mut self) -> &mut Vec<u8> {
        &mut self.payload
    }

    // Generate a new mask for this frame.
    //
    // This method simply generates and stores the mask. It does not change the payload data.
    // Instead, the payload data will be masked with the generated mask when the frame is sent
    // to the other endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn set_mask(&mut self) -> &mut Frame {
        self.mask = Some(generate_mask());
        self
    }

    // This method unmasks the payload and should only be called on frames that are actually
    // masked. In other words, those frames that have just been received from a client endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn remove_mask(&mut self) {
        self.mask.and_then(|mask| {
            Some(apply_mask(&mut self.payload, &mask))
        });
        self.mask = None;
    }

    /// Consume the frame into its payload as binary.
    #[inline]
    pub fn into_data(self) -> Vec<u8> {
        self.payload
    }

    /// Consume the frame into its payload as string.
    #[inline]
    pub fn into_string(self) -> StdResult<String, FromUtf8Error> {
        String::from_utf8(self.payload)
    }

     /// Consume the frame into a closing frame.
    #[inline]
    pub fn into_close(self) -> Result<Option<(CloseCode, String)>> {
        match self.payload.len() {
            0 => Ok(None),
            1 => Err(Error::Protocol("Invalid close sequence".into())),
            _ => {
                let mut data = self.payload;
                let code = NetworkEndian::read_u16(&data[0..2]).into();
                data.drain(0..2);
                let text = String::from_utf8(data)?;
                Ok(Some((code, text)))
            }
        }
    }

    /// Create a new data frame.
    #[inline]
    pub fn message(data: Vec<u8>, code: OpCode, finished: bool) -> Frame {
        debug_assert!(match code {
            OpCode::Data(_) => true,
            _ => false,
        }, "Invalid opcode for data frame.");

        Frame {
            finished: finished,
            opcode: code,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Pong control frame.
    #[inline]
    pub fn pong(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Control(Control::Pong),
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Ping control frame.
    #[inline]
    pub fn ping(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Control(Control::Ping),
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Close control frame.
    #[inline]
    pub fn close(msg: Option<(CloseCode, &str)>) -> Frame {
        let payload = if let Some((code, reason)) = msg {
            let raw: [u8; 2] = unsafe {
                let u: u16 = code.into();
                transmute(u.to_be())
            };
            Vec::from_iter(
                raw[..].iter()
                       .chain(reason.as_bytes().iter())
                       .map(|&b| b))
        } else {
            Vec::new()
        };

        Frame {
            payload: payload,
            .. Frame::default()
        }
    }

    /// Parse the input stream into a frame.
    pub fn parse(cursor: &mut Cursor<Vec<u8>>) -> Result<Option<Frame>> {
        let size = cursor.get_ref().len() as u64 - cursor.position();
        let initial = cursor.position();
        trace!("Position in buffer {}", initial);

        let mut head = [0u8; 2];
        if try!(cursor.read(&mut head)) != 2 {
            cursor.set_position(initial);
            return Ok(None)
        }

        trace!("Parsed headers {:?}", head);

        let first = head[0];
        let second = head[1];
        trace!("First: {:b}", first);
        trace!("Second: {:b}", second);

        let finished = first & 0x80 != 0;

        let rsv1 = first & 0x40 != 0;
        let rsv2 = first & 0x20 != 0;
        let rsv3 = first & 0x10 != 0;

        let opcode = OpCode::from(first & 0x0F);
        trace!("Opcode: {:?}", opcode);

        let masked = second & 0x80 != 0;
        trace!("Masked: {:?}", masked);

        let mut header_length = 2;

        let mut length = (second & 0x7F) as u64;

        if length == 126 {
            let mut length_bytes = [0u8; 2];
            if try!(cursor.read(&mut length_bytes)) != 2 {
                cursor.set_position(initial);
                return Ok(None)
            }

            length = unsafe {
                let mut wide: u16 = transmute(length_bytes);
                wide = u16::from_be(wide);
                wide
            } as u64;
            header_length += 2;
        } else if length == 127 {
            let mut length_bytes = [0u8; 8];
            if try!(cursor.read(&mut length_bytes)) != 8 {
                cursor.set_position(initial);
                return Ok(None)
            }

            unsafe { length = transmute(length_bytes); }
            length = u64::from_be(length);
            header_length += 8;
        }
        trace!("Payload length: {}", length);

        let mask = if masked {
            let mut mask_bytes = [0u8; 4];
            if try!(cursor.read(&mut mask_bytes)) != 4 {
                cursor.set_position(initial);
                return Ok(None)
            } else {
                header_length += 4;
                Some(mask_bytes)
            }
        } else {
            None
        };

        if size < length + header_length {
            cursor.set_position(initial);
            return Ok(None)
        }

        let mut data = Vec::with_capacity(length as usize);
        if length > 0 {
            unsafe {
                try!(cursor.read_exact(data.bytes_mut()));
                data.advance_mut(length as usize);
            }
        }

        // Disallow bad opcode
        match opcode {
            OpCode::Control(Control::Reserved(_)) | OpCode::Data(Data::Reserved(_)) => {
                return Err(Error::Protocol(format!("Encountered invalid opcode: {}", first & 0x0F).into()))
            }
            _ => ()
        }

        let frame = Frame {
            finished: finished,
            rsv1: rsv1,
            rsv2: rsv2,
            rsv3: rsv3,
            opcode: opcode,
            mask: mask,
            payload: data,
        };

        Ok(Some(frame))
    }

    /// Write a frame out to a buffer
    pub fn format<W>(mut self, w: &mut W) -> Result<()>
        where W: Write
    {
        let mut one = 0u8;
        let code: u8 = self.opcode.into();
        if self.is_final() {
            one |= 0x80;
        }
        if self.has_rsv1() {
            one |= 0x40;
        }
        if self.has_rsv2() {
            one |= 0x20;
        }
        if self.has_rsv3() {
            one |= 0x10;
        }
        one |= code;

        let mut two = 0u8;

        if self.is_masked() {
            two |= 0x80;
        }

        if self.payload.len() < 126 {
            two |= self.payload.len() as u8;
            let headers = [one, two];
            try!(w.write(&headers));
        } else if self.payload.len() <= 65535 {
            two |= 126;
            let length_bytes: [u8; 2] = unsafe {
                let short = self.payload.len() as u16;
                transmute(short.to_be())
            };
            let headers = [one, two, length_bytes[0], length_bytes[1]];
            try!(w.write(&headers));
        } else {
            two |= 127;
            let length_bytes: [u8; 8] = unsafe {
                let long = self.payload.len() as u64;
                transmute(long.to_be())
            };
            let headers = [
                one,
                two,
                length_bytes[0],
                length_bytes[1],
                length_bytes[2],
                length_bytes[3],
                length_bytes[4],
                length_bytes[5],
                length_bytes[6],
                length_bytes[7],
            ];
            try!(w.write(&headers));
        }

        if self.is_masked() {
            let mask = self.mask.take().unwrap();
            apply_mask(&mut self.payload, &mask);
            try!(w.write(&mask));
        }

        try!(w.write(&self.payload));
        Ok(())
    }
}

impl Default for Frame {
    fn default() -> Frame {
        Frame {
            finished: true,
            rsv1: false,
            rsv2: false,
            rsv3: false,
            opcode: OpCode::Control(Control::Close),
            mask: None,
            payload: Vec::new(),
        }
    }
}

impl fmt::Display for Frame {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
            "
<FRAME>
final: {}
reserved: {} {} {}
opcode: {}
length: {}
payload length: {}
payload: 0x{}
            ",
            self.finished,
            self.rsv1,
            self.rsv2,
            self.rsv3,
            self.opcode,
            // self.mask.map(|mask| format!("{:?}", mask)).unwrap_or("NONE".into()),
            self.len(),
            self.payload.len(),
            self.payload.iter().map(|byte| format!("{:x}", byte)).collect::<String>())
    }
}

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

    use super::super::coding::{OpCode, Data};
    use std::io::Cursor;

    #[test]
    fn parse() {
        let mut raw: Cursor<Vec<u8>> = Cursor::new(vec![
            0x82, 0x07, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07
        ]);
        let frame = Frame::parse(&mut raw).unwrap().unwrap();
        assert_eq!(frame.into_data(), vec![ 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07 ]);
    }

    #[test]
    fn format() {
        let frame = Frame::ping(vec![0x01, 0x02]);
        let mut buf = Vec::with_capacity(frame.len());
        frame.format(&mut buf).unwrap();
        assert_eq!(buf, vec![0x89, 0x02, 0x01, 0x02]);
    }

    #[test]
    fn display() {
        let f = Frame::message("hi there".into(), OpCode::Data(Data::Text), true);
        let view = format!("{}", f);
        view.contains("payload:");
    }
}
<commit_msg>Remove unnecessary transmute in generate_mask()<commit_after>use std::fmt;
use std::mem::transmute;
use std::io::{Cursor, Read, Write};
use std::default::Default;
use std::iter::FromIterator;
use std::string::{String, FromUtf8Error};
use std::result::Result as StdResult;
use byteorder::{ByteOrder, NetworkEndian};
use bytes::BufMut;

use rand;

use error::{Error, Result};
use super::coding::{OpCode, Control, Data, CloseCode};

fn apply_mask(buf: &mut [u8], mask: &[u8; 4]) {
    let iter = buf.iter_mut().zip(mask.iter().cycle());
    for (byte, &key) in iter {
        *byte ^= key
    }
}

#[inline]
fn generate_mask() -> [u8; 4] {
    rand::random()
}

/// A struct representing a WebSocket frame.
#[derive(Debug, Clone)]
pub struct Frame {
    finished: bool,
    rsv1: bool,
    rsv2: bool,
    rsv3: bool,
    opcode: OpCode,

    mask: Option<[u8; 4]>,

    payload: Vec<u8>,
}

impl Frame {

    /// Get the length of the frame.
    /// This is the length of the header + the length of the payload.
    #[inline]
    pub fn len(&self) -> usize {
        let mut header_length = 2;
        let payload_len = self.payload().len();
        if payload_len > 125 {
            if payload_len <= u16::max_value() as usize {
                header_length += 2;
            } else {
                header_length += 8;
            }
        }

        if self.is_masked() {
            header_length += 4;
        }

        header_length + payload_len
    }

    /// Test whether the frame is a final frame.
    #[inline]
    pub fn is_final(&self) -> bool {
        self.finished
    }

    /// Test whether the first reserved bit is set.
    #[inline]
    pub fn has_rsv1(&self) -> bool {
        self.rsv1
    }

    /// Test whether the second reserved bit is set.
    #[inline]
    pub fn has_rsv2(&self) -> bool {
        self.rsv2
    }

    /// Test whether the third reserved bit is set.
    #[inline]
    pub fn has_rsv3(&self) -> bool {
        self.rsv3
    }

    /// Get the OpCode of the frame.
    #[inline]
    pub fn opcode(&self) -> OpCode {
        self.opcode
    }

    /// Get a reference to the frame's payload.
    #[inline]
    pub fn payload(&self) -> &Vec<u8> {
        &self.payload
    }

    // Test whether the frame is masked.
    #[doc(hidden)]
    #[inline]
    pub fn is_masked(&self) -> bool {
        self.mask.is_some()
    }

    // Get an optional reference to the frame's mask.
    #[doc(hidden)]
    #[allow(dead_code)]
    #[inline]
    pub fn mask(&self) -> Option<&[u8; 4]> {
        self.mask.as_ref()
    }

    /// Make this frame a final frame.
    #[allow(dead_code)]
    #[inline]
    pub fn set_final(&mut self, is_final: bool) -> &mut Frame {
        self.finished = is_final;
        self
    }

    /// Set the first reserved bit.
    #[inline]
    pub fn set_rsv1(&mut self, has_rsv1: bool) -> &mut Frame {
        self.rsv1 = has_rsv1;
        self
    }

    /// Set the second reserved bit.
    #[inline]
    pub fn set_rsv2(&mut self, has_rsv2: bool) -> &mut Frame {
        self.rsv2 = has_rsv2;
        self
    }

    /// Set the third reserved bit.
    #[inline]
    pub fn set_rsv3(&mut self, has_rsv3: bool) -> &mut Frame {
        self.rsv3 = has_rsv3;
        self
    }

    /// Set the OpCode.
    #[allow(dead_code)]
    #[inline]
    pub fn set_opcode(&mut self, opcode: OpCode) -> &mut Frame {
        self.opcode = opcode;
        self
    }

    /// Edit the frame's payload.
    #[allow(dead_code)]
    #[inline]
    pub fn payload_mut(&mut self) -> &mut Vec<u8> {
        &mut self.payload
    }

    // Generate a new mask for this frame.
    //
    // This method simply generates and stores the mask. It does not change the payload data.
    // Instead, the payload data will be masked with the generated mask when the frame is sent
    // to the other endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn set_mask(&mut self) -> &mut Frame {
        self.mask = Some(generate_mask());
        self
    }

    // This method unmasks the payload and should only be called on frames that are actually
    // masked. In other words, those frames that have just been received from a client endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn remove_mask(&mut self) {
        self.mask.and_then(|mask| {
            Some(apply_mask(&mut self.payload, &mask))
        });
        self.mask = None;
    }

    /// Consume the frame into its payload as binary.
    #[inline]
    pub fn into_data(self) -> Vec<u8> {
        self.payload
    }

    /// Consume the frame into its payload as string.
    #[inline]
    pub fn into_string(self) -> StdResult<String, FromUtf8Error> {
        String::from_utf8(self.payload)
    }

     /// Consume the frame into a closing frame.
    #[inline]
    pub fn into_close(self) -> Result<Option<(CloseCode, String)>> {
        match self.payload.len() {
            0 => Ok(None),
            1 => Err(Error::Protocol("Invalid close sequence".into())),
            _ => {
                let mut data = self.payload;
                let code = NetworkEndian::read_u16(&data[0..2]).into();
                data.drain(0..2);
                let text = String::from_utf8(data)?;
                Ok(Some((code, text)))
            }
        }
    }

    /// Create a new data frame.
    #[inline]
    pub fn message(data: Vec<u8>, code: OpCode, finished: bool) -> Frame {
        debug_assert!(match code {
            OpCode::Data(_) => true,
            _ => false,
        }, "Invalid opcode for data frame.");

        Frame {
            finished: finished,
            opcode: code,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Pong control frame.
    #[inline]
    pub fn pong(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Control(Control::Pong),
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Ping control frame.
    #[inline]
    pub fn ping(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Control(Control::Ping),
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Close control frame.
    #[inline]
    pub fn close(msg: Option<(CloseCode, &str)>) -> Frame {
        let payload = if let Some((code, reason)) = msg {
            let raw: [u8; 2] = unsafe {
                let u: u16 = code.into();
                transmute(u.to_be())
            };
            Vec::from_iter(
                raw[..].iter()
                       .chain(reason.as_bytes().iter())
                       .map(|&b| b))
        } else {
            Vec::new()
        };

        Frame {
            payload: payload,
            .. Frame::default()
        }
    }

    /// Parse the input stream into a frame.
    pub fn parse(cursor: &mut Cursor<Vec<u8>>) -> Result<Option<Frame>> {
        let size = cursor.get_ref().len() as u64 - cursor.position();
        let initial = cursor.position();
        trace!("Position in buffer {}", initial);

        let mut head = [0u8; 2];
        if try!(cursor.read(&mut head)) != 2 {
            cursor.set_position(initial);
            return Ok(None)
        }

        trace!("Parsed headers {:?}", head);

        let first = head[0];
        let second = head[1];
        trace!("First: {:b}", first);
        trace!("Second: {:b}", second);

        let finished = first & 0x80 != 0;

        let rsv1 = first & 0x40 != 0;
        let rsv2 = first & 0x20 != 0;
        let rsv3 = first & 0x10 != 0;

        let opcode = OpCode::from(first & 0x0F);
        trace!("Opcode: {:?}", opcode);

        let masked = second & 0x80 != 0;
        trace!("Masked: {:?}", masked);

        let mut header_length = 2;

        let mut length = (second & 0x7F) as u64;

        if length == 126 {
            let mut length_bytes = [0u8; 2];
            if try!(cursor.read(&mut length_bytes)) != 2 {
                cursor.set_position(initial);
                return Ok(None)
            }

            length = unsafe {
                let mut wide: u16 = transmute(length_bytes);
                wide = u16::from_be(wide);
                wide
            } as u64;
            header_length += 2;
        } else if length == 127 {
            let mut length_bytes = [0u8; 8];
            if try!(cursor.read(&mut length_bytes)) != 8 {
                cursor.set_position(initial);
                return Ok(None)
            }

            unsafe { length = transmute(length_bytes); }
            length = u64::from_be(length);
            header_length += 8;
        }
        trace!("Payload length: {}", length);

        let mask = if masked {
            let mut mask_bytes = [0u8; 4];
            if try!(cursor.read(&mut mask_bytes)) != 4 {
                cursor.set_position(initial);
                return Ok(None)
            } else {
                header_length += 4;
                Some(mask_bytes)
            }
        } else {
            None
        };

        if size < length + header_length {
            cursor.set_position(initial);
            return Ok(None)
        }

        let mut data = Vec::with_capacity(length as usize);
        if length > 0 {
            unsafe {
                try!(cursor.read_exact(data.bytes_mut()));
                data.advance_mut(length as usize);
            }
        }

        // Disallow bad opcode
        match opcode {
            OpCode::Control(Control::Reserved(_)) | OpCode::Data(Data::Reserved(_)) => {
                return Err(Error::Protocol(format!("Encountered invalid opcode: {}", first & 0x0F).into()))
            }
            _ => ()
        }

        let frame = Frame {
            finished: finished,
            rsv1: rsv1,
            rsv2: rsv2,
            rsv3: rsv3,
            opcode: opcode,
            mask: mask,
            payload: data,
        };

        Ok(Some(frame))
    }

    /// Write a frame out to a buffer
    pub fn format<W>(mut self, w: &mut W) -> Result<()>
        where W: Write
    {
        let mut one = 0u8;
        let code: u8 = self.opcode.into();
        if self.is_final() {
            one |= 0x80;
        }
        if self.has_rsv1() {
            one |= 0x40;
        }
        if self.has_rsv2() {
            one |= 0x20;
        }
        if self.has_rsv3() {
            one |= 0x10;
        }
        one |= code;

        let mut two = 0u8;

        if self.is_masked() {
            two |= 0x80;
        }

        if self.payload.len() < 126 {
            two |= self.payload.len() as u8;
            let headers = [one, two];
            try!(w.write(&headers));
        } else if self.payload.len() <= 65535 {
            two |= 126;
            let length_bytes: [u8; 2] = unsafe {
                let short = self.payload.len() as u16;
                transmute(short.to_be())
            };
            let headers = [one, two, length_bytes[0], length_bytes[1]];
            try!(w.write(&headers));
        } else {
            two |= 127;
            let length_bytes: [u8; 8] = unsafe {
                let long = self.payload.len() as u64;
                transmute(long.to_be())
            };
            let headers = [
                one,
                two,
                length_bytes[0],
                length_bytes[1],
                length_bytes[2],
                length_bytes[3],
                length_bytes[4],
                length_bytes[5],
                length_bytes[6],
                length_bytes[7],
            ];
            try!(w.write(&headers));
        }

        if self.is_masked() {
            let mask = self.mask.take().unwrap();
            apply_mask(&mut self.payload, &mask);
            try!(w.write(&mask));
        }

        try!(w.write(&self.payload));
        Ok(())
    }
}

impl Default for Frame {
    fn default() -> Frame {
        Frame {
            finished: true,
            rsv1: false,
            rsv2: false,
            rsv3: false,
            opcode: OpCode::Control(Control::Close),
            mask: None,
            payload: Vec::new(),
        }
    }
}

impl fmt::Display for Frame {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
            "
<FRAME>
final: {}
reserved: {} {} {}
opcode: {}
length: {}
payload length: {}
payload: 0x{}
            ",
            self.finished,
            self.rsv1,
            self.rsv2,
            self.rsv3,
            self.opcode,
            // self.mask.map(|mask| format!("{:?}", mask)).unwrap_or("NONE".into()),
            self.len(),
            self.payload.len(),
            self.payload.iter().map(|byte| format!("{:x}", byte)).collect::<String>())
    }
}

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

    use super::super::coding::{OpCode, Data};
    use std::io::Cursor;

    #[test]
    fn parse() {
        let mut raw: Cursor<Vec<u8>> = Cursor::new(vec![
            0x82, 0x07, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07
        ]);
        let frame = Frame::parse(&mut raw).unwrap().unwrap();
        assert_eq!(frame.into_data(), vec![ 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07 ]);
    }

    #[test]
    fn format() {
        let frame = Frame::ping(vec![0x01, 0x02]);
        let mut buf = Vec::with_capacity(frame.len());
        frame.format(&mut buf).unwrap();
        assert_eq!(buf, vec![0x89, 0x02, 0x01, 0x02]);
    }

    #[test]
    fn display() {
        let f = Frame::message("hi there".into(), OpCode::Data(Data::Text), true);
        let view = format!("{}", f);
        view.contains("payload:");
    }
}
<|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.

//! packet metadata fetcher

use metadata::Metadata;

use errors::*;

use retry;
use std::fs::File;
use std::str::FromStr;

use util;

use network;
use network::{Interface,Device,Section,NetworkRoute};
use pnet::util::MacAddr;
use std::net::{IpAddr,Ipv4Addr,Ipv6Addr};
use ipnetwork;
use ipnetwork::{IpNetwork,Ipv4Network,Ipv6Network};

#[derive(Clone,Deserialize)]
struct PacketData {
    id: String,
    hostname: String,
    iqn: String,
    plan: String,
    facility: String,
    tags: Vec<String>,
    ssh_keys: Vec<String>,
    network: PacketNetworkInfo,

    error: Option<String>,
    phone_home_url: String,
}

#[derive(Clone,Deserialize)]
struct PacketNetworkInfo {
    interfaces: Vec<PacketInterfaceInfo>,
    addresses: Vec<PacketAddressInfo>,
    bonding: PacketBondingMode,
}

#[derive(Clone,Deserialize)]
struct PacketBondingMode {
    mode: u32,
}

#[derive(Clone,Deserialize)]
struct PacketInterfaceInfo {
    name: String,
    mac: String,
}

#[derive(Clone,Deserialize)]
struct PacketAddressInfo {
    id: String,
    address_family: i32,
    public: bool,
    management: bool,
    address: IpAddr,
    netmask: IpAddr,
    gateway: IpAddr,
}

pub fn fetch_metadata() -> Result<Metadata> {
    let client = retry::Client::new()?;
    let data: PacketData = client.get(retry::Json, "http://metadata.packet.net/metadata".to_owned()).send()?
        .ok_or("not found")?;

    let (interfaces,network_devices) = parse_network(&data.network)?;

    let attrs = get_attrs(&data)?;

    let mut m = Metadata::builder()
        .add_ssh_keys(data.ssh_keys)?
        .set_hostname(data.hostname);

    for (key,val) in attrs {
        m = m.add_attribute(key,val);
    }
    for iface in interfaces {
        m = m.add_network_interface(iface);
    }
    for netdev in network_devices {
        m = m.add_network_device(netdev);
    }
    Ok(m.build())
}

fn get_dns_servers() -> Result<Vec<IpAddr>> {
    let f = File::open("/run/systemd/netif/state")?;
    let ip_strings = util::key_lookup_reader('=', "DNS", f)?
        .ok_or("DNS not found in netif state file")?;
    let mut addrs = Vec::new();
    for ip_string in ip_strings.split(' ') {
        addrs.push(IpAddr::from_str(ip_string)
            .chain_err(|| "failed to parse IP address")?);
    }
    if addrs.is_empty() {
        return Err("no DNS servers in /run/systemd/netif/state".into());
    }
    Ok(addrs)
}

fn parse_network(netinfo: &PacketNetworkInfo) -> Result<(Vec<Interface>,Vec<Device>)> {
    let mut interfaces = Vec::new();
    for i in netinfo.interfaces.clone() {
        let mac = MacAddr::from_str(&i.mac)
            .map_err(|err| Error::from(format!("{:?}", err)))
            .chain_err(|| format!("failed to parse mac address: '{}'", i.mac))?;
        interfaces.push(Interface {
            mac_address: Some(mac),
            bond: Some("bond0".to_owned()),
            name: None,
            priority: None,
            nameservers: Vec::new(),
            ip_addresses: Vec::new(),
            routes: Vec::new(),
        });
    }
    let mut iface = Interface{
        name: Some("bond0".to_owned()),
        priority: Some(5),
        nameservers: get_dns_servers()?,
        mac_address: None,
        bond: None,
        ip_addresses: Vec::new(),
        routes: Vec::new(),
    };
    for a in netinfo.addresses.clone() {
        let prefix = ipnetwork::ip_mask_to_prefix(a.netmask)
            .chain_err(|| "invalid network mask")?;
        iface.ip_addresses.push(
            match a.address {
                IpAddr::V4(addrv4) => IpNetwork::V4(Ipv4Network::new(addrv4, prefix)
                    .chain_err(|| "invalid IP address or prefix")?),
                IpAddr::V6(addrv6) => IpNetwork::V6(Ipv6Network::new(addrv6, prefix)
                    .chain_err(|| "invalid IP address or prefix")?),
            }
        );
        let dest = match (a.public,a.address) {
            (false,IpAddr::V4(_)) =>
                    IpNetwork::V4(Ipv4Network::new(Ipv4Addr::new(10,0,0,0),8).unwrap()),
            (true,IpAddr::V4(_)) =>
                    IpNetwork::V4(Ipv4Network::new(Ipv4Addr::new(0,0,0,0),0).unwrap()),
            (_,IpAddr::V6(_)) =>
                    IpNetwork::V6(Ipv6Network::new(Ipv6Addr::new(0,0,0,0,0,0,0,0),0).unwrap()),
        };
        iface.routes.push(
            NetworkRoute {
                destination: dest,
                gateway: a.gateway,
            }

        );
    }
    interfaces.push(iface);

    let mut attrs = vec![
		("TransmitHashPolicy".to_owned(), "layer3+4".to_owned()),
		("MIIMonitorSec".to_owned(), ".1".to_owned()),
		("UpDelaySec".to_owned(), ".2".to_owned()),
		("DownDelaySec".to_owned(), ".2".to_owned()),
		("Mode".to_owned(), network::bonding_mode_to_string(&netinfo.bonding.mode)?),
    ];
    if netinfo.bonding.mode == network::BONDING_MODE_LACP {
        attrs.push(("LACPTransmitRate".to_owned(), "fast".to_owned()));
    }
    let network_devices = vec![
        Device{
            name: "bond0".to_owned(),
            kind: "bond".to_owned(),
            mac_address: interfaces[0].mac_address
                .ok_or("first interface doesn't have a mac address, should be impossible")?,
            priority: Some(5),
            sections: vec![
                Section{
                    name: "Bond".to_owned(),
                    attributes: attrs,
                }
            ],
        },
    ];

    Ok((interfaces,network_devices))
}

fn get_attrs(data: &PacketData) -> Result<Vec<(String,String)>> {
    let mut attrs = Vec::new();
    let mut v4_public_counter = 0;
    let mut v4_private_counter = 0;
    let mut v6_public_counter = 0;
    let mut v6_private_counter = 0;
    for a in data.network.addresses.clone() {
        match (a.address,a.public) {
            (IpAddr::V4(a),true) => {
                attrs.push((format!("PACKET_IPV4_PUBLIC_{}", v4_public_counter), format!("{}", a)));
                v4_public_counter += 1;
            }
            (IpAddr::V4(a),false) => {
                attrs.push((format!("PACKET_IPV4_PRIVATE_{}", v4_private_counter), format!("{}", a)));
                v4_private_counter += 1;
            }
            (IpAddr::V6(a),true) => {
                attrs.push((format!("PACKET_IPV6_PUBLIC_{}", v6_public_counter), format!("{}", a)));
                v6_public_counter += 1;
            }
            (IpAddr::V6(a),false) => {
                attrs.push((format!("PACKET_IPV6_PRIVATE_{}", v6_private_counter), format!("{}", a)));
                v6_private_counter += 1;
            }
        }
    }
    attrs.push(("PACKET_HOSTNAME".to_owned(), data.hostname.clone()));
    attrs.push(("PACKET_PHONE_HOME_URL".to_owned(), data.phone_home_url.clone()));
    Ok(attrs)
}
<commit_msg>providers/packet: respect bond interface requests<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.

//! packet metadata fetcher

use metadata::Metadata;

use errors::*;

use retry;
use std::fs::File;
use std::str::FromStr;

use util;

use network;
use network::{Interface,Device,Section,NetworkRoute};
use pnet::util::MacAddr;
use std::net::{IpAddr,Ipv4Addr,Ipv6Addr};
use ipnetwork;
use ipnetwork::{IpNetwork,Ipv4Network,Ipv6Network};

#[derive(Clone,Deserialize)]
struct PacketData {
    id: String,
    hostname: String,
    iqn: String,
    plan: String,
    facility: String,
    tags: Vec<String>,
    ssh_keys: Vec<String>,
    network: PacketNetworkInfo,

    error: Option<String>,
    phone_home_url: String,
}

#[derive(Clone,Deserialize)]
struct PacketNetworkInfo {
    interfaces: Vec<PacketInterfaceInfo>,
    addresses: Vec<PacketAddressInfo>,
    bonding: PacketBondingMode,
}

#[derive(Clone,Deserialize)]
struct PacketBondingMode {
    mode: u32,
}

#[derive(Clone,Deserialize)]
struct PacketInterfaceInfo {
    name: String,
    mac: String,
    bond: Option<String>,
}

#[derive(Clone,Deserialize)]
struct PacketAddressInfo {
    id: String,
    address_family: i32,
    public: bool,
    management: bool,
    address: IpAddr,
    netmask: IpAddr,
    gateway: IpAddr,
}

pub fn fetch_metadata() -> Result<Metadata> {
    let client = retry::Client::new()?;
    let data: PacketData = client.get(retry::Json, "http://metadata.packet.net/metadata".to_owned()).send()?
        .ok_or("not found")?;

    let (interfaces,network_devices) = parse_network(&data.network)?;

    let attrs = get_attrs(&data)?;

    let mut m = Metadata::builder()
        .add_ssh_keys(data.ssh_keys)?
        .set_hostname(data.hostname);

    for (key,val) in attrs {
        m = m.add_attribute(key,val);
    }
    for iface in interfaces {
        m = m.add_network_interface(iface);
    }
    for netdev in network_devices {
        m = m.add_network_device(netdev);
    }
    Ok(m.build())
}

fn get_dns_servers() -> Result<Vec<IpAddr>> {
    let f = File::open("/run/systemd/netif/state")?;
    let ip_strings = util::key_lookup_reader('=', "DNS", f)?
        .ok_or("DNS not found in netif state file")?;
    let mut addrs = Vec::new();
    for ip_string in ip_strings.split(' ') {
        addrs.push(IpAddr::from_str(ip_string)
            .chain_err(|| "failed to parse IP address")?);
    }
    if addrs.is_empty() {
        return Err("no DNS servers in /run/systemd/netif/state".into());
    }
    Ok(addrs)
}

fn parse_network(netinfo: &PacketNetworkInfo) -> Result<(Vec<Interface>,Vec<Device>)> {
    let mut interfaces = Vec::new();
    let mut bonds = Vec::new();
    let dns_servers = get_dns_servers()?;
    for i in netinfo.interfaces.clone() {
        let mac = MacAddr::from_str(&i.mac)
            .map_err(|err| Error::from(format!("{:?}", err)))
            .chain_err(|| format!("failed to parse mac address: '{}'", i.mac))?;
        interfaces.push(Interface {
            mac_address: Some(mac),
            bond: i.bond.clone(),
            name: None,
            priority: None,
            nameservers: Vec::new(),
            ip_addresses: Vec::new(),
            routes: Vec::new(),
        });

        // if there is a bond key, make sure we have a bond device for it
        if let Some(ref bond_name) = i.bond {
            let bond = Interface {
                name: Some(bond_name.clone()),
                priority: Some(5),
                nameservers: dns_servers.clone(),
                mac_address: None,
                bond: None,
                ip_addresses: Vec::new(),
                routes: Vec::new(),
            };
            if !bonds.iter().any(|&(_, ref b): &(MacAddr, Interface)| &bond == b) {
                bonds.push((mac, bond));
            }
        }
    }

    // according to the folks from packet, all the addresses given to us in the
    // network section should be attached to the first bond we find in the list
    // of interfaces. we should always have at least one bond listed, but if we
    // don't find any, we just print out a scary warning and don't attach the
    // addresses to anything.
    if bonds.is_empty() {
        warn!("no bond interfaces. addresses are left unassigned.");
        // the rest of the function operates on bonds, so just return
        return Ok((interfaces, vec![]));
    }

    // remove panics if the index is out of bounds, but we know that there is at
    // least one bond in the vector because we return if it's empty
    let (first_mac, mut first_bond) = bonds.remove(0);
    for a in netinfo.addresses.clone() {
        let prefix = ipnetwork::ip_mask_to_prefix(a.netmask)
            .chain_err(|| "invalid network mask")?;
        first_bond.ip_addresses.push(IpNetwork::new(a.address, prefix)
                                .chain_err(|| "invalid IP address or prefix")?);
        let dest = match (a.public,a.address) {
            (false,IpAddr::V4(_)) =>
                IpNetwork::V4(Ipv4Network::new(Ipv4Addr::new(10,0,0,0),8).unwrap()),
            (true,IpAddr::V4(_)) =>
                IpNetwork::V4(Ipv4Network::new(Ipv4Addr::new(0,0,0,0),0).unwrap()),
            (_,IpAddr::V6(_)) =>
                IpNetwork::V6(Ipv6Network::new(Ipv6Addr::new(0,0,0,0,0,0,0,0),0).unwrap()),
        };
        first_bond.routes.push(NetworkRoute {
            destination: dest,
            gateway: a.gateway,
        });
    }
    bonds.push((first_mac, first_bond));

    let mut attrs = vec![
        ("TransmitHashPolicy".to_owned(), "layer3+4".to_owned()),
        ("MIIMonitorSec".to_owned(), ".1".to_owned()),
        ("UpDelaySec".to_owned(), ".2".to_owned()),
        ("DownDelaySec".to_owned(), ".2".to_owned()),
        ("Mode".to_owned(), network::bonding_mode_to_string(&netinfo.bonding.mode)?),
    ];
    if netinfo.bonding.mode == network::BONDING_MODE_LACP {
        attrs.push(("LACPTransmitRate".to_owned(), "fast".to_owned()));
    }

    let mut network_devices = vec![];
    for (mac, bond) in bonds {
        network_devices.push(Device {
            name: bond.name.clone()
                .ok_or("bond doesn't have a name, should be impossible")?,
            kind: "bond".to_owned(),
            mac_address: mac,
            priority: Some(5),
            sections: vec![
                Section{
                    name: "Bond".to_owned(),
                    attributes: attrs.clone(),
                }
            ],
        });
        // finally, make sure the bond interfaces are in the interface list
        interfaces.push(bond)
    }

    Ok((interfaces,network_devices))
}

fn get_attrs(data: &PacketData) -> Result<Vec<(String,String)>> {
    let mut attrs = Vec::new();
    let mut v4_public_counter = 0;
    let mut v4_private_counter = 0;
    let mut v6_public_counter = 0;
    let mut v6_private_counter = 0;
    for a in data.network.addresses.clone() {
        match (a.address,a.public) {
            (IpAddr::V4(a),true) => {
                attrs.push((format!("PACKET_IPV4_PUBLIC_{}", v4_public_counter), format!("{}", a)));
                v4_public_counter += 1;
            }
            (IpAddr::V4(a),false) => {
                attrs.push((format!("PACKET_IPV4_PRIVATE_{}", v4_private_counter), format!("{}", a)));
                v4_private_counter += 1;
            }
            (IpAddr::V6(a),true) => {
                attrs.push((format!("PACKET_IPV6_PUBLIC_{}", v6_public_counter), format!("{}", a)));
                v6_public_counter += 1;
            }
            (IpAddr::V6(a),false) => {
                attrs.push((format!("PACKET_IPV6_PRIVATE_{}", v6_private_counter), format!("{}", a)));
                v6_private_counter += 1;
            }
        }
    }
    attrs.push(("PACKET_HOSTNAME".to_owned(), data.hostname.clone()));
    attrs.push(("PACKET_PHONE_HOME_URL".to_owned(), data.phone_home_url.clone()));
    Ok(attrs)
}
<|endoftext|>