This is an automated email from the ASF dual-hosted git repository.
He-Pin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/pekko.git
The following commit(s) were added to refs/heads/main by this push:
new f79b53b36e fix(stream): handle deferred TLS output on inbound close
(#3324)
f79b53b36e is described below
commit f79b53b36e32061308c613754e29f59fd15e3121
Author: He-Pin(kerr) <[email protected]>
AuthorDate: Tue Jul 14 10:51:44 2026 +0800
fix(stream): handle deferred TLS output on inbound close (#3324)
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 53390bef9f..1ddc3d78b4 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 65e3a9110c..ab03f2e3ce 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) {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
