Use client code also to handle server messaging in bidirectional transports (#252)

Author: arnetheduck

json_rpc/client.nim

@@ -18,7 +18,7 @@ when (NimMajor, NimMinor, NimPatch) < (2, 2, 6):
     discard Future[string]().value()
 
 import
-  std/[deques, json, tables, macros],
+  std/[deques, hashes, json, tables, macros],
   chronos,
   chronicles,
   stew/byteutils,
@@ -61,11 +61,16 @@ type
       # Client identifier, for logging
     maxMessageSize*: int
 
-    router*: ref RpcRouter
+  RpcConnection* = ref object of RpcClient
+    router*:
+      proc(request: RequestBatchRx): Future[string] {.async: (raises: []).}
       ## Router used for transports that support bidirectional communication
 
   GetJsonRpcRequestHeaders* = proc(): seq[(string, string)] {.gcsafe, raises: [].}
 
+func hash*(v: RpcClient): Hash =
+  cast[Hash](addr v[])
+
 func parseResponse*(payload: openArray[byte], T: type): T {.raises: [JsonRpcError].} =
   try:
     JrpcSys.decode(payload, T)
@@ -111,7 +116,7 @@ proc callOnProcessMessage*(
     ok(true)
 
 proc processMessage*(
-    client: RpcClient, line: sink seq[byte]
+    client: RpcConnection, line: seq[byte]
 ): Future[Result[string, string]] {.async: (raises: []).} =
   if not ?client.callOnProcessMessage(line):
     return ok("")
@@ -142,7 +147,7 @@ proc processMessage*(
       return ok(wrapError(router.INVALID_REQUEST, exc.msg))
 
   if client.router != nil:
-    ok(await client.router[].route(request))
+    ok(await client.router(request))
   else:
     ok("")
 

json_rpc/clients/socketclient.nim

@@ -18,7 +18,7 @@ import
 export client, errors
 
 type
-  RpcSocketClient* = ref object of RpcClient
+  RpcSocketClient* = ref object of RpcConnection
     transport*: StreamTransport
     address*: TransportAddress
     loop*: Future[void]
@@ -28,6 +28,13 @@ proc new*(
     maxMessageSize = defaultMaxMessageSize,
     router = default(ref RpcRouter),
 ): T =
+  let router =
+    if router != nil:
+      proc(request: RequestBatchRx): Future[string] {.async: (raises: [], raw: true).} =
+        router[].route(request)
+    else:
+      nil
+
   T(maxMessageSize: maxMessageSize, router: router)
 
 proc newRpcSocketClient*(
@@ -71,28 +78,28 @@ method request(
 
     await fut
 
-proc processData(client: RpcSocketClient) {.async: (raises: []).} =
+proc processMessages*(client: RpcSocketClient) {.async: (raises: []).} =
+  # Provide backwards compat with consumers that don't set a max message size
+  # for example by constructing RpcWebSocketHandler without going through init
+  let maxMessageSize =
+    if client.maxMessageSize == 0: defaultMaxMessageSize else: client.maxMessageSize
+
   var lastError: ref JsonRpcError
   while true:
-    let data =
-      try:
-        await client.transport.readLine(client.maxMessageSize)
-      except CatchableError as exc:
-        lastError = (ref RpcTransportError)(msg: exc.msg, parent: exc)
+    try:
+      let data = await client.transport.readLine(maxMessageSize)
+      if data == "":
         break
-    if data == "":
-      break
 
-    let resp = await(client.processMessage(data.toBytes())).valueOr:
-      lastError = (ref RequestDecodeError)(msg: error, payload: data.toBytes())
-      break
+      let resp = await(client.processMessage(data.toBytes())).valueOr:
+        lastError = (ref RequestDecodeError)(msg: error, payload: data.toBytes())
+        break
 
-    if resp.len > 0:
-      try:
+      if resp.len > 0:
         discard await client.transport.write(resp & "\r\n")
-      except CatchableError as exc:
-        lastError = (ref RpcTransportError)(msg: exc.msg, parent: exc)
-        break
+    except CatchableError as exc:
+      lastError = (ref RpcTransportError)(msg: exc.msg, parent: exc)
+      break
 
   if lastError == nil:
     lastError = (ref RpcTransportError)(msg: "Connection closed")
@@ -115,7 +122,7 @@ proc connect*(
 
   client.address = address
   client.remote = $client.address
-  client.loop = processData(client)
+  client.loop = processMessages(client)
 
 proc connect*(
     client: RpcSocketClient, address: string, port: Port

json_rpc/clients/websocketclient.nim

@@ -20,7 +20,7 @@ import
 export client, errors
 
 type
-  RpcWebSocketClient* = ref object of RpcClient
+  RpcWebSocketClient* = ref object of RpcConnection
     transport*: WSSession
     uri*: Uri
     loop*: Future[void]
@@ -32,6 +32,13 @@ proc new*(
     maxMessageSize = defaultMaxMessageSize,
     router = default(ref RpcRouter),
 ): T =
+  let router =
+    if router != nil:
+      proc(request: RequestBatchRx): Future[string] {.async: (raises: [], raw: true).} =
+        router[].route(request)
+    else:
+      nil
+
   T(getHeaders: getHeaders, maxMessageSize: maxMessageSize, router: router)
 
 proc newRpcWebSocketClient*(
@@ -82,26 +89,26 @@ method request*(
 
     await fut
 
-proc processData(client: RpcWebSocketClient) {.async: (raises: []).} =
+proc processMessages*(client: RpcWebSocketClient) {.async: (raises: []).} =
+  # Provide backwards compat with consumers that don't set a max message size
+  # for example by constructing RpcWebSocketHandler without going through init
+  let maxMessageSize =
+    if client.maxMessageSize == 0: defaultMaxMessageSize else: client.maxMessageSize
+
   var lastError: ref JsonRpcError
   while client.transport.readyState != ReadyState.Closed:
-    var data =
-      try:
-        await client.transport.recvMsg(client.maxMessageSize)
-      except CatchableError as exc:
-        lastError = (ref RpcTransportError)(msg: exc.msg, parent: exc)
-        break
+    try:
+      let data = await client.transport.recvMsg(maxMessageSize)
 
-    let resp = await(client.processMessage(data)).valueOr:
-      lastError = (ref RequestDecodeError)(msg: error, payload: data)
-      break
+      let resp = await(client.processMessage(data)).valueOr:
+        lastError = (ref RequestDecodeError)(msg: error, payload: data)
+        break
 
-    if resp.len > 0:
-      try:
+      if resp.len > 0:
         await client.transport.send(resp)
-      except CatchableError as exc:
-        lastError = (ref RpcTransportError)(msg: exc.msg, parent: exc)
-        break
+    except CatchableError as exc:
+      lastError = (ref RpcTransportError)(msg: exc.msg, parent: exc)
+      break
 
   if lastError == nil:
     lastError = (ref RpcTransportError)(msg: "Connection closed")
@@ -162,7 +169,7 @@ proc connect*(
   client.transport = ws
   client.uri = uri
   client.remote = uri.hostname & ":" & uri.port
-  client.loop = processData(client)
+  client.loop = processMessages(client)
 
 method close*(client: RpcWebSocketClient) {.async: (raises: []).} =
   await client.loop.cancelAndWait()

json_rpc/server.nim

@@ -10,23 +10,20 @@
 {.push raises: [], gcsafe.}
 
 import
-  std/[json, sequtils],
+  std/[json, sequtils, sets],
   chronos,
   ./[client, errors, jsonmarshal, router],
   ./private/jrpc_sys,
   ./private/shared_wrapper
 
-export
-  chronos,
-  client,
-  jsonmarshal,
-  router
+export chronos, client, jsonmarshal, router, sets
 
 type
   RpcServer* = ref object of RootRef
     router*: RpcRouter
 
-    connections*: seq[RpcClient]
+    # For servers that expose bidirectional connections, keep track of them
+    connections*: HashSet[RpcConnection]
 
 # ------------------------------------------------------------------------------
 # Constructors

json_rpc/servers/socketserver.nim

@@ -16,6 +16,7 @@ import
   json_serialization/std/net as jsnet,
   ../private/utils,
   ../[errors, server],
+  ../private/jrpc_sys,
   ../clients/socketclient
 
 export errors, server, jsnet
@@ -29,44 +30,30 @@ type
     processClientHook: StreamCallback2
     maxMessageSize: int
 
-proc processClient(server: StreamServer, transport: StreamTransport) {.async: (raises: []).} =
+proc processClient(
+    server: StreamServer, transport: StreamTransport
+) {.async: (raises: []).} =
   ## Process transport data to the RPC server
 
   let
     rpc = getUserData[RpcSocketServer](server)
     remote = transport.remoteAddress2().valueOr(default(TransportAddress))
-    c = RpcSocketClient(transport: transport, address: remote, remote: $remote)
-
-  rpc.connections.add(c)
-
-  # Provide backwards compat with consumers that don't set a max message size
-  # for example by constructing RpcWebSocketHandler without going through init
-  let maxMessageSize =
-    if rpc.maxMessageSize == 0: defaultMaxMessageSize else: rpc.maxMessageSize
-
-  try:
-    while true:
-      let req = await transport.readLine(maxMessageSize)
-      if req == "":
-        break
-
-      debug "Received JSON-RPC request",
-        address = transport.remoteAddress(),
-        len = req.len
-
-      let res = await rpc.route(req)
-      if res.len > 0:
-        discard await transport.write(res & "\r\n")
-  except TransportError as ex:
-    error "Transport closed during processing client",
-      remote,
-      msg=ex.msg
-  except CancelledError:
-    debug "JSON-RPC request processing cancelled", remote
-
-  rpc.connections.keepItIf(it != c)
-
-  await transport.closeWait()
+    c = RpcSocketClient(
+      transport: transport,
+      address: remote,
+      remote: $remote,
+      maxMessageSize: rpc.maxMessageSize,
+      router: proc(
+          request: RequestBatchRx
+      ): Future[string] {.async: (raises: [], raw: true).} =
+        rpc.router.route(request),
+    )
+
+  rpc.connections.incl(c)
+
+  await c.processMessages()
+
+  rpc.connections.excl(c)
 
 # Utility functions for setting up servers using stream transport addresses
 

json_rpc/servers/websocketserver.nim

@@ -10,13 +10,13 @@
 {.push raises: [], gcsafe.}
 
 import
-  std/sequtils,
   chronicles,
   chronos,
   websock/[websock, types],
   websock/extensions/compression/deflate,
   json_serialization/std/net as jsnet,
   ../[errors, server],
+  ../private/jrpc_sys,
   ../clients/websocketclient
 
 export errors, server, jsnet
@@ -46,58 +46,39 @@ type
 
 proc serveHTTP*(rpc: RpcWebSocketHandler, request: HttpRequest)
                   {.async: (raises: [CancelledError]).} =
-  try:
+  let ws = try:
     let server = rpc.wsserver
     let ws = await server.handleRequest(request)
     if ws.readyState != ReadyState.Open:
-      error "Failed to open websocket connection",
-        address = $request.uri
+      error "Failed to open websocket connection", address = $request.uri
       return
-
-    trace "Websocket handshake completed"
-    let c = RpcWebSocketClient(transport: ws, remote: $request.uri)
-    rpc.connections.add(c)
-    # Provide backwards compat with consumers that don't set a max message size
-    # for example by constructing RpcWebSocketHandler without going through init
-    let maxMessageSize =
-      if rpc.maxMessageSize == 0: defaultMaxMessageSize else: rpc.maxMessageSize
-    while ws.readyState != ReadyState.Closed:
-
-      let req = await ws.recvMsg(maxMessageSize)
-      debug "Received JSON-RPC request",
-        address = $request.uri,
-        len = req.len
-
-      if ws.readyState == ReadyState.Closed:
-        # if session already terminated by peer,
-        # no need to send response
-        break
-
-      if req.len == 0:
-        await ws.close(
-          reason = "cannot process zero length message"
-        )
-        break
-
-      let data = await rpc.route(req)
-
-      if data.len > 0:
-        await ws.send(data)
-
-    rpc.connections.keepItIf(it != c)
-
+    ws
   except WebSocketError as exc:
-    error "WebSocket error:",
-      address = $request.uri, msg = exc.msg
-
+    error "WebSocket error:", address = $request.uri, msg = exc.msg
+    return
   except CancelledError as exc:
     raise exc
-
   except CatchableError as exc:
-    debug "Internal error while processing JSON-RPC call", msg=exc.msg
-
-proc handleRequest(rpc: RpcWebSocketServer, request: HttpRequest)
-                    {.async: (raises: [CancelledError]).} =
+    debug "Internal error while processing JSON-RPC call", msg = exc.msg
+    return
+
+  trace "Websocket handshake completed"
+  let c = RpcWebSocketClient(
+    transport: ws,
+    remote: $request.uri,
+    maxMessageSize: rpc.maxMessageSize,
+    router: proc(request: RequestBatchRx): Future[string] {.async: (raises: [], raw: true).} =
+      rpc.router.route(request),
+  )
+  rpc.connections.incl(c)
+
+  await c.processMessages()
+
+  rpc.connections.excl(c)
+
+proc handleRequest(
+    rpc: RpcWebSocketServer, request: HttpRequest
+) {.async: (raises: [CancelledError]).} =
   trace "Handling request:", uri = $request.uri
 
   # if hook result is false,