smoltcp/src/socket/tcp.rs

use core::fmt;

use Error;
use Managed;
use wire::{IpProtocol, IpAddress, IpEndpoint};
use wire::{TcpPacket, TcpRepr, TcpControl};
use socket::{Socket, PacketRepr};

/// A TCP stream ring buffer.
#[derive(Debug)]
pub struct SocketBuffer<'a> {
    storage: Managed<'a, [u8]>,
    read_at: usize,
    length:  usize
}

impl<'a> SocketBuffer<'a> {
    /// Create a packet buffer with the given storage.
    pub fn new<T>(storage: T) -> SocketBuffer<'a>
            where T: Into<Managed<'a, [u8]>> {
        SocketBuffer {
            storage: storage.into(),
            read_at: 0,
            length:  0
        }
    }

    /// Return the amount of octets enqueued in the buffer.
    pub fn len(&self) -> usize {
        self.length
    }

    /// Return the maximum amount of octets that can be enqueued in the buffer.
    pub fn capacity(&self) -> usize {
        self.storage.len()
    }

    /// Enqueue a slice of octets up to the given size into the buffer, and return a pointer
    /// to the slice.
    ///
    /// The returned slice may be shorter than requested, as short as an empty slice,
    /// if there is not enough contiguous free space in the buffer.
    pub fn enqueue(&mut self, mut size: usize) -> &mut [u8] {
        let write_at = (self.read_at + self.length) % self.storage.len();
        // We can't enqueue more than there is free space.
        let free = self.storage.len() - self.length;
        if size > free { size = free }
        // We can't contiguously enqueue past the beginning of the storage.
        let until_end = self.storage.len() - write_at;
        if size > until_end { size = until_end }

        self.length += size;
        &mut self.storage[write_at..write_at + size]
    }

    /// Dequeue a slice of octets up to the given size from the buffer, and return a pointer
    /// to the slice.
    ///
    /// The returned slice may be shorter than requested, as short as an empty slice,
    /// if there is not enough contiguous filled space in the buffer.
    pub fn dequeue(&mut self, mut size: usize) -> &[u8] {
        let read_at = self.read_at;
        // We can't dequeue more than was queued.
        if size > self.length { size = self.length }
        // We can't contiguously dequeue past the end of the storage.
        let until_end = self.storage.len() - self.read_at;
        if size > until_end { size = until_end }

        self.read_at = (self.read_at + size) % self.storage.len();
        self.length -= size;
        &self.storage[read_at..read_at + size]
    }
}

impl<'a> Into<SocketBuffer<'a>> for Managed<'a, [u8]> {
    fn into(self) -> SocketBuffer<'a> {
        SocketBuffer::new(self)
    }
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum State {
    Closed,
    Listen,
    SynSent,
    SynReceived,
    Established,
    FinWait1,
    FinWait2,
    CloseWait,
    Closing,
    LastAck,
    TimeWait
}

impl fmt::Display for State {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &State::Closed      => write!(f, "CLOSED"),
            &State::Listen      => write!(f, "LISTEN"),
            &State::SynSent     => write!(f, "SYN_SENT"),
            &State::SynReceived => write!(f, "SYN_RECEIVED"),
            &State::Established => write!(f, "ESTABLISHED"),
            &State::FinWait1    => write!(f, "FIN_WAIT_1"),
            &State::FinWait2    => write!(f, "FIN_WAIT_2"),
            &State::CloseWait   => write!(f, "CLOSE_WAIT"),
            &State::Closing     => write!(f, "CLOSING"),
            &State::LastAck     => write!(f, "LAST_ACK"),
            &State::TimeWait    => write!(f, "TIME_WAIT")
        }
    }
}

#[derive(Debug)]
struct Retransmit {
    sent: bool // FIXME
}

impl Retransmit {
    fn new() -> Retransmit {
        Retransmit { sent: false }
    }

    fn reset(&mut self) {
        self.sent = false
    }

    fn check(&mut self) -> bool {
        let result = !self.sent;
        self.sent = true;
        result
    }
}

/// A Transmission Control Protocol data stream.
#[derive(Debug)]
pub struct TcpSocket<'a> {
    state:         State,
    local_end:     IpEndpoint,
    remote_end:    IpEndpoint,
    local_seq_no:  i32,
    remote_seq_no: i32,
    retransmit:    Retransmit,
    rx_buffer:     SocketBuffer<'a>,
    tx_buffer:     SocketBuffer<'a>
}

impl<'a> TcpSocket<'a> {
    /// Create a socket using the given buffers.
    pub fn new<T>(rx_buffer: T, tx_buffer: T) -> Socket<'a, 'static>
            where T: Into<SocketBuffer<'a>> {
        let rx_buffer = rx_buffer.into();
        if rx_buffer.capacity() > <u16>::max_value() as usize {
            panic!("buffers larger than {} require window scaling, which is not implemented",
                   <u16>::max_value())
        }

        Socket::Tcp(TcpSocket {
            state:         State::Closed,
            local_end:     IpEndpoint::default(),
            remote_end:    IpEndpoint::default(),
            local_seq_no:  0,
            remote_seq_no: 0,
            retransmit:    Retransmit::new(),
            tx_buffer:     tx_buffer.into(),
            rx_buffer:     rx_buffer.into()
        })
    }

    /// Return the connection state.
    #[inline(always)]
    pub fn state(&self) -> State {
        self.state
    }

    /// Return the local endpoint.
    #[inline(always)]
    pub fn local_endpoint(&self) -> IpEndpoint {
        self.local_end
    }

    /// Return the remote endpoint.
    #[inline(always)]
    pub fn remote_endpoint(&self) -> IpEndpoint {
        self.remote_end
    }

    fn set_state(&mut self, state: State) {
        if self.state != state {
            if self.remote_end.addr.is_unspecified() {
                net_trace!("tcp:{}: state={}→{}",
                           self.local_end, self.state, state);
            } else {
                net_trace!("tcp:{}:{}: state={}→{}",
                           self.local_end, self.remote_end, self.state, state);
            }
        }
        self.state = state
    }

    /// Start listening on the given endpoint.
    ///
    /// # Panics
    /// This function will panic if the socket is not in the CLOSED state.
    pub fn listen(&mut self, endpoint: IpEndpoint) {
        assert!(self.state == State::Closed);

        self.local_end  = endpoint;
        self.remote_end = IpEndpoint::default();
        self.set_state(State::Listen);
    }

    /// See [Socket::collect](enum.Socket.html#method.collect).
    pub fn collect(&mut self, src_addr: &IpAddress, dst_addr: &IpAddress,
                   protocol: IpProtocol, payload: &[u8])
            -> Result<(), Error> {
        if protocol != IpProtocol::Tcp { return Err(Error::Rejected) }

        let packet = try!(TcpPacket::new(payload));
        let repr = try!(TcpRepr::parse(&packet, src_addr, dst_addr));

        if self.local_end.port != repr.dst_port { return Err(Error::Rejected) }
        if !self.local_end.addr.is_unspecified() &&
           self.local_end.addr != *dst_addr { return Err(Error::Rejected) }

        if self.remote_end.port != 0 &&
           self.remote_end.port != repr.src_port { return Err(Error::Rejected) }
        if !self.remote_end.addr.is_unspecified() &&
           self.remote_end.addr != *src_addr { return Err(Error::Rejected) }

        match (self.state, repr) {
            (State::Closed, _) => Err(Error::Rejected),

            (State::Listen, TcpRepr {
                src_port, dst_port, control: TcpControl::Syn, seq_number, ack_number: None, ..
            }) => {
                self.local_end     = IpEndpoint::new(*dst_addr, dst_port);
                self.remote_end    = IpEndpoint::new(*src_addr, src_port);
                self.remote_seq_no = seq_number;
                // FIXME: use something more secure
                self.local_seq_no  = !seq_number;
                self.set_state(State::SynReceived);

                // FIXME: queue data from SYN
                self.retransmit.reset();
                Ok(())
            }

            _ => {
                // This will cause the interface to reply with an RST.
                Err(Error::Rejected)
            }
        }
    }

    /// See [Socket::dispatch](enum.Socket.html#method.dispatch).
    pub fn dispatch(&mut self, f: &mut FnMut(&IpAddress, &IpAddress,
                                             IpProtocol, &PacketRepr) -> Result<(), Error>)
            -> Result<(), Error> {
        let mut repr = TcpRepr {
            src_port:   self.local_end.port,
            dst_port:   self.remote_end.port,
            control:    TcpControl::None,
            seq_number: 0,
            ack_number: None,
            window_len: (self.rx_buffer.capacity() - self.rx_buffer.len()) as u16,
            payload:    &[]
        };

        // FIXME: process

        match self.state {
            State::Closed |
            State::Listen => {
                return Err(Error::Exhausted)
            }
            State::SynReceived => {
                if !self.retransmit.check() { return Err(Error::Exhausted) }
                repr.control    = TcpControl::Syn;
                repr.seq_number = self.local_seq_no;
                repr.ack_number = Some(self.remote_seq_no + 1);
                net_trace!("tcp:{}:{}: SYN sent",
                           self.local_end, self.remote_end);
            }
            _ => unreachable!()
        }

        f(&self.local_end.addr, &self.remote_end.addr, IpProtocol::Tcp, &repr)
    }
}

impl<'a> PacketRepr for TcpRepr<'a> {
    fn buffer_len(&self) -> usize {
        self.buffer_len()
    }

    fn emit(&self, src_addr: &IpAddress, dst_addr: &IpAddress, payload: &mut [u8]) {
        let mut packet = TcpPacket::new(payload).expect("undersized payload");
        self.emit(&mut packet, src_addr, dst_addr)
    }
}

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

    #[test]
    fn test_buffer() {
        let mut buffer = SocketBuffer::new(vec![0; 8]); // ........
        buffer.enqueue(6).copy_from_slice(b"foobar");   // foobar..
        assert_eq!(buffer.dequeue(3), b"foo");          // ...bar..
        buffer.enqueue(6).copy_from_slice(b"ba");       // ...barba
        buffer.enqueue(4).copy_from_slice(b"zho");      // zhobarba
        assert_eq!(buffer.dequeue(6), b"barba");        // zho.....
        assert_eq!(buffer.dequeue(8), b"zho");          // ........
        buffer.enqueue(8).copy_from_slice(b"gefug");    // ...gefug
    }
}