chore: Refactor and modularize server code

Signed-off-by: rrouviere <[email protected]>

Author: ruben-etu

topics/p2p-transfer-protocol/src/main.rs

@@ -87,6 +87,7 @@ fn server_mode(bind_addr: &str, file_path: &str) -> io::Result<()> {
     );
 
     // Wait for client
+    // FIXME: We handle multiple clients, but since the output file is shared among them, this is unhelpful.
     for stream in listener.incoming() {
         match stream {
             Ok(mut s) => {

topics/p2p-transfer-protocol/src/protocol/message.rs

@@ -32,9 +32,10 @@ impl FromStr for Message {
         if upper.starts_with("HELLO") {
             let parts: Vec<&str> = s.split_whitespace().collect();
             if parts.len() == 2
-                && let Ok(file_size) = parts[1].parse::<u64>() {
-                    return Ok(Message::Hello { file_size });
-                }
+                && let Ok(file_size) = parts[1].parse::<u64>()
+            {
+                return Ok(Message::Hello { file_size });
+            }
             error!("Invalid HELLO message");
             return Err(());
         };

topics/p2p-transfer-protocol/src/server/connection.rs

@@ -0,0 +1,57 @@
+use log::info;
+use std::fs::File;
+use std::io::{self, Write, copy};
+use std::net::TcpStream;
+
+use crate::protocol::StateMachine;
+use crate::protocol::connection_state::ConnectionState;
+use crate::protocol::message::Message;
+
+use crate::server::events::handle_message_loop;
+
+pub(crate) fn handle_connection(file: &mut File, stream: &mut TcpStream) -> Result<(), io::Error> {
+    // Intialize the state machine
+    let mut state_machine = StateMachine::new();
+    state_machine.transition(ConnectionState::Listening);
+
+    // Start the negociation message loop.
+    handle_message_loop(&mut state_machine, stream)?;
+
+    // Now we're in Established state, receive the file data.
+    let expected_file_size = state_machine.file_size.unwrap_or(0);
+
+    receive_file(file, stream, expected_file_size)
+}
+
+pub fn send_message(stream: &mut TcpStream, msg: &Message) -> Result<(), io::Error> {
+    let msg_str = msg.to_string();
+    stream.write_all(msg_str.as_bytes())?;
+    Ok(())
+}
+
+pub fn receive_file(
+    file: &mut File,
+    stream: &mut TcpStream,
+    expected_file_size: u64,
+) -> io::Result<()> {
+    let mut sized_stream = io::Read::take(stream, expected_file_size);
+
+    let bytes_copied = copy(&mut sized_stream, file)?;
+    file.flush()?;
+
+    if bytes_copied != expected_file_size {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            format!(
+                "Expected {} bytes, but received {} bytes",
+                expected_file_size, bytes_copied
+            ),
+        ));
+    }
+
+    info!(
+        "File received successfully, {} bytes written.",
+        bytes_copied
+    );
+    Ok(())
+}

topics/p2p-transfer-protocol/src/server/events.rs

@@ -0,0 +1,71 @@
+use log::{debug, error, info, warn};
+use std::io::{self, BufRead, BufReader};
+use std::net::TcpStream;
+
+use crate::protocol::StateMachine;
+use crate::protocol::connection_state::ConnectionState;
+use crate::protocol::message::Message;
+use crate::server::connection::send_message;
+
+fn on_hello(file_size: u64, state_machine: &mut StateMachine) -> Message {
+    info!("Client wants to send a file of size {} bytes", file_size);
+    state_machine.transition(ConnectionState::HelloReceived);
+
+    state_machine.file_size = Some(file_size);
+
+    // TODO: Add a basic chec   k to handle the NACK case.
+
+    // send ack
+    state_machine.transition(ConnectionState::ACKSent);
+
+    Message::Ack
+}
+
+fn on_message(msg: Message, state_machine: &mut StateMachine) -> Option<Message> {
+    match msg {
+        Message::Hello { file_size } => Some(on_hello(file_size, state_machine)),
+        Message::Send => {
+            info!("Client is starting to send data.");
+            state_machine.transition(ConnectionState::Established);
+            None
+        }
+
+        _ => {
+            warn!("Unexpected message from client: {:?}", msg);
+            None
+        }
+    }
+}
+pub(crate) fn handle_message_loop(
+    state_machine: &mut StateMachine,
+    stream: &mut TcpStream,
+) -> Result<(), io::Error> {
+    let reader_stream = stream.try_clone().expect("Failed to clone stream");
+    let writer_stream = stream;
+    let reader = BufReader::new(reader_stream);
+
+    for line in reader.lines() {
+        let message = match line {
+            Ok(line) => {
+                debug!("Received: {}", line);
+                line.parse::<Message>()
+            }
+            Err(e) => {
+                error!("Error reading line: {}", e);
+                continue;
+            }
+        };
+
+        // Process message
+        // FIXME: Do not use unwrap, fix the result type mess.
+        if let Some(resp) = on_message(message.unwrap(), state_machine) {
+            send_message(writer_stream, &resp)?
+        }
+
+        // If we are in Established state, stop the line-based logic and receive the data.
+        if let ConnectionState::Established = state_machine.current_state() {
+            break;
+        }
+    }
+    Ok(())
+}

topics/p2p-transfer-protocol/src/server/mod.rs

@@ -1,117 +1,15 @@
-use log::{debug, error, info, warn};
 use std::fs::File;
-use std::io::{self, BufRead, BufReader, Write, copy};
+use std::io;
 use std::net::TcpStream;
 
-use crate::protocol::StateMachine;
-use crate::protocol::connection_state::ConnectionState;
-use crate::protocol::message::Message;
+use crate::server::connection::handle_connection;
 
-fn on_hello(file_size: u64, state_machine: &mut StateMachine) -> Message {
-    info!("Client wants to send a file of size {} bytes", file_size);
-    state_machine.transition(ConnectionState::HelloReceived);
-
-    state_machine.file_size = Some(file_size);
-
-    // TODO: Add a basic check to handle the NACK case.
-
-    // send ack
-    state_machine.transition(ConnectionState::ACKSent);
-
-    Message::Ack
-}
-
-fn on_message(msg: Message, state_machine: &mut StateMachine) -> Option<Message> {
-    match msg {
-        Message::Hello { file_size } => Some(on_hello(file_size, state_machine)),
-        Message::Send => {
-            info!("Client is starting to send data.");
-            state_machine.transition(ConnectionState::Established);
-            None
-        }
-
-        _ => {
-            warn!("Unexpected message from client: {:?}", msg);
-            None
-        }
-    }
-}
-
-fn send_message(stream: &mut TcpStream, msg: &Message) -> Result<(), io::Error> {
-    let msg_str = msg.to_string();
-    stream.write_all(msg_str.as_bytes())?;
-    Ok(())
-}
+mod connection;
+mod events;
 
 pub fn run_server(file: &mut File, stream: &mut TcpStream) -> Result<(), io::Error> {
     let peer_addr = stream.peer_addr()?;
     println!("New connection from {peer_addr}");
 
-    // Intialize the state machine
-    let mut state_machine = StateMachine::new();
-    state_machine.transition(ConnectionState::Listening);
-
-    // Start the negociation message loop.
-    message_loop(&mut state_machine, stream)?;
-
-    // Now we're in Established state, receive the file data.
-    let expected_file_size = state_machine.file_size.unwrap_or(0);
-
-    receive_file(file, stream, expected_file_size)
-}
-
-fn receive_file(
-    file: &mut File,
-    stream: &mut TcpStream,
-    expected_file_size: u64,
-) -> io::Result<()> {
-    let mut sized_stream = io::Read::take(stream, expected_file_size);
-
-    let bytes_copied = copy(&mut sized_stream, file)?;
-    file.flush()?;
-
-    if bytes_copied != expected_file_size {
-        return Err(io::Error::new(
-            io::ErrorKind::InvalidData,
-            format!(
-                "Expected {} bytes, but received {} bytes",
-                expected_file_size, bytes_copied
-            ),
-        ));
-    }
-
-    info!(
-        "File received successfully, {} bytes written.",
-        bytes_copied
-    );
-    Ok(())
-}
-
-fn message_loop(state_machine: &mut StateMachine, stream: &mut TcpStream) -> Result<(), io::Error> {
-    let reader_stream = stream.try_clone().expect("Failed to clone stream");
-    let writer_stream = stream;
-    let reader = BufReader::new(reader_stream);
-
-    for line in reader.lines() {
-        let message = match line {
-            Ok(line) => {
-                debug!("Received: {}", line);
-                line.parse::<Message>()
-            }
-            Err(e) => {
-                error!("Error reading line: {}", e);
-                continue;
-            }
-        };
-
-        // Process message
-        // FIXME: Do not use unwrap, fix the result type mess.
-        if let Some(resp) = on_message(message.unwrap(), state_machine) { send_message(writer_stream, &resp)? }
-
-        // If we are in Established state, stop the line-based logic and receive the data.
-        if let ConnectionState::Established = state_machine.current_state() {
-            break;
-        }
-    }
-    Ok(())
+    handle_connection(file, stream)
 }