Part 7: Code cleanup

Published one month ago


At the end of part 6, we left the server in quite a mess with spaghetti code everywhere and it's not very maintainable at all. Right now we have a list of sockets, two separate dictionaries for client states, and a HandleClient method which leaves much to be desired in terms of maintainability and organisation.

“Mildly complex” is my middle name.

If you were early enough to read the initial draft of part 6, I did mention that we'd be creating our very own client so that we can ditch Telnet once and for all. We will be doing that, but I apologise it won't be this part. I believe this is an important step to complete before we get any further, so let's tidy up the server code so it isn't so abhorrent in nature.

Abstracting the concept of a “client”

Let's create a new Client class to clean up this behemoth. A Client will be created whenever it connects. It should know about the underlying socket to send data to, and this can be private. It also needs to know the State it's in (which will be the enum for either Login or Chat), and the username, both of which we'll expose publicly as properties.

While we're at it, we might as well also have the client encapsulate its NetworkStream and StreamReader/StreamWriter instances since they'll be utilised throughout the entire lifetime of the client's connection.

Lastly, this class should implement IDisposable so that when the client disconnects, the underlying stream and reader/writers can be disposed too. The Client class should now look something like this:

using System.Net.Sockets;

public class Client : IDisposable
{  
    private readonly Socket _socket;
    private readonly NetworkStream _stream;
    private readonly StreamReader _reader;
    private readonly StreamWriter _writer;

    public Client(Socket socket)
    {
	    _socket = socket;

        _stream = new NetworkStream(socket);
        _reader = new StreamReader(_stream);
        _writer = new StreamWriter(_stream); 
    }

	public ClientState State { get; private set; } = ClientState.Login;

	public string Username { get; private set; } = string.Empty;

    public void Dispose()
    {
	    _socket.Dispose();
        _stream.Dispose();
        _reader.Dispose();
        _writer.Dispose();
    }
}

Now we'll need a method which can send data to the socket stream so that we can relay the chat messages to this client. This is simple, it's just a method which calls the WriteLine method on _writer.

public void SendMessage(string message)  
{  
    _writer.WriteLine(message);  
    _writer.Flush();  
}

We can now move most of the logic from HandleClient in the listener into this class. We'll create a ReadMessage method which reads the next line of input from the client and handles it accordingly.

public void ReadMessage()
{
    string? message = _reader.ReadLine();
    if (message is null)
    {
	    // the connection has been closed
	    Console.WriteLine($"Client on endpoint {_socket.RemoteEndPoint} has disconnected");
        return;
    }

    if (State == ClientState.Login)
    {
	    Console.WriteLine($"{_socket.RemoteEndPoint} logged in with username '{message}'");
	    Username = message;
        State = ClientState.Chat;
        return;
    }

    // broadcast message to other clients
    Console.WriteLine($"Client '{Username}' on endpoint {_socket.RemoteEndPoint} sent data: {message}");
}

Now we face a problem. This Client class doesn't know about other Client instances, so how do we get it to broadcast a chat message to everyone else? Since the Listener class will be keeping track of these instances, we need a way to signal to that class that a chat message has been received.

We need an event. Actually we need two, because we also need to tell the listener when this client has disconnected. Now I'm not particularly fond of the .NET event flow, it's becoming a bit of an anti-pattern in my honest opinion. I'd much prefer to use the observer pattern for this situation. However for the sake of simplicity and increasing the odds that most people will be able to follow along, I will be using events here. Just be aware that there are better ways of doing it.

Anyway, we'll create two events. One for when a message has been received, and one for when the client disconnects:

public event EventHandler? Disconnected;

public event EventHandler<string>? MessageReceived;

While it is ideal to create an EventArgs for the message event, frankly I can't be bothered to do that, so I'm just going to use string as the generic argument. But with these two events created, we can now invoke them at the appropriate times:

public void ReadMessage()
{
    string? message = _reader.ReadLine();
    if (message is null)
    {
	    // the connection has been closed
	    Console.WriteLine($"Client on endpoint {_socket.RemoteEndPoint} has disconnected");
        <mark>Disconnected?.Invoke(this, EventArgs.Empty);</mark>
        return;
    }

    if (State == ClientState.Login)
    {
	    Console.WriteLine($"{_socket.RemoteEndPoint} logged in with username '{message}'");
	    Username = message;
        State = ClientState.Chat;
        return;
    }

    // broadcast message to other clients
    Console.WriteLine($"Client '{Username}' on endpoint {_socket.RemoteEndPoint} sent data: {message}");
    <mark>MessageReceived?.Invoke(this, message);</mark>
}

Now the Client class is complete*, and should look like this:

using System.Net.Sockets;

public class Client : IDisposable
{
    private readonly Socket _socket;
    private readonly NetworkStream _stream;
    private readonly StreamReader _reader;
    private readonly StreamWriter _writer;

    public Client(Socket socket)
    {
        _socket = socket;

        _stream = new NetworkStream(socket);
        _reader = new StreamReader(_stream);
        _writer = new StreamWriter(_stream);
    }

    public event EventHandler? Disconnected;

    public event EventHandler<string>? MessageReceived;

    public ClientState State { get; private set; } = ClientState.Login;

    public string Username { get; private set; } = string.Empty;

    public void Dispose()
    {
        _socket.Dispose();
        _stream.Dispose();
        _reader.Dispose();
        _writer.Dispose();
    }

    public void ReadMessage()
    {
        string? message = _reader.ReadLine();
        if (message is null)
        {
            // the connection has been closed
		    Console.WriteLine($"Client on endpoint {_socket.RemoteEndPoint} has disconnected");
            Disconnected?.Invoke(this, EventArgs.Empty);
            return;
        }

        if (State == ClientState.Login)
        {
		    Console.WriteLine($"{_socket.RemoteEndPoint} logged in with username '{message}'");
            Username = message;
            State = ClientState.Chat;
            return;
        }

        // broadcast message to other clients
        Console.WriteLine($"Client '{Username}' on endpoint {_socket.RemoteEndPoint} sent data: {message}");
        MessageReceived?.Invoke(this, message);
    }

    public void SendMessage(string message)
    {
        _writer.WriteLine(message);
        _writer.Flush();
    }
}

Now we can refactor the Listener class so that it no longer tracks client Socket instances but rather instances of Client instead. This means changing the _clients list to a List<Client>, and removing the two dictionaries:

public class Listener
{
    private readonly Socket _socket = new(SocketType.Stream, ProtocolType.Tcp);
    private readonly List<<mark>Client</mark>> _clients = new();

    public bool IsRunning { get; set; }

	// ...

Now we face a new problem. In the Start method, the listener was calling the select function by passing in its client sockets to determine which - if any - had data ready to be received (or if the listener had any pending connections). Since the listener is no longer tracking the raw socket, but the Client abstraction instead, how do we determine which clients have data or want to connect?

This is where I introduce a new socket function! We're going to stop using select and from this point on use the concept of “polling”. This behaviour is extremely similar to “selecting”, so it isn't too difficult to refactor, but this is where - for those in native code land - the API differs slightly. In Linux, the way to poll is by using the poll function (man page), however for Windows, it's called WSAPoll (Winsock docs). The function on both platforms takes the same arguments. It accepts an array of pollfd structures which - surprisingly and fortunately enough - have the same layout on both platforms. It also accepts a timeout, again very similar to select.

The layout for the pollfd structure for Linux is documented on the man page:

struct pollfd {
    int   fd;         /* file descriptor */
    short events;     /* requested events */
    short revents;    /* returned events */
};

And this pretty much aligns with the Windows definition too:

typedef struct pollfd {
  SOCKET fd;
  SHORT  events;
  SHORT  revents;
} WSAPOLLFD, *PWSAPOLLFD, *LPWSAPOLLFD;

But we're not in native code in this guide. In .NET, instead of taking in an array like this, we instead call the Poll method on the individual socket we want to check. This method returns a boolean to indicate whether the polling gave any results, so we can actually just modify the client's ReadMessage to call Poll, and if it tells us no data is ready to be received we'll just skip the rest of the logic:

public void ReadMessage()
{
    <mark>if (!_socket.Poll(TimeSpan.FromMilliseconds(1), SelectMode.SelectRead))</mark>
    <mark>{</mark>
        <mark>return;</mark>
    <mark>}</mark>

    string? message = _reader.ReadLine();
    if (message is null)
    {
        // the connection has been closed
        Disconnected?.Invoke(this, EventArgs.Empty);
        return;
    }

    if (State == ClientState.Login)
    {
		Console.WriteLine($"{_socket.RemoteEndPoint} logged in with username '{message}'");
        Username = message;
        State = ClientState.Chat;
        return;
    }

    // broadcast message to other clients
    Console.WriteLine($"Client '{Username}' on endpoint {_socket.RemoteEndPoint} sent data: {message}");
    MessageReceived?.Invoke(this, message);
}

When I said the Client class was "complete*", now you know why the asterisk was there. Now it's complete*.

* No really, we're done with it for now.

Cleaning up Listener

We can now focus entirely on the Listener class. Most of the logic in this class is going to be purged entirely. We can start by removing both the HandleSockets and HandleClient methods, since the data reading is now implemented in the Client class. All the Listener needs to do now is subscribe to the events we defined to know when a message should be relayed. However before we do that, we have one last problem to solve. The call to Select was including the listener as a way to determine when a new connection was pending. But this is now an easy fix, with the use of the very same Poll method mentioned above!

All we need to do to fix this code is to call Poll on the listener socket. If it returns true, a new connection is pending and we can handle it. To set up the client, we need to:

  • Subscribe to the two events created within the Client class,
  • Add this client to the _clients list,
  • And lastly, send a message which prompts them to enter their username.

After any pending connections have been accepted, we can iterate through the _clients list and call ReadMessage for each of them, allowing each client to poll its socket for data.

public void Start(int port)
{
    var endpoint = new IPEndPoint(IPAddress.Any, port);
    _socket.Bind(endpoint);
    _socket.Listen();

    Console.WriteLine($"Server is starting on endpoint {endpoint}");

    IsRunning = true;

    while (IsRunning)
    {
        if (_socket.Poll(TimeSpan.FromMilliseconds(1), SelectMode.SelectRead))
        {
            // a new connection is pending
            Socket clientSocket = _socket.Accept();
            HandleNewConnection(clientSocket);
        }

        foreach (Client client in _clients)
        {
            client.ReadMessage();
        }
    }
}

private void HandleNewConnection(Socket socket)
{
    Console.WriteLine($"New client connected on endpoint {socket.RemoteEndPoint}");

    var client = new Client(socket);
    client.MessageReceived += OnMessageReceived;
    client.Disconnected += OnDisconnected;

    _clients.Add(client);

    client.SendMessage("Please enter your username: ");
}

Now to implement the event handlers! When the Disconnected event is raised, we need to remove the client from the list and call Dispose on it to dispose the underlying socket, stream, and reader/writer. Since the EventHandler delegate passes the sender as an object, an explicit cast is required here.

private void OnDisconnected(object? sender, EventArgs e)
{
    var client = (Client)sender!;
    client.Dispose();
    _clients.Remove(client);
}

Lastly we implement the MessageReceived handler by broadcasting the chat message to all other clients in the Chat state, except for the client which sent the message to begin with - same as before.

private void OnMessageReceived(object? sender, string message)
{
    var client = (Client)sender!;
    foreach (Client other in _clients)
    {
        if (client == other || client.State != ClientState.Chat)
        {
            continue;
        }

        other.SendMessage($"{client.Username}: {message}");
    }
}

With all of that done, the Listener class now looks like this:

using System.Net;
using System.Net.Sockets;

public class Listener
{
    private readonly Socket _socket = new(SocketType.Stream, ProtocolType.Tcp);
    private readonly List<Client> _clients = new();

    public void Start(int port)
    {
        var endpoint = new IPEndPoint(IPAddress.Any, port);
        _socket.Bind(endpoint);
        _socket.Listen();

        Console.WriteLine($"Server is starting on endpoint {endpoint}");

        IsRunning = true;

        while (IsRunning)
        {
            if (_socket.Poll(TimeSpan.FromMilliseconds(1), SelectMode.SelectRead))
            {
                // a new connection is pending
                Socket clientSocket = _socket.Accept();
                HandleNewConnection(clientSocket);
            }

            foreach (Client client in _clients)
            {
                client.ReadMessage();
            }
        }
    }

    private void HandleNewConnection(Socket socket)
    {
        Console.WriteLine($"New client connected on endpoint {socket.RemoteEndPoint}");

        var client = new Client(socket);
        client.MessageReceived += OnMessageReceived;
        client.Disconnected += OnDisconnected;

        _clients.Add(client);

        client.SendMessage("Please enter your username: ");
    }

    private void OnDisconnected(object? sender, EventArgs e)
    {
        var client = (Client)sender!;
        client.Dispose();
        _clients.Remove(client);
    }

    private void OnMessageReceived(object? sender, string message)
    {
        var client = (Client)sender!;
        foreach (Client other in _clients)
        {
            if (client == other || client.State != ClientState.Chat)
            {
                continue;
            }

            other.SendMessage($"{client.Username}: {message}");
        }
    }
}

We should now be able to run the server and connect with Telnet (one last time) and if you did everything correctly, the same behaviour should present itself:

Alice and Bob: The Sequel: The Musical: On Ice

The only side effect of this is now the username prompt puts the cursor on the next line. This is a simple enough fix, but I leave that to you to resolve. The code is now much clearer and more maintainable, so it should be effortless to do.

In the next part, we're going to really ditch Telnet once and for all. I promise.

Note

The next part is not yet available. Stay tuned.