From b455958d8a66ab6175d4cde4df14e1c3ea86a9f9 Mon Sep 17 00:00:00 2001
From: Dustin Brody <tersec@users.noreply.github.com>
Date: Fri, 4 Mar 2022 19:13:29 +0000
Subject: [PATCH] request header callback

---
 json_rpc/client.nim                  |  2 ++
 json_rpc/clients/httpclient.nim      | 23 ++++++++++++++++++-----
 json_rpc/clients/websocketclient.nim | 14 ++++++++++----
 3 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/json_rpc/client.nim b/json_rpc/client.nim
index 5673fee..a1c27be 100644
--- a/json_rpc/client.nim
+++ b/json_rpc/client.nim
@@ -19,6 +19,8 @@ type
 
   Response* = JsonNode
 
+  GetJsonRpcRequestHeaders* = proc(): seq[(string, string)] {.gcsafe.}
+
 proc getNextId*(client: RpcClient): ClientId =
   client.lastId += 1
   client.lastId
diff --git a/json_rpc/clients/httpclient.nim b/json_rpc/clients/httpclient.nim
index 58adae9..e18f0d1 100644
--- a/json_rpc/clients/httpclient.nim
+++ b/json_rpc/clients/httpclient.nim
@@ -19,11 +19,14 @@ type
     httpSession: HttpSessionRef
     httpAddress: HttpResult[HttpAddress]
     maxBodySize: int
+    getHeaders: GetJsonRpcRequestHeaders
 
 const
   MaxHttpRequestSize = 128 * 1024 * 1024 # maximum size of HTTP body in octets
 
-proc new(T: type RpcHttpClient, maxBodySize = MaxHttpRequestSize, secure = false): T =
+proc new(
+    T: type RpcHttpClient, maxBodySize = MaxHttpRequestSize, secure = false,
+    getHeaders: GetJsonRpcRequestHeaders = nil): T =
   let httpSessionFlags = if secure:
     {
       HttpClientFlag.NoVerifyHost,
@@ -34,11 +37,14 @@ proc new(T: type RpcHttpClient, maxBodySize = MaxHttpRequestSize, secure = false
 
   T(
     maxBodySize: maxBodySize,
-    httpSession: HttpSessionRef.new(flags = httpSessionFlags)
+    httpSession: HttpSessionRef.new(flags = httpSessionFlags),
+    getHeaders: getHeaders
   )
 
-proc newRpcHttpClient*(maxBodySize = MaxHttpRequestSize, secure = false): RpcHttpClient =
-  RpcHttpClient.new(maxBodySize, secure)
+proc newRpcHttpClient*(
+    maxBodySize = MaxHttpRequestSize, secure = false,
+    getHeaders: GetJsonRpcRequestHeaders = nil): RpcHttpClient =
+  RpcHttpClient.new(maxBodySize, secure, getHeaders)
 
 method call*(client: RpcHttpClient, name: string,
              params: JsonNode): Future[Response]
@@ -47,13 +53,20 @@ method call*(client: RpcHttpClient, name: string,
   if client.httpAddress.isErr:
     raise newException(RpcAddressUnresolvableError, client.httpAddress.error)
 
+  var headers =
+    if not isNil(client.getHeaders):
+      client.getHeaders()
+    else:
+      @[]
+  headers.add(("Content-Type", "application/json"))
+
   let
     id = client.getNextId()
     reqBody = $rpcCallNode(name, params, id)
     req = HttpClientRequestRef.post(client.httpSession,
                                     client.httpAddress.get,
                                     body = reqBody.toOpenArrayByte(0, reqBody.len - 1),
-                                    headers = [("Content-Type", "application/json")])
+                                    headers = headers)
     res =
       try:
         await req.send()
diff --git a/json_rpc/clients/websocketclient.nim b/json_rpc/clients/websocketclient.nim
index caf145a..f0083b7 100644
--- a/json_rpc/clients/websocketclient.nim
+++ b/json_rpc/clients/websocketclient.nim
@@ -18,6 +18,7 @@ when useNews:
       transport*: WebSocket
       uri*: string
       loop*: Future[void]
+      getHeaders*: GetJsonRpcRequestHeaders
 
 else:
   import std/[uri, strutils]
@@ -28,13 +29,16 @@ else:
       transport*: WSSession
       uri*: Uri
       loop*: Future[void]
+      getHeaders*: GetJsonRpcRequestHeaders
 
-proc new*(T: type RpcWebSocketClient): T =
-  T()
+proc new*(
+    T: type RpcWebSocketClient, getHeaders: GetJsonRpcRequestHeaders = nil): T =
+  T(getHeaders: getHeaders)
 
-proc newRpcWebSocketClient*: RpcWebSocketClient =
+proc newRpcWebSocketClient*(
+    getHeaders: GetJsonRpcRequestHeaders = nil): RpcWebSocketClient =
   ## Creates a new client instance.
-  RpcWebSocketClient.new()
+  RpcWebSocketClient.new(getHeaders)
 
 method call*(self: RpcWebSocketClient, name: string,
              params: JsonNode): Future[Response] {.
@@ -112,6 +116,8 @@ when useNews:
       # TODO: This is a hack, because the table might be case sensitive. Ideally strtabs module has
       # to be extended with case insensitive accessors.
       headers["Origin"] = "http://localhost"
+    for header in client.getHeaders():
+      headers[header[0]] = header[1]
     client.transport = await newWebSocket(uri, headers)
     client.uri = uri
     client.loop = processData(client)