From 6432c562af6687795df4e826c821207183736fe8 Mon Sep 17 00:00:00 2001 From: PJ Fanning Date: Thu, 29 Oct 2026 19:26:14 +0000 Subject: [PATCH] Keep draining TLS plaintext after a TLS 1.3 peer close_notify Motivation: When the plaintext output is cancelled, the GraphStage TLS engine drains plaintext it has already accepted. If a TLS 1.3 peer sent close_notify during that drain, unwrap returned CLOSED and the stage completed, discarding the accepted plaintext without sending it or its own close_notify, although TLS 1.3 close_notify only closes the peer's sending direction. Modification: - Stay in the inbound-closed drain when unwrap returns CLOSED while accepted plaintext is still being drained and the engine's outbound side is open. The drain then finishes and closes outbound as before. - Document that a TLS 1.3 close_notify does not end the drain, while a TLS 1.2 close_notify (which closes both SSLEngine directions) does. - Add real-SSLEngine tests that send close_notify during the drain for IgnoreComplete and EagerClose, on TLS 1.3 and TLS 1.2. Result: A TLS 1.3 peer that half-closes receives the accepted response followed by close_notify. TLS 1.2 behaviour is unchanged and terminates promptly. Tests: - sbt "stream-tests / Test / testOnly org.apache.pekko.stream.io.TlsGraphStageCancellationSpec": the two TLS 1.3 tests failed before the fix (partial response) and pass after it; 6/6 pass. - sbt "stream-tests / Test / testOnly org.apache.pekko.stream.io.Tls*": 281/281 passed. - Native scalafmt and git diff --check passed. References: Refs #3727 --- docs/src/main/paradox/stream/stream-io.md | 3 + .../io/TlsGraphStageCancellationSpec.scala | 231 +++++++++++++++++- .../pekko/stream/impl/io/TlsGraphStage.scala | 11 +- .../org/apache/pekko/stream/javadsl/TLS.scala | 6 +- .../apache/pekko/stream/scaladsl/TLS.scala | 6 +- 5 files changed, 245 insertions(+), 12 deletions(-) diff --git a/docs/src/main/paradox/stream/stream-io.md b/docs/src/main/paradox/stream/stream-io.md index a3545758ae..86e13b9f8c 100644 --- a/docs/src/main/paradox/stream/stream-io.md +++ b/docs/src/main/paradox/stream/stream-io.md @@ -197,6 +197,9 @@ stops reading. The GraphStage drain also applies when plaintext-output cancellation is observed before the peer's transport input completes. If transport-input completion is observed first, the existing inbound-close contract applies and accepted plaintext is discarded, matching the legacy actor engine. +The drain also continues after a TLS 1.3 peer sends `close_notify`, because that only closes the +peer's sending direction. A TLS 1.2 `close_notify` closes both directions of the `SSLEngine`, so it +ends the drain and accepted plaintext that has not been sent is discarded. ## Streaming File IO diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageCancellationSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageCancellationSpec.scala index 6ee8d641a2..39e3ec9313 100644 --- a/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageCancellationSpec.scala +++ b/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageCancellationSpec.scala @@ -17,7 +17,10 @@ package org.apache.pekko.stream.io -import javax.net.ssl.SSLContext +import java.nio.ByteBuffer +import java.nio.charset.StandardCharsets +import javax.net.ssl.{ SSLContext, SSLEngine, SSLEngineResult } +import javax.net.ssl.SSLEngineResult.HandshakeStatus._ import scala.concurrent.duration._ import scala.util.Success @@ -28,14 +31,14 @@ import pekko.stream.TLSProtocol._ import pekko.stream.impl.io.TlsGraphStage import pekko.stream.scaladsl._ import pekko.stream.stage._ -import pekko.stream.testkit.StreamSpec +import pekko.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber } import pekko.stream.testkit.scaladsl.{ TestSink, TestSource } -import pekko.testkit.WithLogCapturing +import pekko.testkit.{ TestProbe, WithLogCapturing } import pekko.util.ByteString class TlsGraphStageCancellationSpec extends StreamSpec(TlsSpec.configOverrides) with WithLogCapturing { - private def tls(context: SSLContext, client: Boolean) = + private def tls(context: SSLContext, client: Boolean, closing: TLSClosing = IgnoreComplete) = BidiFlow .fromGraph(new TlsGraphStage( () => { @@ -45,7 +48,7 @@ class TlsGraphStageCancellationSpec extends StreamSpec(TlsSpec.configOverrides) engine }, _ => Success(()), - IgnoreComplete)) + closing)) .withAttributes(TlsGraphStage.StreamTlsAttributes) private class EchoThenCancel extends GraphStage[FlowShape[ByteString, SslTlsOutbound]] { @@ -76,6 +79,164 @@ class TlsGraphStageCancellationSpec extends StreamSpec(TlsSpec.configOverrides) } } + /** Echoes the first element, answers the second with `response` and then cancels its input. */ + private class EchoThenRespondAndCancel(response: () => ByteString, drainStarted: TestProbe) + extends GraphStage[FlowShape[ByteString, SslTlsOutbound]] { + val in = Inlet[ByteString]("EchoThenRespondAndCancel.in") + val out = Outlet[SslTlsOutbound]("EchoThenRespondAndCancel.out") + override val shape = FlowShape(in, out) + + override def createLogic(attributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) { + private var first = true + private var cancelOnPull = false + setHandler(in, + new InHandler { + override def onPush(): Unit = { + val bytes = grab(in) + if (first) { + first = false + push(out, SendBytes(bytes)) + } else { + push(out, SendBytes(response())) + cancelOnPull = true + } + } + }) + setHandler(out, + new OutHandler { + override def onPull(): Unit = + if (cancelOnPull) cancel(in) // TLS accepted the response and asks for more + else pull(in) + + override def onDownstreamFinish(cause: Throwable): Unit = { + // TLS stops its plaintext input once it starts draining the accepted response. + drainStarted.ref ! "drain started" + super.onDownstreamFinish(cause) + } + }) + } + } + + /** Drives a plain client `SSLEngine` against the ciphertext ports of a TLS operator. */ + private final class RawTlsClient( + context: SSLContext, + cipherIn: TestPublisher.Probe[ByteString], + cipherOut: TestSubscriber.Probe[ByteString]) { + val engine: SSLEngine = context.createSSLEngine() + engine.setUseClientMode(true) + engine.setEnabledProtocols(Array(context.getProtocol)) + + private var netIn = ByteBuffer.allocate(engine.getSession.getPacketBufferSize) + private var appIn = ByteBuffer.allocate(engine.getSession.getApplicationBufferSize) + private var received = ByteString.empty + private var outstandingDemand = 0L + + /** Ciphertext elements requested from the operator but not received yet. */ + def unusedDemand: Long = outstandingDemand + + def request(elements: Long): Unit = { + cipherOut.request(elements) + outstandingDemand += elements + } + + def handshake(): Unit = { + engine.beginHandshake() + var status = engine.getHandshakeStatus + while (status != NOT_HANDSHAKING && status != FINISHED) { + status match { + case NEED_WRAP => send(ByteBuffer.allocate(0)) + case NEED_UNWRAP => + if (!unwrapBuffered()) receiveCiphertext() + case NEED_TASK => + var task = engine.getDelegatedTask + while (task ne null) { task.run(); task = engine.getDelegatedTask } + case other => fail(s"unexpected handshake status $other") + } + status = engine.getHandshakeStatus + } + } + + def send(plaintext: ByteBuffer): Unit = { + val packet = ByteBuffer.allocate(engine.getSession.getPacketBufferSize) + val result = engine.wrap(plaintext, packet) + result.getStatus shouldBe + (if (engine.isOutboundDone) SSLEngineResult.Status.CLOSED + else SSLEngineResult.Status.OK) + packet.flip() + cipherIn.sendNext(ByteString(packet)) + } + + /** Sends close_notify; TLS 1.3 keeps the inbound direction open. */ + def closeOutbound(): Unit = { + engine.closeOutbound() + send(ByteBuffer.allocate(0)) + } + + /** Reads ciphertext one element at a time until `expected` bytes of plaintext arrived. */ + def readPlaintext(expected: Int): ByteString = { + while (received.size < expected) { + if (!unwrapBuffered()) receiveCiphertext() + } + val result = received + received = ByteString.empty + result + } + + /** Reads until the operator completes its ciphertext output. */ + def readUntilComplete(): ByteString = { + request(Long.MaxValue - outstandingDemand) + var completed = false + while (!completed) { + while (unwrapBuffered()) () + cipherOut.expectNextOrComplete() match { + case Right(bytes) => append(bytes) + case Left(_) => completed = true + } + } + while (unwrapBuffered()) () + val result = received + received = ByteString.empty + result + } + + private def receiveCiphertext(): Unit = { + if (outstandingDemand == 0) request(1) + outstandingDemand -= 1 + append(cipherOut.expectNext()) + } + + private def append(bytes: ByteString): Unit = { + if (netIn.remaining < bytes.size) { + val bigger = ByteBuffer.allocate(netIn.position() + bytes.size) + netIn.flip() + bigger.put(netIn) + netIn = bigger + } + bytes.copyToBuffer(netIn) + } + + /** Unwraps one complete record from the buffered ciphertext, if there is one. */ + private def unwrapBuffered(): Boolean = { + netIn.flip() + val result = engine.unwrap(netIn, appIn) + netIn.compact() + result.getStatus match { + case SSLEngineResult.Status.BUFFER_OVERFLOW => + val bigger = ByteBuffer.allocate(appIn.capacity * 2) + appIn.flip() + bigger.put(appIn) + appIn = bigger + true + case SSLEngineResult.Status.BUFFER_UNDERFLOW => false + case _ => + appIn.flip() + received ++= ByteString(appIn) + appIn.clear() + result.bytesConsumed > 0 || result.bytesProduced > 0 + } + } + } + "TlsGraphStage cancellation ordering" must { List("TLSv1.2", "TLSv1.3").foreach { protocol => s"encrypt plaintext pushed before cancelling its output with $protocol" in { @@ -107,4 +268,64 @@ class TlsGraphStageCancellationSpec extends StreamSpec(TlsSpec.configOverrides) } } } + + /** + * Lets the server accept a response that is larger than its remaining ciphertext demand, cancels the + * server's plaintext output and then sends the client's close_notify. Returns the response and the + * plaintext the client received before the server completed its ciphertext output. + */ + private def halfCloseDuringDrain(protocol: String, closing: TLSClosing): (ByteString, ByteString) = { + val context = TlsSpec.initSslContext(protocol) + val drainStarted = TestProbe() + val transportBufferSize = system.settings.config.getBytes("pekko.io.tcp.direct-buffer-size").toInt + var response = ByteString.empty + val handler = Flow[SslTlsInbound] + .collect { case SessionBytes(_, bytes) => bytes } + .via(new EchoThenRespondAndCancel(() => response, drainStarted)) + + val (cipherIn, cipherOut) = TestSource[ByteString]() + .via(tls(context, client = false, closing).reversed.join(handler)) + .toMat(TestSink[ByteString]())(Keep.both) + .run() + val client = new RawTlsClient(context, cipherIn, cipherOut) + + within(10.seconds) { + // The server must be able to send its whole flight while the client sends handshake records. + client.request(16) + client.handshake() + client.send(ByteBuffer.wrap("ready".getBytes(StandardCharsets.US_ASCII))) + client.readPlaintext(5) shouldBe ByteString("ready") + + // Each ciphertext element holds at most one transport buffer, so the unused demand cannot + // carry the last two buffers of the response before close_notify arrives. + response = ByteString.fromArray( + Array.tabulate[Byte]((client.unusedDemand.toInt + 2) * transportBufferSize)(_.toByte)) + client.send(ByteBuffer.wrap("go".getBytes(StandardCharsets.US_ASCII))) + drainStarted.expectMsg("drain started") + + client.closeOutbound() + val received = client.readUntilComplete() + // The server answered close_notify with its own and stopped. + client.engine.isInboundDone shouldBe true + cipherIn.expectCancellation() + (response, received) + } + } + + "TlsGraphStage cancellation drain" must { + List(IgnoreComplete, EagerClose).foreach { closing => + s"send the accepted response after a TLS 1.3 peer half-closes during the drain with $closing" in { + val (response, received) = halfCloseDuringDrain("TLSv1.3", closing) + received.size shouldBe response.size + (received == response) shouldBe true + } + + s"stop draining when a TLS 1.2 peer closes during the drain with $closing" in { + // A TLS 1.2 close_notify closes both directions of the SSLEngine. + val (response, received) = halfCloseDuringDrain("TLSv1.2", closing) + received.size should be < response.size + (response.take(received.size) == received) shouldBe true + } + } + } } 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 66f48dc4b3..8cd01301d4 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 @@ -401,8 +401,11 @@ import pekko.util.ByteString !closing.ignoreCancel && !canDrainAcceptedUserInput private def drainBeforeInboundClose: Boolean = - transportInput.isDepleted && transportInput.isEmpty && userOutput.isCancelled && - !closing.ignoreCancel && canDrainAcceptedUserInput && (userInBufferHasRemaining || !userInput.isEmpty) + transportInput.isDepleted && transportInput.isEmpty && drainingAcceptedUserInput + + private def drainingAcceptedUserInput: Boolean = + userOutput.isCancelled && !closing.ignoreCancel && canDrainAcceptedUserInput && + (userInBufferHasRemaining || !userInput.isEmpty) private def canDrainAcceptedUserInput: Boolean = (stateBits & PlainDataAllowedFlag) != 0 && mayCloseOutbound @@ -507,7 +510,9 @@ import pekko.util.ByteString case CLOSED => flushToUser() - completeOrFlush() + // A TLS 1.3 close_notify only closes the inbound direction, so keep draining accepted plaintext. + if (phase.toInt != InboundClosedPhase || engine.isOutboundDone || !drainingAcceptedUserInput) + completeOrFlush() case BUFFER_UNDERFLOW => flushToUser() diff --git a/stream/src/main/scala/org/apache/pekko/stream/javadsl/TLS.scala b/stream/src/main/scala/org/apache/pekko/stream/javadsl/TLS.scala index 259a42f531..6b2bd8c523 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/javadsl/TLS.scala +++ b/stream/src/main/scala/org/apache/pekko/stream/javadsl/TLS.scala @@ -54,8 +54,10 @@ import pekko.util.ByteString * drain includes up to 64 queued messages plus any plaintext already buffered for * wrapping; message sizes are unbounded, so cancellation can delay termination * until that data is sent. Data queued after a renegotiation request is discarded, - * and draining stops if a new handshake is required. The engine is selected with - * `pekko.stream.materializer.tls.engine`, which defaults to `legacy-actor`. + * and draining stops if a new handshake is required. The drain continues after a + * TLS 1.3 peer sends `close_notify`, but a TLS 1.2 `close_notify` ends it. The engine + * is selected with `pekko.stream.materializer.tls.engine`, which defaults to + * `legacy-actor`. * * Especially for a client the common idiom of attaching * a finite Source to the plaintext input and transforming the plaintext response diff --git a/stream/src/main/scala/org/apache/pekko/stream/scaladsl/TLS.scala b/stream/src/main/scala/org/apache/pekko/stream/scaladsl/TLS.scala index c95d43c548..f9309da0b4 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/scaladsl/TLS.scala +++ b/stream/src/main/scala/org/apache/pekko/stream/scaladsl/TLS.scala @@ -56,8 +56,10 @@ import com.typesafe.config.{ ConfigException, ConfigFactory } * drain includes up to 64 queued messages plus any plaintext already buffered for * wrapping; message sizes are unbounded, so cancellation can delay termination * until that data is sent. Data queued after a renegotiation request is discarded, - * and draining stops if a new handshake is required. The engine is selected with - * `pekko.stream.materializer.tls.engine`, which defaults to `legacy-actor`. + * and draining stops if a new handshake is required. The drain continues after a + * TLS 1.3 peer sends `close_notify`, but a TLS 1.2 `close_notify` ends it. The engine + * is selected with `pekko.stream.materializer.tls.engine`, which defaults to + * `legacy-actor`. * * Especially for a client the common idiom of attaching * a finite Source to the plaintext input and transforming the plaintext response