Wait for transport demand instead of failing when TLS unwrap needs a wrap #3732

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

@@ -0,0 +1,149 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.pekko.stream.io
+
+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
+
+import org.apache.pekko
+import pekko.NotUsed
+import pekko.stream._
+import pekko.stream.TLSProtocol._
+import pekko.stream.impl.io.{ TlsGraphStage, TlsModule }
+import pekko.stream.scaladsl._
+import pekko.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber }
+import pekko.stream.testkit.scaladsl.{ TestSink, TestSource }
+import pekko.testkit.WithLogCapturing
+import pekko.util.ByteString
+
+class TlsTransportBackpressureSpec extends AbstractTlsTransportBackpressureSpec(useLegacyActor = true)
+class TlsGraphStageTransportBackpressureSpec extends AbstractTlsTransportBackpressureSpec(useLegacyActor = false)
+
+abstract class AbstractTlsTransportBackpressureSpec(useLegacyActor: Boolean)
+    extends StreamSpec(TlsSpec.configOverrides)
+    with WithLogCapturing {
+
+  private def serverTls(
+      context: SSLContext): BidiFlow[SslTlsOutbound, ByteString, ByteString, SslTlsInbound, NotUsed] = {
+    val createEngine = () => {
+      val engine = context.createSSLEngine()
+      engine.setUseClientMode(false)
+      engine.setEnabledProtocols(Array(context.getProtocol))
+      engine
+    }
+    if (useLegacyActor)
+      BidiFlow.fromGraph(TlsModule(Attributes.none, createEngine, _ => Success(()), IgnoreComplete))
+    else
+      BidiFlow
+        .fromGraph(new TlsGraphStage(createEngine, _ => Success(()), IgnoreComplete))
+        .withAttributes(TlsGraphStage.StreamTlsAttributes)
+  }
+
+  /** Drives a plain client `SSLEngine`, requesting server ciphertext only when it needs to unwrap. */
+  private final class LazyTlsClient(
+      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 val appIn = ByteBuffer.allocate(engine.getSession.getApplicationBufferSize * 2)
+    private var received = ByteString.empty
+
+    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)
+      engine.wrap(plaintext, packet).getStatus shouldBe SSLEngineResult.Status.OK
+      packet.flip()
+      cipherIn.sendNext(ByteString(packet))
+    }
+
+    def readPlaintext(expected: Int): ByteString = {
+      while (received.size < expected)
+        if (!unwrapBuffered()) receiveCiphertext()
+      val result = received
+      received = ByteString.empty
+      result
+    }
+
+    private def receiveCiphertext(): Unit = {
+      cipherOut.request(1)
+      val bytes = cipherOut.expectNext()
+      if (netIn.remaining < bytes.size) {
+        val bigger = ByteBuffer.allocate(netIn.position() + bytes.size)
+        netIn.flip()
+        bigger.put(netIn)
+        netIn = bigger
+      }
+      bytes.copyToBuffer(netIn)
+    }
+
+    private def unwrapBuffered(): Boolean = {
+      netIn.flip()
+      val result = engine.unwrap(netIn, appIn)
+      netIn.compact()
+      appIn.flip()
+      received ++= ByteString(appIn)
+      appIn.clear()
+      result.getStatus == SSLEngineResult.Status.OK && (result.bytesConsumed > 0 || result.bytesProduced > 0)
+    }
+  }
+
+  s"TLS server (${if (useLegacyActor) "legacy actor" else "GraphStage"})" must {
+    List("TLSv1.2", "TLSv1.3").foreach { protocol =>
+      s"wait for ciphertext demand when peer data arrives mid-flight with $protocol" in {
+        val context = TlsSpec.initSslContext(protocol)
+        val echo = Flow[SslTlsInbound].collect { case SessionBytes(_, bytes) => SendBytes(bytes): SslTlsOutbound }
+
+        val (cipherIn, cipherOut) = TestSource[ByteString]()
+          .via(serverTls(context).reversed.join(echo))
+          .toMat(TestSink[ByteString]())(Keep.both)
+          .run()
+        val client = new LazyTlsClient(context, cipherIn, cipherOut)
+
+        within(10.seconds) {
+          client.handshake()
+          client.send(ByteBuffer.wrap("hello".getBytes(StandardCharsets.US_ASCII)))
+          client.readPlaintext(5) shouldBe ByteString("hello")
+        }
+      }
+    }
+  }
+}

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

@@ -215,6 +215,13 @@ import pekko.util.ByteString
     def isCompleted = engine.isOutboundDone
   }
 
+  // While the engine needs to wrap, unwrap returns the record unread, so wait for transport demand.
+  val engineNotBlockedOnWrap = new TransferState {
+    def isReady =
+      lastHandshakeStatus != NEED_WRAP || outputBunch.isPending(TransportOut) || outputBunch.isCancelled(TransportOut)
+    def isCompleted = false
+  }
+
   val engineInboundOpen = new TransferState {
     def isReady = true
     def isCompleted = engine.isInboundDone
@@ -232,11 +239,12 @@ import pekko.util.ByteString
 
   // bidirectional case
   val outbound = (userHasData || engineNeedsWrap) && outputBunch.demandAvailableFor(TransportOut)
-  val inbound = (transportInChoppingBlock && outputBunch.demandAvailableFor(UserOut)) || userOutCancelled
+  val inbound =
+    (transportInChoppingBlock && outputBunch.demandAvailableFor(UserOut) && engineNotBlockedOnWrap) || userOutCancelled
 
   // half-closed
   val outboundHalfClosed = engineNeedsWrap && outputBunch.demandAvailableFor(TransportOut)
-  val inboundHalfClosed = transportInChoppingBlock && engineInboundOpen
+  val inboundHalfClosed = transportInChoppingBlock && engineInboundOpen && engineNotBlockedOnWrap
 
   val bidirectional = TransferPhase(outbound || inbound) { () =>
     if (tracing) log.debug("bidirectional")
@@ -423,9 +431,8 @@ import pekko.util.ByteString
         result.getHandshakeStatus match {
           case NEED_WRAP =>
             // https://github.com/apache/pekko/issues/442
-            // A second workaround for an infinite loop we have not been able to reproduce/isolate,
-            // if you see this, and can reproduce consistently, please report back to the Apache Pekko team
-            // with a reproducer or details about the client causing it
+            // Guards against an infinite loop. Peer data arriving while the engine needs to wrap without
+            // transport demand used to hit this; engineNotBlockedOnWrap now waits for demand instead (#3732).
             unwrapPutBackCounter += 1
             if (unwrapPutBackCounter > maxTLSIterations) {
               throw new IllegalStateException(

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

@@ -280,9 +280,14 @@ import pekko.util.ByteString
       private def outboundCompleted: Boolean =
         userOrEngineOutboundCompleted || transportOutput.isClosed
 
-      private def inboundReady: Boolean = (transportInput.isReady && userOutput.isDemandAvailable) ||
+      private def inboundReady: Boolean =
+        (transportInput.isReady && userOutput.isDemandAvailable && !inboundBlockedOnWrap) ||
         userOutput.isCancelled
 
+      // While the engine needs to wrap, unwrap returns the record unread, so wait for ciphertext demand.
+      private def inboundBlockedOnWrap: Boolean =
+        lastHandshakeStatus == NEED_WRAP && !transportOutput.isDemandAvailable && !transportOutput.isClosed
+
       private def inboundCompleted: Boolean = (transportInput.isCompleted || userOutput.isClosed) &&
         (engine.isInboundDone || userOutput.isErrored)
 
@@ -293,7 +298,7 @@ import pekko.util.ByteString
         engineNeedsWrapCompleted || transportOutput.isClosed
 
       private def inboundHalfClosedReady: Boolean =
-        transportInput.isReady && !engine.isInboundDone && !drainBeforeInboundClose
+        transportInput.isReady && !engine.isInboundDone && !drainBeforeInboundClose && !inboundBlockedOnWrap
 
       private def inboundHalfClosedCompleted: Boolean =
         transportInput.isCompleted || engine.isInboundDone