5
\$\begingroup\$

Recently I've been entertaining the idea of making a small encrypted chat app as a way to learn about UI, encryption and networking in C#. After some research/work, I've produced the below protocol class that can reliably and asynchronously communicate over TCP Sockets in a console integration test.

However, I've still got some lingering questions with regards to my use of Tasks, particularly with the AwaitMessageOfType, ContinuousSend, and ContinuousReceive methods. They work, but I have a strong suspicion that they'll fail to scale if I try to create an actual server application that's potentially juggling dozens of protocol instances at once. Is there a better/more proper solution?

Also, feel free to critique my coding style. If something makes your eyes bleed, I'd like to avoid making it a habit. Thanks!

namespace Tightbeam_Common
{
    /// <summary>
    /// Class used to wrap and interface with a Socket instance in a structured manner.
    /// </summary>
    public class TightbeamProtocol
    {
        //Public-readonly enum used to display the general state of the protocol instance
        public ProtocolState State { get; private set; }

        //A pre-connected TCP socket, provided during construction.
        private Socket link;
        //NodeInfo for both ends of the connection.
        private NodeInfo personalInfo;
        private NodeInfo? partnerInfo;

        //Key structs.
        //The partner RSA public key and shared AES private key are received/generated dynamically.
        private RSAKeypair personalKeypair;
        private RSAKeypair partnerKeypair;
        private AESKey sharedPrivateKey;

        //Outgoing message queue.
        private Queue<NetworkMessage> outgoinqQueue = new();
        //Received message queue.
        public Queue<NetworkMessage> ReceivedMessages = new();
        //Event that fires whenever an inbound message is fully received, deserialized and enqueued.
        public event EventHandler? MessageReceived;

        //4-byte buffer used to receive the length of following data streams
        private byte[] lengthBuffer = new byte[sizeof(int)];
        //Buffer used to spool incoming serial data.
        private byte[]? dataBuffer = null;
        //Integer used to hold byte read progress.
        private int bytesReceived = 0;

        public bool anyReceivedMessages
        {
            get
            {
                return (ReceivedMessages.Count > 0);
            }
        }
        private bool anySpooledMessages
        {
            get
            {
                if (outgoinqQueue.Count == 0) { return false; }
                return true;
            }
        }
        private int spooledMessages
        {
            get
            {
                return outgoinqQueue.Count;
            }
        }

        public TightbeamProtocol(Socket socket, NodeInfo info, RSAKeypair personalKeypair)
        {
            State = ProtocolState.StartingUp;

            //Construct instance members
            link = socket;
            personalInfo = info;
            this.personalKeypair = personalKeypair;

            //Spin up continuous I/O tasks
            Task.Run(() => this.ContinuousSend());
            Task.Run(() => this.ContinuousReceive());

            //Start handshake handler
            Task.Run(() => this.PerformHandshake());
            Console.WriteLine(this.GetHashCode() + " Construction complete");
        }

        private async Task PerformHandshake()
        {
            //Send handshake to partner
            Handshake outboundHandshake = new Handshake(personalInfo);
            EnqueueMessage(outboundHandshake, Encryption.None);

            //Await/act upon partner handshake
            Console.WriteLine(this.GetHashCode() + " Handshake sent...");
            NetworkMessage handshakeMessage = await AwaitMessageOfType(typeof(Handshake));
            Console.WriteLine(this.GetHashCode() + " Handshake RECEIVED!");
            Handshake partnerHandshake = Serializer.DeserializeContent<Handshake>(handshakeMessage.Data)!;
            partnerInfo = partnerHandshake.senderInfo;
            partnerKeypair = partnerInfo.PublicKey;

            //Key exchange
            Console.WriteLine(this.GetHashCode() + " Key Exchange sent...");
            AESKey localKey = Cryptographer.GenerateAESKey();
            EnqueueMessage(new KeyExchange(localKey), Encryption.RSA);
            NetworkMessage exchangeMessage = await AwaitMessageOfType(typeof(KeyExchange));
            KeyExchange receivedExchange = Serializer.DeserializeContent<KeyExchange>(exchangeMessage.Data, personalKeypair)!;
            sharedPrivateKey = Cryptographer.CascadingKeySelect(localKey, receivedExchange.AESKey);
            Console.WriteLine(this.GetHashCode() + " " + string.Join(", ", sharedPrivateKey.Key));
            Console.WriteLine(this.GetHashCode() + " Key Exchange COMPLETE!");

            //Unlock protocol instance
            ReceivedMessages.Clear();
            State = ProtocolState.Operational;
        }

        private async Task<NetworkMessage> AwaitMessageOfType(Type type)
        {
            while (true)
            {
                var messageQuery = from msg in ReceivedMessages
                                   where msg.ContentType == type
                                   select msg;

                if (messageQuery.Any()) { return messageQuery.First(); }
                await Task.Delay(100);
            }
        }

        /// <summary>
        /// <br>Starts the graceful shutdown process of a connected protocol instance.</br>
        /// <br>Graceful shutdown waits for all enqueued messages to be sent, and then closes the socket connection.</br>
        /// </summary>
        public async Task StartShutdown()
        {
            State = ProtocolState.ShuttingDown;
            //TODO queue graceful disconnect packet
            while (anySpooledMessages) { await Task.Delay(50); }
            try
            {
                link.Shutdown(SocketShutdown.Both);
            }
            finally
            {
                link.Close();
                State = ProtocolState.Inert;
            }
        }

        /// <summary>
        /// Wraps a given content object and header information into a NetworkMessage instance, and enqueues it accordingly.
        /// </summary>
        public void EnqueueMessage<T>(T content, Encryption encryption)
        {
            if (State > ProtocolState.Operational) { return; }

            byte[] serializedContent;

            switch (encryption)
            {

                case Encryption.AES:
                    serializedContent = Serializer.SerializeContent(content, sharedPrivateKey);
                    break;

                case Encryption.RSA:
                    serializedContent = Serializer.SerializeContent(content, partnerKeypair);
                    break;

                default:
                    serializedContent = Serializer.SerializeContent(content);
                    break;
            }

            var newMessage = new NetworkMessage(serializedContent, typeof(T), encryption);
            outgoinqQueue.Enqueue(newMessage);
        }

        /// <summary>
        /// Persistent Task that continuously attempts to dequeue and send any spooled messages.
        /// </summary>
        private async Task ContinuousSend()
        {
            while (State < ProtocolState.Inert)
            {
                if (anySpooledMessages)
                {
                    NetworkMessage message = outgoinqQueue.Dequeue()!;
                    byte[]? data = Serializer.SerializeMessage(message, true);
                    try
                    {
                        Console.WriteLine(this.GetHashCode() + " Sending message of type: " + message.ContentType.ToString());
                        await link.SendAsync(data, SocketFlags.None);
                    }
                    catch(Exception ex)
                    {
                        Console.WriteLine(ex.ToString());
                    }
                }
            }
        }

        /// <summary>
        /// Persistent Task that continuously attempts to receive inbound bytes.
        /// </summary>
        private async Task ContinuousReceive()
        {
            while (State < ProtocolState.ShuttingDown)
            {
                if (link.Available > 0)
                {
                    byte[] data = new byte[link.Available];
                    try
                    {
                        await link.ReceiveAsync(data, SocketFlags.None);
                        SortData(data);
                    }
                    catch (Exception ex)
                    {
                        Console.WriteLine(ex.ToString());
                    }
                }
                //Brief 100ms delay to allow bytes to accumulate in the socket.
                await Task.Delay(100);
            }
        }

        /// <summary>
        /// Sorts given inbound bytes into the correct buffer.
        /// </summary>
        private void SortData(byte[] data)
        {
            int i = 0;
            while (i != data.Length)
            {
                int availableBytes = data.Length - i;
                if (dataBuffer is not null)
                {
                    //Data buffer has been initialized; assume length buffer has been prepared and read into data buffer.
                    int requestedBytes = dataBuffer.Length - bytesReceived;

                    int transferredBytes = Math.Min(requestedBytes, availableBytes);
                    Array.Copy(data, i, dataBuffer, bytesReceived, transferredBytes);
                    i += transferredBytes;

                    ParseData(transferredBytes);
                }
                else
                {
                    //Data buffer is un-initialized; assume we are reading into the length buffer.
                    int requestedBytes = lengthBuffer.Length - bytesReceived;

                    int transferredBytes = Math.Min(requestedBytes, availableBytes);
                    Array.Copy(data, i, lengthBuffer, bytesReceived, transferredBytes);
                    i += transferredBytes;

                    ParseData(transferredBytes);
                }
            }
        }

        /// <summary>
        /// Acts upon received data once a buffer is full. 
        /// </summary>
        private void ParseData(int count)
        {
            bytesReceived += count;
            if (this.dataBuffer is null)
            {
                //Attempt to parse length buffer.
                if(bytesReceived != sizeof(int)) { /*Pass - awaiting length buffer completion.*/ }
                else
                {
                    //Parse length buffer to determine content stream size and initialize data buffer accordingly.
                    int dataLength = BitConverter.ToInt32(lengthBuffer, 0);
                    if (dataLength < 0 || dataLength > int.MaxValue) { throw new InvalidOperationException("Message length is out of bounds!"); }
                    else { dataBuffer = new byte[dataLength]; bytesReceived = 0; }
                }
            }
            else
            {
                //Attempt to parse contents of data buffer.
                if (bytesReceived != dataBuffer.Length) { /*Pass - awaiting data buffer completion.*/ }
                else
                {
                    //Parse data buffer into a NetworkMessage.
                    NetworkMessage? message = Serializer.DeserializeMessage(dataBuffer);
                    if (message is null) { throw new InvalidDataException("Received corrupt or invalid data - could not deserialize to NetworkMessage."); }
                    //Decrypt data if needed.
                    if (message.Encryption == Encryption.AES) { message.DecryptData(sharedPrivateKey); }
                    ReceivedMessages.Enqueue(message);
                    Console.WriteLine(this.GetHashCode() + " Received message of type: " + ReceivedMessages.Peek().ContentType.ToString());
                    OnMessageReceived();

                    //Reset data buffer and byte reception tally for next message.
                    dataBuffer = null;
                    bytesReceived = 0;
                }
            }
        }

        private void OnMessageReceived()
        {
            //Alert any subscribers that we've received a new message.
            EventHandler? handler = MessageReceived;
            handler?.Invoke(this, EventArgs.Empty);    
        }
    }

    public enum ProtocolState
    {
        StartingUp,
        Operational,
        ShuttingDown,
        Inert,
        Exception
    }
}
```
\$\endgroup\$
12
  • 1
    \$\begingroup\$ Why do you think they will fail at scale? By the way at what scale? :P \$\endgroup\$ Jan 5 at 12:24
  • \$\begingroup\$ Scale-wise, a few dozen instances perhaps? I'm concerned because I know that at some level that async and Tasks are an abstraction of multithreading. When I did my research, the universal recommendation was to use language async features rather than the naive approach of "thread per connection," but it seems to me that the constant checking being done by each Task isn't much better and perhaps actually worse, esp. since the VS console reported several dozen threads being terminated after I let the integration test run for several hours. \$\endgroup\$ Jan 5 at 13:00
  • \$\begingroup\$ Tomorrow or on Monday I will try to put together a review and try to address your concern. It is pretty hard to find some time with a 6 weeks old newborn :) \$\endgroup\$ Jan 6 at 20:19
  • 1
    \$\begingroup\$ If your posted version does work then that's fine. \$\endgroup\$ Jan 8 at 0:40
  • 1
    \$\begingroup\$ I've updated my post please revisit it. Most probably there will be two more updates. \$\endgroup\$ Jan 11 at 8:08
3
+50
\$\begingroup\$

Class members

  • Please try to follow the same naming convention across all your members

    • If you use camel Casing for private members then do for all your private members
    • If you have members that are counterpart of each other then try to use the same naming
      • So, instead of having outgoingQueue and ReceivedMessages
      • Try to name them like : sentMessages and ReceivedMessages
  • Prefer expression bodied property getter over lengthy getter definition

public ProtocolState State { get; private set; }

private Socket link;
private NodeInfo personalInfo;
private NodeInfo? partnerInfo;

private RSAKeypair personalKeypair;
private RSAKeypair partnerKeypair;
private AESKey sharedPrivateKey;

private Queue<NetworkMessage> sentMessages = new();
public Queue<NetworkMessage> ReceivedMessages = new();
public event EventHandler? MessageReceived;

private byte[] lengthBuffer = new byte[sizeof(int)];
private byte[]? dataBuffer = null;
private int bytesReceived = 0;

public bool AnyReceivedMessages => ReceivedMessages.Count > 0;
private bool anySpooledMessages => sentMessages.Count != 0;
private int spooledMessages => sentMessages.Count;
  • If you came from the C++ world then you want to organize your members in the way that public ones come first then the rest of them
    • by not interleaving implementation details with publicly available members

Constructor

  • You can take advantage of ValueTuples (and deconstruction) to assign multiple parameters to class members
  • Instead of repeating the same action (Task.Run(...)) over and over again you can create a collection of async methods and repeat the same action against them
  • Task.Run returns a Task which is not captured by you
    • To make your intent more clear please prefer discard operator (_ = Task.Run(...))
  • Instead of repeating everywhere this piece of code Console.WriteLine(this.GetHashCode() + " ...."); please create a helper method
public TightbeamProtocol(Socket socket, NodeInfo info, RSAKeypair keypair)
{
    State = ProtocolState.StartingUp;
    (link, personalInfo, personalKeypair) = (socket, info, keypair);

    foreach (var action in new Func<Task>[] { () => ContinuousSend(), () => ContinuousReceive(), () => PerformHandshake() })
        _ = Task.Run(action);

    LogDebugInfo("Construction complete");
}

private void LogDebugInfo(string message) => Console.WriteLine($"{GetHashCode()} {message}");

AwaitMessageOfType

  • This piece of code looks really weird to me: await AwaitMessageOfType
    • To avoid stuttering code please rename your method like GetMessageOfTypeByWaitingUntilItArrives
  • The while(true) statements are considered most of the time as code smell
  • You can easily refactor your code to express your intent in a much cleaner way
    • Your version: Perform indefinitely the retrieve for a given message. If it is present then break the loop otherwise wait 100 milliseconds
    • My version: Wait 100 milliseconds periodically until you receive a given message
private async Task<NetworkMessage> GetMessageOfTypeByWaitingUntilItArrives(Type type)
{
    NetworkMessage message = null;
    while (message = (from msg in ReceivedMessages
                    where msg.ContentType == type
                    select msg).FirstOrDefault() == null)
    {
        await Task.Delay(100);
    }
    return message;
}

UPDATE #1 Continuation

StartShutdown

  • In the comment of this method you are stating that this is a graceful shutdown
    • Most of the time this means you are waiting X seconds to finish as many outstanding operations as possible
    • and when X has been passed then abort all the pending actions
    • So there is a time constraint (or in other words there is a guarantee) that the system will terminate after X seconds and will not hang there forever
  • I assume because of this method might be time consuming that's why you have created two different enum values to describe the current state
    • So far (in this and in the previous methods), you have used the State as a Progress tracking mechanism
    • You have used it to inform the consumer of your class about the current situation
    • The thing is most of the time consumers are interested about the state changes
      • They want to know when the given object has transitioned from X to Y
      • You might consider to have an event for this transitions as well
  • while (anySpooledMessages) { await Task.Delay(50); }: This code basically the Donkey in Shrek: "Are we there yet?"
    • Here you have a synchronisation or rendezvous point: One thread waits for the other thread to advance
    • There are lots of built-in signalling primitives which can be used to signal from a thread to another that a certain point has been reached
    • In this particular case an AutoResetEvent could call Set inside the ContinuousSend method whenever the queue is empty and here you could call the Wait method
    • The Wait call is blocking if you want to have a non-blocking version then please visit this SO topic

EnqueueMessage

  • Even though enums are numeric values under the hood I've never used less than or greater than operators on them
    • The really simple reason is that it is fragile
    • Now your enum has the default values and no comments indicate that the ordering is important
    • So, if some future maintainer changes the ordering then your system will break
    • If some future maintainer assigns values to each state it might break the system depending on the values
  • So ,an alternative solution is to use black or whitelisting
    • Check whether or not the current state is listed inside the refusal collection, if so then return early
    • Check whether or not the current state is listed inside the approval collection, if so then carry on
    • Depending on the level of paranoia, the length of each list, future requirements, etc. you should decide which approach you would use
  • Your switch statement can be really easily converted into a switch expression, which is more concise (even the IDE has built-in refactoring mechanism for this)
byte[] serializedContent = encryption switch
{
    Encryption.AES => Serializer.SerializeContent(content, sharedPrivateKey),
    Encryption.RSA => Serializer.SerializeContent(content, partnerKeypair),
    _ => Serializer.SerializeContent(content),
};
  • Let me correct here my suggestion regarding outgoingQueue
    • I've suggested to call it sentMessages which I can see now is misleading
    • A more accurate name would be toBeSentMessages

UPDATE #2 Continuation

ContinuousSend

  • Let's start with the naming. Most of the time in case of C# methods are starting with a verb (like Start, Get, Move, etc.).
    • If you want to follow this guideline then I would suggest SendContinuously or SendUntilInert or TryContinuouslySending, etc.
  • State < ProtocolState.Inert: I've already shared my opinion about enum values and less than or greater than operators
    • By the way this is yet again a rendezvous point: Do something until another thread signals otherwise
    • Here you could use a ManualResetEventSlim, which does not automatically goes back to "normal" state after the signal has reached the interested party (the Wait-er)
      • The MRES does define a property called IsSet, which can be checked whether or not it's time to change the behaviour
      • It works in the same way as CancellationToken's IsCancellationRequested or CountdownEvent's IsSet
    • Here you can learn more about signalling (it is old (so bit outdated) but gold)
  • If you choose to use CancellationToken then you might need to consider to pass that token to the link.SendAsync method to perform collaborative cancellation
  • I have two recommendations regarding your error handling
    • You don't need to explicitly call ToString on the exception, Console.WriteLine will happily do this on your behalf
    • It might make sense for debugging purposes to log some information about the message itself, not just the exception

ContinuousReceive

  • In case of Send the exit condition is Inert, while here it is the Shutdown state
    • I don't know your requirements but for my intuition it does not make too much sense to receive more messages while I try to shutdown the communication
    • So, I do believe here you could use the same signalling object as you do in case of sending
  • Some people (even methodologies like object calisthenics) suggest to try to stick to one level of indentation in your methods
    • In this method your main functionality is at the 3rd level
    • One common way to reduce indentation is the usage of early exit instead of guard expression
    • Another way is to separate responsibilities into their own method
  • Step 1: Use guard expression
if (link.Available <= 0) { await Task.Delay(100); continue; }
                
byte[] data = new byte[link.Available];
try
{
    await link.ReceiveAsync(data, SocketFlags.None);
    SortData(data);
}
catch (Exception ex)
{
    Console.WriteLine(ex);
}
  • Step 2: separate responsibilities
private async Task ReceiveContinuously()
{
    await PerformActionUntilResetEventIsSet(Receive, inertStateSignal);
}

private async Task PerformActionUntilResetEventIsSet(Func<Task> action, ManualResetEventSlim resetEvent)
{
    while (!resetEvent.IsSet) await action();
}

private async Task Receive()
{
    if (link.Available <= 0)
    {
        await Task.Delay(100);
        return;
    }
                
    byte[] data = new byte[link.Available];
    try
    {
        await link.ReceiveAsync(data, SocketFlags.None);
        SortData(data);
    }
    catch (Exception ex)
    {
        Console.WriteLine(ex);
    }          
}

UPDATE #3 Continuation

I forgot to mention in the Receive method that it might make sense to use ArrayPool to avoid allocating bytes arrays frequently.

SortData

  • The i variable name is so meaningless that it ruins the understanding of your intent
    • Please try to use better names like alreadyTransferredBytes or Processed, Consumed, Parsed, whatever ...
  • Inside your while loop the branches execute almost the same code
    • I would suggest to try to minimize your branching by putting there only those calculations that are different
  • Step 1: Extract the common part
int requestedBytes;
Array destination;
if (dataBuffer is not null)
{
    requestedBytes = dataBuffer.Length - bytesReceived;
    destination = dataBuffer;
}
else
{
    requestedBytes = lengthBuffer.Length - bytesReceived;
    destination = lengthBuffer;
}

//The common part
int availableBytes = data.Length - alreadyTransferredBytes;
int transferredBytes = Math.Min(requestedBytes, availableBytes);
Array.Copy(data, alreadyTransferredBytes, destination, bytesReceived, transferredBytes);
alreadyTransferredBytes += transferredBytes;

ParseData(transferredBytes);
  • Step 2: Simplify the branching
Array destination = dataBuffer ?? lengthBuffer;

int requestedBytes = destination.Length - bytesReceived;
int availableBytes = data.Length - alreadyTransferredBytes;
int transferredBytes = Math.Min(requestedBytes, availableBytes);

Array.Copy(data, alreadyTransferredBytes, destination, bytesReceived, transferredBytes);
alreadyTransferredBytes += transferredBytes;
ParseData(transferredBytes);

ParseData

  • Inside the dataBuffer is null branch none of the else are really needed
    • Since return and throw break the normal flow that's why you can use them for early exits
    • The resulting code will be more streamlined and easier to understand IMHO
if (bytesReceived != sizeof(int))
    return;

int dataLength = BitConverter.ToInt32(lengthBuffer, 0);
if (dataLength < 0 || dataLength > int.MaxValue)
    throw new InvalidOperationException("Message length is out of bounds!");

dataBuffer = new byte[dataLength];
bytesReceived = 0;
  • Further simplifying the logic, you can end up a code like this
bytesReceived += count;
if (dataBuffer is null)
{
    if (bytesReceived != sizeof(int))
        return; //early exit

    int dataLength = BitConverter.ToInt32(lengthBuffer, 0);
    if (dataLength < 0 || dataLength > int.MaxValue)
        throw new InvalidOperationException("Message length ..."); //early exit

    dataBuffer = new byte[dataLength];
    bytesReceived = 0;
    return; //early exit
}

if (bytesReceived != dataBuffer.Length)
    return; //early exit

NetworkMessage message = Serializer.DeserializeMessage(dataBuffer)
        ?? throw new InvalidDataException("Received corrupt ..."); //early exit

if (message.Encryption == Encryption.AES)
    message.DecryptData(sharedPrivateKey);

ReceivedMessages.Enqueue(message);
LogDebugInfo($"Received message of type: {ReceivedMessages.Peek().ContentType}");
OnMessageReceived();

dataBuffer = null;
bytesReceived = 0;
  • One can argue against many early exits compared to the single return principle/law
    • If you want to follow that law then try to use smaller functions
\$\endgroup\$
1
  • 1
    \$\begingroup\$ Thank you so much, this is all really really helpful! \$\endgroup\$ Jan 16 at 0:53

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.