Creating functions that send and receive atomic messages in Rust paradigm

Question

I'm converting a library from C to Rust, and I want to see if this is the optimal way to implement network communication. This code is just being used to test the library, but I plan to implement the library's messaging system with the same send_message and receive_message functions.

I want to use only safe Rust.

Some points:

• The Socket type is an enum that can be either TcpStream or UdpSocket and the send and receive functions are wrappers around the respective raw operations for each protocol
• The need to support both protocols is why I'm not using buffered operations like BufReader
• Messages are null terminated because they are unknown length
• Segment size is set to maximum UDP segment payload because TCP can fragment and UDP does not

I've already made some adjustments based on previous review elsewhere, so my original concerns were already addressed, and this is more open-ended.

pub enum Socket {
 Tcp (TcpStream),
 Udp (UdpSocket),
}
impl Socket {
 pub fn send(&mut self, data: &[u8]) -> Result<usize> {
 match self {
 Self::Tcp(ref mut stream) => stream.write(data),
 Self::Udp(socket) => socket.send(data),
 }
 }
 pub fn receive(&mut self, buffer: &mut [u8]) -> Result<usize> {
 match self {
 Self::Tcp(ref mut stream) => stream.read(buffer),
 Self::Udp(socket) => socket.recv(buffer),
 }
 }
}
// Max segment payload based on UDP does not fragment or buffer but TCP does
// (IPv 6 TCP @ 1440 < IPv6 UDP @ 1452 < IPv4 TCP @ 1460 < IPv4 UDP @ 1472)
const SEGMENT_SIZE: usize = 1452; 
fn send_message(socket: &mut Socket, message: &str) -> Result<usize, Error> {
 let mut bytes_sent: usize = 0;
 let message_bytes = message.as_bytes();
 
 loop {
 let chunk = &message_bytes[bytes_sent..(bytes_sent + std::cmp::min(message_bytes.len() - bytes_sent, SEGMENT_SIZE))];
 match socket.send(chunk) {
 Ok(length) => {
 bytes_sent += length;
 if bytes_sent == message_bytes.len() {
 return Ok(bytes_sent)
 }
 },
 Err(_) => todo!(),
 };
 }
}
fn receive_message(socket: &mut Socket) -> Result<String, Error> {
 let mut buffer: Vec<u8> = Vec::with_capacity(SEGMENT_SIZE);
 buffer.resize(SEGMENT_SIZE, 0);
 let mut bytes_received: usize = 0;
 loop {
 match socket.receive(&mut buffer[bytes_received..(bytes_received + SEGMENT_SIZE)]) {
 Ok(length) => {
 bytes_received += length;
 match buffer.get(bytes_received - 1) {
 Some(&0) => {
 match std::str::from_utf8(&buffer[0..bytes_received]) {
 Ok(string) => return Ok(string.to_owned()),
 Err(_) => todo!(),
 };
 },
 Some(_) => buffer.resize(buffer.len() + SEGMENT_SIZE, 0),
 None => todo!(),
 };
 },
 Err(_) => todo!(),
 }
 }
}

Answer 1

Prefer standard library traits over associated functions

Instead of implementing send() and receive() on Socket, consider implementing the Read and Write traits for it.

Prefer associated functions over free functions...

... if their first argument is (a reference to) a concrete type anyway.

Both send_message() and receive_message() receive &mut Socket as a first argument. This is a strong indicator, that they maybe should be associated functions (aka. methods).

Consider adding implementations of conversion traits

Since Socket may be either a TCP or a UDP socket, you may want to add a From<TcpStream> and From<UdpSocket> implementation to simplify crating such a socket from the underlying implementation.