refactor(noise): add comments, restyle code

tests/v2/test_waku_noise.nim

@@ -4,7 +4,9 @@ import
   testutils/unittests,
   std/random,
   stew/byteutils,
+  ../../waku/v2/node/waku_payload,
   ../../waku/v2/protocol/waku_noise/noise,
+  ../../waku/v2/protocol/waku_message,
   ../test_helpers
 
 procSuite "Waku Noise":
@@ -43,7 +45,7 @@ procSuite "Waku Noise":
     check: 
       plaintext.toBytes() == decryptedCiphertext
 
-  test "Encrypt and decrypt Noise public keys":
+  test "Noise public keys: encrypt and decrypt a public key":
 
     let noisePublicKey: NoisePublicKey = genNoisePublicKey(rng[])
 
@@ -55,7 +57,7 @@ procSuite "Waku Noise":
     check: 
       noisePublicKey == decryptedPk
 
-  test "Decrypt unencrypted public key":
+  test "Noise public keys: decrypt an unencrypted public key":
 
     let noisePublicKey: NoisePublicKey = genNoisePublicKey(rng[])
 
@@ -66,7 +68,7 @@ procSuite "Waku Noise":
     check:
       noisePublicKey == decryptedPk
 
-  test "Encrypt encrypted public key":
+  test "Noise public keys: encrypt an encrypted public key":
 
     let noisePublicKey: NoisePublicKey = genNoisePublicKey(rng[])
 
@@ -78,7 +80,7 @@ procSuite "Waku Noise":
     check:
       encryptedPk == encryptedPk2
 
-  test "Encrypt, decrypt and decrypt public key":
+  test "Noise public keys: encrypt, decrypt and decrypt a public key":
 
     let noisePublicKey: NoisePublicKey = genNoisePublicKey(rng[])
 
@@ -91,7 +93,7 @@ procSuite "Waku Noise":
     check: 
       decryptedPk == decryptedPk2
 
-  test "Serialize and deserialize unencrypted public key":
+  test "Noise public keys: serialize and deserialize an unencrypted public key":
 
     let 
       noisePublicKey: NoisePublicKey = genNoisePublicKey(rng[])
@@ -101,7 +103,7 @@ procSuite "Waku Noise":
     check:
       noisePublicKey == deserializedNoisePublicKey
 
-  test "Encrypt, serialize, deserialize and decrypt public key":
+  test "Noise public keys: encrypt, serialize, deserialize and decrypt a public key":
 
     let noisePublicKey: NoisePublicKey = genNoisePublicKey(rng[])
 
@@ -115,41 +117,44 @@ procSuite "Waku Noise":
     check:
       noisePublicKey == decryptedPk
 
-  test "Encode/decode PayloadV2 to byte sequence":
+
+  test "PayloadV2: serialize/deserialize PayloadV2 to byte sequence":
 
     let
-      payload2 = randomPayloadV2(rng[])
+      payload2: PayloadV2 = randomPayloadV2(rng[])
-      encoded_payload = encodeV2(payload2)
+      serializedPayload = serializePayloadV2(payload2)
-      decoded_payload = decodeV2(encoded_payload.get())
 
     check:
-      payload2 == decoded_payload.get()
+      serializedPayload.isOk()
+
+    let deserializedPayload = deserializePayloadV2(serializedPayload.get())
+
+    check:
+      deserializedPayload.isOk()
+      payload2 == deserializedPayload.get()
 
 
-  test "Encode/Decode Waku2 payload (version 2) - ChaChaPoly Keyinfo":
+  test "PayloadV2: Encode/Decode a Waku Message (version 2) to a PayloadV2":
-    # Encoding
+    
+    # We encode to a WakuMessage a random PayloadV2
     let
-      version = 2'u32
+      payload2 = randomPayloadV2(rng[])
-      payload = randomPayloadV2(rng[])
+      msg = encodePayloadV2(payload2)
-      encodedPayload = encodePayloadV2(payload)
 
-    check encodedPayload.isOk()
+    check:
+      msg.isOk()
 
-    let
+    # We create ProtoBuffer from WakuMessage
-      msg = WakuMessage(payload: encodedPayload.get(), version: version)
+    let pb = msg.get().encode()
-      pb =  msg.encode()
 
-    # Decoding
+    # We decode the WakuMessage from the ProtoBuffer
-    let msgDecoded = WakuMessage.init(pb.buffer)
+    let msgFromPb = WakuMessage.init(pb.buffer)
-    check msgDecoded.isOk()
     
-    let
+    check:
-      cipherState = randomChaChaPolyCipherState(rng[])
+      msgFromPb.isOk()
-      keyInfo = KeyInfo(kind: ChaChaPolyEncryption, cs: cipherState)
+
-      decoded = decodePayloadV2(msgDecoded.get(), keyInfo)
+    let decoded = decodePayloadV2(msgFromPb.get())
 
     check:
       decoded.isOk()
-      decoded.get() == payload
+      payload2 == decoded.get()
-
-  #TODO: add encrypt payload with ChaChaPoly

waku/v2/node/waku_payload.nim

@@ -7,16 +7,12 @@ import
   ../protocol/waku_message,
   ../protocol/waku_noise/noise
 
-import libp2p/crypto/[chacha20poly1305, curve25519]
-
-
 export whisper_types, keys, options
 
 type
   KeyKind* = enum
     Symmetric
     Asymmetric
-    ChaChaPolyEncryption
     None
 
   KeyInfo* = object
@@ -25,8 +21,6 @@ type
       symKey*: SymKey
     of Asymmetric:
       privKey*: PrivateKey
-    of ChaChaPolyEncryption:
-      cs*: ChaChaPolyCipherState
     of None:
       discard
 
@@ -86,27 +80,34 @@ proc encode*(payload: Payload, version: uint32, rng: var BrHmacDrbgContext):
     return err("Unsupported WakuMessage version")
 
 
-proc decodePayloadV2*(message: WakuMessage, keyInfo: KeyInfo):
+# Decodes a WakuMessage to a PayloadV2
-    WakuResult[PayloadV2] =
+# Currently, this is just a wrapper over deserializePayloadV2 and encryption/decryption is done on top (no KeyInfo)
+proc decodePayloadV2*(message: WakuMessage): WakuResult[PayloadV2] 
+  {.raises: [Defect, NoiseMalformedHandshake, NoisePublicKeyError].} =
+  # We check message version (only 2 is supported in this proc)
   case message.version
   of 2:
-    case keyInfo.kind
+    # We attempt to decode the WakuMessage payload
-    of ChaChaPolyEncryption:
+    let deserializedPayload2 = deserializePayloadV2(message.payload)
-      let decoded = decodeV2(message.payload)#, keyInfo.cs)
+    if deserializedPayload2.isOk():
-      if decoded.isSome():
+      return ok(deserializedPayload2.get())
-        return ok(decoded.get())
     else:
-        return err("Couldn't decrypt using ChaChaPoly Cipher State")
+      return err("Failed to decode WakuMessage")
   else:
-      discard
+    return err("Wrong message version while decoding payload")
-  else:
-    return err("Key info doesn't match v2 payloads")
 
 
-proc encodePayloadV2*(payload: PayloadV2):
+# Encodes a PayloadV2 to a WakuMessage
-    WakuResult[seq[byte]] =
+# Currently, this is just a wrapper over serializePayloadV2 and encryption/decryption is done on top (no KeyInfo)
-  let encoded = encodeV2(payload)
+proc encodePayloadV2*(payload2: PayloadV2): WakuResult[WakuMessage] 
-  if encoded.isSome():
+  {.raises: [Defect, NoiseMalformedHandshake, NoisePublicKeyError].} =
-    return ok(encoded.get())
+
-  else:
+  # We attempt to encode the PayloadV2
-    return err("Couldn't encode the payload")
+  let serializedPayload2 = serializePayloadV2(payload2)
+  if not serializedPayload2.isOk():
+    return err("Failed to encode PayloadV2")
+
+  # If successful, we create and return a WakuMessage 
+  let msg = WakuMessage(payload: serializedPayload2.get(), version: 2)
+  
+  return ok(msg)

waku/v2/protocol/waku_noise/noise.nim

@@ -7,18 +7,14 @@
 
 {.push raises: [Defect].}
 
-import std/[oids, options, tables]
+import std/[options, tables]
 import chronos
 import chronicles
 import bearssl
 import strutils
-import stew/[endians2]
+import stew/[results, endians2]
 import nimcrypto/[utils, sha2, hmac]
 
-import libp2p/stream/[connection]
-import libp2p/peerid
-import libp2p/peerinfo
-import libp2p/protobuf/minprotobuf
 import libp2p/utility
 import libp2p/errors
 import libp2p/crypto/[crypto, chacha20poly1305, curve25519]
@@ -51,6 +47,7 @@ type
   # This follows https://rfc.vac.dev/spec/35/#public-keys-serialization
   # pk contains the X coordinate of the public key, if unencrypted (this implies flag = 0)
   # or the encryption of the X coordinate concatenated with the authorization tag, if encrypted (this implies flag = 1)
+  # Note: besides encryption, flag can be used to distinguish among multiple supported Elliptic Curves
   NoisePublicKey* = object
     flag: uint8
     pk: seq[byte]
@@ -66,6 +63,15 @@ type
     nonce: ChaChaPolyNonce
     ad: seq[byte]
 
+  # PayloadV2 defines an object for Waku payloads with version 2 as in
+  # https://rfc.vac.dev/spec/35/#public-keys-serialization
+  # It contains a protocol ID field, the handshake message (for Noise handshakes) and 
+  # a transport message (for Noise handshakes and ChaChaPoly encryptions)
+  PayloadV2* = object
+    protocolId: uint8
+    handshakeMessage: seq[NoisePublicKey]
+    transportMessage: seq[byte]
+
   # Some useful error types
   NoiseError* = object of LPError
   NoiseHandshakeError* = object of NoiseError
@@ -86,7 +92,7 @@ proc randomSeqByte*(rng: var BrHmacDrbgContext, size: int): seq[byte] =
   brHmacDrbgGenerate(rng, output)
   return output
 
-# Generate random Curve25519 (public, private) key pairs
+# Generate random (public, private) Elliptic Curve key pairs
 proc genKeyPair*(rng: var BrHmacDrbgContext): KeyPair =
   var keyPair: KeyPair
   keyPair.privateKey = EllipticCurveKey.random(rng)
@@ -253,108 +259,138 @@ proc decryptNoisePublicKey*(cs: ChaChaPolyCipherState, noisePublicKey: NoisePubl
 
 #################################################################
 
-# Payload functions
+# Payload encoding/decoding procedures
-
-type
-  PayloadV2* = object
-    protocol_id: uint8
-    handshake_message: seq[NoisePublicKey]
-    transport_message: seq[byte]
-
 
+# Checks equality between two PayloadsV2 objects
 proc `==`(p1, p2: PayloadV2): bool =
-  result =  (p1.protocol_id == p2.protocol_id) and 
+  return (p1.protocolId == p2.protocolId) and 
-            (p1.handshake_message == p2.handshake_message) and 
+         (p1.handshakeMessage == p2.handshakeMessage) and 
-            (p1.transport_message == p2.transport_message) 
+         (p1.transportMessage == p2.transportMessage) 
-  
   
 
+# Generates a random PayloadV2
 proc randomPayloadV2*(rng: var BrHmacDrbgContext): PayloadV2 =
-  var protocol_id = newSeq[byte](1)
+  var payload2: PayloadV2
-  brHmacDrbgGenerate(rng, protocol_id)
+  # To generate a random protocol id, we generate a random 1-byte long sequence, and we convert the first element to uint8
-  result.protocol_id = protocol_id[0].uint8
+  payload2.protocolId = randomSeqByte(rng, 1)[0].uint8
-  result.handshake_message = @[genNoisePublicKey(rng), genNoisePublicKey(rng), genNoisePublicKey(rng)]
+  # We set the handshake message to three unencrypted random Noise Public Keys
-  result.transport_message = newSeq[byte](128)
+  payload2.handshakeMessage = @[genNoisePublicKey(rng), genNoisePublicKey(rng), genNoisePublicKey(rng)]
-  brHmacDrbgGenerate(rng, result.transport_message)
+  # We set the transport message to a random 128-bytes long sequence
+  payload2.transportMessage = randomSeqByte(rng, 128)
+  return payload2
 
 
-proc encodeV2*(self: PayloadV2): Option[seq[byte]] =
+# Serializes a PayloadV2 object to a byte sequences according to https://rfc.vac.dev/spec/35/.
+# The results can be passed to the payload field of a WakuMessage https://rfc.vac.dev/spec/14/
+proc serializePayloadV2*(self: PayloadV2): Result[seq[byte], cstring] 
+  {.raises: [Defect, NoiseMalformedHandshake, NoisePublicKeyError].} =
 
   #We collect public keys contained in the handshake message
   var
-    ser_handshake_message_len: int = 0
+    # According to https://rfc.vac.dev/spec/35/, the maximum size for the handshake message is 256 bytes, that is 
-    ser_handshake_message = newSeqOfCap[byte](256)
+    # the handshake message length can be represented with 1 byte only. (its length can be stored in 1 byte)
-    ser_pk: seq[byte]
+    # However, to ease public keys length addition operation, we declare it as int and later cast to uit8
-  for pk in self.handshake_message:
+    serializedHandshakeMessageLen: int = 0
-    ser_pk = serializeNoisePublicKey(pk)
+    # This variables will store the concatenation of the serializations of all public keys in the handshake message
-    ser_handshake_message_len +=  ser_pk.len
+    serializedHandshakeMessage = newSeqOfCap[byte](256)
-    ser_handshake_message.add ser_pk
+    # A variable to store the currently processed public key serialization
+    serializedPk: seq[byte]
+  # For each public key in the handshake message
+  for pk in self.handshakeMessage:
+    # We serialize the public key
+    serializedPk = serializeNoisePublicKey(pk)
+    # We sum its serialized length to the total
+    serializedHandshakeMessageLen +=  serializedPk.len
+    # We add its serialization to the concatenation of all serialized public keys in the handshake message 
+    serializedHandshakeMessage.add serializedPk
+    # If we are processing more than 256 byte, we return an error
+    if serializedHandshakeMessageLen > 256:
+      debug "PayloadV2 malformed: too many public keys contained in the handshake message"
+      raise newException(NoiseMalformedHandshake, "Too many public keys in handshake message")
 
+  # We get the transport message byte length
+  let transportMessageLen = self.transportMessage.len
 
-  #RFC: handshake-message-len is 1 byte
+  # The output payload as in https://rfc.vac.dev/spec/35/. We concatenate all the PayloadV2 as 
-  if ser_handshake_message_len > 256:
+  # payload = ( protocolId || serializedHandshakeMessageLen || serializedHandshakeMessage || transportMessageLen || transportMessage)
-    debug "Payload malformed: too many public keys contained in the handshake message"
+  # We declare it as a byte sequence of length accordingly to the PayloadV2 information read 
-    return none(seq[byte])
+  var payload = newSeqOfCap[byte](1 + # 1 byte for protocol ID             
-
+                                  1 + # 1 byte for length of serializedHandshakeMessage field
-  let transport_message_len = self.transport_message.len
+                                  serializedHandshakeMessageLen + # serializedHandshakeMessageLen bytes for serializedHandshakeMessage
-  #let transport_message_len_len = ceil(log(transport_message_len, 8)).int
+                                  8 + # 8 bytes for transportMessageLen
-
+                                  transportMessageLen # transportMessageLen bytes for transportMessage
-  var payload = newSeqOfCap[byte](1 + #self.protocol_id.len +              
-                                  1 + #ser_handshake_message_len 
-                                  ser_handshake_message_len +        
-                                  8 + #transport_message_len
-                                  transport_message_len #self.transport_message
                                   )
   
+  # We concatenate all the data
+  # The protocol ID (1 byte) and handshake message length (1 byte) can be directly casted to byte to allow direct copy to the payload byte sequence
+  payload.add self.protocolId.byte
+  payload.add serializedHandshakeMessageLen.byte
+  payload.add serializedHandshakeMessage
+  # The transport message length is converted from uint64 to bytes in Little-Endian
+  payload.add toBytesLE(transportMessageLen.uint64)
+  payload.add self.transportMessage
 
-  payload.add self.protocol_id.byte
+  return ok(payload)
-  payload.add ser_handshake_message_len.byte
-  payload.add ser_handshake_message
-  payload.add toBytesLE(transport_message_len.uint64)
-  payload.add self.transport_message
-
-  return some(payload)
 
 
 
-#Decode Noise handshake payload
+# Deserializes a byte sequence to a PayloadV2 object according to https://rfc.vac.dev/spec/35/.
-proc decodeV2*(payload: seq[byte]): Option[PayloadV2] =
+proc deserializePayloadV2*(payload: seq[byte]): Result[PayloadV2, cstring] 
-  var res: PayloadV2
+  {.raises: [Defect, NoiseMalformedHandshake, NoisePublicKeyError].} =
 
+  # The output PayloadV2
+  var payload2: PayloadV2
+
+  # i is the read input buffer position index
   var i: uint64 = 0
-  res.protocol_id = payload[i].uint8
+
+  # We start reading the Protocol ID
+  payload2.protocolId = payload[i].uint8
   i+=1
 
-  var handshake_message_len = payload[i].uint64
+  # We read the Handshake Message lenght (1 byte)
+  var handshakeMessageLen = payload[i].uint64
+  if handshakeMessageLen > 256:
+    debug "Payload malformed: too many public keys contained in the handshake message"
+    raise newException(NoiseMalformedHandshake, "Too many public keys in handshake message")
   i+=1
 
-  var handshake_message: seq[NoisePublicKey]
+  # We now read for handshakeMessageLen bytes the buffer and we deserialize each (encrypted/unencrypted) public key read
-
   var
+    # In handshakeMessage we accumulates the deserialized Noise Public keys read
+    handshakeMessage: seq[NoisePublicKey]
     flag: byte
-    pk_len: uint64
+    pkLen: uint64
     written: uint64 = 0
 
-  while written != handshake_message_len:
+  # We read the buffer until handshakeMessageLen are read
-    #Note that flag can be used to add support to multiple Elliptic Curve arithmetics..
+  while written != handshakeMessageLen:
+    # We obtain the current Noise Public key encryption flag
     flag = payload[i]
+    # If the key is unencrypted, we only read the X coordinate of the EC public key and we deserialize into a Noise Public Key
     if flag == 0:
-      pk_len = 1 + Curve25519Key.len
+      pkLen = 1 + EllipticCurveKey.len
-      handshake_message.add intoNoisePublicKey(payload[i..<i+pk_len])
+      handshake_message.add intoNoisePublicKey(payload[i..<i+pkLen])
-      i += pk_len
+      i += pkLen
-      written += pk_len
+      written += pkLen
-    if flag == 1:
+    # If the key is encrypted, we only read the encrypted X coordinate and the authorization tag, and we deserialize into a Noise Public Key
-      pk_len = 1 + Curve25519Key.len + ChaChaPolyTag.len
+    elif flag == 1:
-      handshake_message.add intoNoisePublicKey(payload[i..<i+pk_len])
+      pkLen = 1 + EllipticCurveKey.len + ChaChaPolyTag.len
-      i += pk_len
+      handshakeMessage.add intoNoisePublicKey(payload[i..<i+pkLen])
-      written += pk_len
+      i += pkLen
+      written += pkLen
+    else:
+      raise newException(NoisePublicKeyError, "Invalid flag for Noise public key")
 
-  res.handshake_message = handshake_message
+  # We save in the output PayloadV2 the read handshake message
+  payload2.handshakeMessage = handshakeMessage
 
-  let transport_message_len = fromBytesLE(uint64, payload[i..(i+8-1)])
+  # We read the transport message length (8 bytes) and we convert to uint64 in Little Endian
+  let transportMessageLen = fromBytesLE(uint64, payload[i..(i+8-1)])
   i+=8
 
-  res.transport_message = payload[i..i+transport_message_len-1]
+  # We read the transport message (handshakeMessage bytes) 
-  i+=transport_message_len
+  payload2.transportMessage = payload[i..i+transportMessageLen-1]
+  i+=transportMessageLen
 
-  return some(res)
+  return ok(payload2)