solution: websocket unsubscribe

src/main/kotlin/io/emeraldpay/dshackle/proxy/ProxyServer.kt

@@ -112,6 +112,7 @@ class ProxyServer(
 
     fun setupRoutes(routes: HttpServerRoutes) {
         config.routes.forEach { routeConfig ->
+            // TODO implement a manual handling of the routes and WS upgrade to have a better control over the connection and improve the access logging
             routes.post("/" + routeConfig.id, httpHandler.proxy(routeConfig))
             if (config.websocketEnabled && wsHandler != null) {
                 routes.ws("/" + routeConfig.id, wsHandler.proxy(routeConfig))

src/main/kotlin/io/emeraldpay/dshackle/proxy/WebsocketHandler.kt

@@ -15,6 +15,8 @@
  */
 package io.emeraldpay.dshackle.proxy
 
+import com.google.protobuf.ByteString
+import io.emeraldpay.api.proto.BlockchainOuterClass
 import io.emeraldpay.dshackle.Global
 import io.emeraldpay.dshackle.config.ProxyConfig
 import io.emeraldpay.dshackle.monitoring.accesslog.AccessHandlerHttp
@@ -25,11 +27,11 @@ import io.emeraldpay.etherjar.rpc.json.ResponseJson
 import io.emeraldpay.grpc.Chain
 import io.netty.buffer.ByteBufInputStream
 import io.netty.buffer.Unpooled
-import org.apache.commons.lang3.StringUtils
 import org.reactivestreams.Publisher
 import org.slf4j.LoggerFactory
 import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
+import reactor.core.publisher.Sinks
 import reactor.netty.http.websocket.WebsocketInbound
 import reactor.netty.http.websocket.WebsocketOutbound
 import java.util.concurrent.atomic.AtomicLong
@@ -55,11 +57,14 @@ class WebsocketHandler(
 
     fun nextSubscriptionId(): String {
         val n = subscriptionId.incrementAndGet()
-        return StringUtils.leftPad(n.toString(16), 16, "0")
+        return n.toString(16)
     }
 
     fun proxy(routeConfig: ProxyConfig.Route): BiFunction<WebsocketInbound, WebsocketOutbound, Publisher<Void>> {
         return BiFunction { req, resp ->
+            // each connection keeps a list of subscription controllers
+            val control = HashMap<String, Sinks.One<Boolean>>()
+
             val requests: Flux<RequestJson<Any>> = req.aggregateFrames()
                 .receiveFrames()
                 .map { ByteBufInputStream(it.content()).readAllBytes() }
@@ -67,7 +72,7 @@ class WebsocketHandler(
 
             val eventHandler = accessHandler.start(req, routeConfig.blockchain)
 
-            val responses = respond(routeConfig.blockchain, requests, eventHandler)
+            val responses = respond(routeConfig.blockchain, control, requests, eventHandler)
                 .map { Unpooled.wrappedBuffer(it.toByteArray()) }
 
             resp.send(responses)
@@ -97,7 +102,12 @@ class WebsocketHandler(
         }
     }
 
-    fun respond(blockchain: Chain, requests: Flux<RequestJson<Any>>, eventHandlerFactory: AccessHandlerHttp.WsHandlerFactory): Flux<String> {
+    fun respond(
+        blockchain: Chain,
+        control: MutableMap<String, Sinks.One<Boolean>>,
+        requests: Flux<RequestJson<Any>>,
+        eventHandlerFactory: AccessHandlerHttp.WsHandlerFactory
+    ): Flux<String> {
         return requests.flatMap { call ->
             val method = call.method
 
@@ -113,6 +123,8 @@ class WebsocketHandler(
                             Pair(mp.first, mp.second?.let { Global.objectMapper.writeValueAsBytes(it) })
                         }
                     )
+                    val currentControl = Sinks.one<Boolean>()
+                    control[subscriptionId] = currentControl
                     // first need to respond with ID of the subscription, and the following responses would have it in "subscription" param
                     val start = ResponseJson<String, Any>().also {
                         it.id = call.id
@@ -124,6 +136,7 @@ class WebsocketHandler(
                         .map { event ->
                             WsSubscriptionResponse(params = WsSubscriptionData(event, subscriptionId))
                         }
+                        .takeUntilOther(currentControl.asMono())
                     Flux.concat(Mono.just(start), responses)
                         .map { Global.objectMapper.writeValueAsString(it) }
                         .doOnNext {
@@ -133,6 +146,34 @@ class WebsocketHandler(
                     // TODO should it produce a 404 to the AccessLog?
                     Mono.empty()
                 }
+            } else if (method == "eth_unsubscribe") {
+                val id = call.params?.getOrNull(0) ?: ""
+
+                // put it to the Access Log with fake id=0 (it doesn't matter, except the later reference)
+                val eventHandler: AccessHandlerHttp.RequestHandler = eventHandlerFactory.call()
+                eventHandler.onRequest(
+                    BlockchainOuterClass.NativeCallRequest.newBuilder()
+                        .setChainValue(blockchain.id)
+                        .addItems(
+                            BlockchainOuterClass.NativeCallItem.newBuilder()
+                                .setId(0)
+                                .setMethod("eth_unsubscribe")
+                                .setPayload(ByteString.copyFromUtf8("[\"$id\"]"))
+                                .build()
+                        )
+                        .build()
+                )
+
+                val p = control.remove(id.toString())
+                val success = p?.tryEmitValue(true)?.isSuccess ?: false
+                val response = ResponseJson<Boolean, Any>().also {
+                    it.id = call.id
+                    it.result = success
+                }
+                Mono.just(response)
+                    .map { Global.objectMapper.writeValueAsString(it) }
+                    .doOnNext { eventHandler.onResponse(NativeCall.CallResult.ok(0, it.toByteArray())) }
+                    .doFinally { eventHandler.close() }
             } else {
                 val eventHandler: AccessHandlerHttp.RequestHandler = eventHandlerFactory.call()
                 val proxyCall = readRpcJson.convertToNativeCall(ProxyCall.RpcType.SINGLE, listOf(call))

src/test/groovy/io/emeraldpay/dshackle/proxy/WebsocketHandlerSpec.groovy

@@ -21,6 +21,7 @@ import io.emeraldpay.dshackle.rpc.NativeSubscribe
 import io.emeraldpay.etherjar.rpc.json.RequestJson
 import io.emeraldpay.grpc.Chain
 import reactor.core.publisher.Flux
+import reactor.core.publisher.Sinks
 import spock.lang.Specification
 
 import java.time.Duration
@@ -85,7 +86,7 @@ class WebsocketHandlerSpec extends Specification {
 
         def request = new RequestJson("foo_test", [], 2)
         when:
-        def act = handler.respond(Chain.ETHEREUM, Flux.just(request), requestHandler)
+        def act = handler.respond(Chain.ETHEREUM, new HashMap<String, Sinks.One<Boolean>>(), Flux.just(request), requestHandler)
                 .single()
                 .block(Duration.ofSeconds(1))
         then:
@@ -106,13 +107,51 @@ class WebsocketHandlerSpec extends Specification {
 
         def request = new RequestJson("eth_subscribe", ["foo_test"], 2)
         when:
-        def act = handler.respond(Chain.ETHEREUM, Flux.just(request), requestHandler)
+        def act = handler.respond(Chain.ETHEREUM, new HashMap<String, Sinks.One<Boolean>>(), Flux.just(request), requestHandler)
                 .collectList()
                 .block(Duration.ofSeconds(1))
         then:
-        act[0] == '{"jsonrpc":"2.0","id":2,"result":"0000000000000001"}'
+        act[0] == '{"jsonrpc":"2.0","id":2,"result":"1"}'
-        act[1] == '{"jsonrpc":"2.0","method":"eth_subscription","params":{"result":{"foo":1},"subscription":"0000000000000001"}}'
+        act[1] == '{"jsonrpc":"2.0","method":"eth_subscription","params":{"result":{"foo":1},"subscription":"1"}}'
-        act[2] == '{"jsonrpc":"2.0","method":"eth_subscription","params":{"result":{"foo":2},"subscription":"0000000000000001"}}'
+        act[2] == '{"jsonrpc":"2.0","method":"eth_subscription","params":{"result":{"foo":2},"subscription":"1"}}'
         act.size() == 3
     }
+
+    def "Unsubscribe"() {
+        setup:
+
+        def handler = new WebsocketHandler(
+                new ReadRpcJson(), new WriteRpcJson(), Stub(NativeCall), Stub(NativeSubscribe), requestHandlerFactory, Stub(ProxyServer.RequestMetricsFactory)
+        )
+
+        def control = new HashMap<String, Sinks.One<Boolean>>()
+        Sinks.One<Boolean> sink = Sinks.one();
+        control["5"] = sink
+        def request = new RequestJson("eth_unsubscribe", ["5"], 0)
+        when:
+        def act = handler.respond(Chain.ETHEREUM, control, Flux.just(request), requestHandler)
+                .single()
+                .block(Duration.ofSeconds(1))
+        def sinkResponse = sink.asMono().block()
+        then:
+        act == '{"jsonrpc":"2.0","id":0,"result":true}'
+        sinkResponse != null
+    }
+
+    def "Unsubscribe when no subscription"() {
+        setup:
+
+        def handler = new WebsocketHandler(
+                new ReadRpcJson(), new WriteRpcJson(), Stub(NativeCall), Stub(NativeSubscribe), requestHandlerFactory, Stub(ProxyServer.RequestMetricsFactory)
+        )
+
+        def control = new HashMap<String, Sinks.One<Boolean>>()
+        def request = new RequestJson("eth_unsubscribe", ["5"], 0)
+        when:
+        def act = handler.respond(Chain.ETHEREUM, control, Flux.just(request), requestHandler)
+                .single()
+                .block(Duration.ofSeconds(1))
+        then:
+        act == '{"jsonrpc":"2.0","id":0,"result":false}'
+    }
 }