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

Fix streaming spam (#486)

Browse Source
This commit is contained in:
KirillPamPam
2024-05-28 13:51:25 +04:00
committed by GitHub
parent 4ea965528f
commit 387cd31c98
3 changed files with 95 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
src/main/kotlin/io/emeraldpay/dshackle/upstream/rpcclient/stream/JsonRpcStreamParser.kt
View File

@@ -37,10 +37,7 @@ class JsonRpcStreamParser(
private const val QUOTE: Byte = '"'.code.toByte() private const val QUOTE: Byte = '"'.code.toByte()
} }
fun streamParse( fun streamParse(statusCode: Int, response: Flux<ByteArray>): Mono<out Response> {
statusCode: Int,
response: Flux<ByteArray>,
): Mono<out Response> {
val firstPartSize = AtomicInteger() val firstPartSize = AtomicInteger()
return response.bufferUntil { return response.bufferUntil {
if (firstPartSize.get() > firstChunkMaxSize) { if (firstPartSize.get() > firstChunkMaxSize) {
@@ -110,10 +107,7 @@ class JsonRpcStreamParser(
} }
} }
private fun aggregateResponse( private fun aggregateResponse(response: Flux<ByteArray>, statusCode: Int): Mono<AggregateResponse> {
response: Flux<ByteArray>,
statusCode: Int,
): Mono<AggregateResponse> {
return ByteBufFlux.fromInbound(response).aggregate().asByteArray() return ByteBufFlux.fromInbound(response).aggregate().asByteArray()
.map { AggregateResponse(it, statusCode) } .map { AggregateResponse(it, statusCode) }
} }
@@ -170,7 +164,8 @@ class JsonRpcStreamParser(
var response: SingleResponse? = null var response: SingleResponse? = null
try { try {
jsonFactory.createParser(firstBytes).use { parser -> jsonFactory.createParser(firstBytes).use { parser ->
while (parser.nextToken() != null) { while (true) {
parser.nextToken()
if (firstBytes.size == parser.currentLocation.byteOffset.toInt()) { if (firstBytes.size == parser.currentLocation.byteOffset.toInt()) {
break break
} }
@@ -178,56 +173,51 @@ class JsonRpcStreamParser(
if (parser.currentName == "result") { if (parser.currentName == "result") {
val token = parser.nextToken() val token = parser.nextToken()
val tokenStart = parser.tokenLocation.byteOffset.toInt() val tokenStart = parser.tokenLocation.byteOffset.toInt()
if (token.isScalarValue) { response = if (token.isScalarValue) {
val count = CountQuotesAndSlashes(AtomicInteger(1)) val count = CountQuotesAndSlashes(AtomicInteger(1))
whatCount.set(count) whatCount.set(count)
response = SingleResponse(
SingleResponse( processScalarValue(parser, tokenStart, firstBytes, count, endStream),
processScalarValue(parser, tokenStart, firstBytes, count, endStream), null,
null, )
).merge(response)
} else { } else {
when (token) { when (token) {
JsonToken.START_OBJECT -> { JsonToken.START_OBJECT -> {
val count = val count = CountObjectBrackets(
CountObjectBrackets( AtomicInteger(1),
AtomicInteger(1), CountQuotesAndSlashes(AtomicInteger(0)),
CountQuotesAndSlashes(AtomicInteger(0)), )
)
whatCount.set(count) whatCount.set(count)
response = SingleResponse(
SingleResponse( processAndCountBrackets(
processAndCountBrackets( tokenStart,
tokenStart, firstBytes,
firstBytes, count,
count, endStream,
endStream, OBJECT_OPEN_BRACKET,
OBJECT_OPEN_BRACKET, OBJECT_CLOSE_BRACKET,
OBJECT_CLOSE_BRACKET, ),
), null,
null, )
).merge(response)
} }
JsonToken.START_ARRAY -> { JsonToken.START_ARRAY -> {
val count = val count = CountArrayBrackets(
CountArrayBrackets( AtomicInteger(1),
AtomicInteger(1), CountQuotesAndSlashes(AtomicInteger(0)),
CountQuotesAndSlashes(AtomicInteger(0)), )
)
whatCount.set(count) whatCount.set(count)
response = SingleResponse(
SingleResponse( processAndCountBrackets(
processAndCountBrackets( tokenStart,
tokenStart, firstBytes,
firstBytes, count,
count, endStream,
endStream, ARRAY_OPEN_BRACKET,
ARRAY_OPEN_BRACKET, ARRAY_CLOSE_BRACKET,
ARRAY_CLOSE_BRACKET, ),
), null,
null, )
).merge(response)
} }
else -> { else -> {
@@ -235,15 +225,28 @@ class JsonRpcStreamParser(
} }
} }
} }
if (endStream.get()) {
// we parsed the whole result field, and we can go parse further
parser.skipChildren()
} else {
// otherwise return response assuming there is no error field
return response
}
} else if (parser.currentName == "error") { } else if (parser.currentName == "error") {
return SingleResponse(null, responseRpcParser.readError(parser)).merge(response) return SingleResponse(response?.result, responseRpcParser.readError(parser))
} }
} }
} }
return response return response
} }
} catch (e: Exception) { } catch (e: Exception) {
log.warn("Streaming parsing exception: {}", e.message) if (response == null) {
// something terrible happened when we even don't have a response that means we haven't parsed the first chunk
log.warn("Streaming parsing exception: {}", e.message)
}
// there may be other parsing exceptions that means we have parsed the first chunk, and we have a response
// but when we want to parse further we can get an error if the first chunk is not a finished json
// it doesn't matter, we have the result from the response and can return it
return response return response
} }
} }

15
src/main/kotlin/io/emeraldpay/dshackle/upstream/stream/Responses.kt
View File

@@ -13,21 +13,6 @@ data class SingleResponse(
fun noResponse() = result == null && error == null fun noResponse() = result == null && error == null
fun merge(other: SingleResponse?): SingleResponse {
if (other == null) {
return this
}
var newResult: ByteArray? = result
var newError: ChainCallError? = error
if (newResult == null && other.result != null) {
newResult = other.result
}
if (newError == null && other.error != null) {
newError = other.error
}
return SingleResponse(newResult, newError)
}
override fun equals(other: Any?): Boolean { override fun equals(other: Any?): Boolean {
if (this === other) return true if (this === other) return true
if (other !is SingleResponse) return false if (other !is SingleResponse) return false

56
src/test/kotlin/io/emeraldpay/dshackle/upstream/rpcclient/stream/JsonRpcStreamParserTest.kt
View File

@@ -12,6 +12,7 @@ import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test import org.junit.jupiter.api.Test
import org.junit.jupiter.params.ParameterizedTest import org.junit.jupiter.params.ParameterizedTest
import org.junit.jupiter.params.provider.Arguments import org.junit.jupiter.params.provider.Arguments
import org.junit.jupiter.params.provider.Arguments.of
import org.junit.jupiter.params.provider.MethodSource import org.junit.jupiter.params.provider.MethodSource
import reactor.core.publisher.Flux import reactor.core.publisher.Flux
import reactor.test.StepVerifier import reactor.test.StepVerifier
@@ -60,11 +61,11 @@ class JsonRpcStreamParserTest {
@ParameterizedTest @ParameterizedTest
@MethodSource("data") @MethodSource("data")
fun `if first part has result field then single response`( fun `if first part has result field then single response`(
response: ByteArray, response: List<ByteArray>,
result: ByteArray, result: ByteArray,
) { ) {
val statusCode = 200 val statusCode = 200
val stream: Flux<ByteArray> = Flux.just(response) val stream: Flux<ByteArray> = Flux.fromIterable(response)
StepVerifier.create(streamParser.streamParse(statusCode, stream)) StepVerifier.create(streamParser.streamParse(statusCode, stream))
.expectNext(SingleResponse(result, null)) .expectNext(SingleResponse(result, null))
@@ -116,24 +117,49 @@ class JsonRpcStreamParserTest {
companion object { companion object {
@JvmStatic @JvmStatic
fun data(): List<Arguments> = listOf( fun data(): List<Arguments> = listOf(
Arguments.of("{\"id\": 2,\"result\": \"0x12\"}".toByteArray(), "\"0x12\"".toByteArray()), of(
Arguments.of("{\"id\": 2,\"result\": 11}".toByteArray(), "11".toByteArray()), listOf("{\"id\": 2,\"result\": \"0x12\"}".toByteArray()),
Arguments.of("{\"id\": 2,\"result\": false}".toByteArray(), "false".toByteArray()), "\"0x12\"".toByteArray(),
Arguments.of("{\"id\": 2,\"result\": null}".toByteArray(), "null".toByteArray()), ),
Arguments.of("{\"id\": 2,\"result\": {\"name\": \"value\"}".toByteArray(), "{\"name\": \"value\"}".toByteArray()), of(
Arguments.of("{\"id\": 2,\"result\": [{\"name\": \"value\"}]".toByteArray(), "[{\"name\": \"value\"}]".toByteArray()), listOf("{\"id\": 2,\"result\": 11}".toByteArray()),
"11".toByteArray(),
),
of(
listOf("{\"id\": 2,\"result\": false}".toByteArray()),
"false".toByteArray(),
),
of(
listOf("{\"id\": 2,\"result\": null}".toByteArray()),
"null".toByteArray(),
),
of(
listOf("{\"id\": 2,\"result\": {\"name\": \"value\"}}".toByteArray()),
"{\"name\": \"value\"}".toByteArray(),
),
of(
listOf("{\"id\": 2,\"result\": [{\"name\": \"value\"}]}".toByteArray()),
"[{\"name\": \"value\"}]".toByteArray(),
),
of(
listOf(
"{\"id\": 2,\"result\": [{\"name\": \"value\"}], \"other".toByteArray(),
"\": \"newField\"}".toByteArray(),
),
"[{\"name\": \"value\"}]".toByteArray(),
),
) )
@JvmStatic @JvmStatic
fun dataStream(): List<Arguments> = listOf( fun dataStream(): List<Arguments> = listOf(
Arguments.of( of(
listOf("{\"id\": 2,\"result\": \"0x12".toByteArray(), "222\"}".toByteArray()), listOf("{\"id\": 2,\"result\": \"0x12".toByteArray(), "222\"}".toByteArray()),
listOf( listOf(
Chunk("\"0x12".toByteArray(), false), Chunk("\"0x12".toByteArray(), false),
Chunk("222\"".toByteArray(), true), Chunk("222\"".toByteArray(), true),
), ),
), ),
Arguments.of( of(
listOf( listOf(
"{\"id\": 2,\"result\": \"0x12".toByteArray(), "{\"id\": 2,\"result\": \"0x12".toByteArray(),
"123\\\"".toByteArray(), "123\\\"".toByteArray(),
@@ -145,7 +171,7 @@ class JsonRpcStreamParserTest {
Chunk("222\"".toByteArray(), true), Chunk("222\"".toByteArray(), true),
), ),
), ),
Arguments.of( of(
listOf( listOf(
"{\"id\": 2,\"result\": \"0x12".toByteArray(), "{\"id\": 2,\"result\": \"0x12".toByteArray(),
"1\\n23\\\"".toByteArray(), "1\\n23\\\"".toByteArray(),
@@ -159,7 +185,7 @@ class JsonRpcStreamParserTest {
Chunk("\\222\\\\\\\\\"".toByteArray(), true), Chunk("\\222\\\\\\\\\"".toByteArray(), true),
), ),
), ),
Arguments.of( of(
listOf("{\"id\": 2,\"result\": {\"name\": ".toByteArray(), "\"bigName\"".toByteArray(), "}".toByteArray()), listOf("{\"id\": 2,\"result\": {\"name\": ".toByteArray(), "\"bigName\"".toByteArray(), "}".toByteArray()),
listOf( listOf(
Chunk("{\"name\": ".toByteArray(), false), Chunk("{\"name\": ".toByteArray(), false),
@@ -167,7 +193,7 @@ class JsonRpcStreamParserTest {
Chunk("}".toByteArray(), true), Chunk("}".toByteArray(), true),
), ),
), ),
Arguments.of( of(
listOf( listOf(
"{\"id\": 2,\"result\": [{\"name\": ".toByteArray(), "{\"id\": 2,\"result\": [{\"name\": ".toByteArray(),
"\"bigName\"".toByteArray(), "\"bigName\"".toByteArray(),
@@ -180,7 +206,7 @@ class JsonRpcStreamParserTest {
Chunk("}]".toByteArray(), true), Chunk("}]".toByteArray(), true),
), ),
), ),
Arguments.of( of(
listOf( listOf(
"{\"id\": 2,\"result\": [{\"na]me\": ".toByteArray(), "{\"id\": 2,\"result\": [{\"na]me\": ".toByteArray(),
"\"bigName]".toByteArray(), "\"bigName]".toByteArray(),
@@ -197,7 +223,7 @@ class JsonRpcStreamParserTest {
Chunk("}]".toByteArray(), true), Chunk("}]".toByteArray(), true),
), ),
), ),
Arguments.of( of(
listOf( listOf(
"{\"id\": 2,\"result\": {\"name\": ".toByteArray(), "{\"id\": 2,\"result\": {\"name\": ".toByteArray(),
"\"bigName}".toByteArray(), "\"bigName}".toByteArray(),