diff --git a/src/main/kotlin/io/emeraldpay/dshackle/proxy/ProxyServer.kt b/src/main/kotlin/io/emeraldpay/dshackle/proxy/ProxyServer.kt index 6b10b06a..f0df599d 100644 --- a/src/main/kotlin/io/emeraldpay/dshackle/proxy/ProxyServer.kt +++ b/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)) diff --git a/src/main/kotlin/io/emeraldpay/dshackle/proxy/WebsocketHandler.kt b/src/main/kotlin/io/emeraldpay/dshackle/proxy/WebsocketHandler.kt index 34385553..9b692f43 100644 --- a/src/main/kotlin/io/emeraldpay/dshackle/proxy/WebsocketHandler.kt +++ b/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> { return BiFunction { req, resp -> + // each connection keeps a list of subscription controllers + val control = HashMap>() + val requests: Flux> = 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>, eventHandlerFactory: AccessHandlerHttp.WsHandlerFactory): Flux { + fun respond( + blockchain: Chain, + control: MutableMap>, + requests: Flux>, + eventHandlerFactory: AccessHandlerHttp.WsHandlerFactory + ): Flux { 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() + 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().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().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)) diff --git a/src/test/groovy/io/emeraldpay/dshackle/proxy/WebsocketHandlerSpec.groovy b/src/test/groovy/io/emeraldpay/dshackle/proxy/WebsocketHandlerSpec.groovy index 1ca8415d..35b2b3fb 100644 --- a/src/test/groovy/io/emeraldpay/dshackle/proxy/WebsocketHandlerSpec.groovy +++ b/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>(), 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>(), Flux.just(request), requestHandler) .collectList() .block(Duration.ofSeconds(1)) then: - act[0] == '{"jsonrpc":"2.0","id":2,"result":"0000000000000001"}' - act[1] == '{"jsonrpc":"2.0","method":"eth_subscription","params":{"result":{"foo":1},"subscription":"0000000000000001"}}' - act[2] == '{"jsonrpc":"2.0","method":"eth_subscription","params":{"result":{"foo":2},"subscription":"0000000000000001"}}' + act[0] == '{"jsonrpc":"2.0","id":2,"result":"1"}' + 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":"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>() + Sinks.One 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>() + 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}' + } }