Keep draining TLS plaintext after a TLS 1.3 KeyUpdate (#3738)

* Keep draining TLS plaintext after a TLS 1.3 KeyUpdate

Motivation:
On plaintext-output cancellation the GraphStage TLS engine drains plaintext it
has already accepted, but stops if a new handshake is required. It treated any
NEED_WRAP as a new handshake, so a TLS 1.3 KeyUpdate that requests a response,
received before or during the drain, closed outbound and discarded the
accepted plaintext.

Modification:
- Allow the drain to continue when an established TLS 1.3 session needs to
  wrap. TLS 1.3 has no renegotiation, so this is a post-handshake message such
  as a KeyUpdate response; the response is wrapped before the plaintext.
- TLS 1.2 keeps stopping the drain, because NEED_WRAP may start a
  renegotiation.
- Document the KeyUpdate case in the stream IO guide and both TLS DSLs.

Result:
A TLS 1.3 KeyUpdate no longer discards plaintext accepted before
plaintext-output cancellation.

Tests:
- stream-tests/testOnly TlsGraphStageDeferredCloseSpec: 25 passed; the two new
  TLS 1.3 KeyUpdate tests fail without the fix (ByteString() was not equal to
  ByteString(92, 90)), the TLS 1.2 guard passes both ways
- stream-tests/testOnly org.apache.pekko.stream.io.Tls*: 288 passed
- scalafmt --mode diff-ref=upstream/main: clean

References:
Refs #3727

* Match SunJSSE's FINISHED status for scripted KeyUpdate wraps

Motivation:
Review noted the KeyUpdate tests rely on a scripted engine. A real
SunJSSE TLS 1.3 KeyUpdate (forced with a low jdk.tls.keyLimits) reports
NEED_WRAP after unwrap, and the response wrap returns FINISHED, not
NOT_HANDSHAKING as scripted. FINISHED re-runs session verification.

Modification:
Return FINISHED from the scripted KeyUpdate response wrap, leaving the
engine in NOT_HANDSHAKING afterwards, like SunJSSE.

Result:
The drain tests now cover the FINISHED path a real KeyUpdate takes.

Tests:
- stream-tests/testOnly org.apache.pekko.stream.io.TlsGraphStage*:
  173 passed.
- With the fix disabled, both TLS 1.3 KeyUpdate tests fail.
- Standalone JDK 17 SunJSSE engine pair with low jdk.tls.keyLimits
  confirmed the NEED_WRAP then FINISHED sequence on TLSv1.3.

References:
Review of #3738

Author: pjfanning

docs/src/main/paradox/stream/stream-io.md

@@ -188,7 +188,9 @@ With `pekko.stream.materializer.tls.engine = graph-stage`, the operator stops fu
 and attempts to send data already accepted for the established session, subject to transport
 availability. This includes up to 64 queued messages plus plaintext already buffered for wrapping.
 Message sizes are unbounded, so termination can be delayed until that data is sent. Data queued after
-a renegotiation request is discarded, and draining stops if a new handshake is required. Cancellation
+a renegotiation request is discarded, and draining stops if a new handshake is required. A TLS 1.3
+post-handshake message such as a KeyUpdate is not a new handshake, so the response is sent and draining
+continues; in TLS 1.2 any further handshake stops the drain. Cancellation
 during a handshake or a transport failure can still prevent delivery. The default `legacy-actor`
 engine discards accepted plaintext on cancellation. Use the `IgnoreCancel` closing policy
 (`TLSClosing.ignoreCancel()` in Java) if the handler needs to continue producing a response after it

stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsGraphStageIsolatedSpec.scala

@@ -299,6 +299,7 @@ abstract class TlsGraphStageBufferTestBase(config: Config) extends StreamSpec(co
   protected val TlsRecord = ByteString(Array[Byte](23, 3, 3, 0, 0))
   protected val ApplicationCiphertext = ByteString(0x5A.toByte)
   protected val CloseCiphertext = ByteString(0x5B.toByte)
+  protected val HandshakeCiphertext = ByteString(0x5C.toByte)
 
   protected final class ScriptedSSLEngine(
       delegate: SSLEngine,
@@ -307,7 +308,9 @@ abstract class TlsGraphStageBufferTestBase(config: Config) extends StreamSpec(co
       requestWrapOnSubsequentUnwrap: Boolean,
       @volatile var requiredOutputSpace: Int,
       negotiatedPacketSize: Int,
-      negotiatedApplicationSize: Int) extends SSLEngine {
+      negotiatedApplicationSize: Int,
+      sessionProtocol: String = null,
+      wrapHandshakeOnNeedWrap: Boolean = false) extends SSLEngine {
     private var handshakeStatus: SSLEngineResult.HandshakeStatus = NEED_UNWRAP
     @volatile var unwrapCount = 0
     private var inboundDone = false
@@ -334,6 +337,7 @@ abstract class TlsGraphStageBufferTestBase(config: Config) extends StreamSpec(co
       (_: 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 if (method.getName == "getProtocol" && (sessionProtocol ne null)) sessionProtocol
         else method.invoke(delegate.getSession, Option(args).getOrElse(Array.empty[AnyRef]): _*))
       .asInstanceOf[SSLSession]
 
@@ -356,6 +360,13 @@ abstract class TlsGraphStageBufferTestBase(config: Config) extends StreamSpec(co
         overflowWraps += 1
         return new SSLEngineResult(BUFFER_OVERFLOW, handshakeStatus, 0, 0)
       }
+      if (wrapHandshakeOnNeedWrap && handshakeStatus == NEED_WRAP && !(wrapOnCloseInbound && inboundDone)) {
+        // Like the JDK, a wrap while NEED_WRAP produces only the handshake message and consumes no plaintext.
+        // SunJSSE reports FINISHED for a TLS 1.3 KeyUpdate response, then NOT_HANDSHAKING from the engine.
+        dst.put(HandshakeCiphertext.head)
+        handshakeStatus = NOT_HANDSHAKING
+        return new SSLEngineResult(OK, FINISHED, 0, 1)
+      }
       var consumed = 0
       var index = offset
       while (index < offset + length) {
@@ -440,7 +451,7 @@ abstract class TlsGraphStageBufferTestBase(config: Config) extends StreamSpec(co
     override def getDelegatedTask: Runnable = null
     override def getSession: SSLSession = {
       sessionQueries += 1
-      if (negotiatedPacketSize > 0) reportedSession else delegate.getSession
+      if (negotiatedPacketSize > 0 || (sessionProtocol ne null)) reportedSession else delegate.getSession
     }
     override def getSupportedCipherSuites: Array[String] = delegate.getSupportedCipherSuites
     override def getEnabledCipherSuites: Array[String] = delegate.getEnabledCipherSuites
@@ -474,11 +485,14 @@ abstract class TlsGraphStageBufferTestBase(config: Config) extends StreamSpec(co
       extraAttributes: Attributes = Attributes.none,
       requiredOutputSpace: Int = 0,
       negotiatedPacketSize: Int = 0,
-      negotiatedApplicationSize: Int = 0): DeferredCloseHarness = {
+      negotiatedApplicationSize: Int = 0,
+      sessionProtocol: String = null,
+      wrapHandshakeOnNeedWrap: Boolean = false): DeferredCloseHarness = {
     val engine =
       new ScriptedSSLEngine(TlsSpec.initSslContext("TLSv1.2").createSSLEngine(), wrapOnCloseInbound,
         requestHandshakeAfterApplication, requestWrapOnSubsequentUnwrap,
-        requiredOutputSpace, negotiatedPacketSize, negotiatedApplicationSize)
+        requiredOutputSpace, negotiatedPacketSize, negotiatedApplicationSize, sessionProtocol,
+        wrapHandshakeOnNeedWrap)
     val sessionVerified = Promise[Unit]()
     val probes =
       RunnableGraph
@@ -624,6 +638,48 @@ class TlsGraphStageDeferredCloseSpec extends TlsGraphStageBufferTestBase(TlsGrap
       harness.engine.applicationBytes shouldBe 1
     }
 
+    // The scripted engine answers the first post-handshake unwrap with NEED_WRAP, like a received
+    // TLS 1.3 KeyUpdate that requests an update, and wraps the response before any plaintext.
+    for (keyUpdateBeforeCancel <- Seq(true, false)) {
+      val order = if (keyUpdateBeforeCancel) "before" else "after"
+      s"keep draining accepted plaintext after a TLS 1.3 KeyUpdate received $order output cancellation" in {
+        val harness = newHarness(sessionProtocol = "TLSv1.3", wrapHandshakeOnNeedWrap = true)
+        finishInitialHandshake(harness)
+        harness.plainIn.sendNext(SendBytes(ByteString("a")))
+        harness.plainIn.expectRequest()
+        if (keyUpdateBeforeCancel) {
+          harness.cipherIn.sendNext(TlsRecord)
+          awaitAssert(harness.engine.unwrapCount shouldBe 2)
+        }
+        harness.plainOut.cancel()
+        harness.plainIn.expectCancellation()
+        if (!keyUpdateBeforeCancel) {
+          harness.cipherIn.sendNext(TlsRecord)
+          awaitAssert(harness.engine.unwrapCount shouldBe 2)
+        }
+
+        // A KeyUpdate response is not a new handshake, so it must not discard the accepted plaintext.
+        val sent = harness.cipherOut.toStrict(remainingOrDefault).foldLeft(ByteString.empty)(_ ++ _)
+        sent shouldBe HandshakeCiphertext ++ ApplicationCiphertext
+        harness.engine.applicationBytes shouldBe 1
+      }
+    }
+
+    "stop draining when a TLS 1.2 engine needs to wrap during output cancellation" in {
+      val harness = newHarness(sessionProtocol = "TLSv1.2", wrapHandshakeOnNeedWrap = true)
+      finishInitialHandshake(harness)
+      harness.plainIn.sendNext(SendBytes(ByteString("a")))
+      harness.plainIn.expectRequest()
+      harness.cipherIn.sendNext(TlsRecord)
+      awaitAssert(harness.engine.unwrapCount shouldBe 2)
+      harness.plainOut.cancel()
+      harness.plainIn.expectCancellation()
+
+      // TLS 1.2 NEED_WRAP may start a renegotiation, so draining must stop instead of waiting for it.
+      harness.cipherOut.toStrict(remainingOrDefault) shouldBe empty
+      harness.engine.applicationBytes shouldBe 0
+    }
+
     "stop draining when the transport output is cancelled" in {
       val harness = newHarness()
       finishInitialHandshake(harness)

stream/src/main/scala/org/apache/pekko/stream/impl/io/TlsGraphStage.scala

@@ -412,7 +412,13 @@ import pekko.util.ByteString
         userOutput.isCancelled && !closing.ignoreCancel && canDrainAcceptedUserInput &&
         (userInBufferHasRemaining || !userInput.isEmpty)
 
-      private def canDrainAcceptedUserInput: Boolean = (stateBits & PlainDataAllowedFlag) != 0 && mayCloseOutbound
+      private def canDrainAcceptedUserInput: Boolean = (stateBits & PlainDataAllowedFlag) != 0 &&
+        (mayCloseOutbound || needsPostHandshakeWrap)
+
+      // TLS 1.3 has no renegotiation, so NEED_WRAP in an established session is a post-handshake
+      // message such as a KeyUpdate response. Wrapping it does not start a new handshake.
+      private def needsPostHandshakeWrap: Boolean =
+        lastHandshakeStatus == NEED_WRAP && currentSession.getProtocol == "TLSv1.3"
 
       private def mayCloseOutbound: Boolean =
         lastHandshakeStatus match {

stream/src/main/scala/org/apache/pekko/stream/javadsl/TLS.scala

@@ -54,7 +54,9 @@ 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 drain continues after a
+ * and draining stops if a new handshake is required. A TLS 1.3 post-handshake
+ * message such as a KeyUpdate is not a new handshake and does not stop the drain;
+ * in TLS 1.2 any further handshake stops it. 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`.

stream/src/main/scala/org/apache/pekko/stream/scaladsl/TLS.scala

@@ -56,7 +56,9 @@ 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 drain continues after a
+ * and draining stops if a new handshake is required. A TLS 1.3 post-handshake
+ * message such as a KeyUpdate is not a new handshake and does not stop the drain;
+ * in TLS 1.2 any further handshake stops it. 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`.