Skip to content

fix: bound inbound WebSocket message size #1129

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
devcrocod merged 2 commits into from
Oct 27, 2026
+171 −6
Merged

fix: bound inbound WebSocket message size #1129

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4d46374
fix: bound inbound WebSocket message size
devcrocod Oct 14, 2026
b8dd0d0
Merge branch 'main' into devcrocod/websocket-message-size-limit
devcrocod Oct 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 69 additions & 0 deletions ...lin/io/modelcontextprotocol/kotlin/sdk/integration/WebSocketMessageSizeIntegrationTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
package io.modelcontextprotocol.kotlin.sdk.integration

import io.kotest.matchers.shouldBe
import io.ktor.client.HttpClient
import io.ktor.client.plugins.websocket.WebSockets
import io.ktor.client.plugins.websocket.webSocketSession
import io.ktor.client.request.header
import io.ktor.http.HttpHeaders
import io.ktor.server.engine.embeddedServer
import io.ktor.websocket.Frame
import io.modelcontextprotocol.kotlin.sdk.client.Client
import io.modelcontextprotocol.kotlin.sdk.client.mcpWebSocketTransport
import io.modelcontextprotocol.kotlin.sdk.server.Server
import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions
import io.modelcontextprotocol.kotlin.sdk.server.mcpWebSocket
import io.modelcontextprotocol.kotlin.sdk.shared.MCP_SUBPROTOCOL
import io.modelcontextprotocol.kotlin.sdk.types.Implementation
import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities
import io.modelcontextprotocol.kotlin.test.utils.actualPort
import io.modelcontextprotocol.kotlin.test.utils.runIntegrationTest
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.put
import kotlinx.serialization.json.putJsonObject
import org.junit.jupiter.api.Test
import kotlin.time.Duration.Companion.seconds
import io.ktor.client.engine.cio.CIO as ClientCIO
import io.ktor.server.cio.CIO as ServerCIO

class WebSocketMessageSizeIntegrationTest {

@Test
fun `server closes a connection that sends a message over 16 MiB and keeps serving others`() =
runIntegrationTest(timeout = 30.seconds) {
val server = Server(Implementation("test-server", "1.0.0"), ServerOptions(ServerCapabilities()))
val engine = embeddedServer(ServerCIO, host = HOST, port = 0) { mcpWebSocket { server } }
.startSuspend(wait = false)
val httpClient = HttpClient(ClientCIO) { install(WebSockets) }
val url = "ws://$HOST:${engine.actualPort()}"
try {
val peer = httpClient.webSocketSession(url) {
header(HttpHeaders.SecWebSocketProtocol, MCP_SUBPROTOCOL)
}

// A valid request: without the limit the server would answer it instead of closing.
peer.send(Frame.Text(ping(padding = "a".repeat(16 * 1024 * 1024))))

peer.incoming.receiveCatching().getOrNull() shouldBe null
val client = Client(Implementation("test-client", "1.0.0"))
client.connect(httpClient.mcpWebSocketTransport(url))
client.ping()
client.close()
} finally {
httpClient.close()
engine.stopSuspend(gracePeriodMillis = 0, timeoutMillis = 1_000)
server.close()
}
}

private fun ping(padding: String): String = buildJsonObject {
put("jsonrpc", "2.0")
put("id", 1)
put("method", "ping")
putJsonObject("params") { put("padding", padding) }
}.toString()

private companion object {
const val HOST = "127.0.0.1"
}
}
14 changes: 10 additions & 4 deletions .../src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransport.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,12 +64,19 @@ public abstract class WebSocketMcpTransport : AbstractTransport() {
} catch (_: ClosedReceiveChannelException) {
logger.debug { "Closed receive channel, exiting" }
return@launch
} catch (e: CancellationException) {
throw e
} catch (e: Exception) {
// The session is already failing (e.g. FrameTooBigException) and Ktor closes it.
invokeOnErrorCallback(e)
return@launch
}

// Keep reading after a bad message: an open session nobody reads buffers frames without limit.
if (message !is Frame.Text) {
val e = IllegalArgumentException("Expected text frame, got ${message::class.simpleName}: $message")
_onError.invoke(e)
throw e
invokeOnErrorCallback(e)
continue
}

try {
Expand All @@ -78,8 +85,7 @@ public abstract class WebSocketMcpTransport : AbstractTransport() {
} catch (e: CancellationException) {
throw e
} catch (e: Exception) {
_onError.invoke(e)
throw e
invokeOnErrorCallback(e)
}
}
}
Expand Down
75 changes: 75 additions & 0 deletions ...src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransportTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
package io.modelcontextprotocol.kotlin.sdk.shared

import io.kotest.matchers.collections.shouldHaveSize
import io.kotest.matchers.shouldBe
import io.kotest.matchers.types.shouldBeInstanceOf
import io.ktor.websocket.Frame
import io.ktor.websocket.FrameTooBigException
import io.ktor.websocket.WebSocketExtension
import io.ktor.websocket.WebSocketSession
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.Job
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.ReceiveChannel
import kotlinx.coroutines.channels.SendChannel
import kotlinx.coroutines.test.runTest
import org.junit.jupiter.api.Test
import java.util.concurrent.CopyOnWriteArrayList
import kotlin.coroutines.CoroutineContext

class WebSocketMcpTransportTest {

@Test
fun `should keep reading after messages it cannot parse`() = runTest {
val session = FakeWebSocketSession()
val transport = TestTransport(session)
val errors = CopyOnWriteArrayList<Throwable>()
val received = CompletableDeferred<JSONRPCMessage>()
transport.onError { errors += it }
transport.onMessage { received.complete(it) }
transport.start()

session.incomingFrames.send(Frame.Binary(true, byteArrayOf(1, 2, 3)))
session.incomingFrames.send(Frame.Text("{"))
session.incomingFrames.send(Frame.Text("""{"jsonrpc":"2.0","id":1,"method":"ping"}"""))

received.await().shouldBeInstanceOf<JSONRPCRequest>().method shouldBe "ping"
errors shouldHaveSize 2
session.incomingFrames.close()
}

@Test
fun `should report a session failure through onError`() = runTest {
val session = FakeWebSocketSession()
val transport = TestTransport(session)
val error = CompletableDeferred<Throwable>()
transport.onError { error.complete(it) }
transport.start()

session.incomingFrames.close(FrameTooBigException(42))

error.await().shouldBeInstanceOf<FrameTooBigException>()
}

private class TestTransport(override val session: WebSocketSession) : WebSocketMcpTransport() {
override suspend fun initializeSession() = Unit
}
}

private class FakeWebSocketSession : WebSocketSession {
val incomingFrames = Channel<Frame>(Channel.UNLIMITED)

override val coroutineContext: CoroutineContext = Job()
override var masking: Boolean = false
override var maxFrameSize: Long = Long.MAX_VALUE
override val incoming: ReceiveChannel<Frame> get() = incomingFrames
override val outgoing: SendChannel<Frame> = Channel(Channel.UNLIMITED)
override val extensions: List<WebSocketExtension<*>> = emptyList()

override suspend fun flush() = Unit

@Deprecated("Use cancel() instead.", level = DeprecationLevel.ERROR)
override fun terminate() = Unit
}
19 changes: 17 additions & 2 deletions ...Main/kotlin/io/modelcontextprotocol/kotlin/sdk/server/WebSocketMcpKtorServerExtensions.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,13 @@ import kotlinx.coroutines.awaitCancellation

private val logger = KotlinLogging.logger {}

// Ktor's default maxFrameSize does not limit inbound messages at all.
private const val DEFAULT_MAX_FRAME_SIZE: Long = 16L * 1024 * 1024

@jetbrains-air jetbrains-air Bot Oct 19, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Problem: 16 MiB is hardcoded with no way to raise it on the Application.mcpWebSocket overloads. They call install(WebSockets) internally, and pre-installing the plugin with your own maxFrameSize doesn't layer on top — I checked on this branch and it fails with DuplicatePluginException: … already installed with the same key as 'WebSockets'. So "Breaking Changes: None" doesn't hold for anyone whose peers send larger messages.

Failure scenario: a server hosted via Application.mcpWebSocket whose clients call a tool with a base64-encoded document or image argument above 16 MiB. That works today; after upgrading, Ktor closes the connection with TOO_BIG part-way through the request, and the only remedy is to abandon this overload for install(WebSockets) { maxFrameSize = … } + routing { mcpWebSocket { … } } — which is the path that currently has no limit at all. Between the two, the SDK offers "hardcoded 16 MiB" or "unlimited", with nothing in between.

Suggested fix: take the limit as a parameter so the safe default still ships but stays overridable without switching entry point, e.g. public fun Application.mcpWebSocket(maxFrameSize: Long = DEFAULT_MAX_FRAME_SIZE, block: () -> Server) (and the path variant), passing it through to both the plugin config and the session clamp from the other comment.

@devcrocod devcrocod Oct 22, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DuplicatePluginException isn't new
I'd rather not add a parameter in a security fix, since it needs new public API plus hidden overloads for binary compatibility


/**
* Configures the Ktor Application to handle Model Context Protocol (MCP) over WebSocket.
*
* Set `maxFrameSize` when installing the `WebSockets` plugin: by default Ktor does not limit inbound messages.

@jetbrains-air jetbrains-air Bot Oct 19, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Blocking: the bound is applied only where the SDK itself installs the plugin, so both Route.mcpWebSocket overloads — the ones an app uses when it installs WebSockets itself (for ping config, or because it has other WebSocket routes) — stay completely unbounded after this PR. A KDoc line the user has to read and act on isn't protection, and these overloads are public API listed in kotlin-sdk-server/Module.md.

Failure scenario: install(WebSockets) + routing { mcpWebSocket { server } }. I ran this against this branch: a single 64 MiB ping frame is buffered in full and answered with {"id":1,"result":{},"jsonrpc":"2.0"}. Every open connection can pin arbitrary heap — exactly the DoS this PR exists to close. Note the Application.mcpWebSocket overloads can't be used as the workaround either (see the other comment).

Suggested fix: also clamp on the session itself in mcpWebSocketEndpoint, which covers all four entry points from one place. maxFrameSize is a settable var on WebSocketSession, and Ktor reads it live on every frame — both in the raw reader (readFrame(maxFrameSize, …)) and in DefaultWebSocketSessionImpl.checkMaxFrameSize, which is what caps reassembled fragmented messages. I verified it: adding if (maxFrameSize > DEFAULT_MAX_FRAME_SIZE) maxFrameSize = DEFAULT_MAX_FRAME_SIZE at the top of mcpWebSocketEndpoint turns the probe above into a TOO_BIG close instead of a reply.

@devcrocod devcrocod Oct 22, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Route.mcpWebSocket overload intentionally leaves maxFrameSize to the application. A clamp in mcpWebSocketEndpoint wouldn't prevent early memory allocation because Ktor starts reading frames during WebSocketUpgrade.upgrade(), before the route handler even runs.

*/
@KtorDsl
public fun Route.mcpWebSocket(block: () -> Server) {
Expand All @@ -30,6 +35,8 @@ public fun Route.mcpWebSocket(block: () -> Server) {

/**
* Configures the Ktor Application to handle Model Context Protocol (MCP) over WebSocket.
*
* Set `maxFrameSize` when installing the `WebSockets` plugin: by default Ktor does not limit inbound messages.
*/
@KtorDsl
public fun Route.mcpWebSocket(path: String, block: () -> Server) {
Expand All @@ -40,10 +47,14 @@ public fun Route.mcpWebSocket(path: String, block: () -> Server) {

/**
* Configures the Ktor Application to handle Model Context Protocol (MCP) over WebSocket.
*
* Installs the `WebSockets` plugin with inbound messages limited to 16 MiB.
*/
@KtorDsl
public fun Application.mcpWebSocket(block: () -> Server) {
install(WebSockets)
install(WebSockets) {
maxFrameSize = DEFAULT_MAX_FRAME_SIZE
}

routing {
mcpWebSocket(block)
Expand All @@ -52,10 +63,14 @@ public fun Application.mcpWebSocket(block: () -> Server) {

/**
* Configures the Ktor Application to handle Model Context Protocol (MCP) over WebSocket at the specified path.
*
* Installs the `WebSockets` plugin with inbound messages limited to 16 MiB.
*/
@KtorDsl
public fun Application.mcpWebSocket(path: String, block: () -> Server) {
install(WebSockets)
install(WebSockets) {
maxFrameSize = DEFAULT_MAX_FRAME_SIZE
}

routing {
mcpWebSocket(path, block)
Expand Down
Loading