<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()<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()
{'fraction_non_alphanumeric': 0.1399652978600347, 'fraction_numerical': 0.001662810873337189, 'mean_word_length': 2.0275771503611293, 'pattern_counts': {'":': 0, '<': 94, '<?xml version=': 0, '>': 132, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6226740', 'n_tokens_mistral': 4497, 'n_tokens_neox': 4397, 'n_words': 922}