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