From 45b1e10647497e6fe2b9b9b73f2c961ecd1764c4 Mon Sep 17 00:00:00 2001
From: G <28568419+s1fr0@users.noreply.github.com>
Date: Mon, 4 Apr 2022 17:46:51 +0200
Subject: [PATCH] feat(noise): add support to ChaChaPoly encryption (#930)

* feat(noise): adding ChaChaPoly encryption

* fix(noise): restored waku2 tests

* fix(noise): missing comma

* Fixed header

* fix(noise): made comment about empty ChaChaPolyKey clearer

* refactor(noise): address reviewer's comments

* test(noise): add test for ChaChaPoly Encryption/decryption

* refactor(noise): removed libp2p specific debug tools usage

* style(noise): converted variables to camelCase
---
 tests/all_tests_v2.nim                |  3 +-
 tests/v2/test_waku_noise.nim          | 23 +++++++
 waku/v2/protocol/waku_noise/noise.nim | 94 +++++++++++++++++++++++++++
 3 files changed, 119 insertions(+), 1 deletion(-)
 create mode 100644 tests/v2/test_waku_noise.nim
 create mode 100644 waku/v2/protocol/waku_noise/noise.nim

diff --git a/tests/all_tests_v2.nim b/tests/all_tests_v2.nim
index 3c138777b..54a22119d 100644
--- a/tests/all_tests_v2.nim
+++ b/tests/all_tests_v2.nim
@@ -22,7 +22,8 @@ import
   ./v2/test_enr_utils,
   ./v2/test_waku_store_queue,
   ./v2/test_pagination_utils,
-  ./v2/test_peer_exchange
+  ./v2/test_peer_exchange,
+  ./v2/test_waku_noise
 
 when defined(rln):
   import ./v2/test_waku_rln_relay
diff --git a/tests/v2/test_waku_noise.nim b/tests/v2/test_waku_noise.nim
new file mode 100644
index 000000000..d0b53028c
--- /dev/null
+++ b/tests/v2/test_waku_noise.nim
@@ -0,0 +1,23 @@
+{.used.}
+
+import
+  testutils/unittests,
+  ../../waku/v2/protocol/waku_noise/noise,
+  ../test_helpers
+
+
+procSuite "Waku Noise":
+  
+  let rng = rng()
+
+  test "ChaChaPoly Encryption/Decryption":
+
+    let cipherState = randomChaChaPolyCipherState(rng[])
+
+    let
+      plaintext: seq[byte] = randomSeqByte(rng[], 128)
+      ciphertext: ChaChaPolyCiphertext = encrypt(cipherState, plaintext)
+      decryptedCiphertext: seq[byte] = decrypt(cipherState, ciphertext)
+
+    check: 
+      plaintext == decryptedCiphertext
\ No newline at end of file
diff --git a/waku/v2/protocol/waku_noise/noise.nim b/waku/v2/protocol/waku_noise/noise.nim
new file mode 100644
index 000000000..cdf72330d
--- /dev/null
+++ b/waku/v2/protocol/waku_noise/noise.nim
@@ -0,0 +1,94 @@
+# Waku Noise Protocols for Waku Payload Encryption
+## See spec for more details:
+## https://github.com/vacp2p/rfc/tree/master/content/docs/rfcs/35
+##
+## Implementation partially inspired by noise-libp2p:
+## https://github.com/status-im/nim-libp2p/blob/master/libp2p/protocols/secure/noise.nim
+
+
+{.push raises: [Defect].}
+
+import std/[oids, options]
+import chronos
+import chronicles
+import bearssl
+import nimcrypto/[utils, sha2, hmac]
+
+import libp2p/stream/[connection]
+import libp2p/protobuf/minprotobuf
+import libp2p/utility
+import libp2p/errors
+import libp2p/crypto/[crypto, chacha20poly1305]
+
+logScope:
+  topics = "wakunoise"
+
+const
+  # EmptyKey is a special value which indicates a ChaChaPolyKey has not yet been initialized.
+  EmptyKey = default(ChaChaPolyKey)
+
+type
+  ChaChaPolyCiphertext* = object
+    data: seq[byte]
+    tag: ChaChaPolyTag
+
+  ChaChaPolyCipherState* = object
+    k*: ChaChaPolyKey
+    nonce*: ChaChaPolyNonce
+    ad*: seq[byte]
+
+  NoiseError* = object of LPError
+  NoiseDecryptTagError* = object of NoiseError
+
+#################################################################
+
+# Utilities
+
+# Generates random byte sequences of given size
+proc randomSeqByte*(rng: var BrHmacDrbgContext, size: uint32): seq[byte] =
+  var output = newSeq[byte](size)
+  brHmacDrbgGenerate(rng, output)
+  return output
+
+#################################################################
+
+# ChaChaPoly Symmetric Cipher
+
+# ChaChaPoly encryption
+# It takes a Cipher State (with key, nonce, and associated data) and encrypts a plaintext
+proc encrypt*(
+    state: ChaChaPolyCipherState,
+    plaintext: openArray[byte]): ChaChaPolyCiphertext
+    {.noinit, raises: [Defect].} =
+  #TODO: add padding
+  var ciphertext: ChaChaPolyCiphertext
+  ciphertext.data.add plaintext
+  ChaChaPoly.encrypt(state.k, state.nonce, ciphertext.tag, ciphertext.data, state.ad)
+  return ciphertext
+
+# ChaChaPoly decryption
+# It takes a Cipher State (with key, nonce, and associated data) and decrypts a ciphertext
+proc decrypt*(
+    state: ChaChaPolyCipherState, 
+    ciphertext: ChaChaPolyCiphertext): seq[byte]
+    {.raises: [Defect, NoiseDecryptTagError].} =
+  var
+    tagIn = ciphertext.tag
+    tagOut: ChaChaPolyTag
+  var plaintext = ciphertext.data
+  ChaChaPoly.decrypt(state.k, state.nonce, tagOut, plaintext, state.ad)
+  #TODO: add unpadding
+  trace "decrypt", tagIn = tagIn.shortLog, tagOut = tagOut.shortLog, nonce = state.nonce
+  if tagIn != tagOut:
+    debug "decrypt failed", plaintext = shortLog(plaintext)
+    raise newException(NoiseDecryptTagError, "decrypt tag authentication failed.")
+  return plaintext
+
+# Generates a random Cipher Test for testing encryption/decryption
+proc randomChaChaPolyCipherState*(rng: var BrHmacDrbgContext): ChaChaPolyCipherState =
+  var randomCipherState: ChaChaPolyCipherState
+  brHmacDrbgGenerate(rng, randomCipherState.k)
+  brHmacDrbgGenerate(rng, randomCipherState.nonce)
+  randomCipherState.ad = newSeq[byte](32)
+  brHmacDrbgGenerate(rng, randomCipherState.ad)
+  return randomCipherState