i think the server's register process is done

README.md

@@ -16,6 +16,10 @@
 [ ] Figure out server responses (hopefully manages to be stuck in a 512 bit block as well)
 
 
+## Misc todo:
+
+[ ] Create a Request to String function for easy printing and debugging
+
 ## client todo:
 
 [ ] Check for key when turned on
@@ -26,10 +30,10 @@
 
 ## Server todo:
 
-[ ] Laucnh task for each new connection
-[ ] use RSA key to get first message and extract AES key
+[x] Laucnh task for each new connection
+[x] use RSA key to get first message and extract AES key
 [ ] verify the user using its public RSA key
-	[ ] if it was a register session save the key into the BIG DATA STRUCTURE
+	[x] if it was a register session save the key into the BIG DATA STRUCTURE
 [ ] Keep lists of incoming messages 
 	(doesnt need to know from who, they are just big blobs of shlomp)
 [ ] When user asks for incoming messages, make basic packet and append the incoming messages
@@ -40,3 +44,37 @@
 		to calculate ([128, 200+128=328, 300+328=628, ...])
 		
 
+Register process:
+
+Client			Server
+
+Send AES key (sk)
+Send Register(pub key)
+				Get AES
+				Get Register
+				Send 6 digit code
+Get 6 digit code
+(1) Wait for user to input 6 digit code
+Send 6 digit code (signed)
+				Get 6 digit code and verify (code, then sig)
+				Send OK/NACK
+if NACK goto (1)
+
+Login process:
+
+Client			Server
+
+Send AES key (sk)
+Send Login message (Phone, AES sig)
+				Get AES key
+				Verify AES sig with Phone-pub key
+				if sig is invalid: close connection
+				else: send stored messages because why not
+
+Usual process get messages:
+
+Client 			Server
+
+Send GetMessages
+				Send back messages
+Send GotMessages(amount)

client/Program.cs

@@ -19,7 +19,8 @@ public class Program
             .FirstOrDefault() ?? "0000"; // get the value or deafult if it doesnt exist
 
         // On boot, check if a key is available
-        RSA serverKey = LoadRSAFromFile("key_server.pem")!;
+        RSA? serverKey = LoadRSAFromFile("server_key.pem");
+        if (serverKey == null) { Console.WriteLine("Could not find server key, please run server before clients!"); return; }
         RSA pubKey = LoadRSAFromFile($"pubkey_{user}.pem") ?? RSA.Create(2048);
         RSA privKey = LoadRSAFromFile($"privkey_{user}.pem") ?? pubKey;
 
@@ -33,15 +34,16 @@ public class Program
         Aes sk = Aes.Create(); // creates an AES-256 key
         if (needsRegister)
         {
-            while(true) {
-                try {
+            try
+            {
                 await RegisterClient(user, pubKey, privKey, serverKey, sk, stream);
             }
-                catch (Exception ex) {
+            catch (Exception ex)
+            {
                 Console.WriteLine("Failed registration process");
                 Console.WriteLine("Exception: " + ex.Message);
                 Console.WriteLine("Stack: " + ex.StackTrace);
-                }
+                return;
             }
         }
         else
@@ -49,7 +51,6 @@ public class Program
 
         }
 
-
         var inputTask = Task.Run(async () => await HandleUserInput(client, stream));
         var serverInput = Task.Run(async () => await HandleServerInput(client, stream));
 
@@ -60,28 +61,36 @@ public class Program
     async static Task RegisterClient(string user, RSA pub, RSA priv, RSA server, Aes sk, NetworkStream stream)
     {
         byte counter = 0;
-        // first send the `Register` message
+        // Generate aes key and send it forward
+        byte[] skEnc = server.Encrypt([.. sk.Key, .. sk.IV], RSAEncryptionPadding.OaepSHA256);
+        await stream.WriteAsync(skEnc);
+
+        // Generate the Register msg
         byte[] pubBytes = pub.ExportRSAPublicKey();
-        byte[] data = new byte[16];
+        byte[] data = new byte[12];
         Array.Copy(Utils.NumberToBytes(user), data, 8);
         Array.Copy(BitConverter.GetBytes(pubBytes.Length), 0, data, 8, 4);
-        Array.Copy(BitConverter.GetBytes(sk.Key.Length), 0, data, 12, 4);
         byte[] msg = Request.CreateRequest(RequestType.Register, ref counter, data);
-        byte[] finalPayload = [.. msg, .. pubBytes, .. sk.Key];
-        // signing a hash here is pretty useless tbh
-        // if oscar is changing this message somehow it will just fail the registration process and will restart eventually
-        // so kinda pointless, esp when its signed with the server's key
-        byte[] enc = server.Encrypt(finalPayload, RSAEncryptionPadding.OaepSHA256);
-        await stream.WriteAsync(enc);
-        // get the 6 digit code
+        // Encrypt msg and send it
+        byte[] enc = sk.EncryptCfb(msg, sk.IV, PaddingMode.PKCS7);
+        byte[] payload = [.. enc, .. pubBytes];
+        await stream.WriteAsync(payload);
+
+        // get the 6 digit code (from "secure channel", actually an OK message is expected here but the 6 digit code kinda replaces it)
         byte[] digits = new byte[6];
-        await stream.ReadExactlyAsync(digits, 0, 6);
+        int len = 0;
+        while (len != 6)
+        {
+            len = await stream.ReadAsync(digits);
+        }
         // print the 6 digit code
         Console.WriteLine($"[{DateTime.Now}] 6 digit code: {string.Join(' ', digits.Select(d => d.ToString()))}");
         // get the 6 digit code from the user
-        while(true) {
+        while (true)
+        {
             string? code = Console.ReadLine()?.Trim();
-            if(code == null || code.Take(6).Any(c => !char.IsDigit(c))) {
+            if (code == null || code.Take(6).Any(c => !char.IsDigit(c)))
+            {
                 Console.WriteLine("Invalid code!");
                 continue;
             }
@@ -89,15 +98,21 @@ public class Program
                 .Take(6) // take the first 6 characters
                 .Select(d => byte.Parse(d.ToString())) // parse into bytes
                 .ToArray();
-            msg = Request.CreateRequest(RequestType.ConfirmRegister, ref counter, codeBytes);
-            byte[] signed = priv.SignData(msg, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);
-            enc = sk.EncryptCfb(msg, Array.Empty<byte>(), PaddingMode.None); // no reason to encrpy the signature
-            finalPayload = [.. enc, .. signed]; // should be 128 (enc) + 256 (signed)
-            await stream.WriteAsync(finalPayload);
+            // Debug print the inserted value to see it works :)
+            Console.WriteLine(string.Join(' ', codeBytes.Select(b => b.ToString())));
+
+            // Sign the 6 digit code & Generate ConfirmRegister message
+            byte[] signed = priv.SignData(codeBytes, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);
+            msg = Request.CreateRequest(RequestType.ConfirmRegister, ref counter, [.. codeBytes, .. BitConverter.GetBytes(signed.Length)]);
+            enc = sk.EncryptCfb(msg, sk.IV, PaddingMode.PKCS7); // no reason to encrpy the signature
+            payload = [.. enc, .. signed]; // should be 128 (enc) + 256 (signed)
+            await stream.WriteAsync(payload);
+            // wait for OK/NACK response (anything other than OK is a NACK)
             int incoming = await stream.ReadAsync(enc);
-            msg = sk.DecryptCfb(enc[..incoming], Array.Empty<byte>(), PaddingMode.Zeros);
+            msg = sk.DecryptCfb(enc[..incoming], sk.IV, PaddingMode.PKCS7);
             string r = Encoding.UTF8.GetString(msg);
-            if(r == "OK") {
+            if (r == "OK")
+            {
                 Console.WriteLine("Registration process complete");
                 break;
             }

lib/Request.cs

@@ -12,7 +12,7 @@ public enum RequestType {
 public static class Request {
     public static byte[] CreateRequest(RequestType Type, ref byte counter, byte[] data) {
         if(data.Length > 13) {
-            throw new Exception("extra data is too long");
+            throw new Exception("extra data is too long: " + data.Length.ToString());
         }
 
         byte[] msg = new byte[128];

protocol.md

@@ -49,11 +49,11 @@ to the server.
 ## "Control" requests
 
 - Register:
-	data: Phone - 8 bytes, RSA key size (payload length) - 2 bytes, AES key length - 2 bytes
+	data: Phone - 8 bytes, RSA key size (payload length) - 2 bytes
 - ConfirmRegister (signed & encrypted 6 digit code)
 	data: 6 bytes for the 6 digit code (can use less but it will be padding otherwise)
 - Login (signed hash):
-	data: 8 bytes of user's phone, AES key length - 2 bytes, signed SHA length - 2 bytes
+	data: 8 bytes of user's phone, signed SHA length - 2 bytes
 - GetMessages:
 	data: EMPTY
 - GetUserKey (dont think it needs any signing or encryption technically, as it is very 

server/Data.cs

@@ -0,0 +1,27 @@
+using System.Collections.Generic;
+using System.Security.Cryptography;
+
+namespace server;
+
+public class Data
+{
+    public Dictionary<string, RSA> Keys { set; get; } = [];
+    public Dictionary<string, Queue<byte[]>> Messages { set; get; } = [];
+
+    public RSA? GetKey(string Phone) {
+        return Keys.TryGetValue(Phone, out RSA? value) ? value : null;
+    }
+
+    public Queue<byte[]>? GetMessages(string Phone) {
+        // Check we have a RSA key for the phone and get the messages
+        if(!Keys.ContainsKey(Phone)) { return null; }
+        if(Messages.TryGetValue(Phone, out Queue<byte[]>? value)) {
+            return value;
+        }
+        else {
+            // generate a new queue because one  doesnt already exists
+            Messages[Phone] = new Queue<byte[]>();
+            return Messages[Phone];
+        }
+    }
+}

server/Program.cs

@@ -1,51 +1,226 @@
 using System;
 using System.IO;
+using System.Linq;
 using System.Net;
 using System.Net.Sockets;
 using System.Security.Cryptography;
+using System.Text;
+using System.Threading.Tasks;
+using lib;
+using server;
 
 namespace Server;
 
 public class Program
 {
-    static void Main(string[] args)
+    const int MSG_LEN = 16; // msg len is 128 bits = 16 bytes
+
+    static readonly Data Data = new();
+    static readonly Random Rnd = new((int)DateTime.Now.Ticks);
+
+    static async Task Main()
     {
         // Generally this key would be static but since its not production yet we can generate it every time to make sure
         // the users has the key and could load it from file
-        RSA key = RSA.Create(2048);
+        RSA key = RSA.Create(1024);
         File.WriteAllText("server_key.pem", key.ExportRSAPublicKeyPem());
 
         int port = 12345;
         TcpListener server = new(IPAddress.Parse("0.0.0.0"), port);
-        try {
+        int connectionCounter = 0;
+        try
+        {
             server.Start();
             byte[] buffer = new byte[256];
 
-            while(true) {
+            while (true)
+            {
                 // Currently, every time it gets a block, it will simply send it back but ToUpper
-                using TcpClient client = server.AcceptTcpClient();
-                Console.WriteLine("Got a client!");
-
-                var stream = client.GetStream();
-                int readLen;
-                while((readLen = stream.Read(buffer, 0, buffer.Length)) != 0) {
-                    // for now, lets just read it as an ascii string
-                    string input = System.Text.Encoding.ASCII.GetString(buffer, 0, readLen); 
-                    Console.WriteLine($"Got block: {input}");
-
-                    byte[] ret = System.Text.Encoding.ASCII.GetBytes(input.ToUpper());
-                    stream.Write(ret, 0, ret.Length);
-                }
-
-
+                TcpClient client = await server.AcceptTcpClientAsync();
+                _ = Task.Run(async () => await HandleClient(client, connectionCounter));
+                connectionCounter += 1;
             }
         }
-        catch(Exception ex) {
+        catch (Exception ex)
+        {
             Console.WriteLine($"Server error: {ex.Message}");
             Console.WriteLine("Trace: " + ex.StackTrace);
         }
-        finally {
+        finally
+        {
             server.Stop();
         }
     }
+
+    static async Task HandleClient(TcpClient client, int id)
+    {
+        NetworkStream stream = client.GetStream();
+        byte[] buffer = new byte[1024];
+        byte counter = 0;
+        // Get AES session key
+        int len = await stream.ReadAsync(buffer);
+        Aes sk = Aes.Create();
+        sk.Key = buffer[..256]; // just to make sure no one sends a too big to be true key
+        sk.IV = buffer[256..len];
+        Write(id, "key + iv: " + len.ToString());
+
+        // Get first message (should be either login or )
+        len = await stream.ReadAsync(buffer);
+        byte[] msg = sk.DecryptCfb(buffer[..MSG_LEN], sk.IV, PaddingMode.PKCS7);
+        if (msg[0] != 0)
+        {
+            Write(id, "Invalid session id!");
+            client.Dispose();
+            return;
+        }
+        counter = IncrementCounter(msg[2]); // allow counter to start at a random position
+        if (msg[1] == (byte)RequestType.Register)
+        {
+            // Do register stuff
+            // get phone number
+            string phone = Utils.BytesToNumber(msg[3..11]);
+            Write(id, $"Client wants to register as {phone}");
+            int keyLen = BitConverter.ToInt32(msg, 11);
+            RSA pub = RSA.Create();
+            pub.ImportRSAPublicKey(buffer.AsSpan()[MSG_LEN..], out int bytesRead);
+            Write(id, $"Imported key len: {bytesRead} while client claims it is {keyLen}");
+            // generate the 6 digit code and send it
+            byte[] code = [
+                (byte)Rnd.Next(10),
+                (byte)Rnd.Next(10),
+                (byte)Rnd.Next(10),
+                (byte)Rnd.Next(10),
+                (byte)Rnd.Next(10),
+                (byte)Rnd.Next(10),
+                ];
+            await Send6DigitCodeInSecureChannel(stream, code);
+            // wait for the code to be back with a key
+            int tries = 5; // allow 5 tries before closing the connection and forcing a restart
+            while (tries > 0)
+            {
+                tries -= 1;
+                len = await stream.ReadAsync(buffer);
+                Write(id, $"Got 6 digit code with sig, len: {len}");
+                msg = sk.DecryptCfb(buffer[..MSG_LEN], sk.IV, PaddingMode.PKCS7);
+                byte[] sig = buffer[MSG_LEN..len];
+                if (msg[0] != 0 || msg[1] != (byte)RequestType.ConfirmRegister || msg[2] != counter)
+                {
+                    // invalid or unexpected req, someone might be sending dups
+                    continue;
+                }
+                counter = IncrementCounter(counter);
+                byte[] gottenCode = msg[3..9];
+                int expectedSigLen = BitConverter.ToInt32(msg, 9);
+                if (expectedSigLen != len - MSG_LEN)
+                {
+                    Write(id, $"expected sig len doesnt match read len: {expectedSigLen} / {len - MSG_LEN}");
+                }
+                // check if the codes are equal
+                if (code.Zip(gottenCode).Any(a => a.First != a.Second))
+                {
+                    // codes are not equal, send a nack
+                    msg = sk.EncryptCfb(Encoding.UTF8.GetBytes("BAD CODE"), sk.IV, PaddingMode.PKCS7);
+                    await stream.WriteAsync(msg);
+                }
+                else
+                {
+                    // codes are equal - verify sig
+                    bool sigValid = pub.VerifyData(gottenCode, sig, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);
+                    if (sigValid)
+                    {
+                        msg = sk.EncryptCfb(Encoding.UTF8.GetBytes("OK"), sk.IV, PaddingMode.PKCS7);
+                        await stream.WriteAsync(msg);
+                        Data.Keys[phone] = pub; // save the key
+                        break;
+                    }
+                    else
+                    {
+                        msg = sk.EncryptCfb(Encoding.UTF8.GetBytes("SIG INVALID"), sk.IV, PaddingMode.PKCS7);
+                        await stream.WriteAsync(msg);
+                    }
+                }
+            }
+
+        }
+        else if (msg[1] == (byte)RequestType.Login)
+        {
+            // verify login
+        }
+        else
+        {
+            // invalid connection, quit
+            Write(id, "Client didnt register or login as first message");
+            client.Dispose();
+            return;
+        }
+        // Client registered/logged in, do main messages loop
+        try
+        {
+            while (client.Connected)
+            {
+                // while the client is connected, simply read messages from the client and handle accordingly,
+                // either by getting new messages for other ppl, or sending back keys/pending messages
+                len = await stream.ReadAsync(buffer);
+                msg = sk.DecryptCfb(buffer[..MSG_LEN], sk.IV, PaddingMode.None);
+                // verify that the counter message is correct
+                if (msg[0] != 0 || msg[2] != counter)
+                {
+                    msg = sk.EncryptCfb(Encoding.UTF8.GetBytes("DUPLICATE"), sk.IV, PaddingMode.PKCS7);
+                    await stream.WriteAsync(msg);
+                    continue;
+                }
+                counter = IncrementCounter(counter);
+                switch ((RequestType)msg[1])
+                {
+                    case RequestType.GetMessages:
+                        break;
+                    case RequestType.GetUserKey:
+                        string phone = Utils.BytesToNumber(msg[3..11]);
+                        RSA? key = Data.GetKey(phone);
+                        if (key != null)
+                        {
+                            msg = sk.EncryptCfb(key.ExportRSAPublicKey(), sk.IV, PaddingMode.PKCS7);
+                            await stream.WriteAsync(msg);
+                        }
+                        else
+                        {
+                            msg = sk.EncryptCfb(Encoding.UTF8.GetBytes("USER DOES NOT EXIST"), sk.IV, PaddingMode.PKCS7);
+                            await stream.WriteAsync(msg);
+                        }
+                        break;
+                    case RequestType.SendMessage:
+                        break;
+                    default:
+                        msg = sk.EncryptCfb(Encoding.UTF8.GetBytes("INVALID REQUEST"), sk.IV, PaddingMode.PKCS7);
+                        await stream.WriteAsync(msg);
+                        break;
+                }
+
+            }
+
+        }
+        catch (Exception ex)
+        {
+            Write(id, $"Client failed with error {ex.Message}");
+            Write(id, $"Stack: {ex.StackTrace}");
+        }
+
+        client.Dispose();
+    }
+
+    static byte IncrementCounter(byte counter)
+    {
+        return counter == byte.MaxValue ? (byte)0 : (byte)(counter + 1);
+    }
+
+    static async Task Send6DigitCodeInSecureChannel(NetworkStream stream, byte[] code)
+    {
+        await stream.WriteAsync(code);
+    }
+
+    /// Helper log message so it would print both the time and the id
+    static void Write(int id, string Message)
+    {
+        Console.WriteLine($"[{DateTime.Now:HH:mm:ss}] {id} - {Message}");
+    }
 }