StakeSquid/dshackle
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

solution: websocket unsubscribe

Browse Source
This commit is contained in:
Igor Artamonov
2021-10-29 16:12:54 -04:00
parent e001913acb
commit f1720ad354
3 changed files with 90 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/main/kotlin/io/emeraldpay/dshackle/proxy/ProxyServer.kt
View File

@@ -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))

49
src/main/kotlin/io/emeraldpay/dshackle/proxy/WebsocketHandler.kt
View File

@@ -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))

49
src/test/groovy/io/emeraldpay/dshackle/proxy/WebsocketHandlerSpec.groovy
View File

@@ -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[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<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}'
}
}