replace utils module with direct generate calls (#148)

The `utils` module uses a convoluted mechanism for filling byte arrays with random data. The `generate` function can be used directly, making the `utils` module obsolete.
status-im · Jul 24, 2023 · f8ed9b4 · f8ed9b4
1 parent 2c3ae31
commit f8ed9b4
Show file tree

Hide file tree

Showing 8 changed files with 30 additions and 62 deletions.
diff --git a/tests/extensions/testextflow.nim b/tests/extensions/testextflow.nim
@@ -1,5 +1,5 @@
 ## nim-websock
-## Copyright (c) 2021 Status Research & Development GmbH
+## Copyright (c) 2021-2023 Status Research & Development GmbH
 ## Licensed under either of
 ##  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
 ##  * MIT license ([LICENSE-MIT](LICENSE-MIT))
@@ -118,10 +118,11 @@ suite "Encode frame extensions flow":
       frame.opcode == Opcode.Binary
 
 suite "Decode frame extensions flow":
+  let rng = HmacDrbgContext.new()
   var
     address: TransportAddress
     server: StreamServer
-    maskKey = genMaskKey(HmacDrbgContext.new())
+    maskKey = MaskKey.random(rng[])
     transport: StreamTransport
     reader: AsyncStreamReader
     frame: Frame

diff --git a/tests/helpers.nim b/tests/helpers.nim
@@ -1,5 +1,5 @@
 ## nim-websock
-## Copyright (c) 2021 Status Research & Development GmbH
+## Copyright (c) 2021-2023 Status Research & Development GmbH
 ## Licensed under either of
 ##  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
 ##  * MIT license ([LICENSE-MIT](LICENSE-MIT))
@@ -88,7 +88,7 @@ proc connectClient*(
   onPing: ControlCb = nil,
   onPong: ControlCb = nil,
   onClose: CloseCb = nil,
-  rng: Rng = nil): Future[WSSession] {.async.} =
+  rng = HmacDrbgContext.new()): Future[WSSession] {.async.} =
   let secure = when defined secure: true else: false
   return await WebSocket.connect(
     host = address,

diff --git a/tests/testframes.nim b/tests/testframes.nim
@@ -1,5 +1,5 @@
 ## nim-websock
-## Copyright (c) 2021 Status Research & Development GmbH
+## Copyright (c) 2021-2023 Status Research & Development GmbH
 ## Licensed under either of
 ##  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
 ##  * MIT license ([LICENSE-MIT](LICENSE-MIT))
@@ -12,7 +12,6 @@ import
   pkg/chronos/unittest2/asynctests
 
 include ../websock/frame
-include ../websock/utils
 
 # TODO: Fix Test.
 

diff --git a/tests/testwebsockets.nim b/tests/testwebsockets.nim
@@ -1,5 +1,5 @@
 ## nim-websock
-## Copyright (c) 2021-2022 Status Research & Development GmbH
+## Copyright (c) 2021-2023 Status Research & Development GmbH
 ## Licensed under either of
 ##  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
 ##  * MIT license ([LICENSE-MIT](LICENSE-MIT))
@@ -733,6 +733,7 @@ suite "Test Closing":
 
 suite "Test Payload":
   setup:
+    let rng = HmacDrbgContext.new()
     var
       server: HttpServer
 
@@ -837,7 +838,7 @@ suite "Test Payload":
       address = initTAddress("127.0.0.1:8888"),
       frameSize = maxFrameSize)
 
-    let maskKey = genMaskKey(HmacDrbgContext.new())
+    let maskKey = MaskKey.random(rng[])
     await session.stream.writer.write(
       (await Frame(
         fin: false,
@@ -897,7 +898,7 @@ suite "Test Payload":
         pong = true
     )
 
-    let maskKey = genMaskKey(HmacDrbgContext.new())
+    let maskKey = MaskKey.random(rng[])
     await session.stream.writer.write(
       (await Frame(
         fin: false,

diff --git a/websock/session.nim b/websock/session.nim
@@ -1,5 +1,5 @@
 ## nim-websock
-## Copyright (c) 2021-2022 Status Research & Development GmbH
+## Copyright (c) 2021-2023 Status Research & Development GmbH
 ## Licensed under either of
 ##  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
 ##  * MIT license ([LICENSE-MIT](LICENSE-MIT))
@@ -11,7 +11,7 @@
 
 import std/strformat
 import pkg/[chronos, chronicles, stew/byteutils, stew/endians2]
-import ./types, ./frame, ./utils, ./utf8dfa, ./http
+import ./types, ./frame, ./utf8dfa, ./http
 
 import pkg/chronos/streams/asyncstream
 
@@ -114,8 +114,9 @@ proc nonCancellableSend(
 
   trace "Sending data to remote"
 
-  let maskKey = if ws.masked:
-      genMaskKey(ws.rng)
+  let maskKey =
+    if ws.masked:
+      MaskKey.random(ws.rng[])
     else:
       default(MaskKey)
 

diff --git a/websock/types.nim b/websock/types.nim
@@ -1,5 +1,5 @@
 ## nim-websock
-## Copyright (c) 2021-2022 Status Research & Development GmbH
+## Copyright (c) 2021-2023 Status Research & Development GmbH
 ## Licensed under either of
 ##  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
 ##  * MIT license ([LICENSE-MIT](LICENSE-MIT))
@@ -13,11 +13,11 @@ import std/deques
 import pkg/[chronos,
             chronos/streams/tlsstream,
             chronos/apps/http/httptable,
+            bearssl/rand,
             httputils,
             stew/results]
-import ./utils
 
-export deques
+export deques, rand
 
 const
   SHA1DigestSize* = 20
@@ -55,6 +55,7 @@ type
   HeaderFlags* = set[HeaderFlag]
 
   MaskKey* = array[4, char]
+  WebSecKey* = array[16, byte]
 
   Frame* = ref object
     fin*: bool                 ## Indicates that this is the final fragment in a message.
@@ -89,7 +90,7 @@ type
     masked*: bool             # send masked packets
     binary*: bool             # is payload binary?
     flags*: set[TLSFlags]
-    rng*: Rng
+    rng*: ref HmacDrbgContext
     frameSize*: int           # max frame buffer size
     onPing*: ControlCb
     onPong*: ControlCb
@@ -212,3 +213,6 @@ method encode*(self: Ext, frame: Frame): Future[Frame] {.base, async.} =
 
 method toHttpOptions*(self: Ext): string {.base, gcsafe.} =
   raiseAssert "Not implemented!"
+
+func random*(T: typedesc[MaskKey|WebSecKey], rng: var HmacDrbgContext): T =
+  rng.generate(result)
diff --git a/websock/utils.nim b/websock/utils.nim
diff --git a/websock/websock.nim b/websock/websock.nim
@@ -26,9 +26,9 @@ import pkg/[chronos,
             stew/base10,
             nimcrypto/sha]
 
-import ./utils, ./frame, ./session, /types, ./http, ./extensions/extutils
+import ./frame, ./session, /types, ./http, ./extensions/extutils
 
-export utils, session, frame, types, http, httptable
+export session, frame, types, http, httptable
 
 logScope:
   topics = "websock ws-server"
@@ -117,11 +117,10 @@ proc connect*(
   onPing: ControlCb = nil,
   onPong: ControlCb = nil,
   onClose: CloseCb = nil,
-  rng: Rng = nil): Future[WSSession] {.async.} =
+  rng = HmacDrbgContext.new()): Future[WSSession] {.async.} =
 
   let
-    rng = if isNil(rng): HmacDrbgContext.new() else: rng
-    key = Base64Pad.encode(genWebSecKey(rng))
+    key = Base64Pad.encode(WebSecKey.random(rng[]))
     hostname = if hostName.len > 0: hostName else: $host
 
   let client = if secure:
@@ -216,7 +215,7 @@ proc connect*(
   onPing: ControlCb = nil,
   onPong: ControlCb = nil,
   onClose: CloseCb = nil,
-  rng: Rng = nil): Future[WSSession]
+  rng = HmacDrbgContext.new()): Future[WSSession]
   {.raises: [Defect, WSWrongUriSchemeError].} =
   ## Create a new websockets client
   ## using a Uri
@@ -359,12 +358,12 @@ proc new*(
   onPing: ControlCb = nil,
   onPong: ControlCb = nil,
   onClose: CloseCb = nil,
-  rng: Rng = nil): WSServer =
+  rng = HmacDrbgContext.new()): WSServer =
 
   return WSServer(
     protocols: @protos,
     masked: false,
-    rng: if isNil(rng): HmacDrbgContext.new() else: rng,
+    rng: rng,
     frameSize: frameSize,
     factories: @factories,
     onPing: onPing,