0
\$\begingroup\$

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}
```
\$\endgroup\$

1 Answer 1

2
\$\begingroup\$

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(())
}
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.