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}
```