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<u8>),
}
Structs to hold the state of each connection and the server itself:
#[derive(Debug)]
struct Client {
socket: TcpStream,
status: ClientStatus,
buffer: VecDeque<u8>,
}
impl Client {
fn new(socket: TcpStream) -> Client {
Client {
socket: socket,
status: ClientStatus::Idle,
buffer: VecDeque::new(),
}
}
}
#[derive(Debug)]
struct Server {
socket: TcpListener,
clients: Slab<Client>,
}
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]]);
//The following call will only panic if the cursor doesn't have enough data. That should
// never happen
rdr.read_u16::<LittleEndian>().expect("Header parsing error.")
}
Then, a function that receives the data from the socket until it is empty (or something goes wrong) and manages the state of the client. It works like this:
- Read the socket until it is empty and push all the data to the client buffer
- Once that is done, based on the client status, it tries to gather the whole packet so it can be parsed.
- If the client is idle, wait for two bytes that represent the header and then wait for the rest of the packet.
- Once the packet is done and there is nothing else to consume in the buffer, the client goes back to
Idle.
fn read(client: &mut Client) {
loop {
let mut buf = [0; 2048];
match client.socket.try_read(&mut buf) {
Err(e) => {
println!("Error while reading socket: {:?}", e);
//Should "client" be reset here, or are socket errors unrecoverable and the
// server will just close the connection?
return;
}
Ok(None) => break,
Ok(Some(len)) => {
println!("Received raw: {:?}", &buf[..len]);
client.buffer.reserve(len as usize);
for i in 0..len {
client.buffer.push_back(buf[i]);
}
}
}
}
while client.buffer.len() > 0 {
let buffer_len = client.buffer.len();
if client.status == ClientStatus::Idle {
client.status = ClientStatus::WaitingHeader([(false, 0), (false, 0)]);
}
if let ClientStatus::WaitingHeader(mut h) = client.status {
assert!(h[1].0 == false);
if buffer_len == 1 {
if h[0].0 == false {
h[0].0 = true;
//Length is known, this won't panic:
h[0].1 = client.buffer.pop_front().unwrap();
} else {
h[1].0 = true;
h[1].1 = client.buffer.pop_front().unwrap();
}
} else {
if h[0].0 == false {
h[0].0 = true;
h[0].1 = client.buffer.pop_front().unwrap();
}
h[1].0 = true;
h[1].1 = client.buffer.pop_front().unwrap();
}
if h[0].0 && h[1].0 {
let data_len = convert_header([h[0].1, h[1].1]);
client.status = ClientStatus::ReceivingData(data_len,
Vec::with_capacity(data_len as usize));
} else {
client.status = ClientStatus::WaitingHeader(h);
}
}
let mut parsing_done = false;
if let ClientStatus::ReceivingData(data_len, ref mut data) = client.status {
while let Some(byte) = client.buffer.pop_front() {
data.push(byte);
if data.len() >= (data_len as usize) {
// We have everything, parse it!
println!("Received some data! Size: {}\n\tValue: {}",
data_len,
//Data could be invalid utf8 and it would panic, but for testing
// purposes, unwrapping is enough, though from_utf8_lossy is an
// option.
String::from_utf8(data.clone()).unwrap());
parsing_done = true;
break;
}
}
}
if parsing_done {
client.status = ClientStatus::Idle;
}
}
}
And finally the Handler trait implementation of the server and the main function:
impl Handler for Server {
type Timeout = usize;
type Message = ();
fn ready(&mut self, event_loop: &mut EventLoop<Server>, token: Token, events: EventSet) {
match token {
SERVER_TOKEN => {
let client_socket = match self.socket.accept() {
Err(e) => {
println!("Accept error: {}", e);
return;
}
Ok(None) => {
println!("Socket was not ready.");
return;
}
Ok(Some((sock, _addr))) => sock,
};
let new_token = match self.clients.insert(Client::new(client_socket)) {
Ok(t) => t,
Err(client) => {
println!("Connection could not be accepted, max number reached.");
let _ = client.socket.shutdown(Shutdown::Both);
return;
}
};
println!("New connection! Total connections: {}",
self.clients.count());
match event_loop.register(&self.clients[new_token].socket,
new_token,
EventSet::readable(),
PollOpt::edge() | PollOpt::oneshot()) {
Err(e) => {
println!("IO error when registering a client: {}", e);
self.clients.remove(token);
}
_ => (),
}
}
_ => {
let remove = if !events.is_hup() {
let client: &mut Client = match self.clients.get_mut(token) {
Some(c) => c,
None => {
println!("Received token that was not registered ({}). Ignoring.",
token.0);
return;
}
};
read(client);
match event_loop.reregister(&client.socket,
token,
EventSet::readable() | EventSet::hup(),
PollOpt::edge() | PollOpt::oneshot()) {
Err(e) => {
println!("IO error when reregistering a client: {}", e);
true
}
_ => false,
}
} else {
println!("Removing connection {}", token.0);
true
};
if remove {
self.clients.remove(token);
}
}
}
}
}
fn main() {
let address = "127.0.0.1:10000".parse::<SocketAddr>().unwrap();
let server_socket = TcpListener::bind(&address).unwrap();
let mut event_loop = EventLoop::new().unwrap();
let mut handler = Server {
clients: Slab::new_starting_at(Token(1), 10),
socket: server_socket,
};
event_loop.register(&handler.socket,
SERVER_TOKEN,
EventSet::readable(),
PollOpt::edge())
.unwrap();
event_loop.run(&mut handler).unwrap();
}
Some questions
Should I not call
unwrapwhen I know something can't fail? Is this a sign that I should have structured things differently? For instance, callingpop_frontin theclient.bufferonce I know the length of the buffer is larger than 0.Is my error handling all-right? I feel like I have a lot of
matches trying to catch everything. Would it be better to put my code in another function that returns aResultandtry!the errors?Can I protect the server from errors in the client? If someone starts sending a bunch of random bytes, the server won't parse them into something meaninful, is that ok or can I somehow detect and prevent that?
