TCP Server using Tokio

Question

Solving the following problem on Protohackers: https://protohackers.com/problem/1

The problem:

Each request is a single line containing a JSON object, terminated by a newline character ('\n', or ASCII 10). Each request begets a response, which is also a single line containing a JSON object, terminated by a newline character.

A conforming request object has the required field method, which must always contain the string "isPrime", and the required field number, which must contain a number. Any JSON number is a valid number, including floating-point values.

A conforming response object has the required field method, which must always contain the string "isPrime", and the required field prime, which must contain a boolean value: true if the number in the request was prime, false if it was not.

Whenever you receive a conforming request, send back a correct response, and wait for another request. Whenever you receive a malformed request, send back a single malformed response, and disconnect the client.

main.rs

use clap::Parser;
use two::run;

#[derive(Parser)]
struct Cli {
    /// The port to run the server on
    port: i32,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    env_logger::init();
    let args = Cli::parse();
    run(args.port).await
}

lib.rs

use log::{info, warn};
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};

#[derive(Debug, Default, Deserialize)]
struct Request {
    method: String,
    number: f64,
}

#[derive(Debug)]
enum RequestError {
    InvalidMethod(String),
    InvalidNumber(f64),
    DeserializationError(serde_json::Error),
}

impl From<serde_json::Error> for RequestError {
    fn from(e: serde_json::Error) -> Self {
        Self::DeserializationError(e)
    }
}

impl std::fmt::Display for RequestError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl std::error::Error for RequestError {}

const VALID_METHOD: &str = "isPrime";

impl Request {
    fn from_slice(data: &[u8]) -> Result<Self, RequestError> {
        let request: Self = serde_json::from_slice(data)?;
        request.validate()?;
        Ok(request)
    }

    fn validate(&self) -> Result<(), RequestError> {
        use RequestError::*;
        if self.method != VALID_METHOD {
            return Err(InvalidMethod(self.method.clone()));
        }
        if self.number != (self.number as i64) as f64 {
            return Err(InvalidNumber(self.number));
        }

        Ok(())
    }

    fn number(&self) -> i64 {
        self.number as i64
    }
}

#[derive(Debug, Default, Serialize)]
struct Response {
    method: String,
    prime: bool,
}

impl Response {
    fn new(prime: bool) -> Self {
        Self {
            method: VALID_METHOD.into(),
            prime: prime,
        }
    }

    fn to_bytes(&self) -> Vec<u8> {
        serde_json::to_vec(self).unwrap()
    }
}

/// From https://docs.rs/primes/latest/src/primes/lib.rs.html
fn firstfac(x: i64) -> i64 {
    if x % 2 == 0 {
        return 2;
    };
    
    for n in (1..).map(|m| 2 * m + 1).take_while(|m| m * m <= x) {
        if x % n == 0 {
            return n;
        };
    }
    // No factor found. It must be prime.
    x
}

fn is_prime(n: i64) -> bool {
    if n <= 1 {
        return false;
    }
    firstfac(n) == n
}

async fn handle_connection(mut socket: TcpStream) -> Result<(), Box<dyn std::error::Error>> {
    loop {
        let mut data = vec![];
        let read = socket.read_to_end(&mut data).await?;
        if read == 0 {
            continue;
        }

        let (response, close) = match Request::from_slice(&data) {
            Ok(request) => (Response::new(is_prime(request.number())), false),
            Err(e) => {
                warn!("Received a malformed request. Sending back a malformed response and closing connection: {:?}", e);
                (Response::default(), true)
            }
        };
        let res_bytes = response.to_bytes();
        socket.write_all(&res_bytes).await?;
        if close {
            break;
        }
    }
    Ok(())
}

pub async fn run(port: i32) -> Result<(), Box<dyn std::error::Error>> {
    let addr = format!("127.0.0.1:{}", port);
    info!("Listening on address: {}", addr);

    let listener = TcpListener::bind(addr).await?;
    loop {
        let (socket, _) = listener.accept().await?;
        tokio::spawn(async move {
            _ = handle_connection(socket).await;
        });
    }
}

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

    #[test]
    fn test_request_deserialize() {
        let data = b"{\"method\":\"isPrime\",\"number\":42}";
        let request = Request::from_slice(data).unwrap();

        assert_eq!(request.method, "isPrime");
        assert_eq!(request.number(), 42);
    }

    #[test]
    fn test_response_serialize() {
        let resp = Response::new(true);
        let data = resp.to_bytes();
        let expected = b"{\"method\":\"isPrime\",\"prime\":true}";

        assert_eq!(&data, expected);
    }

    #[test]
    fn test_is_prime_negative() {
        let number = -1;
        assert_eq!(is_prime(number), false);
    }

    #[test]
    fn test_is_prime_zero() {
        let number = 0;
        assert_eq!(is_prime(number), false);
    }

    #[test]
    fn test_is_prime_positive() {
        let number = 13;
        assert_eq!(is_prime(number), true);

        let number = 16;
        assert_eq!(is_prime(number), false);
    }
}

Example usage using netcat

$ echo '{"method":"isPrime","number":42}' | nc localhost 1234                                
{"method":"isPrime","prime":false}

$ echo '{"method":"isPrime","number":13}' | nc localhost 1234
{"method":"isPrime","prime":true}

$ echo '{"method":"isPrime","number":13.43}' | nc localhost 1234
{"method":"","prime":false}

$ echo '{"method":"invalidMethod","number":13}' | nc localhost 1234
{"method":"","prime":false}
```

Answer 1

The handle_connection method is kind of broken:

async fn handle_connection(mut socket: TcpStream) -> Result<(), Box<dyn std::error::Error>> {
    loop {
        let mut data = vec![];
        let read = socket.read_to_end(&mut data).await?;
        if read == 0 {
            continue;
        }

        let (response, close) = match Request::from_slice(&data) {
            Ok(request) => (Response::new(is_prime(request.number())), false),
            Err(e) => {
                warn!("Received a malformed request. Sending back a malformed response and closing connection: {:?}", e);
                (Response::default(), true)
            }
        };
        let res_bytes = response.to_bytes();
        socket.write_all(&res_bytes).await?;
        if close {
            break;
        }
    }
    Ok(())
}

There is no point in using a loop if you are then doing read_to_end. After read_to_end, the connection is closed from the other side. Calling read_to_end again will yield 0 bytes, and as you continue in that case, you will be stuck in an infinite loop.

Other issues:

• your responses are not terminated by a newline
• you cannot query multiple values through the same connection, as you read_to_end() a single query

All of this can be fixed in the handle_connection method. Here is my attempt:

async fn handle_connection(mut socket: TcpStream) -> Result<(), Box<dyn std::error::Error>> {
    let (read_stream, mut write_stream) = socket.split();
    let mut read_stream = BufReader::new(read_stream);

    loop {
        let mut data = String::new();
        let read = read_stream.read_line(&mut data).await?;
        if read == 0 {
            break;
        }

        let (response, close) = match Request::from_str(&data) {
            Ok(request) => (Response::new(is_prime(request.number())), false),
            Err(e) => {
                warn!("Received a malformed request. Sending back a malformed response and closing connection: {:?}", e);
                (Response::default(), true)
            }
        };
        let res_bytes = response.to_bytes();
        write_stream.write_all(&res_bytes).await?;
        write_stream.write_all(b"\n").await?;
        if close {
            break;
        }
    }
    Ok(())
}