<commit_before>/// The tokenizer for the query interpreter

use error::*;

#[derive(Debug, PartialEq, Eq)]
pub enum Token {
    Identifier {
        ident: String,
        next: Option<Box<Token>>
    },

    Index {
        idx: usize,
        next: Option<Box<Token>>
    }
}

impl Token {

    pub fn next(&self) -> Option<&Box<Token>> {
        match self {
            &Token::Identifier { ref next, .. } => next.as_ref(),
            &Token::Index { ref next, .. }      => next.as_ref(),
        }
    }

    /// Convenience function for `token.next().is_some()`
    pub fn has_next(&self) -> bool {
        self.next().is_some()
    }

    pub fn set_next(&mut self, token: Token) {
        match self {
            &mut Token::Identifier { ref mut next, .. } => *next = Some(Box::new(token)),
            &mut Token::Index { ref mut next, .. }      => *next = Some(Box::new(token)),
        }
    }

    #[cfg(test)]
    pub fn identifier(&self) -> &String {
        match self {
            &Token::Identifier { ident: ref ident, .. } => &ident,
            _ => unreachable!(),
        }
    }

    #[cfg(test)]
    pub fn idx(&self) -> usize {
        match self {
            &Token::Index { idx: i, .. } => i,
            _ => unreachable!(),
        }
    }
}

pub fn tokenize_with_seperator(query: &String, seperator: char) -> Result<Token> {
    use std::str::Split;

    /// Creates a Token object from a string
    ///
    /// # Panics
    ///
    /// * If the internal regex does not compile (should never happen)
    /// * If the token is non-valid (that is, a array index with a non-i64)
    /// * If the regex does not find anything
    /// * If the integer in the brackets (`[]`) cannot be parsed to a valid i64
    ///
    /// # Incorrect behaviour
    ///
    /// * If the regex finds multiple captures
    ///
    /// # Returns
    ///
    /// The `Token` object with the correct identifier/index for this token and no next token.
    ///
    fn mk_token_object(s: &str) -> Result<Token> {
        use regex::Regex;
        use std::str::FromStr;

        lazy_static! {
            static ref RE: Regex = Regex::new(r"^\[\d+\]$").unwrap();
        }

        if !has_array_brackets(s) {
            return Ok(Token::Identifier { ident: String::from(s), next: None });
        }

        match RE.captures(s) {
            None => return Err(Error::from(ErrorKind::ArrayAccessWithoutIndex)),
            Some(captures) => {
                match captures.get(0) {
                    None => Ok(Token::Identifier { ident: String::from(s), next: None }),
                    Some(mtch) => {
                        let mtch = without_array_brackets(mtch.as_str());
                        let i : usize = FromStr::from_str(&mtch).unwrap(); // save because regex
                        Ok(Token::Index {
                            idx: i,
                            next: None,
                        })
                    }
                }
            }
        }
    }

    /// Check whether a str begins with '[' and ends with ']'
    fn has_array_brackets(s: &str) -> bool {
        s.as_bytes()[0] == b'[' && s.as_bytes()[s.len() - 1] == b']'
    }

    /// Remove '[' and ']' from a str
    fn without_array_brackets(s: &str) -> String {
        s.replace("[","").replace("]","")
    }

    fn build_token_tree(split: &mut Split<char>, last: &mut Token) -> Result<()> {
        match split.next() {
            None        => { /* No more tokens */ }
            Some(token) => {
                if token.len() == 0 {
                    return Err(Error::from(ErrorKind::EmptyIdentifier));
                }
                let mut token = try!(mk_token_object(token));
                try!(build_token_tree(split, &mut token));
                last.set_next(token);
            }
        }
        Ok(())
    }

    if query.is_empty() {
        return Err(Error::from(ErrorKind::EmptyQueryError));
    }

    let mut tokens = query.split(seperator);

    match tokens.next() {
        None        => Err(Error::from(ErrorKind::EmptyQueryError)),
        Some(token) => {
            if token.len() == 0 {
                return Err(Error::from(ErrorKind::EmptyIdentifier));
            }
            let mut tok = try!(mk_token_object(token));
            let _       = try!(build_token_tree(&mut tokens, &mut tok));
            Ok(tok)
        }
    }
}

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

    use std::ops::Deref;

    #[test]
    fn test_tokenize_empty_query_to_error() {
        let tokens = tokenize_with_seperator(&String::from(""), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::EmptyQueryError { .. }));
    }

    #[test]
    fn test_tokenize_seperator_only() {
        let tokens = tokenize_with_seperator(&String::from("."), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::EmptyIdentifier { .. }));
    }

    #[test]
    fn test_tokenize_array_brackets_only() {
        let tokens = tokenize_with_seperator(&String::from("[]"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::ArrayAccessWithoutIndex { .. }));
    }

    #[test]
    fn test_tokenize_identifiers_with_array_brackets_only() {
        let tokens = tokenize_with_seperator(&String::from("a.b.c.[]"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::ArrayAccessWithoutIndex { .. }));
    }

    #[test]
    fn test_tokenize_identifiers_in_array_brackets() {
        let tokens = tokenize_with_seperator(&String::from("[a]"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::ArrayAccessWithoutIndex { .. }));
    }

    #[test]
    fn test_tokenize_single_token_query() {
        let tokens = tokenize_with_seperator(&String::from("example"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert!(match tokens {
            Token::Identifier {
                ident: ident,
                next: None
            } => {
                assert_eq!(String::from("example"), ident);
                true
            },
            _ => false,
        });
    }

    #[test]
    fn test_tokenize_double_token_query() {
        let tokens = tokenize_with_seperator(&String::from("a.b"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert!(match tokens {
            Token::Identifier { next: Some(ref next), .. } => {
                assert_eq!("b", next.deref().identifier());
                match next.deref() {
                    &Token::Identifier { next: None, .. } => true,
                    _ => false
                }
            },
            _ => false,
        });
        assert_eq!("a", tokens.identifier());
    }

    #[test]
    fn test_tokenize_ident_then_array_query() {
        let tokens = tokenize_with_seperator(&String::from("a.[0]"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert_eq!("a", tokens.identifier());
        assert!(match tokens {
            Token::Identifier { next: Some(ref next), .. } => {
                match next.deref() {
                    &Token::Index { idx: 0, next: None } => true,
                    _ => false
                }
            },
            _ => false,
        });
    }

    #[test]
    fn test_tokenize_many_idents_then_array_query() {
        use std::ops::Deref;

        let tokens = tokenize_with_seperator(&String::from("a.b.c.[1000]"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert_eq!("a", tokens.identifier());

        let expected =
            Token::Identifier {
                ident: String::from("a"),
                next: Some(Box::new(Token::Identifier {
                    ident: String::from("b"),
                    next: Some(Box::new(Token::Identifier {
                        ident: String::from("c"),
                        next: Some(Box::new(Token::Index {
                            idx: 1000,
                            next: None,
                        })),
                    })),
                })),
            };

        assert_eq!(expected, tokens);
    }

    #[test]
    fn test_tokenize_empty_token_after_good_token() {
        let tokens = tokenize_with_seperator(&String::from("a..b"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::EmptyIdentifier { .. }));
    }

    quickcheck! {
        fn test_array_index(i: usize) -> bool {
            match tokenize_with_seperator(&format!("[{}]", i), '.') {
                Ok(Token::Index {
                    idx: i,
                    next: None,
                }) => true,
                _ => false,
            }
        }
    }

}
<commit_msg>Remove unused function<commit_after>/// The tokenizer for the query interpreter

use error::*;

#[derive(Debug, PartialEq, Eq)]
pub enum Token {
    Identifier {
        ident: String,
        next: Option<Box<Token>>
    },

    Index {
        idx: usize,
        next: Option<Box<Token>>
    }
}

impl Token {

    pub fn next(&self) -> Option<&Box<Token>> {
        match self {
            &Token::Identifier { ref next, .. } => next.as_ref(),
            &Token::Index { ref next, .. }      => next.as_ref(),
        }
    }

    pub fn set_next(&mut self, token: Token) {
        match self {
            &mut Token::Identifier { ref mut next, .. } => *next = Some(Box::new(token)),
            &mut Token::Index { ref mut next, .. }      => *next = Some(Box::new(token)),
        }
    }

    #[cfg(test)]
    pub fn identifier(&self) -> &String {
        match self {
            &Token::Identifier { ident: ref ident, .. } => &ident,
            _ => unreachable!(),
        }
    }

    #[cfg(test)]
    pub fn idx(&self) -> usize {
        match self {
            &Token::Index { idx: i, .. } => i,
            _ => unreachable!(),
        }
    }
}

pub fn tokenize_with_seperator(query: &String, seperator: char) -> Result<Token> {
    use std::str::Split;

    /// Creates a Token object from a string
    ///
    /// # Panics
    ///
    /// * If the internal regex does not compile (should never happen)
    /// * If the token is non-valid (that is, a array index with a non-i64)
    /// * If the regex does not find anything
    /// * If the integer in the brackets (`[]`) cannot be parsed to a valid i64
    ///
    /// # Incorrect behaviour
    ///
    /// * If the regex finds multiple captures
    ///
    /// # Returns
    ///
    /// The `Token` object with the correct identifier/index for this token and no next token.
    ///
    fn mk_token_object(s: &str) -> Result<Token> {
        use regex::Regex;
        use std::str::FromStr;

        lazy_static! {
            static ref RE: Regex = Regex::new(r"^\[\d+\]$").unwrap();
        }

        if !has_array_brackets(s) {
            return Ok(Token::Identifier { ident: String::from(s), next: None });
        }

        match RE.captures(s) {
            None => return Err(Error::from(ErrorKind::ArrayAccessWithoutIndex)),
            Some(captures) => {
                match captures.get(0) {
                    None => Ok(Token::Identifier { ident: String::from(s), next: None }),
                    Some(mtch) => {
                        let mtch = without_array_brackets(mtch.as_str());
                        let i : usize = FromStr::from_str(&mtch).unwrap(); // save because regex
                        Ok(Token::Index {
                            idx: i,
                            next: None,
                        })
                    }
                }
            }
        }
    }

    /// Check whether a str begins with '[' and ends with ']'
    fn has_array_brackets(s: &str) -> bool {
        s.as_bytes()[0] == b'[' && s.as_bytes()[s.len() - 1] == b']'
    }

    /// Remove '[' and ']' from a str
    fn without_array_brackets(s: &str) -> String {
        s.replace("[","").replace("]","")
    }

    fn build_token_tree(split: &mut Split<char>, last: &mut Token) -> Result<()> {
        match split.next() {
            None        => { /* No more tokens */ }
            Some(token) => {
                if token.len() == 0 {
                    return Err(Error::from(ErrorKind::EmptyIdentifier));
                }
                let mut token = try!(mk_token_object(token));
                try!(build_token_tree(split, &mut token));
                last.set_next(token);
            }
        }
        Ok(())
    }

    if query.is_empty() {
        return Err(Error::from(ErrorKind::EmptyQueryError));
    }

    let mut tokens = query.split(seperator);

    match tokens.next() {
        None        => Err(Error::from(ErrorKind::EmptyQueryError)),
        Some(token) => {
            if token.len() == 0 {
                return Err(Error::from(ErrorKind::EmptyIdentifier));
            }
            let mut tok = try!(mk_token_object(token));
            let _       = try!(build_token_tree(&mut tokens, &mut tok));
            Ok(tok)
        }
    }
}

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

    use std::ops::Deref;

    #[test]
    fn test_tokenize_empty_query_to_error() {
        let tokens = tokenize_with_seperator(&String::from(""), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::EmptyQueryError { .. }));
    }

    #[test]
    fn test_tokenize_seperator_only() {
        let tokens = tokenize_with_seperator(&String::from("."), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::EmptyIdentifier { .. }));
    }

    #[test]
    fn test_tokenize_array_brackets_only() {
        let tokens = tokenize_with_seperator(&String::from("[]"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::ArrayAccessWithoutIndex { .. }));
    }

    #[test]
    fn test_tokenize_identifiers_with_array_brackets_only() {
        let tokens = tokenize_with_seperator(&String::from("a.b.c.[]"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::ArrayAccessWithoutIndex { .. }));
    }

    #[test]
    fn test_tokenize_identifiers_in_array_brackets() {
        let tokens = tokenize_with_seperator(&String::from("[a]"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::ArrayAccessWithoutIndex { .. }));
    }

    #[test]
    fn test_tokenize_single_token_query() {
        let tokens = tokenize_with_seperator(&String::from("example"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert!(match tokens {
            Token::Identifier {
                ident: ident,
                next: None
            } => {
                assert_eq!(String::from("example"), ident);
                true
            },
            _ => false,
        });
    }

    #[test]
    fn test_tokenize_double_token_query() {
        let tokens = tokenize_with_seperator(&String::from("a.b"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert!(match tokens {
            Token::Identifier { next: Some(ref next), .. } => {
                assert_eq!("b", next.deref().identifier());
                match next.deref() {
                    &Token::Identifier { next: None, .. } => true,
                    _ => false
                }
            },
            _ => false,
        });
        assert_eq!("a", tokens.identifier());
    }

    #[test]
    fn test_tokenize_ident_then_array_query() {
        let tokens = tokenize_with_seperator(&String::from("a.[0]"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert_eq!("a", tokens.identifier());
        assert!(match tokens {
            Token::Identifier { next: Some(ref next), .. } => {
                match next.deref() {
                    &Token::Index { idx: 0, next: None } => true,
                    _ => false
                }
            },
            _ => false,
        });
    }

    #[test]
    fn test_tokenize_many_idents_then_array_query() {
        use std::ops::Deref;

        let tokens = tokenize_with_seperator(&String::from("a.b.c.[1000]"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert_eq!("a", tokens.identifier());

        let expected =
            Token::Identifier {
                ident: String::from("a"),
                next: Some(Box::new(Token::Identifier {
                    ident: String::from("b"),
                    next: Some(Box::new(Token::Identifier {
                        ident: String::from("c"),
                        next: Some(Box::new(Token::Index {
                            idx: 1000,
                            next: None,
                        })),
                    })),
                })),
            };

        assert_eq!(expected, tokens);
    }

    #[test]
    fn test_tokenize_empty_token_after_good_token() {
        let tokens = tokenize_with_seperator(&String::from("a..b"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::EmptyIdentifier { .. }));
    }

    quickcheck! {
        fn test_array_index(i: usize) -> bool {
            match tokenize_with_seperator(&format!("[{}]", i), '.') {
                Ok(Token::Index {
                    idx: i,
                    next: None,
                }) => true,
                _ => false,
            }
        }
    }

}
<|endoftext|>


<commit_before>//! Transport layer.
use bytecodec::io::{BufferedIo, IoDecodeExt, IoEncodeExt};
use bytecodec::{self, Decode, DecodeExt, Encode, EncodeExt, Eos};
use fibers::net::futures::{RecvFrom, SendTo};
use fibers::net::{TcpStream, UdpSocket};
use futures::{Async, Future};
use std::collections::VecDeque;
use std::net::SocketAddr;
use stun_codec::{Attribute, MessageDecoder, MessageEncoder, TransactionId};
use trackable::error::ErrorKindExt;

use constants;
use {Error, ErrorKind, Result};

pub trait Transport {
    type Decoder: Decode;
    type Encoder: Encode;

    fn send(&mut self, peer: SocketAddr, item: <Self::Encoder as Encode>::Item);
    fn recv(&mut self) -> Option<(SocketAddr, <Self::Decoder as Decode>::Item)>;
    fn run_once(&mut self) -> Result<bool>;
}

#[derive(Debug)]
pub enum RecvResult<T> {
    None,
    Some {
        peer: SocketAddr,
        item: T,
    },
    DecodeError {
        peer: SocketAddr,
        error: bytecodec::Error,
    },
}

pub trait UnreliableTransport: Transport {}

pub trait StunTransport<A>:
    Transport<Decoder = MessageDecoder<A>, Encoder = MessageEncoder<A>>
where
    A: Attribute,
{
    fn cancel_retransmission(&mut self, transaction_id: TransactionId);
}

#[derive(Debug)]
pub struct TcpTransporter<D: Decode, E: Encode> {
    stream: BufferedIo<TcpStream>,
    peer: SocketAddr,
    decoder: D,
    encoder: E,
    outgoing_queue: VecDeque<E::Item>,
    last_error: Option<Error>,
}
impl<D, E> TcpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{
    pub fn connect(peer: SocketAddr) -> impl Future<Item = Self, Error = Error> {
        TcpStream::connect(peer)
            .map(move |stream| Self::from((peer, stream)))
            .map_err(|e| track!(Error::from(e)))
    }

    fn start_send(&mut self, item: E::Item) -> Result<()> {
        if self.encoder.is_idle() {
            track!(self.encoder.start_encoding(item))?;
        } else {
            self.outgoing_queue.push_back(item);
        }
        track!(self.poll_send())?;
        Ok(())
    }

    fn poll_send(&mut self) -> Result<()> {
        while !self.stream.is_eos() {
            track!(self.stream.execute_io())?;
            track!(
                self.encoder
                    .encode_to_write_buf(self.stream.write_buf_mut())
            )?;
            if self.encoder.is_idle() {
                if let Some(item) = self.outgoing_queue.pop_front() {
                    track!(self.encoder.start_encoding(item))?;
                } else {
                    break;
                }
            }
            if self.stream.would_block() {
                break;
            }
        }
        Ok(())
    }

    fn poll_recv(&mut self) -> Result<Option<(SocketAddr, D::Item)>> {
        while self.stream.is_eos() {
            track!(self.stream.execute_io())?;
            track!(
                self.decoder
                    .decode_from_read_buf(self.stream.read_buf_mut())
            )?;
            if self.decoder.is_idle() {
                let item = track!(self.decoder.finish_decoding())?;
                return Ok(Some((self.peer, item)));
            }
            if self.stream.would_block() {
                break;
            }
        }
        Ok(None)
    }
}
impl<D, E> From<(SocketAddr, TcpStream)> for TcpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{
    fn from((peer, stream): (SocketAddr, TcpStream)) -> Self {
        let _ = stream.set_nodelay(true);
        TcpTransporter {
            stream: BufferedIo::new(stream, 4096, 4096),
            peer,
            decoder: D::default(),
            encoder: E::default(),
            outgoing_queue: VecDeque::new(),
            last_error: None,
        }
    }
}
impl<D, E> Transport for TcpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{
    type Decoder = D;
    type Encoder = E;

    fn send(&mut self, peer: SocketAddr, item: E::Item) {
        if self.last_error.is_some() {
            return;
        }
        if peer != self.peer {
            let e = ErrorKind::InvalidInput.cause(format!(
                "Unexpected destination peer: actual={}, expected={}",
                peer, self.peer
            ));
            self.last_error = Some(e.into());
            return;
        }
        self.last_error = self.start_send(item).err();
    }

    fn recv(&mut self) -> Option<(SocketAddr, D::Item)> {
        if self.last_error.is_some() {
            return None;
        }
        match self.poll_recv() {
            Err(e) => {
                self.last_error = Some(e);
                None
            }
            Ok(item) => item,
        }
    }

    fn run_once(&mut self) -> Result<bool> {
        if let Some(e) = self.last_error.take() {
            return Err(track!(e));
        }
        track!(self.poll_send())?;
        if self.stream.is_eos() {
            track!(self.decoder.decode(&[], Eos::new(true)))?;
            Ok(true)
        } else {
            Ok(false)
        }
    }
}
impl<A> StunTransport<A> for TcpTransporter<MessageDecoder<A>, MessageEncoder<A>>
where
    A: Attribute,
{
    fn cancel_retransmission(&mut self, _transaction_id: TransactionId) {}
}

#[derive(Debug)]
pub struct UdpTransporter<D: Decode, E: Encode> {
    socket: UdpSocket,
    decoder: D,
    encoder: E,
    outgoing_queue: VecDeque<(SocketAddr, E::Item)>,
    send_to: Option<SendTo<Vec<u8>>>,
    recv_from: RecvFrom<Vec<u8>>, // TODO: parameter
    last_error: Option<Error>,
}
impl<D, E> UdpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{
    pub fn bind(addr: SocketAddr) -> impl Future<Item = Self, Error = Error> {
        UdpSocket::bind(addr)
            .map(Self::from)
            .map_err(|e| track!(Error::from(e)))
    }

    fn poll_send(&mut self) -> Result<()> {
        while track!(
            self.send_to
                .poll()
                .map_err(|(_, _, e)| track!(Error::from(e)))
        )?.is_ready()
        {
            if let Some((peer, item)) = self.outgoing_queue.pop_front() {
                let bytes = track!(self.encoder.encode_into_bytes(item))?;
                self.send_to = Some(self.socket.clone().send_to(bytes, peer));
            } else {
                self.send_to = None;
                break;
            }
        }
        Ok(())
    }

    fn poll_recv(&mut self) -> Result<Option<(SocketAddr, D::Item)>> {
        while let Async::Ready((socket, buf, size, peer)) = self
            .recv_from
            .poll()
            .map_err(|(_, _, e)| track!(Error::from(e)))?
        {
            let item = track!(self.decoder.decode_from_bytes(&buf[..size]))?;
            self.recv_from = socket.recv_from(buf);
            return Ok(Some((peer, item)));
        }
        Ok(None)
    }
}
impl<D, E> From<UdpSocket> for UdpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{
    fn from(f: UdpSocket) -> Self {
        let recv_from = f
            .clone()
            .recv_from(vec![0; constants::DEFAULT_MAX_MESSAGE_SIZE]);
        UdpTransporter {
            socket: f,
            decoder: D::default(),
            encoder: E::default(),
            outgoing_queue: VecDeque::new(),
            send_to: None,
            recv_from,
            last_error: None,
        }
    }
}
impl<D, E> Transport for UdpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{
    type Decoder = D;
    type Encoder = E;

    fn send(&mut self, peer: SocketAddr, item: E::Item) {
        if self.last_error.is_some() {
            return;
        }
        self.outgoing_queue.push_back((peer, item));
        self.last_error = self.poll_send().err();
    }

    fn recv(&mut self) -> Option<(SocketAddr, D::Item)> {
        if self.last_error.is_some() {
            return None;
        }
        match self.poll_recv() {
            Err(e) => {
                self.last_error = Some(e);
                None
            }
            Ok(item) => item,
        }
    }

    fn run_once(&mut self) -> Result<bool> {
        if let Some(e) = self.last_error.take() {
            return Err(track!(e));
        }
        track!(self.poll_send())?;
        Ok(false)
    }
}
impl<D, E> UnreliableTransport for UdpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{}

#[derive(Debug)]
pub struct RetransmitTransporter<A, T> {
    inner: T,
    _phatom: ::std::marker::PhantomData<A>,
}
impl<A, T> RetransmitTransporter<A, T>
where
    A: Attribute,
    T: UnreliableTransport<Decoder = MessageDecoder<A>, Encoder = MessageEncoder<A>>,
{
    pub fn new(inner: T) -> Self {
        RetransmitTransporter {
            inner,
            _phatom: Default::default(),
        }
    }
}
impl<A, T> Transport for RetransmitTransporter<A, T>
where
    A: Attribute,
    T: UnreliableTransport<Decoder = MessageDecoder<A>, Encoder = MessageEncoder<A>>,
{
    type Decoder = MessageDecoder<A>;
    type Encoder = MessageEncoder<A>;

    // TODO:
    fn send(&mut self, peer: SocketAddr, item: <Self::Encoder as Encode>::Item) {
        self.inner.send(peer, item);
    }
    fn recv(&mut self) -> Option<(SocketAddr, <Self::Decoder as Decode>::Item)> {
        self.inner.recv()
    }
    fn run_once(&mut self) -> Result<bool> {
        track!(self.inner.run_once())
    }
}
impl<A, T> StunTransport<A> for RetransmitTransporter<A, T>
where
    A: Attribute,
    T: UnreliableTransport<Decoder = MessageDecoder<A>, Encoder = MessageEncoder<A>>,
{
    fn cancel_retransmission(&mut self, _transaction_id: TransactionId) {
        panic!("TODO")
    }
}
<commit_msg>Add `UdpTransporterBuilder`<commit_after>//! Transport layer.
use bytecodec::io::{BufferedIo, IoDecodeExt, IoEncodeExt};
use bytecodec::{self, Decode, DecodeExt, Encode, EncodeExt, Eos};
use fibers::net::futures::{RecvFrom, SendTo};
use fibers::net::{TcpStream, UdpSocket};
use futures::{Async, Future};
use std::collections::VecDeque;
use std::net::SocketAddr;
use stun_codec::{Attribute, MessageDecoder, MessageEncoder, TransactionId};
use trackable::error::ErrorKindExt;

use constants;
use {Error, ErrorKind, Result};

pub trait Transport {
    type Decoder: Decode;
    type Encoder: Encode;

    fn send(&mut self, peer: SocketAddr, item: <Self::Encoder as Encode>::Item);
    fn recv(&mut self) -> Option<(SocketAddr, <Self::Decoder as Decode>::Item)>;
    fn run_once(&mut self) -> Result<bool>;
}

#[derive(Debug)]
pub enum RecvResult<T> {
    None,
    Some {
        peer: SocketAddr,
        item: T,
    },
    DecodeError {
        peer: SocketAddr,
        error: bytecodec::Error,
    },
}

pub trait UnreliableTransport: Transport {}

pub trait StunTransport<A>:
    Transport<Decoder = MessageDecoder<A>, Encoder = MessageEncoder<A>>
where
    A: Attribute,
{
    fn cancel_retransmission(&mut self, transaction_id: TransactionId);
}

#[derive(Debug)]
pub struct TcpTransporter<D: Decode, E: Encode> {
    stream: BufferedIo<TcpStream>,
    peer: SocketAddr,
    decoder: D,
    encoder: E,
    outgoing_queue: VecDeque<E::Item>,
    last_error: Option<Error>,
}
impl<D, E> TcpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{
    pub fn connect(peer: SocketAddr) -> impl Future<Item = Self, Error = Error> {
        TcpStream::connect(peer)
            .map(move |stream| Self::from((peer, stream)))
            .map_err(|e| track!(Error::from(e)))
    }

    pub fn stream_ref(&self) -> &TcpStream {
        self.stream.stream_ref()
    }

    pub fn stream_mut(&mut self) -> &mut TcpStream {
        self.stream.stream_mut()
    }

    pub fn peer_addr(&self) -> SocketAddr {
        self.peer
    }

    pub fn outgoing_queue_len(&self) -> usize {
        self.outgoing_queue.len() + if self.encoder.is_idle() { 0 } else { 1 }
    }

    pub fn decoder_ref(&self) -> &D {
        &self.decoder
    }

    pub fn decoder_mut(&mut self) -> &mut D {
        &mut self.decoder
    }

    pub fn encoder_ref(&self) -> &E {
        &self.encoder
    }

    pub fn encoder_mut(&mut self) -> &mut E {
        &mut self.encoder
    }

    fn start_send(&mut self, item: E::Item) -> Result<()> {
        if self.encoder.is_idle() {
            track!(self.encoder.start_encoding(item))?;
        } else {
            self.outgoing_queue.push_back(item);
        }
        track!(self.poll_send())?;
        Ok(())
    }

    fn poll_send(&mut self) -> Result<()> {
        while !self.stream.is_eos() {
            track!(self.stream.execute_io())?;
            track!(
                self.encoder
                    .encode_to_write_buf(self.stream.write_buf_mut())
            )?;
            if self.encoder.is_idle() {
                if let Some(item) = self.outgoing_queue.pop_front() {
                    track!(self.encoder.start_encoding(item))?;
                } else {
                    break;
                }
            }
            if self.stream.would_block() {
                break;
            }
        }
        Ok(())
    }

    fn poll_recv(&mut self) -> Result<Option<(SocketAddr, D::Item)>> {
        while self.stream.is_eos() {
            track!(self.stream.execute_io())?;
            track!(
                self.decoder
                    .decode_from_read_buf(self.stream.read_buf_mut())
            )?;
            if self.decoder.is_idle() {
                let item = track!(self.decoder.finish_decoding())?;
                return Ok(Some((self.peer, item)));
            }
            if self.stream.would_block() {
                break;
            }
        }
        Ok(None)
    }
}
impl<D, E> From<(SocketAddr, TcpStream)> for TcpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{
    fn from((peer, stream): (SocketAddr, TcpStream)) -> Self {
        let _ = stream.set_nodelay(true);
        TcpTransporter {
            stream: BufferedIo::new(stream, 8192, 8192),
            peer,
            decoder: D::default(),
            encoder: E::default(),
            outgoing_queue: VecDeque::new(),
            last_error: None,
        }
    }
}
impl<D, E> Transport for TcpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{
    type Decoder = D;
    type Encoder = E;

    fn send(&mut self, peer: SocketAddr, item: E::Item) {
        if self.last_error.is_some() {
            return;
        }
        if peer != self.peer {
            let e = ErrorKind::InvalidInput.cause(format!(
                "Unexpected destination peer: actual={}, expected={}",
                peer, self.peer
            ));
            self.last_error = Some(e.into());
            return;
        }
        self.last_error = self.start_send(item).err();
    }

    fn recv(&mut self) -> Option<(SocketAddr, D::Item)> {
        if self.last_error.is_some() {
            return None;
        }
        match self.poll_recv() {
            Err(e) => {
                self.last_error = Some(e);
                None
            }
            Ok(item) => item,
        }
    }

    fn run_once(&mut self) -> Result<bool> {
        if let Some(e) = self.last_error.take() {
            return Err(track!(e));
        }
        track!(self.poll_send())?;
        if self.stream.is_eos() {
            track!(self.decoder.decode(&[], Eos::new(true)))?;
            Ok(true)
        } else {
            Ok(false)
        }
    }
}
impl<A> StunTransport<A> for TcpTransporter<MessageDecoder<A>, MessageEncoder<A>>
where
    A: Attribute,
{
    fn cancel_retransmission(&mut self, _transaction_id: TransactionId) {}
}

#[derive(Debug, Clone)]
pub struct UdpTransporterBuilder {
    recv_buf_size: usize,
}
impl UdpTransporterBuilder {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn recv_buf_size(&mut self, size: usize) -> &mut Self {
        self.recv_buf_size = size;
        self
    }

    pub fn bind<D, E>(
        &self,
        addr: SocketAddr,
    ) -> impl Future<Item = UdpTransporter<D, E>, Error = Error>
    where
        D: Decode + Default,
        E: Encode + Default,
    {
        let recv_buf = vec![0; self.recv_buf_size];
        UdpSocket::bind(addr)
            .map(|socket| {
                let recv_from = socket.clone().recv_from(recv_buf);
                UdpTransporter {
                    socket,
                    decoder: D::default(),
                    encoder: E::default(),
                    outgoing_queue: VecDeque::new(),
                    send_to: None,
                    recv_from,
                    last_error: None,
                }
            })
            .map_err(|e| track!(Error::from(e)))
    }
}
impl Default for UdpTransporterBuilder {
    fn default() -> Self {
        UdpTransporterBuilder {
            recv_buf_size: constants::DEFAULT_MAX_MESSAGE_SIZE,
        }
    }
}

#[derive(Debug)]
pub struct UdpTransporter<D: Decode, E: Encode> {
    socket: UdpSocket,
    decoder: D,
    encoder: E,
    outgoing_queue: VecDeque<(SocketAddr, E::Item)>,
    send_to: Option<SendTo<Vec<u8>>>,
    recv_from: RecvFrom<Vec<u8>>, // TODO: parameter
    last_error: Option<Error>,
}
impl<D, E> UdpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{
    pub fn bind(addr: SocketAddr) -> impl Future<Item = Self, Error = Error> {
        UdpSocket::bind(addr)
            .map(Self::from)
            .map_err(|e| track!(Error::from(e)))
    }

    fn poll_send(&mut self) -> Result<()> {
        while track!(
            self.send_to
                .poll()
                .map_err(|(_, _, e)| track!(Error::from(e)))
        )?.is_ready()
        {
            if let Some((peer, item)) = self.outgoing_queue.pop_front() {
                let bytes = track!(self.encoder.encode_into_bytes(item))?;
                self.send_to = Some(self.socket.clone().send_to(bytes, peer));
            } else {
                self.send_to = None;
                break;
            }
        }
        Ok(())
    }

    fn poll_recv(&mut self) -> Result<Option<(SocketAddr, D::Item)>> {
        while let Async::Ready((socket, buf, size, peer)) = self
            .recv_from
            .poll()
            .map_err(|(_, _, e)| track!(Error::from(e)))?
        {
            let item = track!(self.decoder.decode_from_bytes(&buf[..size]))?;
            self.recv_from = socket.recv_from(buf);
            return Ok(Some((peer, item)));
        }
        Ok(None)
    }
}
impl<D, E> From<UdpSocket> for UdpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{
    fn from(f: UdpSocket) -> Self {
        let recv_from = f
            .clone()
            .recv_from(vec![0; constants::DEFAULT_MAX_MESSAGE_SIZE]);
        UdpTransporter {
            socket: f,
            decoder: D::default(),
            encoder: E::default(),
            outgoing_queue: VecDeque::new(),
            send_to: None,
            recv_from,
            last_error: None,
        }
    }
}
impl<D, E> Transport for UdpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{
    type Decoder = D;
    type Encoder = E;

    fn send(&mut self, peer: SocketAddr, item: E::Item) {
        if self.last_error.is_some() {
            return;
        }
        self.outgoing_queue.push_back((peer, item));
        self.last_error = self.poll_send().err();
    }

    fn recv(&mut self) -> Option<(SocketAddr, D::Item)> {
        if self.last_error.is_some() {
            return None;
        }
        match self.poll_recv() {
            Err(e) => {
                self.last_error = Some(e);
                None
            }
            Ok(item) => item,
        }
    }

    fn run_once(&mut self) -> Result<bool> {
        if let Some(e) = self.last_error.take() {
            return Err(track!(e));
        }
        track!(self.poll_send())?;
        Ok(false)
    }
}
impl<D, E> UnreliableTransport for UdpTransporter<D, E>
where
    D: Decode + Default,
    E: Encode + Default,
{}

#[derive(Debug)]
pub struct RetransmitTransporter<A, T> {
    inner: T,
    _phatom: ::std::marker::PhantomData<A>,
}
impl<A, T> RetransmitTransporter<A, T>
where
    A: Attribute,
    T: UnreliableTransport<Decoder = MessageDecoder<A>, Encoder = MessageEncoder<A>>,
{
    pub fn new(inner: T) -> Self {
        RetransmitTransporter {
            inner,
            _phatom: Default::default(),
        }
    }
}
impl<A, T> Transport for RetransmitTransporter<A, T>
where
    A: Attribute,
    T: UnreliableTransport<Decoder = MessageDecoder<A>, Encoder = MessageEncoder<A>>,
{
    type Decoder = MessageDecoder<A>;
    type Encoder = MessageEncoder<A>;

    // TODO:
    fn send(&mut self, peer: SocketAddr, item: <Self::Encoder as Encode>::Item) {
        self.inner.send(peer, item);
    }
    fn recv(&mut self) -> Option<(SocketAddr, <Self::Decoder as Decode>::Item)> {
        self.inner.recv()
    }
    fn run_once(&mut self) -> Result<bool> {
        track!(self.inner.run_once())
    }
}
impl<A, T> StunTransport<A> for RetransmitTransporter<A, T>
where
    A: Attribute,
    T: UnreliableTransport<Decoder = MessageDecoder<A>, Encoder = MessageEncoder<A>>,
{
    fn cancel_retransmission(&mut self, _transaction_id: TransactionId) {
        panic!("TODO")
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for implementing transport layer protocols
//!
//! The transport module provides the ability to send and receive packets at
//! the transport layer using IPv4 or IPv6. It also enables layer 3 networking
//! for specific transport protocols, using IPv4 only.
//!
//! Note that this is limited by operating system support - for example, on OS
//! X and FreeBSD, it is impossible to implement protocols which are already
//! implemented in the kernel such as TCP and UDP.

#![macro_use]

extern crate libc;

use self::TransportProtocol::{Ipv4, Ipv6};
use self::TransportChannelType::{Layer3, Layer4};

use std::io;
use std::io::Error;
use std::iter::repeat;
use std::net;
use std::mem;
use std::sync::Arc;

use packet::Packet;
use packet::ip::IpNextHeaderProtocol;
use packet::ipv4::Ipv4Packet;
use packet::udp::UdpPacket;

use internal;
use util;

/// Represents a transport layer protocol
#[derive(Clone, Copy)]
pub enum TransportProtocol {
    /// Represents a transport protocol built on top of IPv4
    Ipv4(IpNextHeaderProtocol),
    /// Represents a transport protocol built on top of IPv6
    Ipv6(IpNextHeaderProtocol),
}

/// Type of transport channel to present
#[derive(Clone, Copy)]
pub enum TransportChannelType {
    /// The application will send and receive transport layer packets
    Layer4(TransportProtocol),
    /// The application will send and receive IPv4 packets, with the specified transport protocol
    Layer3(IpNextHeaderProtocol),
}

/// Structure used for sending at the transport layer. Should be created with transport_channel()
pub struct TransportSender {
    socket: Arc<internal::FileDesc>,
    _channel_type: TransportChannelType,
}

/// Structure used for sending at the transport layer. Should be created with transport_channel()
pub struct TransportReceiver {
    socket: Arc<internal::FileDesc>,
    buffer: Vec<u8>,
    channel_type: TransportChannelType,
}

#[cfg(windows)]
const INVALID_SOCKET: libc::SOCKET = libc::INVALID_SOCKET;

#[cfg(not(windows))]
const INVALID_SOCKET: libc::c_int = -1;

/// Create a new (TransportSender, TransportReceiver) pair
///
/// This allows for sending and receiving packets at the transport layer. The buffer size should be
/// large enough to handle the largest packet you wish to receive.
///
/// The channel type specifies what layer to send and receive packets at, and the transport
/// protocol you wish to implement. For example, `Layer4(Ipv4(IpNextHeaderProtocols::Udp))` would
/// allow sending and receiving UDP packets using IPv4; whereas Layer3(IpNextHeaderProtocols::Udp)
/// would include the IPv4 Header in received values, and require manual construction of an IP
/// header when sending.
pub fn transport_channel(buffer_size: usize,
                         channel_type: TransportChannelType)
    -> io::Result<(TransportSender, TransportReceiver)> {
    use std::net;

    // This hack makes sure that winsock is initialised
    let _ = {
        let ip = net::Ipv4Addr::new(255, 255, 255, 255);
        let sockaddr = net::SocketAddr::V4(net::SocketAddrV4::new(ip, 0));

        net::UdpSocket::bind(sockaddr)
    };

    let socket = unsafe {
        match channel_type {
            Layer4(Ipv4(IpNextHeaderProtocol(proto))) | Layer3(IpNextHeaderProtocol(proto)) =>
                libc::socket(libc::AF_INET, libc::SOCK_RAW, proto as libc::c_int),
            Layer4(Ipv6(IpNextHeaderProtocol(proto))) =>
                libc::socket(libc::AF_INET6, libc::SOCK_RAW, proto as libc::c_int),
        }
    };
    if socket != INVALID_SOCKET {
        if match channel_type {
            Layer3(_) | Layer4(Ipv4(_)) => true,
            _ => false,
        } {
            let hincl: libc::c_int = match channel_type {
                Layer4(..) => 0,
                _ => 1,
            };
            let res = unsafe {
                libc::setsockopt(socket,
                                 libc::IPPROTO_IP,
                                 libc::IP_HDRINCL,
                                 (&hincl as *const libc::c_int) as *const libc::c_void,
                                 mem::size_of::<libc::c_int>() as libc::socklen_t)
            };
            if res == -1 {
                let err = Error::last_os_error();
                unsafe {
                    internal::close(socket);
                }
                return Err(err);
            }
        }

        let sock = Arc::new(internal::FileDesc { fd: socket });
        let sender = TransportSender {
            socket: sock.clone(),
            _channel_type: channel_type,
        };
        let receiver = TransportReceiver {
            socket: sock,
            buffer: repeat(0u8).take(buffer_size).collect(),
            channel_type: channel_type,
        };

        Ok((sender, receiver))
    } else {
        Err(Error::last_os_error())
    }
}

impl TransportSender {
    fn send<T: Packet>(&mut self, packet: T, dst: util::IpAddr) -> io::Result<usize> {
        let mut caddr = unsafe { mem::zeroed() };
        let sockaddr = match dst {
            util::IpAddr::V4(ip_addr) =>
                net::SocketAddr::V4(net::SocketAddrV4::new(ip_addr, 0)),
            util::IpAddr::V6(ip_addr) =>
                net::SocketAddr::V6(net::SocketAddrV6::new(ip_addr, 0, 0, 0)),
        };
        let slen = internal::addr_to_sockaddr(sockaddr, &mut caddr);
        let caddr_ptr = (&caddr as *const libc::sockaddr_storage) as *const libc::sockaddr;

        internal::send_to(self.socket.fd, packet.packet(), caddr_ptr, slen)
    }

    /// Send a packet to the provided desination
    #[inline]
    pub fn send_to<T: Packet>(&mut self, packet: T, destination: util::IpAddr) -> io::Result<usize> {
        self.send_to_impl(packet, destination)
    }

    #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: util::IpAddr) -> io::Result<usize> {
        self.send(packet, dst)
    }

    #[cfg(any(target_os = "freebsd", target_os = "macos"))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: util::IpAddr) -> io::Result<usize> {
        use packet::MutablePacket;
        use packet::ipv4::MutableIpv4Packet;

        // FreeBSD and OS X expect total length and fragment offset fields of IPv4
        // packets to be in
        // host byte order rather than network byte order (man 4 ip/Raw IP Sockets)
        if match self._channel_type {
            Layer3(..) => true,
            _ => false,
        } {
            let mut mut_slice: Vec<u8> = repeat(0u8).take(packet.packet().len()).collect();

            let mut new_packet = MutableIpv4Packet::new(&mut mut_slice[..]).unwrap();
            new_packet.clone_from(&packet);
            let length = new_packet.get_total_length().to_be();
            new_packet.set_total_length(length);
            let offset = new_packet.get_fragment_offset().to_be();
            new_packet.set_fragment_offset(offset);

            return self.send(new_packet, dst);
        }

        self.send(packet, dst)
    }
}

/// Create an iterator for some packet type.
///
/// Usage:
/// ```
/// transport_channel_iterator!(Ipv4Packet, // Type to iterate over
///                             Ipv4TransportChannelIterator, // Name for iterator struct
///                             ipv4_packet_iter) // Name of function to create iterator
/// ```
#[macro_export]
macro_rules! transport_channel_iterator {
    ($ty:ident, $iter:ident, $func:ident) => (
        /// An iterator over packets of type $ty
        pub struct $iter<'a> {
            tr: &'a mut TransportReceiver
        }
        /// Return a packet iterator with packets of type $ty for some transport receiver
        pub fn $func(tr: &mut TransportReceiver) -> $iter {
            $iter {
                tr: tr
            }
        }
        impl<'a> $iter<'a> {
            /// Get the next ($ty, IpAddr) pair for the given channel
            pub fn next(&mut self) -> io::Result<($ty, util::IpAddr)> {
                let mut caddr: libc::sockaddr_storage = unsafe { mem::zeroed() };
                let res = internal::recv_from(self.tr.socket.fd,
                                              &mut self.tr.buffer[..],
                                              &mut caddr);

                let offset = match self.tr.channel_type {
                    Layer4(Ipv4(_)) => {
                        let ip_header = Ipv4Packet::new(&self.tr.buffer[..]).unwrap();

                        ip_header.get_header_length() as usize * 4usize
                    },
                    Layer3(_) => {
                        fixup_packet(&mut self.tr.buffer[..]);

                        0
                    },
                    _ => 0
                };
                return match res {
                    Ok(len) => {
                        let packet = $ty::new(&self.tr.buffer[offset..len]).unwrap();
                        let addr = internal::sockaddr_to_addr(
                                        &caddr,
                                        mem::size_of::<libc::sockaddr_storage>()
                                   );
                        let ip = match addr.unwrap() {
                            net::SocketAddr::V4(sa) => util::IpAddr::V4(*sa.ip()),
                            net::SocketAddr::V6(sa) => util::IpAddr::V6(*sa.ip()),
                        };
                        Ok((packet, ip))
                    },
                    Err(e) => Err(e),
                };

                #[cfg(any(target_os = "freebsd", target_os = "macos"))]
                fn fixup_packet(buffer: &mut [u8]) {
                    use packet::ipv4::MutableIpv4Packet;

                    let buflen = buffer.len();
                    let mut new_packet = MutableIpv4Packet::new(buffer).unwrap();

                    let length = u16::from_be(new_packet.get_total_length());
                    new_packet.set_total_length(length);

                    // OS X does this awesome thing where it removes the header length
                    // from the total length sometimes.
                    let length = new_packet.get_total_length() as usize +
                                 (new_packet.get_header_length() as usize * 4usize);
                    if length == buflen {
                        new_packet.set_total_length(length as u16)
                    }

                    let offset = u16::from_be(new_packet.get_fragment_offset());
                    new_packet.set_fragment_offset(offset);
                }

                #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
                fn fixup_packet(_buffer: &mut [u8]) {}
            }
        }
    )
}

transport_channel_iterator!(Ipv4Packet,
                            Ipv4TransportChannelIterator,
                            ipv4_packet_iter);

transport_channel_iterator!(UdpPacket,
                            UdpTransportChannelIterator,
                            udp_packet_iter);
<commit_msg>Auto merge of #165 - faern:fix_osx_flags, r=mrmonday<commit_after>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for implementing transport layer protocols
//!
//! The transport module provides the ability to send and receive packets at
//! the transport layer using IPv4 or IPv6. It also enables layer 3 networking
//! for specific transport protocols, using IPv4 only.
//!
//! Note that this is limited by operating system support - for example, on OS
//! X and FreeBSD, it is impossible to implement protocols which are already
//! implemented in the kernel such as TCP and UDP.

#![macro_use]

extern crate libc;

use self::TransportProtocol::{Ipv4, Ipv6};
use self::TransportChannelType::{Layer3, Layer4};

use std::io;
use std::io::Error;
use std::iter::repeat;
use std::net;
use std::mem;
use std::sync::Arc;

use packet::Packet;
use packet::ip::IpNextHeaderProtocol;
use packet::ipv4::Ipv4Packet;
use packet::udp::UdpPacket;

use internal;
use util;

/// Represents a transport layer protocol
#[derive(Clone, Copy)]
pub enum TransportProtocol {
    /// Represents a transport protocol built on top of IPv4
    Ipv4(IpNextHeaderProtocol),
    /// Represents a transport protocol built on top of IPv6
    Ipv6(IpNextHeaderProtocol),
}

/// Type of transport channel to present
#[derive(Clone, Copy)]
pub enum TransportChannelType {
    /// The application will send and receive transport layer packets
    Layer4(TransportProtocol),
    /// The application will send and receive IPv4 packets, with the specified transport protocol
    Layer3(IpNextHeaderProtocol),
}

/// Structure used for sending at the transport layer. Should be created with transport_channel()
pub struct TransportSender {
    socket: Arc<internal::FileDesc>,
    _channel_type: TransportChannelType,
}

/// Structure used for sending at the transport layer. Should be created with transport_channel()
pub struct TransportReceiver {
    socket: Arc<internal::FileDesc>,
    buffer: Vec<u8>,
    channel_type: TransportChannelType,
}

#[cfg(windows)]
const INVALID_SOCKET: libc::SOCKET = libc::INVALID_SOCKET;

#[cfg(not(windows))]
const INVALID_SOCKET: libc::c_int = -1;

/// Create a new (TransportSender, TransportReceiver) pair
///
/// This allows for sending and receiving packets at the transport layer. The buffer size should be
/// large enough to handle the largest packet you wish to receive.
///
/// The channel type specifies what layer to send and receive packets at, and the transport
/// protocol you wish to implement. For example, `Layer4(Ipv4(IpNextHeaderProtocols::Udp))` would
/// allow sending and receiving UDP packets using IPv4; whereas Layer3(IpNextHeaderProtocols::Udp)
/// would include the IPv4 Header in received values, and require manual construction of an IP
/// header when sending.
pub fn transport_channel(buffer_size: usize,
                         channel_type: TransportChannelType)
    -> io::Result<(TransportSender, TransportReceiver)> {
    use std::net;

    // This hack makes sure that winsock is initialised
    let _ = {
        let ip = net::Ipv4Addr::new(255, 255, 255, 255);
        let sockaddr = net::SocketAddr::V4(net::SocketAddrV4::new(ip, 0));

        net::UdpSocket::bind(sockaddr)
    };

    let socket = unsafe {
        match channel_type {
            Layer4(Ipv4(IpNextHeaderProtocol(proto))) | Layer3(IpNextHeaderProtocol(proto)) =>
                libc::socket(libc::AF_INET, libc::SOCK_RAW, proto as libc::c_int),
            Layer4(Ipv6(IpNextHeaderProtocol(proto))) =>
                libc::socket(libc::AF_INET6, libc::SOCK_RAW, proto as libc::c_int),
        }
    };
    if socket != INVALID_SOCKET {
        if match channel_type {
            Layer3(_) | Layer4(Ipv4(_)) => true,
            _ => false,
        } {
            let hincl: libc::c_int = match channel_type {
                Layer4(..) => 0,
                _ => 1,
            };
            let res = unsafe {
                libc::setsockopt(socket,
                                 libc::IPPROTO_IP,
                                 libc::IP_HDRINCL,
                                 (&hincl as *const libc::c_int) as *const libc::c_void,
                                 mem::size_of::<libc::c_int>() as libc::socklen_t)
            };
            if res == -1 {
                let err = Error::last_os_error();
                unsafe {
                    internal::close(socket);
                }
                return Err(err);
            }
        }

        let sock = Arc::new(internal::FileDesc { fd: socket });
        let sender = TransportSender {
            socket: sock.clone(),
            _channel_type: channel_type,
        };
        let receiver = TransportReceiver {
            socket: sock,
            buffer: repeat(0u8).take(buffer_size).collect(),
            channel_type: channel_type,
        };

        Ok((sender, receiver))
    } else {
        Err(Error::last_os_error())
    }
}

impl TransportSender {
    fn send<T: Packet>(&mut self, packet: T, dst: util::IpAddr) -> io::Result<usize> {
        let mut caddr = unsafe { mem::zeroed() };
        let sockaddr = match dst {
            util::IpAddr::V4(ip_addr) =>
                net::SocketAddr::V4(net::SocketAddrV4::new(ip_addr, 0)),
            util::IpAddr::V6(ip_addr) =>
                net::SocketAddr::V6(net::SocketAddrV6::new(ip_addr, 0, 0, 0)),
        };
        let slen = internal::addr_to_sockaddr(sockaddr, &mut caddr);
        let caddr_ptr = (&caddr as *const libc::sockaddr_storage) as *const libc::sockaddr;

        internal::send_to(self.socket.fd, packet.packet(), caddr_ptr, slen)
    }

    /// Send a packet to the provided desination
    #[inline]
    pub fn send_to<T: Packet>(&mut self, packet: T, destination: util::IpAddr) -> io::Result<usize> {
        self.send_to_impl(packet, destination)
    }

    #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: util::IpAddr) -> io::Result<usize> {
        self.send(packet, dst)
    }

    #[cfg(any(target_os = "freebsd", target_os = "macos"))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: util::IpAddr) -> io::Result<usize> {
        use packet::MutablePacket;
        use packet::ipv4::MutableIpv4Packet;

        // FreeBSD and OS X expect total length and fragment offset fields of IPv4
        // packets to be in host byte order rather than network byte order. Fragment offset is the
        // ip_off field in the ip struct and contains both the offset and the three flag bits.
        // See `man 4 ip`/Raw IP Sockets)
        if let Layer3(_) = self._channel_type {
            let mut mut_slice: Vec<u8> = repeat(0u8).take(packet.packet().len()).collect();

            let mut new_packet = MutableIpv4Packet::new(&mut mut_slice[..]).unwrap();
            new_packet.clone_from(&packet);
            let length = new_packet.get_total_length().to_be();
            new_packet.set_total_length(length);
            {
                // Turn fragment offset into host order
                let d = new_packet.packet_mut();
                let host_order = u16::from_be((d[6] as u16) << 8 | d[7] as u16);
                d[6] = (host_order >> 8) as u8;
                d[7] = host_order as u8;
            }
            return self.send(new_packet, dst);
        }

        self.send(packet, dst)
    }
}

/// Create an iterator for some packet type.
///
/// Usage:
/// ```
/// transport_channel_iterator!(Ipv4Packet, // Type to iterate over
///                             Ipv4TransportChannelIterator, // Name for iterator struct
///                             ipv4_packet_iter) // Name of function to create iterator
/// ```
#[macro_export]
macro_rules! transport_channel_iterator {
    ($ty:ident, $iter:ident, $func:ident) => (
        /// An iterator over packets of type $ty
        pub struct $iter<'a> {
            tr: &'a mut TransportReceiver
        }
        /// Return a packet iterator with packets of type $ty for some transport receiver
        pub fn $func(tr: &mut TransportReceiver) -> $iter {
            $iter {
                tr: tr
            }
        }
        impl<'a> $iter<'a> {
            /// Get the next ($ty, IpAddr) pair for the given channel
            pub fn next(&mut self) -> io::Result<($ty, util::IpAddr)> {
                let mut caddr: libc::sockaddr_storage = unsafe { mem::zeroed() };
                let res = internal::recv_from(self.tr.socket.fd,
                                              &mut self.tr.buffer[..],
                                              &mut caddr);

                let offset = match self.tr.channel_type {
                    Layer4(Ipv4(_)) => {
                        let ip_header = Ipv4Packet::new(&self.tr.buffer[..]).unwrap();

                        ip_header.get_header_length() as usize * 4usize
                    },
                    Layer3(_) => {
                        fixup_packet(&mut self.tr.buffer[..]);

                        0
                    },
                    _ => 0
                };
                return match res {
                    Ok(len) => {
                        let packet = $ty::new(&self.tr.buffer[offset..len]).unwrap();
                        let addr = internal::sockaddr_to_addr(
                                        &caddr,
                                        mem::size_of::<libc::sockaddr_storage>()
                                   );
                        let ip = match addr.unwrap() {
                            net::SocketAddr::V4(sa) => util::IpAddr::V4(*sa.ip()),
                            net::SocketAddr::V6(sa) => util::IpAddr::V6(*sa.ip()),
                        };
                        Ok((packet, ip))
                    },
                    Err(e) => Err(e),
                };

                #[cfg(any(target_os = "freebsd", target_os = "macos"))]
                fn fixup_packet(buffer: &mut [u8]) {
                    use packet::ipv4::MutableIpv4Packet;

                    let buflen = buffer.len();
                    let mut new_packet = MutableIpv4Packet::new(buffer).unwrap();

                    let length = u16::from_be(new_packet.get_total_length());
                    new_packet.set_total_length(length);

                    // OS X does this awesome thing where it removes the header length
                    // from the total length sometimes.
                    let length = new_packet.get_total_length() as usize +
                                 (new_packet.get_header_length() as usize * 4usize);
                    if length == buflen {
                        new_packet.set_total_length(length as u16)
                    }

                    let offset = u16::from_be(new_packet.get_fragment_offset());
                    new_packet.set_fragment_offset(offset);
                }

                #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
                fn fixup_packet(_buffer: &mut [u8]) {}
            }
        }
    )
}

transport_channel_iterator!(Ipv4Packet,
                            Ipv4TransportChannelIterator,
                            ipv4_packet_iter);

transport_channel_iterator!(UdpPacket,
                            UdpTransportChannelIterator,
                            udp_packet_iter);
<|endoftext|>


<commit_before>//! This module contains the core algorithms.

use std::borrow::Borrow;
use {TrieKey, NibbleVec};
use trie_node::TrieNode;
use keys::{match_keys, KeyMatch};

use self::DescendantResult::*;

impl<K, V> TrieNode<K, V>
    where K: TrieKey
{
    pub fn get(&self, nv: &NibbleVec) -> Option<&TrieNode<K, V>> {
        iterative_get(self, nv)
    }

    pub fn get_mut(&mut self, nv: &NibbleVec) -> Option<&mut TrieNode<K, V>> {
        iterative_get_mut(self, nv)
    }

    pub fn insert(&mut self, key: K, value: V, nv: NibbleVec) -> Option<V> {
        iterative_insert(self, key, value, nv)
    }

    pub fn remove<Q: ?Sized>(&mut self, key: &Q) -> Option<V>
        where K: Borrow<Q>, Q: TrieKey {
        recursive_remove(self, key)
    }

    pub fn get_ancestor(&self, nv: &NibbleVec) -> Option<(&TrieNode<K, V>, usize)> {
        get_ancestor(self, nv)
    }

    pub fn get_raw_ancestor(&self, nv: &NibbleVec) -> (&TrieNode<K, V>, usize) {
        get_raw_ancestor(self, nv)
    }

    pub fn get_raw_descendant<'a>(&'a self, nv: &NibbleVec) -> Option<DescendantResult<'a, K, V>> {
        get_raw_descendant(self, nv)
    }
}

macro_rules! get_func {
    (
        name: $name:ident,
        trie_type: $trie_type:ty,
        mutability: $($mut_:tt)*
    ) => {id!{
        fn $name<'a, K, V>(trie: $trie_type, nv: &NibbleVec) -> Option<$trie_type> {
            if nv.len() == 0 {
                return Some(trie);
            }

            let mut prev = trie;
            let mut depth = 0;

            loop {
                let bucket = nv.get(depth) as usize;
                let current = prev;
                if let Some(ref $($mut_)* child) = current.children[bucket] {
                    match match_keys(depth, nv, &child.key) {
                        KeyMatch::Full => {
                            return Some(child);
                        }
                        KeyMatch::SecondPrefix => {
                            depth += child.key.len();
                            prev = child;
                        }
                        _ => {
                            return None;
                        }
                    }
                } else {
                    return None;
                }
            }
        }
    }}
}

get_func!(name: iterative_get, trie_type: &'a TrieNode<K, V>, mutability: );
get_func!(name: iterative_get_mut, trie_type: &'a mut TrieNode<K, V>, mutability: mut);

fn iterative_insert<K, V>(trie: &mut TrieNode<K, V>,
                          key: K,
                          value: V,
                          mut nv: NibbleVec)
                          -> Option<V>
    where K: TrieKey
{
    if nv.len() == 0 {
        return trie.replace_value(key, value);
    }

    let mut prev = trie;
    let mut depth = 0;

    loop {
        let bucket = nv.get(depth) as usize;
        let current = prev;
        if let Some(ref mut child) = current.children[bucket] {
            match match_keys(depth, &nv, &child.key) {
                KeyMatch::Full => {
                    return child.replace_value(key, value);
                }
                KeyMatch::Partial(idx) => {
                    // Split the existing child.
                    child.split(idx);

                    // Insert the new key below the prefix node.
                    let new_key = nv.split(depth + idx);
                    let new_key_bucket = new_key.get(0) as usize;

                    child.add_child(new_key_bucket,
                                    Box::new(TrieNode::with_key_value(new_key, key, value)));

                    return None;
                }
                KeyMatch::FirstPrefix => {
                    child.split(nv.len() - depth);
                    child.add_key_value(key, value);
                    return None;
                }
                KeyMatch::SecondPrefix => {
                    depth += child.key.len();
                    prev = child;
                }
            }
        } else {
            let node_key = nv.split(depth);
            current.add_child(bucket,
                              Box::new(TrieNode::with_key_value(node_key, key, value)));
            return None;
        }
    }
}

// TODO: clean this up and make it iterative.
fn recursive_remove<K, Q: ?Sized, V>(trie: &mut TrieNode<K, V>, key: &Q) -> Option<V>
    where K: TrieKey, K: Borrow<Q>, Q: TrieKey
{
    let nv = key.encode();

    if nv.len() == 0 {
        return trie.take_value(key);
    }

    let bucket = nv.get(0) as usize;

    let child = trie.take_child(bucket);

    match child {
        Some(mut child) => {
            let depth = child.key.len();
            if depth == nv.len() {
                let result = child.take_value(key);
                if child.child_count != 0 {
                    // If removing this node's value has made it a value-less node with a
                    // single child, then merge its child.
                    let repl = if child.child_count == 1 {
                        get_merge_child(&mut child)
                    } else {
                        child
                    };
                    trie.add_child(bucket, repl);
                }
                result
            } else {
                rec_remove(trie, child, bucket, key, depth, &nv)
            }
        }
        None => None,
    }
}

fn get_merge_child<K, V>(trie: &mut TrieNode<K, V>) -> Box<TrieNode<K, V>>
    where K: TrieKey
{
    let mut child = trie.take_only_child();

    // Join the child's key onto the existing one.
    child.key = trie.key.clone().join(&child.key);

    child
}

// Tail-recursive remove function used by `recursive_remove`.
fn rec_remove<K, Q: ?Sized, V>(parent: &mut TrieNode<K, V>,
                    mut middle: Box<TrieNode<K, V>>,
                    prev_bucket: usize,
                    key: &Q,
                    depth: usize,
                    nv: &NibbleVec)
                    -> Option<V>
    where K: TrieKey, K: Borrow<Q>, Q: TrieKey
{
    let bucket = nv.get(depth) as usize;

    let child = middle.take_child(bucket);
    parent.add_child(prev_bucket, middle);

    match child {
        Some(mut child) => {
            let middle = parent.children[prev_bucket].as_mut().unwrap();
            match match_keys(depth, nv, &child.key) {
                KeyMatch::Full => {
                    let result = child.take_value(key);

                    // If this node has children, keep it.
                    if child.child_count != 0 {
                        // If removing this node's value has made it a value-less node with a
                        // single child, then merge its child.
                        let repl = if child.child_count == 1 {
                            get_merge_child(&mut *child)
                        } else {
                            child
                        };
                        middle.add_child(bucket, repl);
                    }
                    // Otherwise, if the parent node now only has a single child, merge it.
                    else if middle.child_count == 1 && middle.key_value.is_none() {
                        let repl = get_merge_child(middle);
                        *middle = repl;
                    }

                    result
                }
                KeyMatch::SecondPrefix => {
                    let new_depth = depth + child.key.len();
                    rec_remove(middle, child, bucket, key, new_depth, nv)
                }
                _ => None,
            }
        }
        None => None,
    }
}

fn get_ancestor<'a, K, V>(trie: &'a TrieNode<K, V>,
                          nv: &NibbleVec)
                          -> Option<(&'a TrieNode<K, V>, usize)>
    where K: TrieKey
{
    if nv.len() == 0 {
        return trie.as_value_node().map(|node| (node, 0));
    }

    let mut prev = trie;
    // The ancestor is such that all nodes upto and including `prev` have
    // already been considered.
    let mut ancestor = prev.as_value_node();
    let mut depth = 0;

    loop {
        let bucket = nv.get(depth) as usize;
        let current = prev;
        if let Some(ref child) = current.children[bucket] {
            match match_keys(depth, nv, &child.key) {
                KeyMatch::Full => {
                    return child.as_value_node()
                        .map(|node| (node, depth + node.key.len()))
                        .or_else(|| ancestor.map(|anc| (anc, depth)));
                }
                KeyMatch::FirstPrefix |
                KeyMatch::Partial(_) => {
                    return ancestor.map(|anc| (anc, depth));
                }
                KeyMatch::SecondPrefix => {
                    depth += child.key.len();
                    ancestor = child.as_value_node().or(ancestor);
                    prev = child;
                }
            }
        } else {
            return ancestor.map(|anc| (anc, depth));
        }
    }
}

fn get_raw_ancestor<'a, K, V>(trie: &'a TrieNode<K, V>,
                              nv: &NibbleVec)
                              -> (&'a TrieNode<K, V>, usize)
    where K: TrieKey
{
    if nv.len() == 0 {
        return (trie, 0);
    }

    let mut prev = trie;
    // The ancestor is such that all nodes upto and including `prev` have
    // already been considered.
    let mut ancestor = prev;
    let mut depth = 0;

    loop {
        let bucket = nv.get(depth) as usize;
        let current = prev;
        if let Some(ref child) = current.children[bucket] {
            match match_keys(depth, nv, &child.key) {
                KeyMatch::Full => {
                    return (child, depth + child.key.len());
                }
                KeyMatch::FirstPrefix |
                KeyMatch::Partial(_) => {
                    return (ancestor, depth);
                }
                KeyMatch::SecondPrefix => {
                    depth += child.key.len();
                    ancestor = child;
                    prev = child;
                }
            }
        } else {
            return (ancestor, depth);
        }
    }
}

// Type used to propogate subtrie construction instructions to the top-level `get_raw_descendant`
// method.
pub enum DescendantResult<'a, K: 'a, V: 'a> {
    NoModification(&'a TrieNode<K, V>),
    ExtendKey(&'a TrieNode<K, V>, usize, &'a NibbleVec),
}

fn get_raw_descendant<'a, K, V>(trie: &'a TrieNode<K, V>,
                                nv: &NibbleVec)
                                -> Option<DescendantResult<'a, K, V>> {
    if nv.len() == 0 {
        return Some(NoModification(trie));
    }

    let mut prev = trie;
    let mut depth = 0;

    loop {
        let bucket = nv.get(depth) as usize;
        let current = prev;
        if let Some(ref child) = current.children[bucket] {
            match match_keys(depth, nv, &child.key) {
                KeyMatch::Full => {
                    return Some(NoModification(child));
                }
                KeyMatch::FirstPrefix => {
                    return Some(ExtendKey(child, depth, &child.key));
                }
                KeyMatch::SecondPrefix => {
                    depth += child.key.len();
                    prev = child;
                }
                _ => {
                    return None;
                }
            }
        } else {
            return None;
        }
    }
}
<commit_msg>Prevent panics when removing non-existent keys<commit_after>//! This module contains the core algorithms.

use std::borrow::Borrow;
use {TrieKey, NibbleVec};
use trie_node::TrieNode;
use keys::{match_keys, KeyMatch};

use self::DescendantResult::*;

impl<K, V> TrieNode<K, V>
    where K: TrieKey
{
    pub fn get(&self, nv: &NibbleVec) -> Option<&TrieNode<K, V>> {
        iterative_get(self, nv)
    }

    pub fn get_mut(&mut self, nv: &NibbleVec) -> Option<&mut TrieNode<K, V>> {
        iterative_get_mut(self, nv)
    }

    pub fn insert(&mut self, key: K, value: V, nv: NibbleVec) -> Option<V> {
        iterative_insert(self, key, value, nv)
    }

    pub fn remove<Q: ?Sized>(&mut self, key: &Q) -> Option<V>
        where K: Borrow<Q>, Q: TrieKey {
        recursive_remove(self, key)
    }

    pub fn get_ancestor(&self, nv: &NibbleVec) -> Option<(&TrieNode<K, V>, usize)> {
        get_ancestor(self, nv)
    }

    pub fn get_raw_ancestor(&self, nv: &NibbleVec) -> (&TrieNode<K, V>, usize) {
        get_raw_ancestor(self, nv)
    }

    pub fn get_raw_descendant<'a>(&'a self, nv: &NibbleVec) -> Option<DescendantResult<'a, K, V>> {
        get_raw_descendant(self, nv)
    }
}

macro_rules! get_func {
    (
        name: $name:ident,
        trie_type: $trie_type:ty,
        mutability: $($mut_:tt)*
    ) => {id!{
        fn $name<'a, K, V>(trie: $trie_type, nv: &NibbleVec) -> Option<$trie_type> {
            if nv.len() == 0 {
                return Some(trie);
            }

            let mut prev = trie;
            let mut depth = 0;

            loop {
                let bucket = nv.get(depth) as usize;
                let current = prev;
                if let Some(ref $($mut_)* child) = current.children[bucket] {
                    match match_keys(depth, nv, &child.key) {
                        KeyMatch::Full => {
                            return Some(child);
                        }
                        KeyMatch::SecondPrefix => {
                            depth += child.key.len();
                            prev = child;
                        }
                        _ => {
                            return None;
                        }
                    }
                } else {
                    return None;
                }
            }
        }
    }}
}

get_func!(name: iterative_get, trie_type: &'a TrieNode<K, V>, mutability: );
get_func!(name: iterative_get_mut, trie_type: &'a mut TrieNode<K, V>, mutability: mut);

fn iterative_insert<K, V>(trie: &mut TrieNode<K, V>,
                          key: K,
                          value: V,
                          mut nv: NibbleVec)
                          -> Option<V>
    where K: TrieKey
{
    if nv.len() == 0 {
        return trie.replace_value(key, value);
    }

    let mut prev = trie;
    let mut depth = 0;

    loop {
        let bucket = nv.get(depth) as usize;
        let current = prev;
        if let Some(ref mut child) = current.children[bucket] {
            match match_keys(depth, &nv, &child.key) {
                KeyMatch::Full => {
                    return child.replace_value(key, value);
                }
                KeyMatch::Partial(idx) => {
                    // Split the existing child.
                    child.split(idx);

                    // Insert the new key below the prefix node.
                    let new_key = nv.split(depth + idx);
                    let new_key_bucket = new_key.get(0) as usize;

                    child.add_child(new_key_bucket,
                                    Box::new(TrieNode::with_key_value(new_key, key, value)));

                    return None;
                }
                KeyMatch::FirstPrefix => {
                    child.split(nv.len() - depth);
                    child.add_key_value(key, value);
                    return None;
                }
                KeyMatch::SecondPrefix => {
                    depth += child.key.len();
                    prev = child;
                }
            }
        } else {
            let node_key = nv.split(depth);
            current.add_child(bucket,
                              Box::new(TrieNode::with_key_value(node_key, key, value)));
            return None;
        }
    }
}

// TODO: clean this up and make it iterative.
fn recursive_remove<K, Q: ?Sized, V>(trie: &mut TrieNode<K, V>, key: &Q) -> Option<V>
    where K: TrieKey, K: Borrow<Q>, Q: TrieKey
{
    let nv = key.encode();

    if nv.len() == 0 {
        return trie.take_value(key);
    }

    let bucket = nv.get(0) as usize;

    let child = trie.take_child(bucket);

    match child {
        Some(mut child) => {
            match match_keys(0, &nv, &child.key) {
                KeyMatch::Full => {
                    let result = child.take_value(key);
                    if child.child_count != 0 {
                        // If removing this node's value has made it a value-less node with a
                        // single child, then merge its child.
                        let repl = if child.child_count == 1 {
                            get_merge_child(&mut child)
                        } else {
                            child
                        };
                        trie.add_child(bucket, repl);
                    }
                    result
                }
                KeyMatch::SecondPrefix => {
                    let depth = child.key.len();
                    rec_remove(trie, child, bucket, key, depth, &nv)
                }
                _ => None,
            }
        }
        None => None,
    }
}

fn get_merge_child<K, V>(trie: &mut TrieNode<K, V>) -> Box<TrieNode<K, V>>
    where K: TrieKey
{
    let mut child = trie.take_only_child();

    // Join the child's key onto the existing one.
    child.key = trie.key.clone().join(&child.key);

    child
}

// Tail-recursive remove function used by `recursive_remove`.
fn rec_remove<K, Q: ?Sized, V>(parent: &mut TrieNode<K, V>,
                    mut middle: Box<TrieNode<K, V>>,
                    prev_bucket: usize,
                    key: &Q,
                    depth: usize,
                    nv: &NibbleVec)
                    -> Option<V>
    where K: TrieKey, K: Borrow<Q>, Q: TrieKey
{
    let bucket = nv.get(depth) as usize;

    let child = middle.take_child(bucket);
    parent.add_child(prev_bucket, middle);

    match child {
        Some(mut child) => {
            let middle = parent.children[prev_bucket].as_mut().unwrap();
            match match_keys(depth, nv, &child.key) {
                KeyMatch::Full => {
                    let result = child.take_value(key);

                    // If this node has children, keep it.
                    if child.child_count != 0 {
                        // If removing this node's value has made it a value-less node with a
                        // single child, then merge its child.
                        let repl = if child.child_count == 1 {
                            get_merge_child(&mut *child)
                        } else {
                            child
                        };
                        middle.add_child(bucket, repl);
                    }
                    // Otherwise, if the parent node now only has a single child, merge it.
                    else if middle.child_count == 1 && middle.key_value.is_none() {
                        let repl = get_merge_child(middle);
                        *middle = repl;
                    }

                    result
                }
                KeyMatch::SecondPrefix => {
                    let new_depth = depth + child.key.len();
                    rec_remove(middle, child, bucket, key, new_depth, nv)
                }
                _ => None,
            }
        }
        None => None,
    }
}

fn get_ancestor<'a, K, V>(trie: &'a TrieNode<K, V>,
                          nv: &NibbleVec)
                          -> Option<(&'a TrieNode<K, V>, usize)>
    where K: TrieKey
{
    if nv.len() == 0 {
        return trie.as_value_node().map(|node| (node, 0));
    }

    let mut prev = trie;
    // The ancestor is such that all nodes upto and including `prev` have
    // already been considered.
    let mut ancestor = prev.as_value_node();
    let mut depth = 0;

    loop {
        let bucket = nv.get(depth) as usize;
        let current = prev;
        if let Some(ref child) = current.children[bucket] {
            match match_keys(depth, nv, &child.key) {
                KeyMatch::Full => {
                    return child.as_value_node()
                        .map(|node| (node, depth + node.key.len()))
                        .or_else(|| ancestor.map(|anc| (anc, depth)));
                }
                KeyMatch::FirstPrefix |
                KeyMatch::Partial(_) => {
                    return ancestor.map(|anc| (anc, depth));
                }
                KeyMatch::SecondPrefix => {
                    depth += child.key.len();
                    ancestor = child.as_value_node().or(ancestor);
                    prev = child;
                }
            }
        } else {
            return ancestor.map(|anc| (anc, depth));
        }
    }
}

fn get_raw_ancestor<'a, K, V>(trie: &'a TrieNode<K, V>,
                              nv: &NibbleVec)
                              -> (&'a TrieNode<K, V>, usize)
    where K: TrieKey
{
    if nv.len() == 0 {
        return (trie, 0);
    }

    let mut prev = trie;
    // The ancestor is such that all nodes upto and including `prev` have
    // already been considered.
    let mut ancestor = prev;
    let mut depth = 0;

    loop {
        let bucket = nv.get(depth) as usize;
        let current = prev;
        if let Some(ref child) = current.children[bucket] {
            match match_keys(depth, nv, &child.key) {
                KeyMatch::Full => {
                    return (child, depth + child.key.len());
                }
                KeyMatch::FirstPrefix |
                KeyMatch::Partial(_) => {
                    return (ancestor, depth);
                }
                KeyMatch::SecondPrefix => {
                    depth += child.key.len();
                    ancestor = child;
                    prev = child;
                }
            }
        } else {
            return (ancestor, depth);
        }
    }
}

// Type used to propogate subtrie construction instructions to the top-level `get_raw_descendant`
// method.
pub enum DescendantResult<'a, K: 'a, V: 'a> {
    NoModification(&'a TrieNode<K, V>),
    ExtendKey(&'a TrieNode<K, V>, usize, &'a NibbleVec),
}

fn get_raw_descendant<'a, K, V>(trie: &'a TrieNode<K, V>,
                                nv: &NibbleVec)
                                -> Option<DescendantResult<'a, K, V>> {
    if nv.len() == 0 {
        return Some(NoModification(trie));
    }

    let mut prev = trie;
    let mut depth = 0;

    loop {
        let bucket = nv.get(depth) as usize;
        let current = prev;
        if let Some(ref child) = current.children[bucket] {
            match match_keys(depth, nv, &child.key) {
                KeyMatch::Full => {
                    return Some(NoModification(child));
                }
                KeyMatch::FirstPrefix => {
                    return Some(ExtendKey(child, depth, &child.key));
                }
                KeyMatch::SecondPrefix => {
                    depth += child.key.len();
                    prev = child;
                }
                _ => {
                    return None;
                }
            }
        } else {
            return None;
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use parser::{Program, Decl, Stmt, Expr};
use types::Type;
use error::Error;

type Symtable = HashMap<Expr, Type>;

pub fn tc_program(p: &Program) -> Result<(), Error> {
    let mut init_symtable = try!(tc_decls(&p.decls));
    tc_stmts(&p.stmts, &mut init_symtable)
}

fn tc_decls(decls: &Vec<Decl>) -> Result<Symtable, Error> {
    let mut symtable: Symtable = HashMap::new();
    for decl in decls {
        try!(tc_decl(&decl, &mut symtable));
    }
    println!("{:?}", symtable);
    Ok(symtable)
}

fn tc_decl(decl: &Decl, symtable: &mut Symtable) -> Result<(), Error> {
    match *decl {
        Decl::Decl(ref id, ty) => {
            let id_expr = Expr::Id(id.clone());
            if symtable.contains_key(&id_expr) {
                return Err(Error::GenericError);
            } else {
                symtable.insert(id_expr, ty);
            }
        }
    }
    Ok(())
}

fn tc_stmts(stmts: &Vec<Stmt>, symtable: &mut Symtable) -> Result<(), Error> {
    for stmt in stmts {
        try!(tc_stmt(&stmt, symtable));
    }
    Ok(())
}

fn tc_stmt(stmt: &Stmt, symtable: &mut Symtable) -> Result<(), Error> {
    match *stmt {
        Stmt::Assign(ref id, ref e) => tc_stmt_assign(id, e, symtable),
        Stmt::Read(ref id) => tc_stmt_read(id, symtable),
        Stmt::Print(ref e) => tc_stmt_print(e, symtable),
        Stmt::If(ref e, ref then_stmts, ref else_stmts) => tc_stmt_if(e, then_stmts, else_stmts, symtable),
        Stmt::While(ref e, ref stmts) => tc_stmt_while(e, stmts, symtable),
    }
}

fn tc_stmt_assign(id: &String, expr: &Expr, symtable: &mut Symtable) -> Result<(), Error> {
    let expr_id = Expr::Id(id.clone());
    let t = try!(tc_expr(expr, symtable));
    match symtable.get(&expr_id) {
        Some(&ty) => {
            if t == ty {
                Ok(())
            } else {
                Err(Error::GenericError)
            }
        }
        None => Err(Error::GenericError)
    }
}

fn tc_stmt_read(id: &String, symtable: &Symtable) -> Result<(), Error> {
    let expr_id = Expr::Id(id.clone());
    if symtable.contains_key(&expr_id) {
        Ok(())
    } else {
        Err(Error::GenericError)
    }
}

fn tc_stmt_print(expr: &Expr, symtable: &mut Symtable) -> Result<(), Error> {
    try!(tc_expr(expr, symtable));
    Ok(())
}

fn tc_stmt_if(expr: &Expr, then_stmts: &Vec<Stmt>, else_stmts: &Vec<Stmt>, symtable: &mut Symtable) -> Result<(), Error> {
    let t = try!(tc_expr(expr, symtable));
    match t {
        Type::Int => {
            try!(tc_stmts(then_stmts, symtable));
            try!(tc_stmts(else_stmts, symtable));
            Ok(())
        }
        _ => {
            Err(Error::GenericError)
        }
    }
}

fn tc_stmt_while(expr: &Expr, stmts: &Vec<Stmt>, symtable: &mut Symtable) -> Result<(), Error> {
    let t = try!(tc_expr(expr, symtable));
    match t {
        Type::Int => {
            tc_stmts(stmts, symtable)
        }
        _ => {
            Err(Error::GenericError)
        }
    }
}

fn tc_expr(expr: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    let ty = try!(match *expr {
        Expr::Int(_) => Ok(Type::Int),
        Expr::Float(_) => Ok(Type::Float),
        Expr::Id(_) => tc_expr_id(expr, symtable),
        Expr::Negate(ref e) => tc_expr_negate(e, symtable),
        Expr::Add(ref e1, ref e2) |
        Expr::Sub(ref e1, ref e2) |
        Expr::Mul(ref e1, ref e2) |
        Expr::Div(ref e1, ref e2) => tc_expr_binop(e1, e2, symtable),
    });
    // TODO(vfoley): insert expr -> ty into symtable
    Ok(ty)
}

fn tc_expr_id(expr: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    match symtable.get(expr) {
        Some(ty) => Ok(*ty),
        None => Err(Error::GenericError),
    }
}

fn tc_expr_negate(e: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    tc_expr(e, symtable)
}

fn tc_expr_binop(e1: &Expr, e2: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    let t1 = try!(tc_expr(e1, symtable));
    let t2 = try!(tc_expr(e2, symtable));

    match (t1, t2) {
        (Type::Int   , Type::Int)   => Ok(Type::Int),
        (Type::Int   , Type::Float) => Ok(Type::Float),
        (Type::Float , Type::Int)   => Ok(Type::Float),
        (Type::Float , Type::Float) => Ok(Type::Float),
    }
}
<commit_msg>Allow assigning ints into floats<commit_after>use std::collections::HashMap;

use parser::{Program, Decl, Stmt, Expr};
use types::Type;
use error::Error;

type Symtable = HashMap<Expr, Type>;

pub fn tc_program(p: &Program) -> Result<(), Error> {
    let mut init_symtable = try!(tc_decls(&p.decls));
    tc_stmts(&p.stmts, &mut init_symtable)
}

fn tc_decls(decls: &Vec<Decl>) -> Result<Symtable, Error> {
    let mut symtable: Symtable = HashMap::new();
    for decl in decls {
        try!(tc_decl(&decl, &mut symtable));
    }
    println!("{:?}", symtable);
    Ok(symtable)
}

fn tc_decl(decl: &Decl, symtable: &mut Symtable) -> Result<(), Error> {
    match *decl {
        Decl::Decl(ref id, ty) => {
            let id_expr = Expr::Id(id.clone());
            if symtable.contains_key(&id_expr) {
                return Err(Error::GenericError);
            } else {
                symtable.insert(id_expr, ty);
            }
        }
    }
    Ok(())
}

fn tc_stmts(stmts: &Vec<Stmt>, symtable: &mut Symtable) -> Result<(), Error> {
    for stmt in stmts {
        try!(tc_stmt(&stmt, symtable));
    }
    Ok(())
}

fn tc_stmt(stmt: &Stmt, symtable: &mut Symtable) -> Result<(), Error> {
    match *stmt {
        Stmt::Assign(ref id, ref e) => tc_stmt_assign(id, e, symtable),
        Stmt::Read(ref id) => tc_stmt_read(id, symtable),
        Stmt::Print(ref e) => tc_stmt_print(e, symtable),
        Stmt::If(ref e, ref then_stmts, ref else_stmts) => tc_stmt_if(e, then_stmts, else_stmts, symtable),
        Stmt::While(ref e, ref stmts) => tc_stmt_while(e, stmts, symtable),
    }
}

fn tc_stmt_assign(id: &String, expr: &Expr, symtable: &mut Symtable) -> Result<(), Error> {
    let expr_id = Expr::Id(id.clone());
    let expr_ty = try!(tc_expr(expr, symtable));
    match symtable.get(&expr_id) {
        Some(&id_ty) => {
            match (id_ty, expr_ty) {
                (Type::Int, Type::Int) => Ok(()),
                (Type::Int, Type::Float) => Err(Error::GenericError),
                (Type::Float, Type::Int) => Ok(()),
                (Type::Float, Type::Float) => Ok(()),
            }
        }
        None => Err(Error::GenericError)
    }
}

fn tc_stmt_read(id: &String, symtable: &Symtable) -> Result<(), Error> {
    let expr_id = Expr::Id(id.clone());
    if symtable.contains_key(&expr_id) {
        Ok(())
    } else {
        Err(Error::GenericError)
    }
}

fn tc_stmt_print(expr: &Expr, symtable: &mut Symtable) -> Result<(), Error> {
    try!(tc_expr(expr, symtable));
    Ok(())
}

fn tc_stmt_if(expr: &Expr, then_stmts: &Vec<Stmt>, else_stmts: &Vec<Stmt>, symtable: &mut Symtable) -> Result<(), Error> {
    let t = try!(tc_expr(expr, symtable));
    match t {
        Type::Int => {
            try!(tc_stmts(then_stmts, symtable));
            try!(tc_stmts(else_stmts, symtable));
            Ok(())
        }
        _ => {
            Err(Error::GenericError)
        }
    }
}

fn tc_stmt_while(expr: &Expr, stmts: &Vec<Stmt>, symtable: &mut Symtable) -> Result<(), Error> {
    let t = try!(tc_expr(expr, symtable));
    match t {
        Type::Int => {
            tc_stmts(stmts, symtable)
        }
        _ => {
            Err(Error::GenericError)
        }
    }
}

fn tc_expr(expr: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    let ty = try!(match *expr {
        Expr::Int(_) => Ok(Type::Int),
        Expr::Float(_) => Ok(Type::Float),
        Expr::Id(_) => tc_expr_id(expr, symtable),
        Expr::Negate(ref e) => tc_expr_negate(e, symtable),
        Expr::Add(ref e1, ref e2) |
        Expr::Sub(ref e1, ref e2) |
        Expr::Mul(ref e1, ref e2) |
        Expr::Div(ref e1, ref e2) => tc_expr_binop(e1, e2, symtable),
    });
    // TODO(vfoley): insert expr -> ty into symtable
    Ok(ty)
}

fn tc_expr_id(expr: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    match symtable.get(expr) {
        Some(ty) => Ok(*ty),
        None => Err(Error::GenericError),
    }
}

fn tc_expr_negate(e: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    tc_expr(e, symtable)
}

fn tc_expr_binop(e1: &Expr, e2: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    let t1 = try!(tc_expr(e1, symtable));
    let t2 = try!(tc_expr(e2, symtable));

    match (t1, t2) {
        (Type::Int   , Type::Int)   => Ok(Type::Int),
        (Type::Int   , Type::Float) => Ok(Type::Float),
        (Type::Float , Type::Int)   => Ok(Type::Float),
        (Type::Float , Type::Float) => Ok(Type::Float),
    }
}
<|endoftext|>