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