From 6e6f80512fb8fcc9f3c031cfd693ccbcf9c4aaab Mon Sep 17 00:00:00 2001 From: Pavel Gorgulov Date: Tue, 2 Jun 2026 12:31:28 +0200 Subject: [PATCH] Merge commit from fork --- kotlin-sdk-core/api/kotlin-sdk-core.api | 6 +++ .../kotlin/sdk/shared/ReadBuffer.kt | 40 ++++++++++++++++++- .../kotlin/sdk/shared/ReadBufferTest.kt | 33 +++++++++++++++ .../kotlin/sdk/server/StdioServerTransport.kt | 3 ++ 4 files changed, 80 insertions(+), 2 deletions(-) diff --git a/kotlin-sdk-core/api/kotlin-sdk-core.api b/kotlin-sdk-core/api/kotlin-sdk-core.api index b588eef..0895722 100644 --- a/kotlin-sdk-core/api/kotlin-sdk-core.api +++ b/kotlin-sdk-core/api/kotlin-sdk-core.api @@ -125,6 +125,12 @@ public final class io/modelcontextprotocol/kotlin/sdk/shared/RequestOptions : io public fun toString ()Ljava/lang/String; } +public final class io/modelcontextprotocol/kotlin/sdk/shared/TooLongFrameException : java/io/IOException { + public fun (JI)V + public final fun getFrameSize ()J + public final fun getMaxFrameSize ()I +} + public abstract interface class io/modelcontextprotocol/kotlin/sdk/shared/Transport { public abstract fun close (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public abstract fun onClose (Lkotlin/jvm/functions/Function0;)V diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBuffer.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBuffer.kt index c0c853f..8102cf6 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBuffer.kt +++ b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBuffer.kt @@ -4,13 +4,24 @@ import io.github.oshai.kotlinlogging.KotlinLogging import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage import io.modelcontextprotocol.kotlin.sdk.types.McpJson import kotlinx.io.Buffer +import kotlinx.io.IOException import kotlinx.io.indexOf import kotlinx.io.readString /** * Buffers a continuous stdio stream into discrete JSON-RPC messages. + * + * A single newline-terminated frame may not exceed an internal size cap (16 MiB): a peer that streams + * bytes without ever sending a newline triggers a [TooLongFrameException] from [readMessage] instead + * of growing the buffer without bound. A non-positive cap disables the check. */ -public class ReadBuffer { +public class ReadBuffer internal constructor(private val maxFrameSize: Int) { + + public constructor() : this(DEFAULT_MAX_FRAME_SIZE) + + internal companion object { + internal const val DEFAULT_MAX_FRAME_SIZE: Int = 16 * 1024 * 1024 + } private val logger = KotlinLogging.logger { } @@ -53,7 +64,11 @@ public class ReadBuffer { private fun readNextLine(): String? { val lfIndex = if (buffer.exhausted()) -1L else buffer.indexOf('\n'.code.toByte()) - if (lfIndex == -1L) return null + if (lfIndex == -1L) { + failIfFrameTooLong(buffer.size) + return null + } + failIfFrameTooLong(lfIndex) return if (lfIndex == 0L) { buffer.skip(1) @@ -71,6 +86,13 @@ public class ReadBuffer { } } + private fun failIfFrameTooLong(frameSize: Long) { + if (maxFrameSize in 1..(line) /** Serializes a [JSONRPCMessage] to its JSON string representation with a trailing newline. */ diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBufferTest.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBufferTest.kt index 7b4f053..e5e50f8 100644 --- a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBufferTest.kt +++ b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBufferTest.kt @@ -6,7 +6,9 @@ import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification import kotlinx.serialization.json.Json import kotlin.test.Test +import kotlin.test.assertContains import kotlin.test.assertEquals +import kotlin.test.assertFailsWith import kotlin.test.assertNull class ReadBufferTest { @@ -75,4 +77,35 @@ class ReadBufferTest { val message = readBuffer.readMessage() assertEquals(testMessage, message) } + + @Test + fun `should fail when an unframed blob exceeds the cap`() { + val readBuffer = ReadBuffer(maxFrameSize = 64) + // No newline ever arrives: the memory-exhaustion vector. + readBuffer.append(ByteArray(100) { 'a'.code.toByte() }) + val ex = assertFailsWith { readBuffer.readMessage() } + assertContains(ex.message.orEmpty(), "maximum size") + } + + @Test + fun `should fail when a completed line exceeds the cap`() { + val readBuffer = ReadBuffer(maxFrameSize = 64) + readBuffer.append(ByteArray(100) { 'a'.code.toByte() } + '\n'.code.toByte()) + assertFailsWith { readBuffer.readMessage() } + } + + @Test + fun `should not enforce a cap when maxFrameSize is non-positive`() { + val readBuffer = ReadBuffer(maxFrameSize = 0) + // Well beyond any small cap and still no newline — must not throw when disabled. + readBuffer.append(ByteArray(8192) { 'a'.code.toByte() }) + assertNull(readBuffer.readMessage()) + } + + @Test + fun `should parse a message that fits under the cap`() { + val readBuffer = ReadBuffer(maxFrameSize = 1024) + readBuffer.append(serializeMessage(testMessage).encodeToByteArray()) + assertEquals(testMessage, readBuffer.readMessage()) + } } diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport.kt index 2b76089..9261be2 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport.kt @@ -4,6 +4,7 @@ import io.github.oshai.kotlinlogging.KotlinLogging import io.modelcontextprotocol.kotlin.sdk.internal.IODispatcher import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer +import io.modelcontextprotocol.kotlin.sdk.shared.TooLongFrameException import io.modelcontextprotocol.kotlin.sdk.shared.TransportSendOptions import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage @@ -208,6 +209,8 @@ public class StdioServerTransport private constructor( readBuffer.readMessage() } catch (e: CancellationException) { throw e + } catch (e: TooLongFrameException) { + throw e } catch (e: Throwable) { _onError(e) null