From 66208055bc393e7ff7a4f525768632a0c7589164 Mon Sep 17 00:00:00 2001
From: jangko <jangko128@gmail.com>
Date: Fri, 12 Jan 2024 17:05:55 +0700
Subject: [PATCH] Both http server and client now can handle chunked transfer

---
 json_rpc/servers/httpserver.nim | 32 ++++++++++++++++++++++++++------
 tests/testhttp.nim              | 29 ++++++++++++++++++++++++-----
 2 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/json_rpc/servers/httpserver.nim b/json_rpc/servers/httpserver.nim
index 6cf3cc6..114fa71 100644
--- a/json_rpc/servers/httpserver.nim
+++ b/json_rpc/servers/httpserver.nim
@@ -37,6 +37,7 @@ type
   RpcHttpServer* = ref object of RpcServer
     httpServers: seq[HttpServerRef]
     authHooks: seq[HttpAuthHook]
+    maxChunkSize: int
 
 proc processClientRpc(rpcServer: RpcHttpServer): HttpProcessCallback2 =
   return proc (req: RequestFence): Future[HttpResponseRef] {.async: (raises: [CancelledError]).} =
@@ -64,15 +65,31 @@ proc processClientRpc(rpcServer: RpcHttpServer): HttpProcessCallback2 =
     let
       headers = HttpTable.init([("Content-Type",
                                  "application/json; charset=utf-8")])
+      chunkSize = rpcServer.maxChunkSize
+
     try:
       let
         body = await request.getBody()
-
         data = await rpcServer.route(string.fromBytes(body))
-        res = await request.respond(Http200, data, headers)
 
-      trace "JSON-RPC result has been sent"
-      return res
+      if data.len <= chunkSize:
+        let res = await request.respond(Http200, data, headers)
+        trace "JSON-RPC result has been sent"
+        return res
+
+      let response = request.getResponse()
+      response.status = Http200
+      response.addHeader("Content-Type", "application/json; charset=utf-8")
+      await response.prepare()
+      let maxLen = data.len
+      var len = data.len
+      while len > chunkSize:
+        await response.sendChunk(data[maxLen - len].unsafeAddr, chunkSize)
+        len -= chunkSize
+      if len > 0:
+        await response.sendChunk(data[maxLen - len].unsafeAddr, len)
+      await response.finish()
+
     except CancelledError as exc:
       raise exc
     except CatchableError as exc:
@@ -243,10 +260,10 @@ proc addSecureHttpServer*(server: RpcHttpServer,
     server.addSecureHttpServer(a, tlsPrivateKey, tlsCertificate)
 
 proc new*(T: type RpcHttpServer, authHooks: seq[HttpAuthHook] = @[]): T =
-  T(router: RpcRouter.init(), httpServers: @[], authHooks: authHooks)
+  T(router: RpcRouter.init(), httpServers: @[], authHooks: authHooks, maxChunkSize: 8192)
 
 proc new*(T: type RpcHttpServer, router: RpcRouter, authHooks: seq[HttpAuthHook] = @[]): T =
-  T(router: router, httpServers: @[], authHooks: authHooks)
+  T(router: router, httpServers: @[], authHooks: authHooks, maxChunkSize: 8192)
 
 proc newRpcHttpServer*(authHooks: seq[HttpAuthHook] = @[]): RpcHttpServer =
   RpcHttpServer.new(authHooks)
@@ -295,3 +312,6 @@ proc closeWait*(server: RpcHttpServer) {.async.} =
 proc localAddress*(server: RpcHttpServer): seq[TransportAddress] =
   for item in server.httpServers:
     result.add item.instance.localAddress()
+
+proc setMaxChunkSize*(server: RpcHttpServer, maxChunkSize: int) =
+  server.maxChunkSize = maxChunkSize
diff --git a/tests/testhttp.nim b/tests/testhttp.nim
index 23b29b5..2bd0c82 100644
--- a/tests/testhttp.nim
+++ b/tests/testhttp.nim
@@ -7,8 +7,9 @@
 # This file may not be copied, modified, or distributed except according to
 # those terms.
 
-import unittest2
-import ../json_rpc/[rpcserver, rpcclient]
+import
+  unittest2,
+  ../json_rpc/[rpcserver, rpcclient, jsonmarshal]
 
 const TestsCount = 100
 
@@ -46,6 +47,15 @@ proc invalidTest(address: string): Future[bool] {.async.} =
   if invalidA and invalidB:
     result = true
 
+const bigChunkSize = 4 * 8192
+
+proc chunkedTest(address: string): Future[bool] {.async.} =
+  var client = newRpcHttpClient()
+  await client.connect("http://" & address)
+  let r = await client.call("bigchunkMethod", %[])
+  let data = JrpcConv.decode(r.string, seq[byte])
+  return data.len == bigChunkSize
+
 var httpsrv = newRpcHttpServer(["127.0.0.1:0"])
 
 # Create RPC on server
@@ -54,15 +64,24 @@ httpsrv.rpc("myProc") do(input: string, data: array[0..3, int]):
 httpsrv.rpc("noParamsProc") do():
   result = %("Hello world")
 
+httpsrv.rpc("bigchunkMethod") do() -> seq[byte]:
+  result = newSeq[byte](bigChunkSize)
+  for i in 0..<result.len:
+    result[i] = byte(i mod 255)
+
+httpsrv.setMaxChunkSize(8192)
 httpsrv.start()
+let serverAddress = $httpsrv.localAddress()[0]
 
 suite "JSON-RPC test suite":
   test "Simple RPC call":
-    check waitFor(simpleTest($httpsrv.localAddress()[0])) == true
+    check waitFor(simpleTest(serverAddress)) == true
   test "Continuous RPC calls (" & $TestsCount & " messages)":
-    check waitFor(continuousTest($httpsrv.localAddress()[0])) == TestsCount
+    check waitFor(continuousTest(serverAddress)) == TestsCount
   test "Invalid RPC calls":
-    check waitFor(invalidTest($httpsrv.localAddress()[0])) == true
+    check waitFor(invalidTest(serverAddress)) == true
+  test "Http client can handle chunked transfer encoding":
+    check waitFor(chunkedTest(serverAddress)) == true
 
 waitFor httpsrv.stop()
 waitFor httpsrv.closeWait()