From 211ec6967463751dc35109904d83783c4e69c062 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=99=8E=E9=B8=A3?= Date: Sat, 11 Jul 2026 15:21:23 +0800 Subject: [PATCH] fix(stream): handle deferred TLS output on inbound close Motivation: The GraphStage TLS engine can receive inbound completion while plaintext batching or ciphertext flushing is deferred. Its close state can then stall or complete before pending protocol output is emitted. Modification: Flush buffered ciphertext before choosing the close phase and base completion on whether the SSLEngine still requires a wrap. Add deterministic tests that hold zero-delay timers and exercise pending plaintext, buffered ciphertext, and engine wrap output. Result: Inbound close no longer strands the stage in FlushingOutbound, buffered ciphertext is emitted before completion, and pending engine output is wrapped before shutdown. Tests: - rtk sbt "stream-tests / Test / testOnly org.apache.pekko.stream.io.TlsGraphStageDeferredCloseSpec" (3 passed) - rtk sbt "stream-tests / Test / testOnly org.apache.pekko.stream.io.TlsGraphStageIsolatedSpec" (7 passed) - rtk scalafmt --list --mode diff-ref=origin/main - rtk sbt "headerCreateAll" "+headerCheckAll" - rtk git diff --check References: Refs #2860 --- .../stream/io/TlsGraphStageIsolatedSpec.scala | 221 +++++++++++++++++- .../pekko/stream/impl/io/TlsGraphStage.scala | 8 +- 2 files changed, 225 insertions(+), 4 deletions(-) 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 53390bef9fa..1ddc3d78b4f 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 @@ -17,21 +17,31 @@ package org.apache.pekko.stream.io +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 scala.concurrent.Await +import scala.concurrent.ExecutionContext import scala.concurrent.duration._ import scala.util.Success +import com.typesafe.config.{ Config, ConfigFactory } + import org.apache.pekko import pekko.NotUsed +import pekko.actor.Cancellable +import pekko.event.LoggingAdapter import pekko.stream._ import pekko.stream.TLSProtocol._ import pekko.stream.impl.io.TlsGraphStage import pekko.stream.scaladsl._ import pekko.stream.scaladsl.GraphDSL.Implicits._ -import pekko.stream.testkit.{ StreamSpec, TestPublisher } -import pekko.testkit.WithLogCapturing +import pekko.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber } +import pekko.stream.testkit.scaladsl.{ TestSink, TestSource } +import pekko.testkit.{ ExplicitlyTriggeredScheduler, WithLogCapturing } import pekko.util.ByteString /** @@ -236,3 +246,210 @@ class TlsGraphStageIsolatedSpec extends StreamSpec(TlsSpec.configOverrides) with } } } + +class TlsGraphStageDeferredCloseSpec extends StreamSpec(TlsGraphStageDeferredCloseSpec.config) { + + private val TlsRecord = ByteString(Array[Byte](23, 3, 3, 0, 0)) + private val ApplicationCiphertext = ByteString(0x5A.toByte) + private val CloseCiphertext = ByteString(0x5B.toByte) + + private final class ScriptedSSLEngine(delegate: SSLEngine, wrapOnCloseInbound: Boolean) extends SSLEngine { + private var handshakeStatus: SSLEngineResult.HandshakeStatus = NEED_UNWRAP + @volatile var unwrapCount = 0 + private var inboundDone = false + private var outboundDone = false + + @volatile var applicationWraps = 0 + + override def beginHandshake(): Unit = handshakeStatus = NEED_UNWRAP + + override def getHandshakeStatus: SSLEngineResult.HandshakeStatus = handshakeStatus + + override def wrap( + srcs: Array[ByteBuffer], + offset: Int, + length: Int, + dst: ByteBuffer): SSLEngineResult = { + var consumed = 0 + var index = offset + while (index < offset + length) { + val src = srcs(index) + consumed += src.remaining() + src.position(src.limit()) + index += 1 + } + + val closeWrap = wrapOnCloseInbound && inboundDone && handshakeStatus == NEED_WRAP + val produced = + if (closeWrap) { + dst.put(CloseCiphertext.head) + outboundDone = true + 1 + } else if (consumed > 0) { + applicationWraps += 1 + dst.put(ApplicationCiphertext.head) + 1 + } else 0 + + handshakeStatus = NOT_HANDSHAKING + new SSLEngineResult(if (closeWrap) CLOSED else OK, NOT_HANDSHAKING, consumed, produced) + } + + override def unwrap( + src: ByteBuffer, + dsts: Array[ByteBuffer], + offset: Int, + length: Int): SSLEngineResult = { + val consumed = src.remaining() + src.position(src.limit()) + unwrapCount += 1 + + if (unwrapCount == 1) { + handshakeStatus = NOT_HANDSHAKING + new SSLEngineResult(OK, FINISHED, consumed, 0) + } else { + handshakeStatus = NEED_WRAP + new SSLEngineResult(OK, NEED_WRAP, consumed, 0) + } + } + + override def closeInbound(): Unit = { + inboundDone = true + handshakeStatus = if (wrapOnCloseInbound) NEED_WRAP else NOT_HANDSHAKING + } + + override def closeOutbound(): Unit = { + outboundDone = true + handshakeStatus = NOT_HANDSHAKING + } + + override def isInboundDone: Boolean = inboundDone + override def isOutboundDone: Boolean = outboundDone + override def getDelegatedTask: Runnable = null + override def getSession: SSLSession = 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) + override def getSupportedProtocols: Array[String] = delegate.getSupportedProtocols + override def getEnabledProtocols: Array[String] = delegate.getEnabledProtocols + override def setEnabledProtocols(protocols: Array[String]): Unit = delegate.setEnabledProtocols(protocols) + override def setUseClientMode(mode: Boolean): Unit = delegate.setUseClientMode(mode) + override def getUseClientMode: Boolean = delegate.getUseClientMode + override def setNeedClientAuth(need: Boolean): Unit = delegate.setNeedClientAuth(need) + override def getNeedClientAuth: Boolean = delegate.getNeedClientAuth + override def setWantClientAuth(want: Boolean): Unit = delegate.setWantClientAuth(want) + override def getWantClientAuth: Boolean = delegate.getWantClientAuth + override def setEnableSessionCreation(flag: Boolean): Unit = delegate.setEnableSessionCreation(flag) + override def getEnableSessionCreation: Boolean = delegate.getEnableSessionCreation + } + + private final class DeferredCloseHarness( + val engine: ScriptedSSLEngine, + val plainIn: TestPublisher.Probe[SslTlsOutbound], + val cipherIn: TestPublisher.Probe[ByteString], + val cipherOut: TestSubscriber.Probe[ByteString], + val plainOut: TestSubscriber.Probe[SslTlsInbound]) + + private def newHarness(wrapOnCloseInbound: Boolean = false): DeferredCloseHarness = { + val engine = + new ScriptedSSLEngine(TlsSpec.initSslContext("TLSv1.2").createSSLEngine(), wrapOnCloseInbound) + val probes = + RunnableGraph + .fromGraph( + GraphDSL.createGraph( + TestSource[SslTlsOutbound](), + TestSource[ByteString](), + TestSink[ByteString](), + TestSink[SslTlsInbound]())(Tuple4.apply) { implicit b => (plainIn, cipherIn, cipherOut, plainOut) => + val tls = b.add( + BidiFlow + .fromGraph(new TlsGraphStage(() => engine, _ => Success(()), IgnoreComplete)) + .withAttributes(Attributes.inputBuffer(1, 1))) + plainIn ~> tls.in1 + tls.out1 ~> cipherOut + cipherIn ~> tls.in2 + tls.out2 ~> plainOut + ClosedShape + }) + .run() + + val harness = new DeferredCloseHarness(engine, probes._1, probes._2, probes._3, probes._4) + harness.plainIn.expectRequest() + harness.cipherIn.expectRequest() + harness.cipherOut.ensureSubscription() + harness.plainOut.ensureSubscription() + harness + } + + private def finishInitialHandshake(harness: DeferredCloseHarness): Unit = { + harness.plainOut.request(1) + harness.cipherIn.sendNext(TlsRecord) + harness.cipherIn.expectRequest() + awaitAssert(harness.engine.unwrapCount shouldBe 1) + } + + "TlsGraphStage deferred close handling" must { + "complete when inbound closes before a deferred user batch runs" in { + val harness = newHarness() + finishInitialHandshake(harness) + + harness.plainIn.sendNext(SendBytes(ByteString("a"))) + harness.plainIn.expectRequest() // proves onPush ran before cipherIn completion + harness.cipherIn.sendComplete() + + harness.cipherOut.expectComplete() + harness.plainOut.expectComplete() + harness.plainIn.expectCancellation() + } + + "flush deferred cipher bytes before completing after inbound close" in { + val harness = newHarness() + finishInitialHandshake(harness) + + harness.cipherIn.sendNext(TlsRecord) // make the small user input bypass the user batch timer + 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) + harness.cipherOut.expectNoMessage(100.millis) + + harness.cipherIn.sendComplete() + + harness.cipherOut.expectNext(ApplicationCiphertext) + harness.cipherOut.expectComplete() + harness.plainOut.expectComplete() + harness.plainIn.expectCancellation() + } + + "wrap pending engine output before completing after inbound close" in { + val harness = newHarness(wrapOnCloseInbound = true) + finishInitialHandshake(harness) + + harness.cipherOut.request(1) + harness.cipherIn.sendComplete() + + harness.cipherOut.expectNext(CloseCiphertext) + harness.cipherOut.expectComplete() + harness.plainOut.expectComplete() + harness.plainIn.expectCancellation() + } + } +} + +class DeferredZeroTimerScheduler(config: Config, log: LoggingAdapter, threadFactory: ThreadFactory) + extends ExplicitlyTriggeredScheduler(config, log, threadFactory) { + + override def scheduleOnce(delay: FiniteDuration, runnable: Runnable)(implicit + executor: ExecutionContext): Cancellable = + super.scheduleOnce(if (delay <= Duration.Zero) 1.millis else delay, runnable) +} + +object TlsGraphStageDeferredCloseSpec { + val config: Config = + ConfigFactory + .parseString( + s"""pekko.scheduler.implementation = "${classOf[DeferredZeroTimerScheduler].getName}"""") + .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 65e3a9110cb..ab03f2e3cef 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 @@ -300,10 +300,14 @@ import pekko.util.ByteString transportInput.pullIfNeeded(transportInput.isEmpty) } - private def completeOrFlush(): Unit = - if (engine.isOutboundDone || (engine.isInboundDone && userInput.isEmpty && !userInBufferHasRemaining)) + private def completeOrFlush(): Unit = { + if (transportOutBuffer.position() > 0) + flushToTransport() + + if (engine.isOutboundDone || (engine.isInboundDone && !engineNeedsWrapReady)) nextPhase(CompletedPhase) else nextPhase(FlushingOutboundPhase) + } private def doInbound(isOutboundClosed: Boolean, inboundHalfClosedMode: Boolean): Boolean = if (transportInput.isDepleted && transportInput.isEmpty) {