Cleaned up parts of the serialization by removing redundant code.

src/common/crypto.nim

@@ -44,6 +44,15 @@ proc decrypt*(key: Key, iv: Iv, encData: seq[byte], sequenceNumber: uint64): (se
     
     return (data, tag)
 
+proc validateDecryption*(key: Key, iv: Iv, encData: seq[byte], sequenceNumber: uint64, header: Header): seq[byte] = 
+
+    let (decData, gmac) = decrypt(key, iv, encData, sequenceNumber)
+
+    if gmac != header.gmac: 
+        raise newException(CatchableError, "Invalid authentication tag.")
+
+    return decData
+
 #[
     Key exchange using X25519 and Blake2b
     Elliptic curve cryptography ensures that the actual session key is never sent over the network
@@ -147,5 +156,4 @@ proc loadKeyPair*(keyFile: string): KeyPair =
     except IOError: 
         let keyPair = generateKeyPair() 
         writeKeyToDisk(keyFile, keyPair.privateKey)
-        return keyPair
-
+        return keyPair

src/common/sequence.nim

@@ -7,7 +7,7 @@ proc nextSequence*(agentId: uint32): uint64 =
     sequenceTable[agentId] = sequenceTable.getOrDefault(agentId, 0'u64) + 1
     return sequenceTable[agentId]
 
-proc validateSequence*(agentId: uint32, seqNr: uint64, packetType: uint8): bool = 
+proc validateSequence(agentId: uint32, seqNr: uint64, packetType: uint8): bool = 
     let lastSeqNr = sequenceTable.getOrDefault(agentId, 0'u64)
 
     # Heartbeat messages are not used for sequence tracking
@@ -26,3 +26,17 @@ proc validateSequence*(agentId: uint32, seqNr: uint64, packetType: uint8): bool
     # Update sequence number
     sequenceTable[agentId] = seqNr
     return true
+
+proc validatePacket*(header: Header, expectedType: uint8) = 
+    
+    # Validate magic number
+    if header.magic != MAGIC:
+        raise newException(CatchableError, "Invalid magic bytes.")
+
+    # Validate packet type
+    if header.packetType != expectedType: 
+        raise newException(CatchableError, "Invalid packet type.")
+
+    # Validate sequence number 
+    if not validateSequence(header.agentId, header.seqNr, header.packetType): 
+        raise newException(CatchableError, "Invalid sequence number.")

src/common/serialize.nim

@@ -1,5 +1,5 @@
-import streams, strutils
-import ./[types, utils]
+import streams, strutils, tables
+import ./[types, utils, crypto, sequence]
 type 
     Packer* = ref object 
         stream: StringStream
@@ -17,9 +17,8 @@ proc addData*(packer: Packer, data: openArray[byte]): Packer {.discardable.} =
     return packer
 
 proc addArgument*(packer: Packer, arg: TaskArg): Packer {.discardable.} = 
-    
+    # Optional argument was passed as "", ignore
     if arg.data.len <= 0: 
-        # Optional argument was passed as "", ignore
         return
 
     packer.add(arg.argType)
@@ -34,7 +33,6 @@ proc addArgument*(packer: Packer, arg: TaskArg): Packer {.discardable.} =
     return packer
 
 proc addVarLengthMetadata*(packer: Packer, metadata: seq[byte]): Packer {.discardable.} = 
-
     # Add length of metadata field
     packer.add(cast[uint32](metadata.len))
 
@@ -160,7 +158,8 @@ proc getVarLengthMetadata*(unpacker: Unpacker): string =
     # Read content
     return unpacker.getBytes(int(length)).toString()
 
-proc packHeader*(packer: Packer, header: Header, bodySize: uint32): seq[byte] = 
+# Serialization & Deserialization functions
+proc serializeHeader*(packer: Packer, header: Header, bodySize: uint32): seq[byte] = 
     packer
         .add(header.magic)
         .add(header.version)
@@ -174,7 +173,7 @@ proc packHeader*(packer: Packer, header: Header, bodySize: uint32): seq[byte] =
 
     return packer.pack()
 
-proc unpackHeader*(unpacker: Unpacker): Header=
+proc deserializeHeader*(unpacker: Unpacker): Header=
     return Header(
         magic: unpacker.getUint32(),
         version: unpacker.getUint8(),
@@ -185,4 +184,5 @@ proc unpackHeader*(unpacker: Unpacker): Header=
         seqNr: unpacker.getUint64(),
         iv: unpacker.getIv(),
         gmac: unpacker.getAuthenticationTag()
-    )
+    )
+