<commit_before>//! WebSocket protocol implementation.
//!
//! To setup a WebSocket, first perform the WebSocket handshake then on success convert `Payload` into a
//! `WsStream` stream and then use `WsWriter` to communicate with the peer.

use std::io;

use derive_more::{Display, Error, From};
use http::{header, Method, StatusCode};

use crate::{
    body::AnyBody, header::HeaderValue, message::RequestHead, response::Response,
    ResponseBuilder,
};

mod codec;
mod dispatcher;
mod frame;
mod mask;
mod proto;

pub use self::codec::{Codec, Frame, Item, Message};
pub use self::dispatcher::Dispatcher;
pub use self::frame::Parser;
pub use self::proto::{hash_key, CloseCode, CloseReason, OpCode};

/// WebSocket protocol errors.
#[derive(Debug, Display, Error, From)]
pub enum ProtocolError {
    /// Received an unmasked frame from client.
    #[display(fmt = "Received an unmasked frame from client.")]
    UnmaskedFrame,

    /// Received a masked frame from server.
    #[display(fmt = "Received a masked frame from server.")]
    MaskedFrame,

    /// Encountered invalid opcode.
    #[display(fmt = "Invalid opcode: {}.", _0)]
    InvalidOpcode(#[error(not(source))] u8),

    /// Invalid control frame length
    #[display(fmt = "Invalid control frame length: {}.", _0)]
    InvalidLength(#[error(not(source))] usize),

    /// Bad opcode.
    #[display(fmt = "Bad opcode.")]
    BadOpCode,

    /// A payload reached size limit.
    #[display(fmt = "A payload reached size limit.")]
    Overflow,

    /// Continuation is not started.
    #[display(fmt = "Continuation is not started.")]
    ContinuationNotStarted,

    /// Received new continuation but it is already started.
    #[display(fmt = "Received new continuation but it is already started.")]
    ContinuationStarted,

    /// Unknown continuation fragment.
    #[display(fmt = "Unknown continuation fragment: {}.", _0)]
    ContinuationFragment(#[error(not(source))] OpCode),

    /// I/O error.
    #[display(fmt = "I/O error: {}", _0)]
    Io(io::Error),
}

/// WebSocket handshake errors
#[derive(Debug, PartialEq, Display, Error)]
pub enum HandshakeError {
    /// Only get method is allowed.
    #[display(fmt = "Method not allowed.")]
    GetMethodRequired,

    /// Upgrade header if not set to WebSocket.
    #[display(fmt = "WebSocket upgrade is expected.")]
    NoWebsocketUpgrade,

    /// Connection header is not set to upgrade.
    #[display(fmt = "Connection upgrade is expected.")]
    NoConnectionUpgrade,

    /// WebSocket version header is not set.
    #[display(fmt = "WebSocket version header is required.")]
    NoVersionHeader,

    /// Unsupported WebSocket version.
    #[display(fmt = "Unsupported WebSocket version.")]
    UnsupportedVersion,

    /// WebSocket key is not set or wrong.
    #[display(fmt = "Unknown websocket key.")]
    BadWebsocketKey,
}

impl From<&HandshakeError> for Response<AnyBody> {
    fn from(err: &HandshakeError) -> Self {
        match err {
            HandshakeError::GetMethodRequired => {
                let mut res = Response::new(StatusCode::METHOD_NOT_ALLOWED);
                res.headers_mut()
                    .insert(header::ALLOW, HeaderValue::from_static("GET"));
                res
            }

            HandshakeError::NoWebsocketUpgrade => {
                let mut res = Response::bad_request();
                res.head_mut().reason = Some("No WebSocket Upgrade header found");
                res
            }

            HandshakeError::NoConnectionUpgrade => {
                let mut res = Response::bad_request();
                res.head_mut().reason = Some("No Connection upgrade");
                res
            }

            HandshakeError::NoVersionHeader => {
                let mut res = Response::bad_request();
                res.head_mut().reason = Some("WebSocket version header is required");
                res
            }

            HandshakeError::UnsupportedVersion => {
                let mut res = Response::bad_request();
                res.head_mut().reason = Some("Unsupported WebSocket version");
                res
            }

            HandshakeError::BadWebsocketKey => {
                let mut res = Response::bad_request();
                res.head_mut().reason = Some("Handshake error");
                res
            }
        }
    }
}

impl From<HandshakeError> for Response<AnyBody> {
    fn from(err: HandshakeError) -> Self {
        (&err).into()
    }
}

/// Verify WebSocket handshake request and create handshake response.
pub fn handshake(req: &RequestHead) -> Result<ResponseBuilder, HandshakeError> {
    verify_handshake(req)?;
    Ok(handshake_response(req))
}

/// Verify WebSocket handshake request.
pub fn verify_handshake(req: &RequestHead) -> Result<(), HandshakeError> {
    // WebSocket accepts only GET
    if req.method != Method::GET {
        return Err(HandshakeError::GetMethodRequired);
    }

    // Check for "UPGRADE" to WebSocket header
    let has_hdr = if let Some(hdr) = req.headers().get(header::UPGRADE) {
        if let Ok(s) = hdr.to_str() {
            s.to_ascii_lowercase().contains("websocket")
        } else {
            false
        }
    } else {
        false
    };
    if !has_hdr {
        return Err(HandshakeError::NoWebsocketUpgrade);
    }

    // Upgrade connection
    if !req.upgrade() {
        return Err(HandshakeError::NoConnectionUpgrade);
    }

    // check supported version
    if !req.headers().contains_key(header::SEC_WEBSOCKET_VERSION) {
        return Err(HandshakeError::NoVersionHeader);
    }
    let supported_ver = {
        if let Some(hdr) = req.headers().get(header::SEC_WEBSOCKET_VERSION) {
            hdr == "13" || hdr == "8" || hdr == "7"
        } else {
            false
        }
    };
    if !supported_ver {
        return Err(HandshakeError::UnsupportedVersion);
    }

    // check client handshake for validity
    if !req.headers().contains_key(header::SEC_WEBSOCKET_KEY) {
        return Err(HandshakeError::BadWebsocketKey);
    }
    Ok(())
}

/// Create WebSocket handshake response.
///
/// This function returns handshake `Response`, ready to send to peer.
pub fn handshake_response(req: &RequestHead) -> ResponseBuilder {
    let key = {
        let key = req.headers().get(header::SEC_WEBSOCKET_KEY).unwrap();
        proto::hash_key(key.as_ref())
    };

    Response::build(StatusCode::SWITCHING_PROTOCOLS)
        .upgrade("websocket")
        .insert_header((header::TRANSFER_ENCODING, "chunked"))
        .insert_header((
            header::SEC_WEBSOCKET_ACCEPT,
            // key is known to be header value safe ascii
            HeaderValue::from_bytes(&key).unwrap(),
        ))
        .take()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{body::AnyBody, test::TestRequest};
    use http::{header, Method};

    #[test]
    fn test_handshake() {
        let req = TestRequest::default().method(Method::POST).finish();
        assert_eq!(
            HandshakeError::GetMethodRequired,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default().finish();
        assert_eq!(
            HandshakeError::NoWebsocketUpgrade,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default()
            .insert_header((header::UPGRADE, header::HeaderValue::from_static("test")))
            .finish();
        assert_eq!(
            HandshakeError::NoWebsocketUpgrade,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default()
            .insert_header((
                header::UPGRADE,
                header::HeaderValue::from_static("websocket"),
            ))
            .finish();
        assert_eq!(
            HandshakeError::NoConnectionUpgrade,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default()
            .insert_header((
                header::UPGRADE,
                header::HeaderValue::from_static("websocket"),
            ))
            .insert_header((
                header::CONNECTION,
                header::HeaderValue::from_static("upgrade"),
            ))
            .finish();
        assert_eq!(
            HandshakeError::NoVersionHeader,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default()
            .insert_header((
                header::UPGRADE,
                header::HeaderValue::from_static("websocket"),
            ))
            .insert_header((
                header::CONNECTION,
                header::HeaderValue::from_static("upgrade"),
            ))
            .insert_header((
                header::SEC_WEBSOCKET_VERSION,
                header::HeaderValue::from_static("5"),
            ))
            .finish();
        assert_eq!(
            HandshakeError::UnsupportedVersion,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default()
            .insert_header((
                header::UPGRADE,
                header::HeaderValue::from_static("websocket"),
            ))
            .insert_header((
                header::CONNECTION,
                header::HeaderValue::from_static("upgrade"),
            ))
            .insert_header((
                header::SEC_WEBSOCKET_VERSION,
                header::HeaderValue::from_static("13"),
            ))
            .finish();
        assert_eq!(
            HandshakeError::BadWebsocketKey,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default()
            .insert_header((
                header::UPGRADE,
                header::HeaderValue::from_static("websocket"),
            ))
            .insert_header((
                header::CONNECTION,
                header::HeaderValue::from_static("upgrade"),
            ))
            .insert_header((
                header::SEC_WEBSOCKET_VERSION,
                header::HeaderValue::from_static("13"),
            ))
            .insert_header((
                header::SEC_WEBSOCKET_KEY,
                header::HeaderValue::from_static("13"),
            ))
            .finish();
        assert_eq!(
            StatusCode::SWITCHING_PROTOCOLS,
            handshake_response(req.head()).finish().status()
        );
    }

    #[test]
    fn test_ws_error_http_response() {
        let resp: Response<AnyBody> = HandshakeError::GetMethodRequired.into();
        assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);
        let resp: Response<AnyBody> = HandshakeError::NoWebsocketUpgrade.into();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: Response<AnyBody> = HandshakeError::NoConnectionUpgrade.into();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: Response<AnyBody> = HandshakeError::NoVersionHeader.into();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: Response<AnyBody> = HandshakeError::UnsupportedVersion.into();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: Response<AnyBody> = HandshakeError::BadWebsocketKey.into();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
    }
}
<commit_msg>remove chunked encoding header for websockets<commit_after>//! WebSocket protocol implementation.
//!
//! To setup a WebSocket, first perform the WebSocket handshake then on success convert `Payload` into a
//! `WsStream` stream and then use `WsWriter` to communicate with the peer.

use std::io;

use derive_more::{Display, Error, From};
use http::{header, Method, StatusCode};

use crate::{
    body::AnyBody, header::HeaderValue, message::RequestHead, response::Response,
    ResponseBuilder,
};

mod codec;
mod dispatcher;
mod frame;
mod mask;
mod proto;

pub use self::codec::{Codec, Frame, Item, Message};
pub use self::dispatcher::Dispatcher;
pub use self::frame::Parser;
pub use self::proto::{hash_key, CloseCode, CloseReason, OpCode};

/// WebSocket protocol errors.
#[derive(Debug, Display, Error, From)]
pub enum ProtocolError {
    /// Received an unmasked frame from client.
    #[display(fmt = "Received an unmasked frame from client.")]
    UnmaskedFrame,

    /// Received a masked frame from server.
    #[display(fmt = "Received a masked frame from server.")]
    MaskedFrame,

    /// Encountered invalid opcode.
    #[display(fmt = "Invalid opcode: {}.", _0)]
    InvalidOpcode(#[error(not(source))] u8),

    /// Invalid control frame length
    #[display(fmt = "Invalid control frame length: {}.", _0)]
    InvalidLength(#[error(not(source))] usize),

    /// Bad opcode.
    #[display(fmt = "Bad opcode.")]
    BadOpCode,

    /// A payload reached size limit.
    #[display(fmt = "A payload reached size limit.")]
    Overflow,

    /// Continuation is not started.
    #[display(fmt = "Continuation is not started.")]
    ContinuationNotStarted,

    /// Received new continuation but it is already started.
    #[display(fmt = "Received new continuation but it is already started.")]
    ContinuationStarted,

    /// Unknown continuation fragment.
    #[display(fmt = "Unknown continuation fragment: {}.", _0)]
    ContinuationFragment(#[error(not(source))] OpCode),

    /// I/O error.
    #[display(fmt = "I/O error: {}", _0)]
    Io(io::Error),
}

/// WebSocket handshake errors
#[derive(Debug, PartialEq, Display, Error)]
pub enum HandshakeError {
    /// Only get method is allowed.
    #[display(fmt = "Method not allowed.")]
    GetMethodRequired,

    /// Upgrade header if not set to WebSocket.
    #[display(fmt = "WebSocket upgrade is expected.")]
    NoWebsocketUpgrade,

    /// Connection header is not set to upgrade.
    #[display(fmt = "Connection upgrade is expected.")]
    NoConnectionUpgrade,

    /// WebSocket version header is not set.
    #[display(fmt = "WebSocket version header is required.")]
    NoVersionHeader,

    /// Unsupported WebSocket version.
    #[display(fmt = "Unsupported WebSocket version.")]
    UnsupportedVersion,

    /// WebSocket key is not set or wrong.
    #[display(fmt = "Unknown websocket key.")]
    BadWebsocketKey,
}

impl From<&HandshakeError> for Response<AnyBody> {
    fn from(err: &HandshakeError) -> Self {
        match err {
            HandshakeError::GetMethodRequired => {
                let mut res = Response::new(StatusCode::METHOD_NOT_ALLOWED);
                res.headers_mut()
                    .insert(header::ALLOW, HeaderValue::from_static("GET"));
                res
            }

            HandshakeError::NoWebsocketUpgrade => {
                let mut res = Response::bad_request();
                res.head_mut().reason = Some("No WebSocket Upgrade header found");
                res
            }

            HandshakeError::NoConnectionUpgrade => {
                let mut res = Response::bad_request();
                res.head_mut().reason = Some("No Connection upgrade");
                res
            }

            HandshakeError::NoVersionHeader => {
                let mut res = Response::bad_request();
                res.head_mut().reason = Some("WebSocket version header is required");
                res
            }

            HandshakeError::UnsupportedVersion => {
                let mut res = Response::bad_request();
                res.head_mut().reason = Some("Unsupported WebSocket version");
                res
            }

            HandshakeError::BadWebsocketKey => {
                let mut res = Response::bad_request();
                res.head_mut().reason = Some("Handshake error");
                res
            }
        }
    }
}

impl From<HandshakeError> for Response<AnyBody> {
    fn from(err: HandshakeError) -> Self {
        (&err).into()
    }
}

/// Verify WebSocket handshake request and create handshake response.
pub fn handshake(req: &RequestHead) -> Result<ResponseBuilder, HandshakeError> {
    verify_handshake(req)?;
    Ok(handshake_response(req))
}

/// Verify WebSocket handshake request.
pub fn verify_handshake(req: &RequestHead) -> Result<(), HandshakeError> {
    // WebSocket accepts only GET
    if req.method != Method::GET {
        return Err(HandshakeError::GetMethodRequired);
    }

    // Check for "UPGRADE" to WebSocket header
    let has_hdr = if let Some(hdr) = req.headers().get(header::UPGRADE) {
        if let Ok(s) = hdr.to_str() {
            s.to_ascii_lowercase().contains("websocket")
        } else {
            false
        }
    } else {
        false
    };
    if !has_hdr {
        return Err(HandshakeError::NoWebsocketUpgrade);
    }

    // Upgrade connection
    if !req.upgrade() {
        return Err(HandshakeError::NoConnectionUpgrade);
    }

    // check supported version
    if !req.headers().contains_key(header::SEC_WEBSOCKET_VERSION) {
        return Err(HandshakeError::NoVersionHeader);
    }
    let supported_ver = {
        if let Some(hdr) = req.headers().get(header::SEC_WEBSOCKET_VERSION) {
            hdr == "13" || hdr == "8" || hdr == "7"
        } else {
            false
        }
    };
    if !supported_ver {
        return Err(HandshakeError::UnsupportedVersion);
    }

    // check client handshake for validity
    if !req.headers().contains_key(header::SEC_WEBSOCKET_KEY) {
        return Err(HandshakeError::BadWebsocketKey);
    }
    Ok(())
}

/// Create WebSocket handshake response.
///
/// This function returns handshake `Response`, ready to send to peer.
pub fn handshake_response(req: &RequestHead) -> ResponseBuilder {
    let key = {
        let key = req.headers().get(header::SEC_WEBSOCKET_KEY).unwrap();
        proto::hash_key(key.as_ref())
    };

    Response::build(StatusCode::SWITCHING_PROTOCOLS)
        .upgrade("websocket")
        .insert_header((
            header::SEC_WEBSOCKET_ACCEPT,
            // key is known to be header value safe ascii
            HeaderValue::from_bytes(&key).unwrap(),
        ))
        .take()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{body::AnyBody, test::TestRequest};
    use http::{header, Method};

    #[test]
    fn test_handshake() {
        let req = TestRequest::default().method(Method::POST).finish();
        assert_eq!(
            HandshakeError::GetMethodRequired,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default().finish();
        assert_eq!(
            HandshakeError::NoWebsocketUpgrade,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default()
            .insert_header((header::UPGRADE, header::HeaderValue::from_static("test")))
            .finish();
        assert_eq!(
            HandshakeError::NoWebsocketUpgrade,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default()
            .insert_header((
                header::UPGRADE,
                header::HeaderValue::from_static("websocket"),
            ))
            .finish();
        assert_eq!(
            HandshakeError::NoConnectionUpgrade,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default()
            .insert_header((
                header::UPGRADE,
                header::HeaderValue::from_static("websocket"),
            ))
            .insert_header((
                header::CONNECTION,
                header::HeaderValue::from_static("upgrade"),
            ))
            .finish();
        assert_eq!(
            HandshakeError::NoVersionHeader,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default()
            .insert_header((
                header::UPGRADE,
                header::HeaderValue::from_static("websocket"),
            ))
            .insert_header((
                header::CONNECTION,
                header::HeaderValue::from_static("upgrade"),
            ))
            .insert_header((
                header::SEC_WEBSOCKET_VERSION,
                header::HeaderValue::from_static("5"),
            ))
            .finish();
        assert_eq!(
            HandshakeError::UnsupportedVersion,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default()
            .insert_header((
                header::UPGRADE,
                header::HeaderValue::from_static("websocket"),
            ))
            .insert_header((
                header::CONNECTION,
                header::HeaderValue::from_static("upgrade"),
            ))
            .insert_header((
                header::SEC_WEBSOCKET_VERSION,
                header::HeaderValue::from_static("13"),
            ))
            .finish();
        assert_eq!(
            HandshakeError::BadWebsocketKey,
            verify_handshake(req.head()).unwrap_err(),
        );

        let req = TestRequest::default()
            .insert_header((
                header::UPGRADE,
                header::HeaderValue::from_static("websocket"),
            ))
            .insert_header((
                header::CONNECTION,
                header::HeaderValue::from_static("upgrade"),
            ))
            .insert_header((
                header::SEC_WEBSOCKET_VERSION,
                header::HeaderValue::from_static("13"),
            ))
            .insert_header((
                header::SEC_WEBSOCKET_KEY,
                header::HeaderValue::from_static("13"),
            ))
            .finish();
        assert_eq!(
            StatusCode::SWITCHING_PROTOCOLS,
            handshake_response(req.head()).finish().status()
        );
    }

    #[test]
    fn test_ws_error_http_response() {
        let resp: Response<AnyBody> = HandshakeError::GetMethodRequired.into();
        assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);
        let resp: Response<AnyBody> = HandshakeError::NoWebsocketUpgrade.into();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: Response<AnyBody> = HandshakeError::NoConnectionUpgrade.into();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: Response<AnyBody> = HandshakeError::NoVersionHeader.into();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: Response<AnyBody> = HandshakeError::UnsupportedVersion.into();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        let resp: Response<AnyBody> = HandshakeError::BadWebsocketKey.into();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
    }
}
<|endoftext|>


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

use self::Node::*;

/// A `Node` in the `Rope`'s tree.
///
/// A `Node` is either a `Leaf` holding a `String`, or a
/// a `Branch` concatenating together two `Node`s.
#[derive(Debug, Clone)]
pub enum Node {
    /// A leaf node
    Leaf(String)
  , /// A branch concatenating together `l`eft and `r`ight nodes.
    Branch(BranchNode)
}

#[derive(Debug, Clone)]
pub struct BranchNode {
    /// The length of this node
    len: usize
  , /// The weight of a node is the summed weight of its left subtree
    weight: usize
  , /// The left branch node
    pub left: Box<Node>
  , /// The right branch node
    pub right: Box<Node>
}

impl Default for Node {
    fn default() -> Self { Node::empty() }
}

/// Returns the _n_th fibonacci number.
// TODO: replace with an iterative implementation and/or lookup table?
fn fibonacci(n: usize) -> usize {
    match n { 0 => 1
            , 1 => 1
            , _ => fibonacci(n - 1) + fibonacci(n - 2)
            }
}

impl BranchNode {

    #[inline]
    fn new(left: Node, right: Node) -> Self {
        BranchNode { len: left.len() + right.len()
                   , weight: left.subtree_weight()
                   , left: box left
                   , right: box right
                   }
    }

    #[inline]
    fn take_left(&mut self) -> Node {
        mem::replace(self.left.as_mut(), Node::empty())
    }

    #[inline]
    fn take_right(&mut self) -> Node {
        mem::replace(self.right.as_mut(), Node::empty())
    }

    fn split(&mut self, index: usize) -> &mut Node {
        if index < self.weight {
            // split the left node
            let mut left = self.take_left();
            left.split(index);
            // replacing *self with a new BranchNode will update the
            // node's weight automagically
            *self = BranchNode::new(left, self.take_right());
            self.left.as_mut()
        } else {
            // split the right node
            let mut right = self.take_right();
            right.split(index);
            // replacing *self with a new BranchNode will update the
            // node's weight automagically
            *self = BranchNode::new(self.take_left(), right);
            self.left.as_mut()
        }
    }

}

impl Node {
    fn split_leaf(&mut self, index: usize) -> bool {
        // we mutably borrow `self` here. This precludes us from changing it
        // directly as in `*self = ...`, because the borrow checker won't allow
        // it. Therefore, the assignment to `self` must be outside the `if let`
        // clause.
        *self = if let Leaf(ref mut s) = *self {

            // if this node is a Leaf, take the String out of it
            // (note that empty strings don't allocate).
            let string = mem::replace(s, String::new());
            // split the string into left and right parts...
            let left = Leaf(string[index..].to_string());
            let right = Leaf(string[..index].to_string());
            // construct the new Branch node that will be reassigned to `self`
            Node::new_branch(left, right)
        } else {
            // if this node is not a Leaf, we return immediately, thus skipping
            // the assignment
            return false
        };
        return true
    }

    pub fn split(&mut self, index: usize) -> &mut Node {
        if self.split_leaf(index) == true {
            self
        } else {
            if let &mut Branch(ref mut node) = self {
                node.split(index)
            } else {
                unreachable!()
            }
        }
    }

    pub fn empty() -> Self {
        Leaf(String::new())
    }

    /// Concatenate two `Node`s to return a new `Branch` node.
    #[inline]
    pub fn new_branch(left: Self, right: Self) -> Self {
        Branch(BranchNode::new(left, right))
    }

    #[inline]
    pub fn concat(&mut self, right: Self) {
        *self = Node::new_branch(mem::replace(self, Node::empty()), right)
    }

    #[inline]
    pub const fn new_leaf(string: String) -> Self {
        Leaf(string)
    }

    /// Returns true if this node is balanced
    ///
    /// "We define the depth of a leaf to be 0, and the depth of a
    /// concatenation to be one plus the maximum depth of its children. Let
    /// _Fn_ be the _n_th Fibonacci number. A rope of depth _n_ is balanced if
    /// its length is at least _Fn_+2, e.g. a balanced rope of depth 1 must
    /// have length at least 2. Note that balanced ropes may contain unbalanced
    /// subropes."
    /// – from "Ropes: An Alternative to Strings"
    #[inline]
    pub fn is_balanced(&self) -> bool {
        self.len() >= fibonacci(self.depth() + 2)
    }

    /// Returns the depth in the tree of a node
    #[inline]
    fn depth(&self) -> usize {
        use std::cmp::max;

        match self {
            &Node::Leaf(_) => 0
          , &Node::Branch(BranchNode { ref left, ref right, .. }) =>
                max(left.depth(), right.depth()) + 1
            }
    }

    /// Returns the length of a node
    //  TODO: do we want to cache this?
    pub fn len(&self) -> usize {
        match self { &Leaf(ref s) => s.len()
                   , &Branch(BranchNode { ref left, ref right, .. }) =>
                        left.len() + right.len()
                    }
    }

    /// Calculates the weight of a node
    #[inline]
    fn subtree_weight (&self) -> usize {
        match self { &Leaf(ref s) => s.len()
                   , &Branch(BranchNode { ref left, .. }) =>
                        left.subtree_weight()
                    }
    }

    /// Rebalance the subrope starting at this `Node`, returning a new `Node`
    ///
    /// From "Ropes: An Alternative to Strings":
    /// > "The rebalancing operation maintains an ordered sequence of (empty
    /// > or) balanced ropes, one for each length interval [_Fn_, _Fn_+1), for
    /// > _n_ >= 2. We traverse the rope from left to right, inserting each
    /// > leaf into the appropriate sequence position, depending on its length.
    ///
    /// > The concatenation of the sequence of ropes in order of decreasing
    /// > length is equivalent to the prefix of the rope we have traversed so
    /// > far. Each new leaf _x_ is with_insert_ropeed into the appropriate entry of the
    /// > sequence. Assume that _x_’s length is in the interval [_Fn_, _Fn_+1),
    /// > and thus it should be put in slot _n_ (which also corresponds to
    /// > maximum depth _n_ − 2). If all lower and equal numbered levels are
    /// > empty, then this can be done directly. If not, then we concatenate
    /// > ropes in slots 2, ... ,(_n_ − 1) (concatenating onto the left), and
    /// > concatenate _x_ to the right of the result. We then continue to
    /// > concatenate ropes from the sequence in increasing order to the left
    /// > of this result, until the result fits into an empty slot in the
    /// > sequence."
    pub fn rebalance(self) -> Self {
        // TODO: this is a huge mess, I based it on the IBM Java implementation
        //       please refactor until it stops being ugly!
        //        - eliza, 12/17/2016

        if self.is_balanced() {
            // the subrope is already balanced, do nothing
            self
        } else {
            let mut leaves: Vec<Option<Node>> =
                self.into_leaves().map(Option::Some).collect();
            let len = leaves.len();
            fn _rebalance(l: &mut Vec<Option<Node>>, start: usize, end: usize)
                          -> Node {
                match end - start {
                    1 => l[start].take().unwrap()
                  , 2 => l[start].take().unwrap() + l[start + 1].take().unwrap()
                  , n => {
                        let mid = start + (n / 2);
                        _rebalance(l, start, mid) + _rebalance(l, mid, end)

                    }
                }
            };
            _rebalance(&mut leaves, 0, len)
        }
    }

    /// Returns an iterator over all leaf nodes in this `Node`'s subrope
    fn leaves<'a>(&'a self) -> Leaves<'a> {
        Leaves(vec![self])
    }

    /// Returns a move iterator over all leaf nodes in this `Node`'s subrope
    fn into_leaves(self) -> IntoLeaves {
        IntoLeaves(vec![self])
    }

    /// Returns an iterator over all the strings in this `Node`s subrope'
    #[inline]
    pub fn strings<'a>(&'a self) -> impl Iterator<Item=&'a str> {
        box self.leaves().map(|n| match n {
            &Leaf(ref s) => s.as_ref()
          , _ => unreachable!("Node.leaves() iterator contained something \
                               that wasn't a leaf. Barring _force majeure_, \
                               this should be impossible. Something's broken.")
        })
    }

    /// Returns a move iterator over all the strings in this `Node`s subrope'
    ///
    /// Consumes `self`.
    #[inline]
    pub fn into_strings(self) -> impl Iterator<Item=String> {
        box self.into_leaves().map(|n| match n {
            Leaf(s) => s
            , _ => unreachable!("Node.into_leaves() iterator contained \
                                 something  that wasn't a leaf. Barring _force \
                                 majeure_, this should be impossible. \
                                 Something's broken.")
        })
    }

    str_iters! {
        #[doc="Returns an iterator over all the bytes in this `Node`'s \
               subrope \n\

               \nAs a Rope consists of a sequence of bytes, we can iterate \
               through a rope by byte. This method returns such an iterator."]
        #[inline]
        impl bytes<u8> for Node {}
        #[doc="Returns an iterator over all the characters in this `Node`'s \
               subrope \n\

               \nAs a `Rope` consists of valid UTF-8, we can iterate through a \
               `Rope` by `char`. This method returns such an iterator. \n\

               \nIt's important to remember that `char` represents a Unicode \
               Scalar Value, and may not match your idea of what a \
               'character' is. Iteration over grapheme clusters may be what \
               you actually want."]
        #[inline]
        impl chars<char> for Node {}
        // TODO: this is actually Wrong, the indices will wrap around once we
        //       iterate into the next leaf node. we'll need to write our own
        //       char_indices iterator that tracks the character's index in the
        //       global Rope. shouldn't be too hard, just a fold...
        //          - eliza, 12/18/2016
        // #[inline]
        // impl char_indices<(usize, char)> for Node {}
        #[inline]
        impl split_whitespace<&'a str> for Node {}
        #[inline]
        impl lines<&'a str> for Node {}
    }

    // /// Returns n iterator over the bytes of this `Node`'s subrope
    // ///
    // ///
    // #[inline]
    // pub fn bytes<'a>(&'a self) -> impl Iterator<Item=u8> + 'a {
    //     self.strings().flat_map(str::bytes)
    // }

    /// Returns an iterator over the grapheme clusters of this `Node`'s subrope'
    ///
    /// This is the iterator returned by `Node::into_iter`.
    #[inline]
    pub fn graphemes<'a>(&'a self) -> impl Iterator<Item=&'a str> {
        // the compiler won't let me mark this as unimplemented using the
        // unimplemented!() macro, due to Reasons (i suspect relating to
        // returning `impl Trait`)
        //  - eliza, 12/18/2016
        panic!("Unimplemented!");
        self.strings()
    }

}

/// An iterator over a series of leaf `Node`s
struct Leaves<'a>(Vec<&'a Node>);

impl<'a> Iterator for Leaves<'a> {
    type Item = &'a Node;
    fn next(&mut self) -> Option<Self::Item> {
        match self.0.pop() {
            None => None
          , Some(leaf @ &Leaf(_)) => if leaf.len() == 0 { None }
                                     else { Some(leaf) }
          , Some(&Branch(BranchNode { ref left, ref right, .. })) => {
                self.0.push(left);
                self.0.push(right);
                self.next()
            }
        }
    }
}

/// A move iterator over a series of leaf `Node`s
struct IntoLeaves(Vec<Node>);

impl Iterator for IntoLeaves {
    type Item = Node;
    fn next(&mut self) -> Option<Self::Item> {
        match self.0.pop() {
            None => None
          , Some(leaf @ Leaf(_)) => if leaf.len() == 0 { None }
                                    else { Some(leaf) }
          , Some(Branch(BranchNode { box left, box right, .. })) => {
                self.0.push(right);
                self.0.push(left);
                self.next()
            }
        }
    }
}

impl ops::Add for Node {
    type Output = Self;
    /// Concatenate two `Node`s, returning a `Branch` node.
    fn add(self, right: Self) -> Self { Node::new_branch(self, right) }
}
<commit_msg>rope: fix broken iterators<commit_after>use std::ops;
use std::mem;

use self::Node::*;

/// A `Node` in the `Rope`'s tree.
///
/// A `Node` is either a `Leaf` holding a `String`, or a
/// a `Branch` concatenating together two `Node`s.
#[derive(Debug, Clone)]
pub enum Node {
    /// A leaf node
    Leaf(String)
  , /// A branch concatenating together `l`eft and `r`ight nodes.
    Branch(BranchNode)
}

#[derive(Debug, Clone)]
pub struct BranchNode {
    /// The length of this node
    len: usize
  , /// The weight of a node is the summed weight of its left subtree
    weight: usize
  , /// The left branch node
    pub left: Box<Node>
  , /// The right branch node
    pub right: Box<Node>
}

impl Default for Node {
    fn default() -> Self { Node::empty() }
}

/// Returns the _n_th fibonacci number.
// TODO: replace with an iterative implementation and/or lookup table?
fn fibonacci(n: usize) -> usize {
    match n { 0 => 1
            , 1 => 1
            , _ => fibonacci(n - 1) + fibonacci(n - 2)
            }
}

impl BranchNode {

    #[inline]
    fn new(left: Node, right: Node) -> Self {
        BranchNode { len: left.len() + right.len()
                   , weight: left.subtree_weight()
                   , left: box left
                   , right: box right
                   }
    }

    #[inline]
    fn take_left(&mut self) -> Node {
        mem::replace(self.left.as_mut(), Node::empty())
    }

    #[inline]
    fn take_right(&mut self) -> Node {
        mem::replace(self.right.as_mut(), Node::empty())
    }

    fn split(&mut self, index: usize) -> &mut Node {
        if index < self.weight {
            // split the left node
            let mut left = self.take_left();
            left.split(index);
            // replacing *self with a new BranchNode will update the
            // node's weight automagically
            *self = BranchNode::new(left, self.take_right());
            self.left.as_mut()
        } else {
            // split the right node
            let mut right = self.take_right();
            right.split(index);
            // replacing *self with a new BranchNode will update the
            // node's weight automagically
            *self = BranchNode::new(self.take_left(), right);
            self.left.as_mut()
        }
    }

}

impl Node {
    fn split_leaf(&mut self, index: usize) -> bool {
        // we mutably borrow `self` here. This precludes us from changing it
        // directly as in `*self = ...`, because the borrow checker won't allow
        // it. Therefore, the assignment to `self` must be outside the `if let`
        // clause.
        *self = if let Leaf(ref mut s) = *self {

            // if this node is a Leaf, take the String out of it
            // (note that empty strings don't allocate).
            let string = mem::replace(s, String::new());
            // split the string into left and right parts...
            let left = Leaf(string[index..].to_string());
            let right = Leaf(string[..index].to_string());
            // construct the new Branch node that will be reassigned to `self`
            Node::new_branch(left, right)
        } else {
            // if this node is not a Leaf, we return immediately, thus skipping
            // the assignment
            return false
        };
        return true
    }

    pub fn split(&mut self, index: usize) -> &mut Node {
        if self.split_leaf(index) == true {
            self
        } else {
            if let &mut Branch(ref mut node) = self {
                node.split(index)
            } else {
                unreachable!()
            }
        }
    }

    pub fn empty() -> Self {
        Leaf(String::new())
    }

    /// Concatenate two `Node`s to return a new `Branch` node.
    #[inline]
    pub fn new_branch(left: Self, right: Self) -> Self {
        Branch(BranchNode::new(left, right))
    }

    #[inline]
    pub fn concat(&mut self, right: Self) {
        *self = Node::new_branch(mem::replace(self, Node::empty()), right)
    }

    #[inline]
    pub const fn new_leaf(string: String) -> Self {
        Leaf(string)
    }

    /// Returns true if this node is balanced
    ///
    /// "We define the depth of a leaf to be 0, and the depth of a
    /// concatenation to be one plus the maximum depth of its children. Let
    /// _Fn_ be the _n_th Fibonacci number. A rope of depth _n_ is balanced if
    /// its length is at least _Fn_+2, e.g. a balanced rope of depth 1 must
    /// have length at least 2. Note that balanced ropes may contain unbalanced
    /// subropes."
    /// – from "Ropes: An Alternative to Strings"
    #[inline]
    pub fn is_balanced(&self) -> bool {
        self.len() >= fibonacci(self.depth() + 2)
    }

    /// Returns the depth in the tree of a node
    #[inline]
    fn depth(&self) -> usize {
        use std::cmp::max;

        match self {
            &Node::Leaf(_) => 0
          , &Node::Branch(BranchNode { ref left, ref right, .. }) =>
                max(left.depth(), right.depth()) + 1
            }
    }

    /// Returns the length of a node
    //  TODO: do we want to cache this?
    pub fn len(&self) -> usize {
        match self { &Leaf(ref s) => s.len()
                   , &Branch(BranchNode { ref left, ref right, .. }) =>
                        left.len() + right.len()
                    }
    }

    /// Calculates the weight of a node
    #[inline]
    fn subtree_weight (&self) -> usize {
        match self { &Leaf(ref s) => s.len()
                   , &Branch(BranchNode { ref left, .. }) =>
                        left.subtree_weight()
                    }
    }

    /// Rebalance the subrope starting at this `Node`, returning a new `Node`
    ///
    /// From "Ropes: An Alternative to Strings":
    /// > "The rebalancing operation maintains an ordered sequence of (empty
    /// > or) balanced ropes, one for each length interval [_Fn_, _Fn_+1), for
    /// > _n_ >= 2. We traverse the rope from left to right, inserting each
    /// > leaf into the appropriate sequence position, depending on its length.
    ///
    /// > The concatenation of the sequence of ropes in order of decreasing
    /// > length is equivalent to the prefix of the rope we have traversed so
    /// > far. Each new leaf _x_ is with_insert_ropeed into the appropriate entry of the
    /// > sequence. Assume that _x_’s length is in the interval [_Fn_, _Fn_+1),
    /// > and thus it should be put in slot _n_ (which also corresponds to
    /// > maximum depth _n_ − 2). If all lower and equal numbered levels are
    /// > empty, then this can be done directly. If not, then we concatenate
    /// > ropes in slots 2, ... ,(_n_ − 1) (concatenating onto the left), and
    /// > concatenate _x_ to the right of the result. We then continue to
    /// > concatenate ropes from the sequence in increasing order to the left
    /// > of this result, until the result fits into an empty slot in the
    /// > sequence."
    pub fn rebalance(self) -> Self {
        // TODO: this is a huge mess, I based it on the IBM Java implementation
        //       please refactor until it stops being ugly!
        //        - eliza, 12/17/2016

        if self.is_balanced() {
            // the subrope is already balanced, do nothing
            self
        } else {
            let mut leaves: Vec<Option<Node>> =
                self.into_leaves().map(Option::Some).collect();
            let len = leaves.len();
            fn _rebalance(l: &mut Vec<Option<Node>>, start: usize, end: usize)
                          -> Node {
                match end - start {
                    1 => l[start].take().unwrap()
                  , 2 => l[start].take().unwrap() + l[start + 1].take().unwrap()
                  , n => {
                        let mid = start + (n / 2);
                        _rebalance(l, start, mid) + _rebalance(l, mid, end)

                    }
                }
            };
            _rebalance(&mut leaves, 0, len)
        }
    }

    /// Returns an iterator over all leaf nodes in this `Node`'s subrope
    fn leaves<'a>(&'a self) -> Leaves<'a> {
        Leaves(vec![self])
    }

    /// Returns a move iterator over all leaf nodes in this `Node`'s subrope
    fn into_leaves(self) -> IntoLeaves {
        IntoLeaves(vec![self])
    }

    /// Returns an iterator over all the strings in this `Node`s subrope'
    #[inline]
    pub fn strings<'a>(&'a self) -> impl Iterator<Item=&'a str> {
        box self.leaves().map(|n| match n {
            &Leaf(ref s) => s.as_ref()
          , _ => unreachable!("Node.leaves() iterator contained something \
                               that wasn't a leaf. Barring _force majeure_, \
                               this should be impossible. Something's broken.")
        })
    }

    /// Returns a move iterator over all the strings in this `Node`s subrope'
    ///
    /// Consumes `self`.
    #[inline]
    pub fn into_strings(self) -> impl Iterator<Item=String> {
        box self.into_leaves().map(|n| match n {
            Leaf(s) => s
            , _ => unreachable!("Node.into_leaves() iterator contained \
                                 something  that wasn't a leaf. Barring _force \
                                 majeure_, this should be impossible. \
                                 Something's broken.")
        })
    }

    str_iters! {
        #[doc="Returns an iterator over all the bytes in this `Node`'s \
               subrope \n\

               \nAs a Rope consists of a sequence of bytes, we can iterate \
               through a rope by byte. This method returns such an iterator."]
        #[inline]
        impl bytes<u8> for Node {}
        #[doc="Returns an iterator over all the characters in this `Node`'s \
               subrope \n\

               \nAs a `Rope` consists of valid UTF-8, we can iterate through a \
               `Rope` by `char`. This method returns such an iterator. \n\

               \nIt's important to remember that `char` represents a Unicode \
               Scalar Value, and may not match your idea of what a \
               'character' is. Iteration over grapheme clusters may be what \
               you actually want."]
        #[inline]
        impl chars<char> for Node {}
        // TODO: this is actually Wrong, the indices will wrap around once we
        //       iterate into the next leaf node. we'll need to write our own
        //       char_indices iterator that tracks the character's index in the
        //       global Rope. shouldn't be too hard, just a fold...
        //          - eliza, 12/18/2016
        // #[inline]
        // impl char_indices<(usize, char)> for Node {}
        #[inline]
        impl split_whitespace<&'a str> for Node {}
        #[inline]
        impl lines<&'a str> for Node {}
    }

    // /// Returns n iterator over the bytes of this `Node`'s subrope
    // ///
    // ///
    // #[inline]
    // pub fn bytes<'a>(&'a self) -> impl Iterator<Item=u8> + 'a {
    //     self.strings().flat_map(str::bytes)
    // }

    /// Returns an iterator over the grapheme clusters of this `Node`'s subrope'
    ///
    /// This is the iterator returned by `Node::into_iter`.
    #[inline]
    pub fn graphemes<'a>(&'a self) -> impl Iterator<Item=&'a str> {
        // the compiler won't let me mark this as unimplemented using the
        // unimplemented!() macro, due to Reasons (i suspect relating to
        // returning `impl Trait`)
        //  - eliza, 12/18/2016
        panic!("Unimplemented!");
        self.strings()
    }

}

/// An iterator over a series of leaf `Node`s
struct Leaves<'a>(Vec<&'a Node>);

impl<'a> Iterator for Leaves<'a> {
    type Item = &'a Node;
    fn next(&mut self) -> Option<Self::Item> {
        match self.0.pop() {
            None => None
          , Some(leaf @ &Leaf(_)) => Some(leaf)
          , Some(&Branch(BranchNode { ref left, ref right, .. })) => {
                self.0.push(left);
                self.0.push(right);
                self.next()
            }
        }
    }
}

/// A move iterator over a series of leaf `Node`s
struct IntoLeaves(Vec<Node>);

impl Iterator for IntoLeaves {
    type Item = Node;
    fn next(&mut self) -> Option<Self::Item> {
        match self.0.pop() {
            None => None
          , Some(leaf @ Leaf(_)) => Some(leaf)
          , Some(Branch(BranchNode { box left, box right, .. })) => {
                self.0.push(right);
                self.0.push(left);
                self.next()
            }
        }
    }
}

impl ops::Add for Node {
    type Output = Self;
    /// Concatenate two `Node`s, returning a `Branch` node.
    fn add(self, right: Self) -> Self { Node::new_branch(self, right) }
}
<|endoftext|>


<commit_before>use aster;
use syntax::ast::{self, MetaItem, TyPath};
use syntax::attr::AttrMetaMethods;
use syntax::codemap::Span;
use syntax::ext::base::{Annotatable, ExtCtxt};
use syntax::ptr::P;
use syntax::parse::token::{InternedString, intern_and_get_ident, str_to_ident};

use attr::Attr;
use model::Model;

pub fn expand_changeset_for(
    cx: &mut ExtCtxt,
    span: Span,
    meta_item: &MetaItem,
    annotatable: &Annotatable,
    push: &mut FnMut(Annotatable),
) {
    let builder = aster::AstBuilder::new().span(span);

    if let Some(model) = Model::from_annotable(cx, span, annotatable) {
        let options = changeset_options(cx, meta_item).unwrap();
        push(Annotatable::Item(changeset_impl(cx, builder, &options, &model).unwrap()));
        if let Some(item) = save_changes_impl(cx, builder, &options, &model) {
            push(Annotatable::Item(item));
        }
    } else {
        cx.span_err(meta_item.span,
            "`changeset_for` may only be apllied to enums and structs");
    }
}

struct ChangesetOptions {
    table_name: ast::Ident,
    skip_visibility: bool,
    treat_none_as_null: bool,
}

fn changeset_options(cx: &mut ExtCtxt, meta_item: &MetaItem) -> Result<ChangesetOptions, ()> {
    match meta_item.node {
        ast::MetaList(_, ref meta_items) => {
            let table_name = try!(table_name(cx, &meta_items[0]));
            let skip_visibility = try!(boolean_option(cx, &meta_items[1..], "__skip_visibility"))
                .unwrap_or(false);
            let treat_none_as_null = try!(boolean_option(cx, &meta_items[1..], "treat_none_as_null"))
                .unwrap_or(false);
            Ok(ChangesetOptions {
                table_name: str_to_ident(&table_name),
                skip_visibility: skip_visibility,
                treat_none_as_null: treat_none_as_null,
            })
        }
        _ => usage_error(cx, meta_item),
    }
}

fn table_name(cx: &mut ExtCtxt, meta_item: &MetaItem) -> Result<InternedString, ()> {
    match meta_item.node {
        ast::MetaWord(ref word) => Ok(word.clone()),
        _ => usage_error(cx, meta_item),
    }
}

fn boolean_option(cx: &mut ExtCtxt, meta_items: &[P<MetaItem>], option_name: &str)
    -> Result<Option<bool>, ()>
{
    if let Some(item) = meta_items.iter().find(|item| item.name() == option_name) {
        match item.value_str() {
            Some(ref s) if *s == "true" => Ok(Some(true)),
            Some(ref s) if *s == "false" => Ok(Some(false)),
            _ => {
                cx.span_err(item.span,
                    &format!("Expected {} to be in the form `option=\"true\"` or \
                            option=\"false\"", option_name));
                Err(())
            }
        }
    } else {
        Ok(None)
    }
}

fn usage_error<T>(cx: &mut ExtCtxt, meta_item: &MetaItem) -> Result<T, ()> {
    cx.span_err(meta_item.span,
        "`changeset_for` must be used in the form `#[changeset_for(table1)]`");
    Err(())
}

fn changeset_impl(
    cx: &mut ExtCtxt,
    builder: aster::AstBuilder,
    options: &ChangesetOptions,
    model: &Model,
) -> Option<P<ast::Item>> {
    let ref struct_name = model.ty;
    let pk = model.primary_key_name();
    let attrs_for_changeset = model.attrs.iter().filter(|a| a.column_name != pk)
        .collect::<Vec<_>>();
    let changeset_ty = builder.ty().tuple()
        .with_tys(attrs_for_changeset.iter()
                  .map(|a| changeset_ty(cx, builder, &options, a)))
        .build();
    let changeset_body = builder.expr().tuple()
        .with_exprs(attrs_for_changeset.iter()
                    .map(|a| changeset_expr(cx, builder, &options, a)))
        .build();
    quote_item!(cx,
        impl<'a: 'update, 'update> ::diesel::query_builder::AsChangeset for
            &'update $struct_name
        {
            type Changeset = $changeset_ty;

            fn as_changeset(self) -> Self::Changeset {
                $changeset_body
            }
        }
    )
}

#[allow(unused_mut)]
fn save_changes_impl(
    cx: &mut ExtCtxt,
    builder: aster::AstBuilder,
    options: &ChangesetOptions,
    model: &Model,
) -> Option<P<ast::Item>> {
    let ref struct_name = model.ty;
    let pk = model.primary_key_name();
    let sql_type = builder.path()
        .segment(&options.table_name).build()
        .segment("SqlType").build()
        .build();
    let table = builder.path()
        .segment(&options.table_name).build()
        .segment("table").build()
        .build();
    let _pub = if options.skip_visibility {
        quote_tokens!(cx, )
    } else {
        quote_tokens!(cx, pub)
    };
    model.attrs.iter().find(|a| a.column_name == pk).and_then(|pk| {
        let pk_field = pk.field_name.unwrap();
        quote_item!(cx,
            impl<'a> $struct_name {
                $_pub fn save_changes<T>(&self, connection: &::diesel::Connection)
                    -> ::diesel::QueryResult<T> where
                    T: Queryable<$sql_type>,
                {
                    use ::diesel::update;
                    update($table.filter($table.primary_key().eq(&self.$pk_field)))
                        .set(self)
                        .get_result(&connection)
                }
            }
        )
    })
}

fn changeset_ty(
    cx: &mut ExtCtxt,
    builder: aster::AstBuilder,
    options: &ChangesetOptions,
    attr: &Attr,
) -> P<ast::Ty> {
    let column = builder.path()
        .segment(options.table_name).build()
        .segment(attr.column_name).build()
        .build();
    match (options.treat_none_as_null, ty_param_of_option(&attr.ty)) {
        (false, Some(ty)) => {
            let inner_ty = inner_changeset_ty(cx, column, &ty);
            quote_ty!(cx, Option<$inner_ty>)
        }
        _ => inner_changeset_ty(cx, column, &attr.ty),
    }
}

fn inner_changeset_ty(
    cx: &mut ExtCtxt,
    column: ast::Path,
    field_ty: &ast::Ty,
) -> P<ast::Ty> {
    quote_ty!(cx,
        ::diesel::expression::predicates::Eq<
            $column,
            ::diesel::expression::bound::Bound<
                <$column as ::diesel::expression::Expression>::SqlType,
                &'update $field_ty,
            >,
        >
    )
}

fn changeset_expr(
    cx: &mut ExtCtxt,
    builder: aster::AstBuilder,
    options: &ChangesetOptions,
    attr: &Attr,
) -> P<ast::Expr> {
    let column = builder.path()
        .segment(options.table_name).build()
        .segment(attr.column_name).build()
        .build();
    let field_name = &attr.field_name.unwrap();
    if !options.treat_none_as_null && is_option_ty(&attr.ty) {
        quote_expr!(cx, self.$field_name.as_ref().map(|f| $column.eq(f)))
    } else {
        quote_expr!(cx, $column.eq(&self.$field_name))
    }
}

fn ty_param_of_option(ty: &ast::Ty) -> Option<&P<ast::Ty>> {
    match ty.node {
        TyPath(_, ref path) => {
            path.segments.first().iter()
                .filter(|s| s.identifier.name.as_str() == intern_and_get_ident("Option"))
                .flat_map(|s| s.parameters.types().first().map(|p| *p))
                .next()
        }
        _ => None,
    }
}

fn is_option_ty(ty: &ast::Ty) -> bool {
    ty_param_of_option(ty).is_some()
}
<commit_msg>Remove aster from `update` in `codegen`<commit_after>use syntax::ast::{self, MetaItem, TyPath};
use syntax::attr::AttrMetaMethods;
use syntax::codemap::Span;
use syntax::ext::base::{Annotatable, ExtCtxt};
use syntax::ext::build::AstBuilder;
use syntax::ptr::P;
use syntax::parse::token::{InternedString, intern_and_get_ident, str_to_ident};

use attr::Attr;
use model::Model;

pub fn expand_changeset_for(
    cx: &mut ExtCtxt,
    span: Span,
    meta_item: &MetaItem,
    annotatable: &Annotatable,
    push: &mut FnMut(Annotatable),
) {
    if let Some(model) = Model::from_annotable(cx, span, annotatable) {
        let options = changeset_options(cx, meta_item).unwrap();
        push(Annotatable::Item(changeset_impl(cx, span, &options, &model).unwrap()));
        if let Some(item) = save_changes_impl(cx, span, &options, &model) {
            push(Annotatable::Item(item));
        }
    } else {
        cx.span_err(meta_item.span,
            "`changeset_for` may only be apllied to enums and structs");
    }
}

struct ChangesetOptions {
    table_name: ast::Ident,
    skip_visibility: bool,
    treat_none_as_null: bool,
}

fn changeset_options(cx: &mut ExtCtxt, meta_item: &MetaItem) -> Result<ChangesetOptions, ()> {
    match meta_item.node {
        ast::MetaList(_, ref meta_items) => {
            let table_name = try!(table_name(cx, &meta_items[0]));
            let skip_visibility = try!(boolean_option(cx, &meta_items[1..], "__skip_visibility"))
                .unwrap_or(false);
            let treat_none_as_null = try!(boolean_option(cx, &meta_items[1..], "treat_none_as_null"))
                .unwrap_or(false);
            Ok(ChangesetOptions {
                table_name: str_to_ident(&table_name),
                skip_visibility: skip_visibility,
                treat_none_as_null: treat_none_as_null,
            })
        }
        _ => usage_error(cx, meta_item),
    }
}

fn table_name(cx: &mut ExtCtxt, meta_item: &MetaItem) -> Result<InternedString, ()> {
    match meta_item.node {
        ast::MetaWord(ref word) => Ok(word.clone()),
        _ => usage_error(cx, meta_item),
    }
}

fn boolean_option(cx: &mut ExtCtxt, meta_items: &[P<MetaItem>], option_name: &str)
    -> Result<Option<bool>, ()>
{
    if let Some(item) = meta_items.iter().find(|item| item.name() == option_name) {
        match item.value_str() {
            Some(ref s) if *s == "true" => Ok(Some(true)),
            Some(ref s) if *s == "false" => Ok(Some(false)),
            _ => {
                cx.span_err(item.span,
                    &format!("Expected {} to be in the form `option=\"true\"` or \
                            option=\"false\"", option_name));
                Err(())
            }
        }
    } else {
        Ok(None)
    }
}

fn usage_error<T>(cx: &mut ExtCtxt, meta_item: &MetaItem) -> Result<T, ()> {
    cx.span_err(meta_item.span,
        "`changeset_for` must be used in the form `#[changeset_for(table1)]`");
    Err(())
}

fn changeset_impl(
    cx: &mut ExtCtxt,
    span: Span,
    options: &ChangesetOptions,
    model: &Model,
) -> Option<P<ast::Item>> {
    let ref struct_name = model.ty;
    let pk = model.primary_key_name();
    let attrs_for_changeset = model.attrs.iter().filter(|a| a.column_name != pk)
        .collect::<Vec<_>>();
    let changeset_ty = cx.ty(span, ast::TyTup(
        attrs_for_changeset.iter()
              .map(|a| changeset_ty(cx, span, &options, a))
              .collect()
    ));
    let changeset_body = cx.expr_tuple(span, attrs_for_changeset.iter()
        .map(|a| changeset_expr(cx, span, &options, a))
        .collect());
    quote_item!(cx,
        impl<'a: 'update, 'update> ::diesel::query_builder::AsChangeset for
            &'update $struct_name
        {
            type Changeset = $changeset_ty;

            fn as_changeset(self) -> Self::Changeset {
                $changeset_body
            }
        }
    )
}

#[allow(unused_mut)]
fn save_changes_impl(
    cx: &mut ExtCtxt,
    span: Span,
    options: &ChangesetOptions,
    model: &Model,
) -> Option<P<ast::Item>> {
    let ref struct_name = model.ty;
    let pk = model.primary_key_name();
    let sql_type = cx.path(span, vec![options.table_name, str_to_ident("SqlType")]);
    let table = cx.path(span, vec![options.table_name, str_to_ident("table")]);
    let _pub = if options.skip_visibility {
        quote_tokens!(cx, )
    } else {
        quote_tokens!(cx, pub)
    };
    model.attrs.iter().find(|a| a.column_name == pk).and_then(|pk| {
        let pk_field = pk.field_name.unwrap();
        quote_item!(cx,
            impl<'a> $struct_name {
                $_pub fn save_changes<T>(&self, connection: &::diesel::Connection)
                    -> ::diesel::QueryResult<T> where
                    T: Queryable<$sql_type>,
                {
                    use ::diesel::update;
                    update($table.filter($table.primary_key().eq(&self.$pk_field)))
                        .set(self)
                        .get_result(&connection)
                }
            }
        )
    })
}

fn changeset_ty(
    cx: &ExtCtxt,
    span: Span,
    options: &ChangesetOptions,
    attr: &Attr,
) -> P<ast::Ty> {
    let column = cx.path(span, vec![options.table_name, attr.column_name]);
    match (options.treat_none_as_null, ty_param_of_option(&attr.ty)) {
        (false, Some(ty)) => {
            let inner_ty = inner_changeset_ty(cx, column, &ty);
            quote_ty!(cx, Option<$inner_ty>)
        }
        _ => inner_changeset_ty(cx, column, &attr.ty),
    }
}

fn inner_changeset_ty(
    cx: &ExtCtxt,
    column: ast::Path,
    field_ty: &ast::Ty,
) -> P<ast::Ty> {
    quote_ty!(cx,
        ::diesel::expression::predicates::Eq<
            $column,
            ::diesel::expression::bound::Bound<
                <$column as ::diesel::expression::Expression>::SqlType,
                &'update $field_ty,
            >,
        >
    )
}

fn changeset_expr(
    cx: &ExtCtxt,
    span: Span,
    options: &ChangesetOptions,
    attr: &Attr,
) -> P<ast::Expr> {
    let column = cx.path(span, vec![options.table_name, attr.column_name]);
    let field_name = &attr.field_name.unwrap();
    if !options.treat_none_as_null && is_option_ty(&attr.ty) {
        quote_expr!(cx, self.$field_name.as_ref().map(|f| $column.eq(f)))
    } else {
        quote_expr!(cx, $column.eq(&self.$field_name))
    }
}

fn ty_param_of_option(ty: &ast::Ty) -> Option<&P<ast::Ty>> {
    match ty.node {
        TyPath(_, ref path) => {
            path.segments.first().iter()
                .filter(|s| s.identifier.name.as_str() == intern_and_get_ident("Option"))
                .flat_map(|s| s.parameters.types().first().map(|p| *p))
                .next()
        }
        _ => None,
    }
}

fn is_option_ty(ty: &ast::Ty) -> bool {
    ty_param_of_option(ty).is_some()
}
<|endoftext|>


<commit_before>//! `i686`'s Streaming SIMD Extensions 4.1 (SSE4.1)

use crate::{
    core_arch::{simd_llvm::*, x86::*},
    mem::transmute,
};

#[cfg(test)]
use stdarch_test::assert_instr;

/// Extracts an 64-bit integer from `a` selected with `imm8`
///
/// [Intel's documentation](https://software.intel.com/sites/landingpage/IntrinsicsGuide/#text=_mm_extract_epi64)
#[inline]
#[target_feature(enable = "sse4.1")]
#[cfg_attr(all(test, not(target_os = "windows")), assert_instr(pextrq, imm8 = 1))]
#[rustc_args_required_const(1)]
#[stable(feature = "simd_x86", since = "1.27.0")]
pub unsafe fn _mm_extract_epi64(a: __m128i, imm8: i32) -> i64 {
    let a = a.as_i64x2();
    match imm8 & 1 {
        0 => simd_extract(a, 0),
        _ => simd_extract(a, 1),
    }
}

/// Returns a copy of `a` with the 64-bit integer from `i` inserted at a
/// location specified by `imm8`.
///
/// [Intel's documentation](https://software.intel.com/sites/landingpage/IntrinsicsGuide/#text=_mm_insert_epi64)
#[inline]
#[target_feature(enable = "sse4.1")]
#[cfg_attr(test, assert_instr(pinsrq, imm8 = 0))]
#[rustc_args_required_const(2)]
#[stable(feature = "simd_x86", since = "1.27.0")]
pub unsafe fn _mm_insert_epi64(a: __m128i, i: i64, imm8: i32) -> __m128i {
    let a = a.as_i64x2();
    match imm8 & 1 {
        0 => transmute(simd_insert(a, 0, i)),
        _ => transmute(simd_insert(a, 1, i)),
    }
}

#[cfg(test)]
mod tests {
    use crate::core_arch::arch::x86_64::*;
    use stdarch_test::simd_test;

    #[simd_test(enable = "sse4.1")]
    unsafe fn test_mm_extract_epi64() {
        let a = _mm_setr_epi64x(0, 1);
        let r = _mm_extract_epi64(a, 1);
        assert_eq!(r, 1);
        let r = _mm_extract_epi64(a, 3);
        assert_eq!(r, 1);
    }

    #[simd_test(enable = "sse4.1")]
    unsafe fn test_mm_insert_epi64() {
        let a = _mm_set1_epi64x(0);
        let e = _mm_setr_epi64x(0, 32);
        let r = _mm_insert_epi64(a, 32, 1);
        assert_eq_m128i(r, e);
        let r = _mm_insert_epi64(a, 32, 3);
        assert_eq_m128i(r, e);
    }
}
<commit_msg>convert `_mm_extract_epi64` to const generics<commit_after>//! `i686`'s Streaming SIMD Extensions 4.1 (SSE4.1)

use crate::{
    core_arch::{simd_llvm::*, x86::*},
    mem::transmute,
};

#[cfg(test)]
use stdarch_test::assert_instr;

/// Extracts an 64-bit integer from `a` selected with `IMM1`
///
/// [Intel's documentation](https://software.intel.com/sites/landingpage/IntrinsicsGuide/#text=_mm_extract_epi64)
#[inline]
#[target_feature(enable = "sse4.1")]
#[cfg_attr(all(test, not(target_os = "windows")), assert_instr(pextrq, IMM1 = 1))]
#[rustc_legacy_const_generics(1)]
#[stable(feature = "simd_x86", since = "1.27.0")]
pub unsafe fn _mm_extract_epi64<const IMM1: i32>(a: __m128i) -> i64 {
    static_assert_imm1!(IMM1);
    simd_extract(a.as_i64x2(), IMM1 as u32)
}

/// Returns a copy of `a` with the 64-bit integer from `i` inserted at a
/// location specified by `imm8`.
///
/// [Intel's documentation](https://software.intel.com/sites/landingpage/IntrinsicsGuide/#text=_mm_insert_epi64)
#[inline]
#[target_feature(enable = "sse4.1")]
#[cfg_attr(test, assert_instr(pinsrq, imm8 = 0))]
#[rustc_args_required_const(2)]
#[stable(feature = "simd_x86", since = "1.27.0")]
pub unsafe fn _mm_insert_epi64(a: __m128i, i: i64, imm8: i32) -> __m128i {
    let a = a.as_i64x2();
    match imm8 & 1 {
        0 => transmute(simd_insert(a, 0, i)),
        _ => transmute(simd_insert(a, 1, i)),
    }
}

#[cfg(test)]
mod tests {
    use crate::core_arch::arch::x86_64::*;
    use stdarch_test::simd_test;

    #[simd_test(enable = "sse4.1")]
    unsafe fn test_mm_extract_epi64() {
        let a = _mm_setr_epi64x(0, 1);
        let r = _mm_extract_epi64::<1>(a);
        assert_eq!(r, 1);
        let r = _mm_extract_epi64::<0>(a);
        assert_eq!(r, 0);
    }

    #[simd_test(enable = "sse4.1")]
    unsafe fn test_mm_insert_epi64() {
        let a = _mm_set1_epi64x(0);
        let e = _mm_setr_epi64x(0, 32);
        let r = _mm_insert_epi64(a, 32, 1);
        assert_eq_m128i(r, e);
        let r = _mm_insert_epi64(a, 32, 3);
        assert_eq_m128i(r, e);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
extern crate modbus;
use modbus::{Client, Coil};
use modbus::tcp;
use clap::{App};

fn main() {
    let matches = App::new("client")
                    .author("Falco Hirschenberger <falco.hirschenberger@gmail.com>")
                    .version(&crate_version!()[..])
                    .about("Modbus Tcp client")
                    .args_from_usage(
                        "<SERVER> 'The IP address or hostname of the server'
                        --read-coils=[ADDR] [QUANTITY] 'Read QUANTITY coils from ADDR'
                        --read-discrete-inputs=[ADDR] [QUANTITY] 'Read QUANTITY inputs from ADDR'
                        --write-single-coil=[ADDR] [On,Off] 'Write the coil value (On or Off) to \
                          ADDR'
                        --write-multiple-coils=[ADDR] [On,Off..] 'Write multiple coil values \
                          (On or Off) to ADDR (use \"..\" without spaces to group them like \
                          \"On,Off,On,Off\")'
                        --read-input-registers=[ADDR], [QUANTITY] 'Read QUANTITY input registers\
                          from ADDR'
                        --read-holding-registers=[ADDR], [QUANTITY] 'Read QUANTITY holding \
                          registers from ADDR'
                        --write-single-register=[ADDR] [VALUE] 'Write VALUE to register ADDR'
                        --write-multiple-registers=[ADDR] [V1,V2...] 'Write multiple register \
                          values to ADDR (use \"..\" without spaces to group them like \
                          \"23, 24, 25\")'")
                    .get_matches();

    let mut client = tcp::Transport::new(matches.value_of("SERVER").unwrap()).unwrap();

    if let Some(args) = matches.values_of("read-coils") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let qtty: u16 = args[1].parse().expect(matches.usage());
        println!("{:?}", client.read_coils(addr, qtty).expect("IO Error"));
    }
    else if let Some(args) = matches.values_of("read-discrete-inputs") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let qtty: u16 = args[1].parse().expect(matches.usage());
        println!("{:?}", client.read_discrete_inputs(addr, qtty).expect("IO Error"));
    }
    else if let Some(args) = matches.values_of("write-single-coil") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let value: Coil = args[1].parse().expect(matches.usage());
        client.write_single_coil(addr, value).expect("IO Error");
    }
    else if let Some(args) = matches.values_of("write-multiple-coils") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let values: Vec<Coil> = args[1].split(',').map(|s| s.parse().expect(matches.usage())).collect();
        client.write_multiple_coils(addr, &values).expect("IO Error");
    }
    else if let Some(args) = matches.values_of("read-holding-registers") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let qtty: u16 = args[1].parse().expect(matches.usage());
        println!("{:?}", client.read_holding_registers(addr, qtty).expect("IO Error"));
    }
    else if let Some(args) = matches.values_of("write-single-register") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let value: u16 = args[1].parse().expect(matches.usage());
        client.write_single_register(addr, value).expect("IO Error");
    }
    else if let Some(args) = matches.values_of("write-multiple-registers") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let values: Vec<u16> = args[1].split(',').map(|s| s.parse().expect(matches.usage())).collect();
        client.write_multiple_registers(addr, &values).expect("IO Error");
    }

}
<commit_msg>Trim the list values before parsing for more intuitive CLI behaviour<commit_after>#[macro_use]
extern crate clap;
extern crate modbus;
use modbus::{Client, Coil};
use modbus::tcp;
use clap::App;

fn main() {
    let matches = App::new("client")
                      .author("Falco Hirschenberger <falco.hirschenberger@gmail.com>")
                      .version(&crate_version!()[..])
                      .about("Modbus Tcp client")
                      .args_from_usage("<SERVER> 'The IP address or hostname of the server'
                        \
                                        --read-coils=[ADDR] [QUANTITY] 'Read QUANTITY coils from \
                                        ADDR'
                        \
                                        --read-discrete-inputs=[ADDR] [QUANTITY] 'Read QUANTITY \
                                        inputs from ADDR'
                        \
                                        --write-single-coil=[ADDR] [On,Off] 'Write the coil \
                                        value (On or Off) to ADDR'
                        \
                                        --write-multiple-coils=[ADDR] [On,Off..] 'Write multiple \
                                        coil values (On or Off) to ADDR (use \"..\" without \
                                        spaces to group them e.g. \"On, Off, On, Off\")'
                        \
                                        --read-input-registers=[ADDR], [QUANTITY] 'Read QUANTITY \
                                        input registersfrom ADDR'
                        \
                                        --read-holding-registers=[ADDR], [QUANTITY] 'Read \
                                        QUANTITY holding registers from ADDR'
                        \
                                        --write-single-register=[ADDR] [VALUE] 'Write VALUE to \
                                        register ADDR'
                        \
                                        --write-multiple-registers=[ADDR] [V1,V2...] 'Write \
                                        multiple register values to ADDR (use \"..\" to group \
                                        them e.g. \"23, 24, 25\")'")
                      .get_matches();

    let mut client = tcp::Transport::new(matches.value_of("SERVER").unwrap()).unwrap();

    if let Some(args) = matches.values_of("read-coils") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let qtty: u16 = args[1].parse().expect(matches.usage());
        println!("{:?}", client.read_coils(addr, qtty).expect("IO Error"));
    } else if let Some(args) = matches.values_of("read-discrete-inputs") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let qtty: u16 = args[1].parse().expect(matches.usage());
        println!("{:?}",
                 client.read_discrete_inputs(addr, qtty).expect("IO Error"));
    } else if let Some(args) = matches.values_of("write-single-coil") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let value: Coil = args[1].parse().expect(matches.usage());
        client.write_single_coil(addr, value).expect("IO Error");
    } else if let Some(args) = matches.values_of("write-multiple-coils") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let values: Vec<Coil> = args[1]
                                    .split(',')
                                    .map(|s| s.trim().parse().expect(matches.usage()))
                                    .collect();
        client.write_multiple_coils(addr, &values).expect("IO Error");
    } else if let Some(args) = matches.values_of("read-holding-registers") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let qtty: u16 = args[1].parse().expect(matches.usage());
        println!("{:?}",
                 client.read_holding_registers(addr, qtty).expect("IO Error"));
    } else if let Some(args) = matches.values_of("write-single-register") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let value: u16 = args[1].parse().expect(matches.usage());
        client.write_single_register(addr, value).expect("IO Error");
    } else if let Some(args) = matches.values_of("write-multiple-registers") {
        let addr: u16 = args[0].parse().expect(matches.usage());
        let values: Vec<u16> = args[1]
                                   .split(',')
                                   .map(|s| s.trim().parse().expect(matches.usage()))
                                   .collect();
        client.write_multiple_registers(addr, &values).expect("IO Error");
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2018, Cloudflare, Inc.
// Copyright (C) 2018, Alessandro Ghedini
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * 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 THE COPYRIGHT HOLDERS AND CONTRIBUTORS "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 THE COPYRIGHT HOLDER 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.

#[macro_use]
extern crate log;

use std::net;

use std::collections::HashMap;

use ring::rand::*;

const LOCAL_CONN_ID_LEN: usize = 16;

const MAX_DATAGRAM_SIZE: usize = 1350;

const USAGE: &str = "Usage:
  server [options]
  server -h | --help

Options:
  --listen <addr>   Listen on the given IP:port [default: 127.0.0.1:4433]
  --cert <file>     TLS certificate path [default: examples/cert.crt]
  --key <file>      TLS certificate key path [default: examples/cert.key]
  --root <dir>      Root directory [default: examples/root/]
  --name <str>      Name of the server [default: quic.tech]
  --no-retry        Disable stateless retry.
  -h --help         Show this screen.
";

type ConnMap = HashMap<Vec<u8>, (net::SocketAddr, Box<quiche::Connection>)>;

fn main() -> Result<(), Box<std::error::Error>> {
    let mut buf = [0; 65535];
    let mut out = [0; MAX_DATAGRAM_SIZE];

    env_logger::init();

    let args = docopt::Docopt::new(USAGE)
        .and_then(|dopt| dopt.parse())
        .unwrap_or_else(|e| e.exit());

    let socket = net::UdpSocket::bind(args.get_str("--listen"))?;

    let poll = mio::Poll::new()?;
    let mut events = mio::Events::with_capacity(1024);

    let socket = mio::net::UdpSocket::from_socket(socket)?;
    poll.register(
        &socket,
        mio::Token(0),
        mio::Ready::readable(),
        mio::PollOpt::edge(),
    )?;

    let mut connections = ConnMap::new();

    let mut config = quiche::Config::new(quiche::VERSION_DRAFT18)?;

    config.load_cert_chain_from_pem_file(args.get_str("--cert"))?;
    config.load_priv_key_from_pem_file(args.get_str("--key"))?;

    config.set_application_protos(&[b"h3-18", b"hq-18", b"http/0.9"])?;

    config.set_idle_timeout(30);
    config.set_max_packet_size(MAX_DATAGRAM_SIZE as u64);
    config.set_initial_max_data(10_000_000);
    config.set_initial_max_stream_data_bidi_local(1_000_000);
    config.set_initial_max_stream_data_bidi_remote(1_000_000);
    config.set_initial_max_stream_data_uni(1_000_000);
    config.set_initial_max_streams_bidi(100);
    config.set_initial_max_streams_uni(5);
    config.set_disable_migration(true);

    loop {
        // TODO: use event loop that properly supports timers
        let timeout = connections.values().filter_map(|(_, c)| c.timeout()).min();

        poll.poll(&mut events, timeout)?;

        'read: loop {
            if events.is_empty() {
                debug!("timed out");

                connections.values_mut().for_each(|(_, c)| c.on_timeout());

                break 'read;
            }

            let (len, src) = match socket.recv_from(&mut buf) {
                Ok(v) => v,

                Err(e) => {
                    if e.kind() == std::io::ErrorKind::WouldBlock {
                        debug!("recv() would block");
                        break 'read;
                    }

                    panic!("recv() failed: {:?}", e);
                },
            };

            debug!("got {} bytes", len);

            let pkt_buf = &mut buf[..len];

            let hdr = match quiche::Header::from_slice(pkt_buf, LOCAL_CONN_ID_LEN)
            {
                Ok(v) => v,

                Err(e) => {
                    error!("Parsing packet header failed: {:?}", e);
                    continue;
                },
            };

            trace!("got packet {:?}", hdr);

            if hdr.ty == quiche::Type::VersionNegotiation {
                error!("Version negotiation invalid on the server");
                continue;
            }

            let (_, conn) = if !connections.contains_key(&hdr.dcid) {
                if hdr.ty != quiche::Type::Initial {
                    error!("Packet is not Initial");
                    continue;
                }

                if hdr.version != quiche::VERSION_DRAFT18 {
                    warn!("Doing version negotiation");

                    let len = quiche::negotiate_version(
                        &hdr.scid, &hdr.dcid, &mut out,
                    )?;

                    let out = &out[..len];

                    socket.send_to(out, &src)?;
                    continue;
                }

                let mut scid = [0; LOCAL_CONN_ID_LEN];
                SystemRandom::new().fill(&mut scid[..])?;

                let mut odcid = None;

                if !args.get_bool("--no-retry") {
                    // Token is always present in Initial packets.
                    let token = hdr.token.as_ref().unwrap();

                    if token.is_empty() {
                        warn!("Doing stateless retry");

                        let new_token = mint_token(&hdr, &src);

                        let len = quiche::retry(
                            &hdr.scid, &hdr.dcid, &scid, &new_token, &mut out,
                        )?;

                        let out = &out[..len];

                        socket.send_to(out, &src)?;
                        continue;
                    }

                    odcid = validate_token(&src, token);

                    if odcid == None {
                        error!("Invalid address validation token");
                        continue;
                    }

                    scid.copy_from_slice(&hdr.dcid);
                }

                debug!(
                    "New connection: dcid={} scid={}",
                    hex_dump(&hdr.dcid),
                    hex_dump(&scid)
                );

                let conn = quiche::accept(&scid, odcid, &mut config)?;

                connections.insert(scid.to_vec(), (src, conn));

                connections.get_mut(&scid[..]).unwrap()
            } else {
                connections.get_mut(&hdr.dcid).unwrap()
            };

            // Process potentially coalesced packets.
            let read = match conn.recv(pkt_buf) {
                Ok(v) => v,

                Err(quiche::Error::Done) => {
                    debug!("{} done reading", conn.trace_id());
                    break;
                },

                Err(e) => {
                    error!("{} recv failed: {:?}", conn.trace_id(), e);
                    conn.close(false, e.to_wire(), b"fail")?;
                    break 'read;
                },
            };

            debug!("{} processed {} bytes", conn.trace_id(), read);

            let streams: Vec<u64> = conn.readable().collect();
            for s in streams {
                while let Ok((read, fin)) = conn.stream_recv(s, &mut buf) {
                    debug!("{} received {} bytes", conn.trace_id(), read);

                    let stream_buf = &buf[..read];

                    debug!(
                        "{} stream {} has {} bytes (fin? {})",
                        conn.trace_id(),
                        s,
                        stream_buf.len(),
                        fin
                    );

                    handle_stream(conn, s, stream_buf, args.get_str("--root"));
                }
            }
        }

        for (peer, conn) in connections.values_mut() {
            loop {
                let write = match conn.send(&mut out) {
                    Ok(v) => v,

                    Err(quiche::Error::Done) => {
                        debug!("{} done writing", conn.trace_id());
                        break;
                    },

                    Err(e) => {
                        error!("{} send failed: {:?}", conn.trace_id(), e);
                        conn.close(false, e.to_wire(), b"fail")?;
                        break;
                    },
                };

                // TODO: coalesce packets.
                socket.send_to(&out[..write], &peer)?;

                debug!("{} written {} bytes", conn.trace_id(), write);
            }
        }

        // Garbage collect closed connections.
        connections.retain(|_, (_, ref mut c)| {
            debug!("Collecting garbage");

            if c.is_closed() {
                info!("{} connection collected {:?}", c.trace_id(), c.stats());
            }

            !c.is_closed()
        });
    }
}

fn handle_stream(
    conn: &mut quiche::Connection, stream: u64, buf: &[u8], root: &str,
) {
    if buf.len() > 4 && &buf[..4] == b"GET " {
        let uri = &buf[4..buf.len()];
        let uri = String::from_utf8(uri.to_vec()).unwrap();
        let uri = String::from(uri.lines().next().unwrap());
        let uri = std::path::Path::new(&uri);
        let mut path = std::path::PathBuf::from(root);

        for c in uri.components() {
            if let std::path::Component::Normal(v) = c {
                path.push(v)
            }
        }

        info!(
            "{} got GET request for {:?} on stream {}",
            conn.trace_id(),
            path,
            stream
        );

        let data = std::fs::read(path.as_path())
            .unwrap_or_else(|_| b"Not Found!\r\n".to_vec());

        info!(
            "{} sending response of size {} on stream {}",
            conn.trace_id(),
            data.len(),
            stream
        );

        if let Err(e) = conn.stream_send(stream, &data, true) {
            error!("{} stream send failed {:?}", conn.trace_id(), e);
        }
    }
}

fn mint_token(hdr: &quiche::Header, src: &net::SocketAddr) -> Vec<u8> {
    let mut token = Vec::new();

    token.extend_from_slice(b"quiche");

    let addr = match src.ip() {
        std::net::IpAddr::V4(a) => a.octets().to_vec(),
        std::net::IpAddr::V6(a) => a.octets().to_vec(),
    };

    token.extend_from_slice(&addr);
    token.extend_from_slice(&hdr.dcid);

    token
}

fn validate_token<'a>(
    src: &net::SocketAddr, token: &'a [u8],
) -> Option<&'a [u8]> {
    if token.len() < 6 {
        return None;
    }

    if &token[..6] != b"quiche" {
        return None;
    }

    let token = &token[6..];

    let addr = match src.ip() {
        std::net::IpAddr::V4(a) => a.octets().to_vec(),
        std::net::IpAddr::V6(a) => a.octets().to_vec(),
    };

    if token.len() < addr.len() || &token[..addr.len()] != addr.as_slice() {
        return None;
    }

    let token = &token[addr.len()..];

    Some(&token[..])
}

fn hex_dump(buf: &[u8]) -> String {
    let vec: Vec<String> = buf.iter().map(|b| format!("{:02x}", b)).collect();

    vec.join("")
}
<commit_msg>server: add support for enabling key logging<commit_after>// Copyright (C) 2018, Cloudflare, Inc.
// Copyright (C) 2018, Alessandro Ghedini
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * 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 THE COPYRIGHT HOLDERS AND CONTRIBUTORS "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 THE COPYRIGHT HOLDER 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.

#[macro_use]
extern crate log;

use std::net;

use std::collections::HashMap;

use ring::rand::*;

const LOCAL_CONN_ID_LEN: usize = 16;

const MAX_DATAGRAM_SIZE: usize = 1350;

const USAGE: &str = "Usage:
  server [options]
  server -h | --help

Options:
  --listen <addr>   Listen on the given IP:port [default: 127.0.0.1:4433]
  --cert <file>     TLS certificate path [default: examples/cert.crt]
  --key <file>      TLS certificate key path [default: examples/cert.key]
  --root <dir>      Root directory [default: examples/root/]
  --name <str>      Name of the server [default: quic.tech]
  --no-retry        Disable stateless retry.
  -h --help         Show this screen.
";

type ConnMap = HashMap<Vec<u8>, (net::SocketAddr, Box<quiche::Connection>)>;

fn main() -> Result<(), Box<std::error::Error>> {
    let mut buf = [0; 65535];
    let mut out = [0; MAX_DATAGRAM_SIZE];

    env_logger::init();

    let args = docopt::Docopt::new(USAGE)
        .and_then(|dopt| dopt.parse())
        .unwrap_or_else(|e| e.exit());

    let socket = net::UdpSocket::bind(args.get_str("--listen"))?;

    let poll = mio::Poll::new()?;
    let mut events = mio::Events::with_capacity(1024);

    let socket = mio::net::UdpSocket::from_socket(socket)?;
    poll.register(
        &socket,
        mio::Token(0),
        mio::Ready::readable(),
        mio::PollOpt::edge(),
    )?;

    let mut connections = ConnMap::new();

    let mut config = quiche::Config::new(quiche::VERSION_DRAFT18)?;

    config.load_cert_chain_from_pem_file(args.get_str("--cert"))?;
    config.load_priv_key_from_pem_file(args.get_str("--key"))?;

    config.set_application_protos(&[b"h3-18", b"hq-18", b"http/0.9"])?;

    config.set_idle_timeout(30);
    config.set_max_packet_size(MAX_DATAGRAM_SIZE as u64);
    config.set_initial_max_data(10_000_000);
    config.set_initial_max_stream_data_bidi_local(1_000_000);
    config.set_initial_max_stream_data_bidi_remote(1_000_000);
    config.set_initial_max_stream_data_uni(1_000_000);
    config.set_initial_max_streams_bidi(100);
    config.set_initial_max_streams_uni(5);
    config.set_disable_migration(true);

    if std::env::var_os("SSLKEYLOGFILE").is_some() {
        config.log_keys();
    }

    loop {
        // TODO: use event loop that properly supports timers
        let timeout = connections.values().filter_map(|(_, c)| c.timeout()).min();

        poll.poll(&mut events, timeout)?;

        'read: loop {
            if events.is_empty() {
                debug!("timed out");

                connections.values_mut().for_each(|(_, c)| c.on_timeout());

                break 'read;
            }

            let (len, src) = match socket.recv_from(&mut buf) {
                Ok(v) => v,

                Err(e) => {
                    if e.kind() == std::io::ErrorKind::WouldBlock {
                        debug!("recv() would block");
                        break 'read;
                    }

                    panic!("recv() failed: {:?}", e);
                },
            };

            debug!("got {} bytes", len);

            let pkt_buf = &mut buf[..len];

            let hdr = match quiche::Header::from_slice(pkt_buf, LOCAL_CONN_ID_LEN)
            {
                Ok(v) => v,

                Err(e) => {
                    error!("Parsing packet header failed: {:?}", e);
                    continue;
                },
            };

            trace!("got packet {:?}", hdr);

            if hdr.ty == quiche::Type::VersionNegotiation {
                error!("Version negotiation invalid on the server");
                continue;
            }

            let (_, conn) = if !connections.contains_key(&hdr.dcid) {
                if hdr.ty != quiche::Type::Initial {
                    error!("Packet is not Initial");
                    continue;
                }

                if hdr.version != quiche::VERSION_DRAFT18 {
                    warn!("Doing version negotiation");

                    let len = quiche::negotiate_version(
                        &hdr.scid, &hdr.dcid, &mut out,
                    )?;

                    let out = &out[..len];

                    socket.send_to(out, &src)?;
                    continue;
                }

                let mut scid = [0; LOCAL_CONN_ID_LEN];
                SystemRandom::new().fill(&mut scid[..])?;

                let mut odcid = None;

                if !args.get_bool("--no-retry") {
                    // Token is always present in Initial packets.
                    let token = hdr.token.as_ref().unwrap();

                    if token.is_empty() {
                        warn!("Doing stateless retry");

                        let new_token = mint_token(&hdr, &src);

                        let len = quiche::retry(
                            &hdr.scid, &hdr.dcid, &scid, &new_token, &mut out,
                        )?;

                        let out = &out[..len];

                        socket.send_to(out, &src)?;
                        continue;
                    }

                    odcid = validate_token(&src, token);

                    if odcid == None {
                        error!("Invalid address validation token");
                        continue;
                    }

                    scid.copy_from_slice(&hdr.dcid);
                }

                debug!(
                    "New connection: dcid={} scid={}",
                    hex_dump(&hdr.dcid),
                    hex_dump(&scid)
                );

                let conn = quiche::accept(&scid, odcid, &mut config)?;

                connections.insert(scid.to_vec(), (src, conn));

                connections.get_mut(&scid[..]).unwrap()
            } else {
                connections.get_mut(&hdr.dcid).unwrap()
            };

            // Process potentially coalesced packets.
            let read = match conn.recv(pkt_buf) {
                Ok(v) => v,

                Err(quiche::Error::Done) => {
                    debug!("{} done reading", conn.trace_id());
                    break;
                },

                Err(e) => {
                    error!("{} recv failed: {:?}", conn.trace_id(), e);
                    conn.close(false, e.to_wire(), b"fail")?;
                    break 'read;
                },
            };

            debug!("{} processed {} bytes", conn.trace_id(), read);

            let streams: Vec<u64> = conn.readable().collect();
            for s in streams {
                while let Ok((read, fin)) = conn.stream_recv(s, &mut buf) {
                    debug!("{} received {} bytes", conn.trace_id(), read);

                    let stream_buf = &buf[..read];

                    debug!(
                        "{} stream {} has {} bytes (fin? {})",
                        conn.trace_id(),
                        s,
                        stream_buf.len(),
                        fin
                    );

                    handle_stream(conn, s, stream_buf, args.get_str("--root"));
                }
            }
        }

        for (peer, conn) in connections.values_mut() {
            loop {
                let write = match conn.send(&mut out) {
                    Ok(v) => v,

                    Err(quiche::Error::Done) => {
                        debug!("{} done writing", conn.trace_id());
                        break;
                    },

                    Err(e) => {
                        error!("{} send failed: {:?}", conn.trace_id(), e);
                        conn.close(false, e.to_wire(), b"fail")?;
                        break;
                    },
                };

                // TODO: coalesce packets.
                socket.send_to(&out[..write], &peer)?;

                debug!("{} written {} bytes", conn.trace_id(), write);
            }
        }

        // Garbage collect closed connections.
        connections.retain(|_, (_, ref mut c)| {
            debug!("Collecting garbage");

            if c.is_closed() {
                info!("{} connection collected {:?}", c.trace_id(), c.stats());
            }

            !c.is_closed()
        });
    }
}

fn handle_stream(
    conn: &mut quiche::Connection, stream: u64, buf: &[u8], root: &str,
) {
    if buf.len() > 4 && &buf[..4] == b"GET " {
        let uri = &buf[4..buf.len()];
        let uri = String::from_utf8(uri.to_vec()).unwrap();
        let uri = String::from(uri.lines().next().unwrap());
        let uri = std::path::Path::new(&uri);
        let mut path = std::path::PathBuf::from(root);

        for c in uri.components() {
            if let std::path::Component::Normal(v) = c {
                path.push(v)
            }
        }

        info!(
            "{} got GET request for {:?} on stream {}",
            conn.trace_id(),
            path,
            stream
        );

        let data = std::fs::read(path.as_path())
            .unwrap_or_else(|_| b"Not Found!\r\n".to_vec());

        info!(
            "{} sending response of size {} on stream {}",
            conn.trace_id(),
            data.len(),
            stream
        );

        if let Err(e) = conn.stream_send(stream, &data, true) {
            error!("{} stream send failed {:?}", conn.trace_id(), e);
        }
    }
}

fn mint_token(hdr: &quiche::Header, src: &net::SocketAddr) -> Vec<u8> {
    let mut token = Vec::new();

    token.extend_from_slice(b"quiche");

    let addr = match src.ip() {
        std::net::IpAddr::V4(a) => a.octets().to_vec(),
        std::net::IpAddr::V6(a) => a.octets().to_vec(),
    };

    token.extend_from_slice(&addr);
    token.extend_from_slice(&hdr.dcid);

    token
}

fn validate_token<'a>(
    src: &net::SocketAddr, token: &'a [u8],
) -> Option<&'a [u8]> {
    if token.len() < 6 {
        return None;
    }

    if &token[..6] != b"quiche" {
        return None;
    }

    let token = &token[6..];

    let addr = match src.ip() {
        std::net::IpAddr::V4(a) => a.octets().to_vec(),
        std::net::IpAddr::V6(a) => a.octets().to_vec(),
    };

    if token.len() < addr.len() || &token[..addr.len()] != addr.as_slice() {
        return None;
    }

    let token = &token[addr.len()..];

    Some(&token[..])
}

fn hex_dump(buf: &[u8]) -> String {
    let vec: Vec<String> = buf.iter().map(|b| format!("{:02x}", b)).collect();

    vec.join("")
}
<|endoftext|>