Editing A TCP Multi-Client Chat Server In Rust

I started this project simply to learn Rust, this is to say, I’m relatively new to rust as a programming language.

[[file:A_TCP_Multi-Client_Chat.jpg|500px]]

Photo by Pavan Trikutam on Unsplash



At the end of this project, the aim is to send and receive messages across devices (clients built with flutter), admit users into chatroom, list current chat rooms e.t.c. I would like to re-emphasise, this is not a full blown chat server. I will, however, explain things as best as I could, despite my limited knowledge on some concepts. For concepts I feel needs extra reading, I will attach helpful resource. This post assumes basic familiarity with async Rust, but not with network programming.



Now, what even is “TCP (Transmission Control Protocol)” ?

It is the communication standard for delivering data and messages through computers. It is one of the main building blocks holding together the internet we’ve come to know. TCP is designed to send packets across the internet, with the assurance of accurate and correct transmission of data(“packets”) between devices. Read more

How do we get our hands one of these “bad boys” ? You might ask
For the course of this, I leveraged the tokio TCP api. It’s the one of the viable option I came across.
In theory, our transmission protocol works by first listening, I mean, the only reasonable way to transmit anything is to somehow get what to transmit and and the best way to do that is by listening for data.

Setting up the listener


use anyhow::Result;
use tokio::{
    net::{
        TcpListener
   },
};

async fn main() -> Result<()> {
let port = "0.0.0.0:8080"
    // Create a listener at Port 8080 on our local machine
    let listener = TcpListener::bind("0.0.0.0:8080").await?;
    println!("Server running on {}", port);
}

The listener doesn’t transmit data itself — it simply waits for connections and hands you a TcpStream for each client. That stream is what we use to send and receive bytes.



This leads to another interesting concept — The 3-way handshake. It’s a process TCP uses to establish reliable connection between a client and a server by exchanging 3 packets of data.

How TCP connections are established

* 		Client → Server: sends a SYN packet (says “I’d like to talk on this port”).
* 		Server → Client: replies with SYN-ACK (says “Okay, I heard you, and I’m ready”).
* 		Client → Server: sends ACK (says “Cool, we’re connected”).
      let (tcp_stream, socket_addr) = match listener.accept().await {
            Ok(result) => result,
            Err(e) => {
                eprintln!("Failed to accept connection: {}", e);
                continue;
            }
        };
        println!("New connection from {}", socket_addr);

        let (read, mut write) = tcp_stream.into_split();
        let mut reader: BufReader<OwnedReadHalf> = BufReader::new(read);
        let mut first_line = String::new();
        let n = match reader.read_line(&mut first_line).await {
            Ok(n) => n,
            Err(e) => {
                eprintln!("Failed to read handshake from {}: {}", socket_addr, e);
                continue;
            }
        };
        println!("Read {} bytes for handshake from {}", n, socket_addr);
        if n == 0 {
            println!("Client {} disconnected before handshake", socket_addr);
            continue;
        }
        println!("This is first line , {}", first_line);
        let msg = first_line.trim_end().to_string();
        println!("Handshake received from {}: {}", socket_addr, msg);

A complete connection gives us two Objects, a tcp_stream and the socket_address. We split so that one async task can handle reads while another handles writes concurrently Read more
We are now able to accept the “speaker” — a client who speaks to our listener.
Our main focus from here is the tcp_stream, if you noticed we have
let (read, mut write) = tcp_stream.into_split();

The .into_split() method we invoked on the tcp_stream object returns two useful objects.
Splitting the stream allows two async tasks to run independently: one reads messages while the other writes responses. Without splitting, one task would block the other.
The read to read message from the client
reader.read_line(&mut line).await
The write to write messages back to the client
write.write_all(msg.as_bytes()).await
So far, we’ve indirectly learned how a TCP echo server works.

From Echo Server to Multi-Client Server

Simply connecting and sending exact messages back to a client isn’t the exact scope of this post.
The next part is to figure out how to accept multiple speakers (clients) to our listener. I was able to stumble upon mpsc — multiple producer, single consumer. A channel for sending messages across async tasks.
This channel is what extends our project from a simple TCP echo server to a multi — client chat server.

It is important to get the correct mental model before we proceed into the complicated bits.

So far, we’ve been able to write a program that gets a message from a client and send it back (echo server), but we need to accept multiple client, with the mpsc channel we created earlier, this is possible.
The channel is best viewed as a pipe, now you might ask, isn’t the tcp_stream also a pipe/channel, yes you would be correct, but it is a channel between two computers, unlike the mpsc channel that connects background task inside the server.

If this is not immediately clear, think of it this way, a client can connect to our current server, in fact, multiple client can connect, but that’s all it is without an mpsc channel. The clients need a way to communicate to each other. The channel can be best viewed as a mobile phone/mail box registered on the server, when a message is dropped by any client, our server should look up all mailboxes, send this message to them somehow and delivered to them (the client).

The TCP stream is the communication line between two computers; the mpsc channel is an in-memory communication line between tasks inside the same process

Now that we can relay messages internally, we need to track who’s who. A session ties a connected client’s TCP writer (tx) to their ID, and a Publisher manages chat rooms as groups of sessions.

A simple user session would look like this

use std::collections::HashMap;
use tokio::sync::mpsc::Sender;

use super::subscriber::Subscriber;

#[derive(Debug)]
pub struct Session {
    pub sub: Subscriber,
    pub tx: Sender<String>,
}

#[derive(Debug)]
pub struct SessionStore {
    sessions: HashMap<i32, Session>,
}

impl SessionStore {
    pub fn new() -> Self {
        SessionStore {
            sessions: HashMap::new(),
        }
    }
    pub fn get(&self, id: &i32) -> Option<&Session> {
        self.sessions.get(id)
    }
    pub fn insert(&mut self, id: i32, session: Session) {
        self.sessions.insert(id, session);
    }

    pub fn remove(&mut self, id: &i32) {
        self.sessions.remove(id);
    }

    pub fn iter(&self) -> impl Iterator<Item = (&i32, &Session)> {
        self.sessions.iter()
    }
}

A session currently hold a tx, the part of the pipe responsible for broadcasting the message we get from any client into the rx which will then receive messages sent transmitted by tx and relay it back to clients with

while let Some(msg) = rx.recv().await {
    write.write_all(msg.as_bytes()).await?;
}

Now that our server can relay messages internally, we need a way to group clients into chat rooms. This is where the Publisher struct comes in — it tracks which clients belong to which chat rooms and routes messages accordingly.

use tokio::sync::Mutex;

use crate::rust_models::session::SessionStore;
use crate::rust_models::sq_lite::SqLite;
use crate::rust_models::{conversation::ChatRoom, message::Message, subscriber::Subscriber};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

// Chat service

pub struct Publisher {
    pub conversations: HashMap<i32, ChatRoom>, // registry of subscribers/listeners
    pub db: SqLite,
}

impl Publisher {
    pub fn create_chat_room(&mut self, chat_room: ChatRoom) {
        println!("Chat Room Created For {:?} ", chat_room.convo_id);
        // Save in-memory
        self.conversations
            .insert(chat_room.convo_id.to_owned(), chat_room.to_owned());

    }
    pub fn subscribe_to_chat_room(
        &mut self,
        conversation_id: i32,
        subscribers: HashSet<Subscriber>,
    ) {
        if let Some(exisiting_convo) = self.conversations.get_mut(&conversation_id) {
            // conversation exists → push subscriber

            for s in &subscribers {
                exisiting_convo.subscribers.insert(s.to_owned());
            }
        } else {
            let conversation = ChatRoom {
                convo_id: conversation_id,
                subscribers: subscribers,
                messages: Vec::new(),
            };
            // conversation doesn’t exist → create chat room instead
            self.create_chat_room(conversation);
        }
    }

    pub fn unsubscribe_from_chat_service(
        &mut self,
        conversation_id: i32,
        subscribers: HashSet<Subscriber>,
    ) {
        if let Some(existing_members) = self.conversations.get_mut(&conversation_id) {
            for s in &subscribers {
                existing_members.subscribers.remove(&s);
            }
            println!(
                "Removed the following subscribers {:?} from this Chat {}",
                subscribers, conversation_id
            )
        }
    }

    pub async fn dispatch_messages(
        &mut self,
        conversation_id: i32,
        message: &Message,
        store: Arc<Mutex<SessionStore>>,
    ) {
        if let Some(convo) = self.conversations.get_mut(&conversation_id) {
            println!("Got here to dispatch");
            convo.messages.push(message.clone());
            let sessions = store.lock().await;
            for sub in &convo.subscribers {
                println!("Dsipatching messages to {}", sub.subscriber_id);
                if let Some(sess) = sessions.get(&sub.subscriber_id) {
                    if let Err(e) = sess.tx.send(message.content.clone()).await {
                        eprintln!("Failed to send to {}: {}", sub.subscriber_id, e);
                    }
                }
            }
        }
    }

    pub fn delete_chat_room(&mut self, conversation_id: i32) {
        self.conversations.remove(&conversation_id);

        println!("Chat Room No: {} deleted successfully", conversation_id);
    }
    pub fn list_chat_rooms(&mut self) {
        for (convo_id, convo) in self.conversations.iter() {
            println!(
                "Chat Room {}  -> subscribers {:?} messages {:?} ",
                convo_id, convo.subscribers, convo.messages
            )
        }
    }
}

This experiment was a great deep dive into async Rust, socket programming, and concurrent systems design.
 While the implementation is simple, it covers real-world primitives — connections, message routing, and state management — that form the foundation of most distributed systems.
You can explore the complete code, including the Flutter client, on my GitHub repository.

Read the full article here: https://medium.com/@simplicityspaces/a-tcp-multi-client-chat-server-in-rust-eb00ba042c7e