From 3c4452d18aa44c60419f8d940aa0fd64d0183ac7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=99=8E=E9=B8=A3?= Date: Thu, 29 Oct 2026 22:36:07 +0800 Subject: [PATCH] Fix TLS GraphStage buffer allocation and growth #3728 Motivation: TLS GraphStage cannot initialize when the TCP pool is smaller than the SSLEngine packet requirement, and fixed-size pooled buffers cannot satisfy growth. The partial-record regression test also races handshake processing. Modification: Use stage-owned heap buffers when needed, sharing the initial allocation between three disjoint transport slices. Flush before growing, preserve pending data, refresh session requirements and batching, and shrink drained output after renegotiation. Keep foreign buffers out of the TCP pool. Add buffer and cancellation regressions and document allocation behavior. Synchronize the test harness with session verification and check query counts after cancellation acknowledges completion of unwrap processing. Result: Small TCP buffers no longer prevent TLS progress. Buffer growth and recycling preserve data and ownership, and the partial-record test checks session query behavior without racing the stage. Tests: - JDK 17: 312 TLS/TCP tests passed, including 28 focused buffer/close cases. - Negative control: unconditional underflow session queries fail with 3 != 2; the temporary mutation was restored. - Native scalafmt, headerCreateAll, checkCodeStyle and git diff --check passed. - +mimaReportBinaryIssues passed on Scala 2.13.18 and 3.3.9. - sortImports encountered a Scalafix/Scalameta NoSuchMethodError; its unrelated edits were reverted and the native formatter check passed. - docs/paradox stopped at user request; validatePullRequest and the remaining Scala test matrix are delegated to remote CI. References: Refs #2860, Refs #3727, Refs #3728. --- docs/src/main/paradox/stream/stream-io.md | 13 + .../stream/io/TlsGraphStageIsolatedSpec.scala | 350 +++++++++++++++++- .../pekko/stream/impl/io/TlsGraphStage.scala | 121 ++++-- 3 files changed, 441 insertions(+), 43 deletions(-) diff --git a/docs/src/main/paradox/stream/stream-io.md b/docs/src/main/paradox/stream/stream-io.md index 8053747840..a3545758ae 100644 --- a/docs/src/main/paradox/stream/stream-io.md +++ b/docs/src/main/paradox/stream/stream-io.md @@ -150,6 +150,19 @@ Similar factories as shown above for raw TCP but where the data is encrypted usi through `outgoingConnectionWithTls`, `bindWithTls` and `bindAndHandleWithTls`, see the @scala[@scaladoc[`Tcp Scaladoc`](pekko.stream.scaladsl.Tcp)]@java[@javadoc[`Tcp Javadoc`](pekko.stream.javadsl.Tcp)] for details. +When using `pekko.stream.materializer.tls.engine = graph-stage`, TLS packet buffers are sized to +the requirements of the `SSLEngine`. The `pekko.io.tcp.direct-buffer-size` setting controls the +size of reusable TCP buffers; it does not limit the TLS record size. If a TLS buffer needs more +space than the TCP pool provides, the TLS stage uses heap buffers for that connection, initially +sharing one allocation between its three transport buffers. Packet buffers grow if session requirements +increase, preserving pending ciphertext. After a handshake reduces the required output capacity, the +stage shrinks its output buffer once pending ciphertext has been sent, returning to a pooled buffer +when it fits. A capacity required by an empty-buffer overflow is retained until the next handshake +to avoid repeated allocation. Heap fallback buffers are never returned to the shared TCP pool. +Small transport buffers can restrict plaintext +batches to one TLS record per engine call; batch limits are recalculated from the actual output +capacity as it grows. Measure allocation and throughput for your workload before tuning this setting. + Using TLS requires a keystore and a truststore and then a somewhat involved dance of configuring the SSLEngine and the details for how the session should be negotiated: Scala diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageIsolatedSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageIsolatedSpec.scala index 544ff63fbe..2102ff03b1 100644 --- a/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageIsolatedSpec.scala +++ b/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageIsolatedSpec.scala @@ -21,10 +21,11 @@ import java.nio.ByteBuffer import java.util.concurrent.ThreadFactory import javax.net.ssl._ import javax.net.ssl.SSLEngineResult.HandshakeStatus.{ FINISHED, NEED_UNWRAP, NEED_WRAP, NOT_HANDSHAKING } -import javax.net.ssl.SSLEngineResult.Status.{ CLOSED, OK } +import javax.net.ssl.SSLEngineResult.Status.{ BUFFER_OVERFLOW, BUFFER_UNDERFLOW, CLOSED, OK } import scala.concurrent.Await import scala.concurrent.ExecutionContext +import scala.concurrent.Promise import scala.concurrent.duration._ import scala.util.Success @@ -41,6 +42,7 @@ import pekko.stream.scaladsl._ import pekko.stream.scaladsl.GraphDSL.Implicits._ import pekko.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber } import pekko.stream.testkit.scaladsl.{ TestSink, TestSource } +import pekko.stream.testkit.scaladsl.StreamTestKit.assertAllStagesStopped import pekko.testkit.{ ExplicitlyTriggeredScheduler, WithLogCapturing } import pekko.util.ByteString @@ -247,28 +249,97 @@ class TlsGraphStageIsolatedSpec extends StreamSpec(TlsSpec.configOverrides) with } } -class TlsGraphStageDeferredCloseSpec extends StreamSpec(TlsGraphStageDeferredCloseSpec.config) { +class TlsGraphStageSmallBufferSpec + extends StreamSpec( + ConfigFactory.parseString("pekko.io.tcp.direct-buffer-size = 8192") + .withFallback(ConfigFactory.parseString(TlsSpec.configOverrides))) { + + "TlsGraphStage" must { + for (protocol <- Seq("TLSv1.2", "TLSv1.3")) { + s"round-trip $protocol records larger than the TCP pool buffers" in { + assertAllStagesStopped { + val context = TlsSpec.initSslContext(protocol) + def tls(client: Boolean) = + BidiFlow.fromGraph(new TlsGraphStage( + () => { + val engine = context.createSSLEngine() + engine.setUseClientMode(client) + engine.setEnabledProtocols(Array(protocol)) + engine + }, + _ => Success(()), + IgnoreComplete)).withAttributes(TlsGraphStage.StreamTlsAttributes) + + val payload = ByteString.fromArray(Array.tabulate[Byte](32768)(_.toByte)) + val echo = Flow[SslTlsInbound].collect { case SessionBytes(_, bytes) => SendBytes(bytes) } + val received = Source.single[SslTlsOutbound](SendBytes(payload)) + .via(tls(client = true).atop(tls(client = false).reversed).join(echo)) + .collect { case SessionBytes(_, bytes) => bytes } + .scan(ByteString.empty)(_ ++ _) + .filter(_.size >= payload.size) + .runWith(Sink.head) + + Await.result(received, remainingOrDefault) shouldBe payload + } + + // TLS fallback buffers must not enter the shared, fixed-size TCP pool. + val pool = pekko.io.Tcp(system).bufferPool + val buffer = pool.acquire() + try { + buffer.isDirect shouldBe true + buffer.capacity() shouldBe 8192 + } finally pool.release(buffer) + } + } + } +} - private val TlsRecord = ByteString(Array[Byte](23, 3, 3, 0, 0)) - private val ApplicationCiphertext = ByteString(0x5A.toByte) - private val CloseCiphertext = ByteString(0x5B.toByte) +abstract class TlsGraphStageBufferTestBase(config: Config) extends StreamSpec(config) { - private final class ScriptedSSLEngine( + protected val TlsRecord = ByteString(Array[Byte](23, 3, 3, 0, 0)) + protected val ApplicationCiphertext = ByteString(0x5A.toByte) + protected val CloseCiphertext = ByteString(0x5B.toByte) + + protected final class ScriptedSSLEngine( delegate: SSLEngine, wrapOnCloseInbound: Boolean, requestHandshakeAfterApplication: Boolean, - requestWrapOnSubsequentUnwrap: Boolean) extends SSLEngine { + requestWrapOnSubsequentUnwrap: Boolean, + @volatile var requiredOutputSpace: Int, + negotiatedPacketSize: Int, + negotiatedApplicationSize: Int) extends SSLEngine { private var handshakeStatus: SSLEngineResult.HandshakeStatus = NEED_UNWRAP @volatile var unwrapCount = 0 private var inboundDone = false private var outboundDone = false + private var handshakePending = true @volatile var applicationWraps = 0 @volatile var applicationBytes = 0 @volatile var handshakesStarted = 0 + @volatile var overflowWraps = 0 + @volatile var largestOutputCapacity = 0 + @volatile var inputArray: Array[Byte] = _ + @volatile var outputArray: Array[Byte] = _ + @volatile var underflowUnwraps = 0 + @volatile var overflowUnwraps = 0 + // Keep the scripted session's initial requirements independent of the JDK provider. + private var reportedPacketSize = 8192 + private var reportedApplicationSize = 4096 + private var consecutiveBufferRetries = 0 + @volatile var sessionQueries = 0 + private lazy val reportedSession = java.lang.reflect.Proxy.newProxyInstance( + classOf[SSLSession].getClassLoader, + Array[Class[?]](classOf[SSLSession]), + (_: AnyRef, method: java.lang.reflect.Method, args: Array[AnyRef]) => + if (method.getName == "getPacketBufferSize") Int.box(reportedPacketSize) + else if (method.getName == "getApplicationBufferSize") Int.box(reportedApplicationSize) + else method.invoke(delegate.getSession, Option(args).getOrElse(Array.empty[AnyRef]): _*)) + .asInstanceOf[SSLSession] override def beginHandshake(): Unit = { handshakesStarted += 1 + handshakePending = true handshakeStatus = NEED_UNWRAP } @@ -279,6 +350,12 @@ class TlsGraphStageDeferredCloseSpec extends StreamSpec(TlsGraphStageDeferredClo offset: Int, length: Int, dst: ByteBuffer): SSLEngineResult = { + largestOutputCapacity = math.max(largestOutputCapacity, dst.capacity()) + if (dst.hasArray) outputArray = dst.array() + if (dst.remaining() < requiredOutputSpace) { + overflowWraps += 1 + return new SSLEngineResult(BUFFER_OVERFLOW, handshakeStatus, 0, 0) + } var consumed = 0 var index = offset while (index < offset + length) { @@ -311,13 +388,34 @@ class TlsGraphStageDeferredCloseSpec extends StreamSpec(TlsGraphStageDeferredClo dsts: Array[ByteBuffer], offset: Int, length: Int): SSLEngineResult = { + if (src.hasArray) inputArray = src.array() + if (!handshakePending && negotiatedPacketSize > 0) { + reportedPacketSize = negotiatedPacketSize + if (negotiatedApplicationSize > 0) reportedApplicationSize = negotiatedApplicationSize + consecutiveBufferRetries += 1 + if (consecutiveBufferRetries > 8) throw new SSLException("TLS buffers were not resized after repeated retries") + if (dsts(offset).remaining() < negotiatedApplicationSize) { + overflowUnwraps += 1 + return new SSLEngineResult(BUFFER_OVERFLOW, NOT_HANDSHAKING, 0, 0) + } + if (src.remaining() < negotiatedPacketSize) { + underflowUnwraps += 1 + return new SSLEngineResult(BUFFER_UNDERFLOW, NOT_HANDSHAKING, 0, 0) + } + } + consecutiveBufferRetries = 0 val consumed = src.remaining() src.position(src.limit()) unwrapCount += 1 - if (unwrapCount == 1) { + if (handshakePending) { + handshakePending = false handshakeStatus = NOT_HANDSHAKING new SSLEngineResult(OK, FINISHED, consumed, 0) + } else if (negotiatedPacketSize > 0) { + val produced = math.max(1, negotiatedApplicationSize) + dsts(offset).put(Array.fill[Byte](produced)(0x41.toByte)) + new SSLEngineResult(OK, NOT_HANDSHAKING, consumed, produced) } else if (requestWrapOnSubsequentUnwrap) { handshakeStatus = NEED_WRAP new SSLEngineResult(OK, NEED_WRAP, consumed, 0) @@ -340,7 +438,10 @@ class TlsGraphStageDeferredCloseSpec extends StreamSpec(TlsGraphStageDeferredClo override def isInboundDone: Boolean = inboundDone override def isOutboundDone: Boolean = outboundDone override def getDelegatedTask: Runnable = null - override def getSession: SSLSession = delegate.getSession + override def getSession: SSLSession = { + sessionQueries += 1 + if (negotiatedPacketSize > 0) reportedSession else delegate.getSession + } override def getSupportedCipherSuites: Array[String] = delegate.getSupportedCipherSuites override def getEnabledCipherSuites: Array[String] = delegate.getEnabledCipherSuites override def setEnabledCipherSuites(suites: Array[String]): Unit = delegate.setEnabledCipherSuites(suites) @@ -357,22 +458,28 @@ class TlsGraphStageDeferredCloseSpec extends StreamSpec(TlsGraphStageDeferredClo override def getEnableSessionCreation: Boolean = delegate.getEnableSessionCreation } - private final class DeferredCloseHarness( + protected final class DeferredCloseHarness( val engine: ScriptedSSLEngine, + val sessionVerified: Promise[Unit], val plainIn: TestPublisher.Probe[SslTlsOutbound], val cipherIn: TestPublisher.Probe[ByteString], val cipherOut: TestSubscriber.Probe[ByteString], val plainOut: TestSubscriber.Probe[SslTlsInbound]) - private def newHarness( + protected def newHarness( wrapOnCloseInbound: Boolean = false, closing: TLSClosing = IgnoreComplete, requestHandshakeAfterApplication: Boolean = false, requestWrapOnSubsequentUnwrap: Boolean = true, - extraAttributes: Attributes = Attributes.none): DeferredCloseHarness = { + extraAttributes: Attributes = Attributes.none, + requiredOutputSpace: Int = 0, + negotiatedPacketSize: Int = 0, + negotiatedApplicationSize: Int = 0): DeferredCloseHarness = { val engine = new ScriptedSSLEngine(TlsSpec.initSslContext("TLSv1.2").createSSLEngine(), wrapOnCloseInbound, - requestHandshakeAfterApplication, requestWrapOnSubsequentUnwrap) + requestHandshakeAfterApplication, requestWrapOnSubsequentUnwrap, + requiredOutputSpace, negotiatedPacketSize, negotiatedApplicationSize) + val sessionVerified = Promise[Unit]() val probes = RunnableGraph .fromGraph( @@ -383,7 +490,13 @@ class TlsGraphStageDeferredCloseSpec extends StreamSpec(TlsGraphStageDeferredClo TestSink[SslTlsInbound]())(Tuple4.apply) { implicit b => (plainIn, cipherIn, cipherOut, plainOut) => val tls = b.add( BidiFlow - .fromGraph(new TlsGraphStage(() => engine, _ => Success(()), closing)) + .fromGraph(new TlsGraphStage( + () => engine, + _ => { + sessionVerified.trySuccess(()) + Success(()) + }, + closing)) .withAttributes(Attributes.inputBuffer(1, 1).and(extraAttributes))) plainIn ~> tls.in1 tls.out1 ~> cipherOut @@ -397,7 +510,7 @@ class TlsGraphStageDeferredCloseSpec extends StreamSpec(TlsGraphStageDeferredClo }) .run() - val harness = new DeferredCloseHarness(engine, probes._1, probes._2, probes._3, probes._4) + val harness = new DeferredCloseHarness(engine, sessionVerified, probes._1, probes._2, probes._3, probes._4) harness.plainIn.expectRequest() harness.cipherIn.expectRequest() harness.cipherOut.ensureSubscription() @@ -405,13 +518,19 @@ class TlsGraphStageDeferredCloseSpec extends StreamSpec(TlsGraphStageDeferredClo harness } - private def finishInitialHandshake(harness: DeferredCloseHarness): Unit = { + protected def finishInitialHandshake(harness: DeferredCloseHarness): Unit = { harness.plainOut.request(1) harness.cipherIn.sendNext(TlsRecord) harness.cipherIn.expectRequest() awaitAssert(harness.engine.unwrapCount shouldBe 1) + // unwrapCount is incremented before the stage processes FINISHED and queries the session. + Await.result(harness.sessionVerified.future, remainingOrDefault) } +} + +class TlsGraphStageDeferredCloseSpec extends TlsGraphStageBufferTestBase(TlsGraphStageDeferredCloseSpec.config) { + "TlsGraphStage deferred close handling" must { for (closing <- Seq(IgnoreComplete, EagerClose)) { @@ -433,6 +552,25 @@ class TlsGraphStageDeferredCloseSpec extends StreamSpec(TlsGraphStageDeferredClo harness.engine.applicationBytes shouldBe 1 } + s"grow the transport buffer while draining accepted plaintext on output cancellation with $closing" in { + val poolBufferSize = pekko.io.Tcp(system).Settings.DirectBufferSize + val harness = newHarness(closing = closing, requiredOutputSpace = poolBufferSize + 1) + finishInitialHandshake(harness) + harness.plainIn.sendNext(SendBytes(ByteString("a"))) + harness.plainIn.expectRequest() + harness.plainOut.cancel() + harness.plainIn.expectCancellation() + + harness.engine.applicationWraps shouldBe 0 + harness.cipherOut.request(1) + harness.cipherOut.expectNext(ApplicationCiphertext) + harness.cipherOut.expectComplete() + harness.cipherIn.expectCancellation() + harness.engine.applicationBytes shouldBe 1 + harness.engine.overflowWraps shouldBe 1 + harness.engine.largestOutputCapacity should be > poolBufferSize + } + s"drain a partially sent large message on output cancellation with $closing" in { val harness = newHarness(closing = closing) finishInitialHandshake(harness) @@ -646,6 +784,182 @@ class TlsGraphStageDeferredCloseSpec extends StreamSpec(TlsGraphStageDeferredClo harness.plainOut.expectComplete() harness.plainIn.expectCancellation() } + + "grow an empty transport buffer and return only unique pooled buffers after termination" in { + val tcp = pekko.io.Tcp(system) + val poolBufferSize = tcp.Settings.DirectBufferSize + assertAllStagesStopped { + val harness = newHarness(requiredOutputSpace = poolBufferSize + 1) + finishInitialHandshake(harness) + + harness.plainIn.sendNext(SendBytes(ByteString.fromArray(Array.fill[Byte](2048)(1)))) + harness.cipherOut.request(1) + harness.cipherOut.expectNext(ApplicationCiphertext) + harness.engine.overflowWraps shouldBe 1 + harness.engine.largestOutputCapacity should be > poolBufferSize + + harness.cipherIn.sendComplete() + harness.cipherOut.expectComplete() + harness.plainOut.expectComplete() + harness.plainIn.expectCancellation() + } + + // Hold the entire free list at once to detect oversized buffers and duplicate releases. + val buffers = Vector.fill(tcp.Settings.MaxDirectBufferPoolSize)(tcp.bufferPool.acquire()) + try { + buffers.foreach { buffer => + buffer.isDirect shouldBe true + buffer.capacity() shouldBe poolBufferSize + } + for (i <- buffers.indices; j <- 0 until i) + (buffers(i) eq buffers(j)) shouldBe false + } finally buffers.foreach(tcp.bufferPool.release) + } + + "return to a pooled output buffer after a handshake lowers the capacity requirement" in { + val poolBufferSize = pekko.io.Tcp(system).Settings.DirectBufferSize + val harness = newHarness(requiredOutputSpace = poolBufferSize + 1) + finishInitialHandshake(harness) + val payload = SendBytes(ByteString.fromArray(Array.fill[Byte](2048)(1))) + + for (_ <- 1 to 2) { + harness.plainIn.sendNext(payload) + harness.cipherOut.request(1) + harness.cipherOut.expectNext(ApplicationCiphertext) + } + // A real empty-buffer overflow sets a lower bound until the next completed handshake. + harness.engine.overflowWraps shouldBe 1 + harness.engine.largestOutputCapacity should be > poolBufferSize + harness.engine.requiredOutputSpace = 0 + harness.engine.largestOutputCapacity = 0 + + harness.cipherOut.request(1) + harness.plainIn.sendNext(NegotiateNewSession) + awaitAssert(harness.engine.handshakesStarted shouldBe 2) + harness.cipherIn.sendNext(TlsRecord) + harness.cipherIn.expectRequest() + awaitAssert(harness.engine.unwrapCount shouldBe 2) + harness.engine.largestOutputCapacity = 0 + harness.plainIn.sendNext(payload) + harness.cipherOut.expectNext(ApplicationCiphertext) + harness.engine.largestOutputCapacity shouldBe poolBufferSize + harness.engine.overflowWraps shouldBe 1 + + harness.cipherIn.sendComplete() + harness.cipherOut.expectComplete() + harness.plainOut.expectComplete() + harness.plainIn.expectCancellation() + } + + "flush deferred cipher bytes before retrying an overflow without growing the buffer" in { + val poolBufferSize = system.settings.config.getBytes("pekko.io.tcp.direct-buffer-size").toInt + val harness = newHarness(requiredOutputSpace = poolBufferSize) + finishInitialHandshake(harness) + + harness.cipherIn.sendNext(TlsRecord) // bypass the user batch timer for the first write + harness.cipherIn.expectRequest() + awaitAssert(harness.engine.unwrapCount shouldBe 2) + harness.plainIn.sendNext(SendBytes(ByteString("a"))) + harness.plainIn.expectRequest() + harness.cipherOut.request(1) + awaitAssert(harness.engine.applicationWraps shouldBe 1) + + // The overflow is caused by buffered ciphertext, not an undersized empty buffer. + harness.plainIn.sendNext(SendBytes(ByteString.fromArray(Array.fill[Byte](2048)(1)))) + harness.cipherOut.expectNext(ApplicationCiphertext) + harness.cipherOut.request(1) + harness.cipherOut.expectNext(ApplicationCiphertext) + harness.engine.overflowWraps shouldBe 1 + harness.engine.applicationWraps shouldBe 2 + harness.engine.largestOutputCapacity shouldBe poolBufferSize + + harness.cipherIn.sendComplete() + harness.cipherOut.expectComplete() + harness.plainOut.expectComplete() + harness.plainIn.expectCancellation() + } + } +} + +class TlsGraphStageHeapBufferSpec + extends TlsGraphStageBufferTestBase( + ConfigFactory.parseString("pekko.io.tcp.direct-buffer-size = 4096") + .withFallback(TlsGraphStageDeferredCloseSpec.config)) { + + "TlsGraphStage heap buffers" must { + "share one allocation between the three initial transport buffers" in { + val harness = newHarness(negotiatedPacketSize = 32768) + finishInitialHandshake(harness) + harness.plainIn.sendNext(SendBytes(ByteString.fromArray(Array.fill[Byte](2048)(1)))) + harness.cipherOut.request(1) + harness.cipherOut.expectNext(ApplicationCiphertext) + + harness.engine.inputArray.length shouldBe 8192 * 3 + (harness.engine.inputArray eq harness.engine.outputArray) shouldBe true + harness.cipherIn.sendComplete() + harness.cipherOut.expectComplete() + harness.plainOut.expectComplete() + harness.plainIn.expectCancellation() + } + } +} + +class TlsGraphStageBufferGrowthSpec + extends TlsGraphStageBufferTestBase( + ConfigFactory.parseString("pekko.io.tcp.direct-buffer-size = 18432") + .withFallback(TlsGraphStageDeferredCloseSpec.config)) { + + "TlsGraphStage buffer growth" must { + "avoid querying session sizes for an ordinary partial record" in { + val harness = newHarness(negotiatedPacketSize = 32768) + finishInitialHandshake(harness) + val sessionQueries = harness.engine.sessionQueries + harness.cipherIn.sendNext(TlsRecord) + harness.cipherIn.expectRequest() + awaitAssert(harness.engine.underflowUnwraps shouldBe 1) + harness.plainOut.cancel() + harness.cipherOut.cancel() + harness.cipherIn.expectCancellation() + harness.plainIn.expectCancellation() + // Cancellation is processed after the stage handles the BUFFER_UNDERFLOW result. + harness.engine.sessionQueries shouldBe sessionQueries + } + + for (applicationSize <- Seq(0, 32700)) { + s"grow inbound buffers after session sizes increase (application size $applicationSize)" in { + val packetSize = 32768 + val harness = newHarness(negotiatedPacketSize = packetSize, negotiatedApplicationSize = applicationSize) + finishInitialHandshake(harness) + val bodySize = packetSize - 5 + val record = ByteString(23.toByte, 3.toByte, 3.toByte, (bodySize >> 8).toByte, bodySize.toByte) ++ + ByteString.fromArray(Array.fill[Byte](bodySize)(1)) + + harness.cipherIn.sendNext(record) + harness.plainOut.expectNext().asInstanceOf[SessionBytes].bytes shouldBe + ByteString.fromArray(Array.fill[Byte](math.max(1, applicationSize))(0x41.toByte)) + harness.engine.unwrapCount shouldBe 2 + if (applicationSize == 0) harness.engine.underflowUnwraps should be > 0 + else harness.engine.overflowUnwraps should be > 0 + + harness.plainIn.sendNext(SendBytes(ByteString.fromArray(Array.fill[Byte](65536)(1)))) + harness.cipherOut.request(1) + harness.cipherOut.expectNext(ApplicationCiphertext) + harness.engine.largestOutputCapacity should be >= packetSize + + // Successful records reset the retry guard; the connection can process more than eight records. + for (_ <- 1 to 10) { + harness.plainOut.request(1) + harness.cipherIn.sendNext(record) + harness.plainOut.expectNext().asInstanceOf[SessionBytes].bytes.size shouldBe math.max(1, applicationSize) + } + harness.plainOut.request(1) + harness.cipherIn.sendComplete() + harness.cipherOut.expectComplete() + harness.plainOut.expectComplete() + harness.plainIn.expectCancellation() + } + } + } } @@ -661,6 +975,8 @@ object TlsGraphStageDeferredCloseSpec { val config: Config = ConfigFactory .parseString( - s"""pekko.scheduler.implementation = "${classOf[DeferredZeroTimerScheduler].getName}"""") + s"""pekko.scheduler.implementation = "${classOf[DeferredZeroTimerScheduler].getName}" + |pekko.io.tcp.direct-buffer-pool-limit = 8 + |""".stripMargin) .withFallback(ConfigFactory.parseString(TlsSpec.configOverrides)) } diff --git a/stream/src/main/scala/org/apache/pekko/stream/impl/io/TlsGraphStage.scala b/stream/src/main/scala/org/apache/pekko/stream/impl/io/TlsGraphStage.scala index 79a5dd5ea9..66f48dc4b3 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/impl/io/TlsGraphStage.scala +++ b/stream/src/main/scala/org/apache/pekko/stream/impl/io/TlsGraphStage.scala @@ -82,6 +82,7 @@ import pekko.util.ByteString private var bufferPool: BufferPool = _ private var pooledBufferSize = 0 private var transportPacketSize = 0 + private var minimumTransportOutBufferSize = 0 private var applicationRecordSize = 0 private var maxUserInBufferSize = 0 private var maxUserOutBufferSize = 0 @@ -136,15 +137,21 @@ import pekko.util.ByteString private def allocateBuffers(session: SSLSession): Unit = { val packetSize = math.max(1, session.getPacketBufferSize) val applicationSize = math.max(1, session.getApplicationBufferSize) - val applicationBatchSize = applicationBufferSize(applicationSize, packetSize) transportPacketSize = packetSize + minimumTransportOutBufferSize = packetSize applicationRecordSize = applicationSize - maxUserInBufferSize = applicationBatchSize - maxUserOutBufferSize = applicationBatchSize - - transportOutBuffer = acquireTransportBuffer(packetSize) - transportInBuffer = acquireTransportBuffer(packetSize) - transportUnreadBuffer = acquireTransportBuffer(packetSize) + if (packetSize > pooledBufferSize && packetSize <= Int.MaxValue / 3) { + // Share one heap allocation without allowing the three buffers to overwrite each other. + val transportBuffers = ByteBuffer.allocate(packetSize * 3) + transportOutBuffer = transportBuffers.slice(0, packetSize) + transportInBuffer = transportBuffers.slice(packetSize, packetSize) + transportUnreadBuffer = transportBuffers.slice(packetSize * 2, packetSize) + } else { + transportOutBuffer = acquireTransportBuffer(packetSize) + transportInBuffer = acquireTransportBuffer(packetSize) + transportUnreadBuffer = acquireTransportBuffer(packetSize) + } + updateApplicationBufferSizes() userInBuffer = ByteBuffer.allocate(applicationSize) userOutBuffer = ByteBuffer.allocate(applicationSize) @@ -154,11 +161,17 @@ import pekko.util.ByteString } private def applicationBufferSize(applicationSize: Int, packetSize: Int): Int = { - val recordsThatFitTransportBuffer = math.max(1, pooledBufferSize / packetSize) + val recordsThatFitTransportBuffer = math.max(1, transportOutBuffer.capacity() / packetSize) val recordsPerEngineCall = math.min(MaxApplicationRecordsPerEngineCall, recordsThatFitTransportBuffer) applicationSize * recordsPerEngineCall } + private def updateApplicationBufferSizes(): Unit = { + val applicationBatchSize = applicationBufferSize(applicationRecordSize, transportPacketSize) + maxUserInBufferSize = applicationBatchSize + maxUserOutBufferSize = applicationBatchSize + } + private def runPump(): Unit = if ((stateBits & (StageClosedFlag | StageCompletionPendingFlag)) == 0 && (stateBits & InitialPumpEnabledFlag) != 0) { @@ -421,7 +434,8 @@ import pekko.util.ByteString else nextPhase(AwaitingClosePhase) case BUFFER_OVERFLOW => - growTransportOutBuffer() + if (transportOutBuffer.position() > 0) flushToTransport() + else growTransportOutBuffer() case status => failTls(new IllegalStateException(s"unexpected status $status in doWrap()")) @@ -456,6 +470,7 @@ import pekko.util.ByteString growUserOutBuffer() val oldInPosition = transportInBuffer.position() + val emptyUserOutBuffer = userOutBuffer.position() == 0 val result = engine.unwrap(transportInBuffer, userOutBuffer) if (ignoreOutput) userOutBuffer.clear() lastHandshakeStatus = result.getHandshakeStatus @@ -496,9 +511,11 @@ import pekko.util.ByteString case BUFFER_UNDERFLOW => flushToUser() + if (transportInBuffer.remaining() == transportInBuffer.capacity()) ensureInboundBufferSizes() case BUFFER_OVERFLOW => flushToUser() + if (emptyUserOutBuffer) ensureInboundBufferSizes() transportInput.putBackUnreadBuffer(transportInBuffer) case null => @@ -506,6 +523,33 @@ import pekko.util.ByteString } } + private def ensureInboundBufferSizes(): Unit = ensureInboundBufferSizes(engine.getSession) + + private def ensureInboundBufferSizes(session: SSLSession): Unit = { + // Session requirements can increase while an incoming record is being inspected. + val packetSize = math.max(1, session.getPacketBufferSize) + transportInBuffer = ensureReadBufferCapacity(transportInBuffer, packetSize) + transportUnreadBuffer = ensureReadBufferCapacity(transportUnreadBuffer, packetSize) + transportPacketSize = math.max(transportPacketSize, packetSize) + minimumTransportOutBufferSize = math.max(minimumTransportOutBufferSize, transportPacketSize) + ensureTransportOutBufferCapacity(transportPacketSize) + + val applicationSize = math.max(1, session.getApplicationBufferSize) + applicationRecordSize = math.max(applicationRecordSize, applicationSize) + updateApplicationBufferSizes() + if (userOutBuffer.capacity() < applicationSize) growUserOutBuffer() + } + + private def ensureReadBufferCapacity(buffer: ByteBuffer, requiredCapacity: Int): ByteBuffer = + if (buffer.capacity() >= requiredCapacity) buffer + else { + val expanded = acquireTransportBuffer(requiredCapacity) + expanded.put(buffer) + expanded.flip() + releaseTransportBuffer(buffer) + expanded + } + private def runDelegatedTasks(): Unit = if (TlsEngineHelpers.runDelegatedTasks(engine) > 0 || lastHandshakeStatus == NEED_TASK) lastHandshakeStatus = engine.getHandshakeStatus @@ -516,6 +560,12 @@ import pekko.util.ByteString case Success(()) => currentSession = session stateBits |= PlainDataAllowedFlag + transportPacketSize = math.max(1, session.getPacketBufferSize) + applicationRecordSize = math.max(1, session.getApplicationBufferSize) + // An empty-buffer overflow remains a lower bound until a new session is established. + minimumTransportOutBufferSize = transportPacketSize + ensureInboundBufferSizes(session) + shrinkTransportOutBufferIfNeeded() flushToUser() true @@ -545,6 +595,7 @@ import pekko.util.ByteString if (transportOutBuffer.hasRemaining && !transportOutput.isClosed) transportOutput.enqueue(ByteString(transportOutBuffer)) transportOutBuffer.clear() + shrinkTransportOutBufferIfNeeded() } private def flushToTransportIfNeeded(result: SSLEngineResult): Unit = @@ -599,18 +650,35 @@ import pekko.util.ByteString } private def growTransportOutBuffer(): Unit = { - val oldBuffer = transportOutBuffer - val oldCapacity = oldBuffer.capacity() + val oldCapacity = transportOutBuffer.capacity() if (oldCapacity > Int.MaxValue / 2) throw new IllegalStateException(s"Cannot grow TLS transport output buffer beyond $oldCapacity bytes") - val bigger = acquireTransportBuffer(oldCapacity * 2) - transportOutBuffer.flip() - bigger.put(transportOutBuffer) - releaseTransportBuffer(oldBuffer) - transportOutBuffer = bigger + minimumTransportOutBufferSize = oldCapacity * 2 + ensureTransportOutBufferCapacity(minimumTransportOutBufferSize) + updateApplicationBufferSizes() + } + + private def shrinkTransportOutBufferIfNeeded(): Unit = { + val requiredCapacity = math.max(pooledBufferSize, minimumTransportOutBufferSize) + if (transportOutBuffer.position() == 0 && transportOutBuffer.capacity() > requiredCapacity) { + val oldBuffer = transportOutBuffer + transportOutBuffer = acquireTransportBuffer(requiredCapacity) + releaseTransportBuffer(oldBuffer) + updateApplicationBufferSizes() + } } + private def ensureTransportOutBufferCapacity(requiredCapacity: Int): Unit = + if (transportOutBuffer.capacity() < requiredCapacity) { + val oldBuffer = transportOutBuffer + val bigger = acquireTransportBuffer(requiredCapacity) + oldBuffer.flip() + bigger.put(oldBuffer) + transportOutBuffer = bigger + releaseTransportBuffer(oldBuffer) + } + private def failTls(ex: Throwable): Unit = failTls(ex, closeTransport = true) private def failTls(ex: Throwable, closeTransport: Boolean): Unit = @@ -702,18 +770,19 @@ import pekko.util.ByteString userInBuffer = null } - private def acquireTransportBuffer(requiredCapacity: Int): ByteBuffer = { - val buffer = bufferPool.acquire() - if (buffer.capacity >= requiredCapacity) buffer - else { - releaseTransportBuffer(buffer) - throw new IllegalStateException( - s"TLS packet buffer requires $requiredCapacity bytes but pekko.io.tcp.direct-buffer-size is $pooledBufferSize") - } - } + private def acquireTransportBuffer(requiredCapacity: Int): ByteBuffer = + if (requiredCapacity <= pooledBufferSize) { + val buffer = bufferPool.acquire() + if (buffer.capacity() >= requiredCapacity) buffer + else { + releaseTransportBuffer(buffer) + ByteBuffer.allocate(requiredCapacity) + } + } else ByteBuffer.allocate(requiredCapacity) private def releaseTransportBuffer(buffer: ByteBuffer): Unit = - if (buffer ne null) { + // Stage-owned heap buffers and buffers with another size must not enter the fixed-size TCP pool. + if ((buffer ne null) && buffer.isDirect && buffer.capacity() == pooledBufferSize) { buffer.clear() bufferPool.release(buffer) }