Fix deferred TLS output handling on inbound close

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))
+}

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) {