HiveBrain v1.2.0
Get Started
← Back to all entries
patternrustMinor

A socket server that receives strings and prints them

Submitted by: @import:stackexchange-codereview··
0
Viewed 0 times
printsandstringsthatreceivesthemserversocket

Problem

Summary: Is my server handling every error and corner case possible? If not, what did I miss?

Details

I'm trying to write the most basic socket server I can, so I can learn TCP and the pitfalls of parsing network data. I'm also worried about security (DDOS mitigation) and making sure the server never panics (unless it is an exceptional error, of course). I'm using mio to manage my event loop and sockets.

This is what I have so far:

My imports:

extern crate mio;
extern crate byteorder;

use mio::*;
use mio::tcp::*;
use mio::util::Slab;
use std::net::SocketAddr;
use std::collections::VecDeque;


An enum that represents the status of a client connection. Idle means the connection is not sending anything. WaitingHeader means the server received something but doesn't have enough data to parse the header (the header is 2 u8s, which represent a single u16). The header is the size of the message. Finally ReceivingData has the now-parsed header and a Vec that is filled as the network sends bytes, once this vector is full, the contents are parsed and the client goes back to Idle

#[derive(Debug, PartialEq, Eq)]
enum ClientStatus {
    Idle,
    WaitingHeader([(bool, u8); 2]),
    ReceivingData(u16, Vec),
}


Structs to hold the state of each connection and the server itself:

#[derive(Debug)]
struct Client {
    socket: TcpStream,
    status: ClientStatus,
    buffer: VecDeque,
}

impl Client {
    fn new(socket: TcpStream) -> Client {
        Client {
            socket: socket,
            status: ClientStatus::Idle,
            buffer: VecDeque::new(),
        }
    }
}

#[derive(Debug)]
struct Server {
    socket: TcpListener,
    clients: Slab,
}


The server token and a function to convert [u8; 2]s into u16s:

```
const SERVER_TOKEN: Token = Token(0);

fn convert_header(header: [u8; 2]) -> u16 {
use std::io::Cursor;
use byteorder::{LittleEndian, ReadBytesExt};

let mut rdr = Cursor::new([header[0], header[1]]);

Solution

Disclaimer you didn't provide any suggestions on how to run or test your code. I was able to compile and run the server, but wasn't sure how to properly send a message — my netcat-fu wasn't strong enough.

I also don't have extensive experience with MIO and haven't ever had to defend against a DOS, much less a DDOS. My best advice (which isn't saying much) is to always have limits. Don't allow buffers to grow forever, don't allow infinite amounts of clients.

The code was mostly understandable. My eyes glazed over a bit at the page full of h[1].0 and I was surprised to see VecDeque used.

-
Having a (bool, u8) is quite confusing. Giving names to the members would have been better. This would have shown me that the code is implementing an array that can be partially full. That functionality exists in a crate called arrayvec.

-
A Cursor isn't needed when parsing out the size field as the code will always read from the beginning of the slice.

-
I inlined the header parsing as it basically was only one line after other changes.

-
There's a few spots that are less-than-efficient.

-
Anytime you process byte-by-byte is likely to be inefficient, as are any extra copies of data.

-
Use extend instead of pushing bytes one-by-one. This can be combined with drain. Beware that this is still copying data.

-
Move buf up one level. I think if it is inside the loop it will be reinitialized each time.

-
Instead of reading from the socket into a temporary buffer at all, directly read to the final buffer.

-
When printing the response, convert the byte array to a &str, avoiding the need to clone and make a String.

-
Instead of unwrapping the response, print the Debug form of the result. This prevents a panic case.

-
Try to place comments into expect calls, making the code more self-documenting.

-
Don't include no trailing punctuation in expect. A colon will be added automatically.

-
I'd prefer to see more methods instead of free functions.

-
Splitting up read would help with readability.

-
Definitely use Result. It even helps when refactoring. For example, when I extracted drain_socket, I didn't originally handle the failure case correctly. I just returned from the drain_socket function.

-
A match with one useful arm should be written as if let.

-
Looping while the message state machine can continue allows processing multiple exchanges that have all been buffered up. This should help avoid a case where many small messages fill up the buffer.

-
After moving the state machine to ClientStatus, I generalized the concept of deferring the change of state (the old done_processing flag is an example). To do this, I used mem::replace to switch out the old state while processing it.

-
Prefer expect over unwrap. This helps you and users understand what failed when it does.

```
#[macro_use]
extern crate quick_error;
extern crate mio;
extern crate byteorder;
extern crate arrayvec;

use arrayvec::ArrayVec;
use byteorder::{LittleEndian, ReadBytesExt};
use mio::*;
use mio::tcp::*;
use mio::util::Slab;
use std::net::SocketAddr;
use std::{cmp, mem, io};

quick_error! {
#[derive(Debug)]
pub enum Error {
ReadFromSocket(err: io::Error) {
description("could not read from socket")
display("could not read from socket: {}", err)
cause(err)
}
}
}

#[derive(Debug, PartialEq, Eq)]
enum ClientStatus {
Idle,
WaitingHeader(ArrayVec),
ReceivingData(u16, Vec),
}

#[derive(Debug)]
struct Client {
socket: TcpStream,
status: ClientStatus,
buffer: Vec,
}

impl Client {
fn new(socket: TcpStream) -> Client {
Client {
socket: socket,
status: ClientStatus::Idle,
buffer: Vec::new(),
}
}
}

#[derive(Debug)]
struct Server {
socket: TcpListener,
clients: Slab,
}

const SERVER_TOKEN: Token = Token(0);

impl ClientStatus {
fn next(self, buffer: &mut Vec) -> (Self, bool) {
use ClientStatus::*;

match self {
Idle => {
(WaitingHeader(Default::default()), true)
},
WaitingHeader(mut h) => {
let remaining = h.capacity() - h.len();
assert!(remaining > 0);
let available = cmp::min(remaining, buffer.len());
h.extend(buffer.drain(..available));

if h.len() == h.capacity() {
let data_len = h.as_ref().read_u16::().expect("Not enough data to parse header");
(ReceivingData(data_len, Vec::with_capacity(data_len as usize)), true)
} else {
(WaitingHeader(h), false)
}
},
ReceivingData(data_len, mut data) => {
let remaining = (data_len as usize) - data.len();
if buffer.len() >= remaining {
data.extend(buffer.drain(0..remaining));

Code Snippets

#[macro_use]
extern crate quick_error;
extern crate mio;
extern crate byteorder;
extern crate arrayvec;

use arrayvec::ArrayVec;
use byteorder::{LittleEndian, ReadBytesExt};
use mio::*;
use mio::tcp::*;
use mio::util::Slab;
use std::net::SocketAddr;
use std::{cmp, mem, io};

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        ReadFromSocket(err: io::Error) {
            description("could not read from socket")
            display("could not read from socket: {}", err)
            cause(err)
        }
    }
}

#[derive(Debug, PartialEq, Eq)]
enum ClientStatus {
    Idle,
    WaitingHeader(ArrayVec<[u8; 2]>),
    ReceivingData(u16, Vec<u8>),
}

#[derive(Debug)]
struct Client {
    socket: TcpStream,
    status: ClientStatus,
    buffer: Vec<u8>,
}

impl Client {
    fn new(socket: TcpStream) -> Client {
        Client {
            socket: socket,
            status: ClientStatus::Idle,
            buffer: Vec::new(),
        }
    }
}

#[derive(Debug)]
struct Server {
    socket: TcpListener,
    clients: Slab<Client>,
}

const SERVER_TOKEN: Token = Token(0);

impl ClientStatus {
    fn next(self, buffer: &mut Vec<u8>) -> (Self, bool) {
        use ClientStatus::*;

        match self {
            Idle => {
                (WaitingHeader(Default::default()), true)
            },
            WaitingHeader(mut h) => {
                let remaining = h.capacity() - h.len();
                assert!(remaining > 0);
                let available = cmp::min(remaining, buffer.len());
                h.extend(buffer.drain(..available));

                if h.len() == h.capacity() {
                    let data_len = h.as_ref().read_u16::<LittleEndian>().expect("Not enough data to parse header");
                    (ReceivingData(data_len, Vec::with_capacity(data_len as usize)), true)
                } else {
                    (WaitingHeader(h), false)
                }
            },
            ReceivingData(data_len, mut data) => {
                let remaining = (data_len as usize) - data.len();
                if buffer.len() >= remaining {
                    data.extend(buffer.drain(0..remaining));

                    // We have everything, parse it!
                    println!("Received some data! Size: {}\n\tValue: {:?}",
                             data.len(),
                             std::str::from_utf8(&data));

                    (Idle, true)
                } else {
                    data.extend(buffer.drain(..));
                    (ReceivingData(data_len, data), false)
                }
            },
        }
    }
}

impl Client {
    fn drain_socket(&mut self) -> Result<(), Error> {
        loop {
            let bytes = try!(self.socket.try_read_buf(&mut self.buffer).map_err(Error::ReadFromSocket));
            // Should "self" be reset on failure, or are socket errors
            // unrecoverable and the server will just close the
            // connection?

            match bytes {
                None =>
if let Some(v) = client.buffer.pop_front()

Context

StackExchange Code Review Q#129528, answer score: 2

Revisions (0)

No revisions yet.