diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsTransportBackpressureSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/io/TlsTransportBackpressureSpec.scala new file mode 100644 index 0000000000..de74f35892 --- /dev/null +++ b/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") + } + } + } + } +} diff --git a/stream/src/main/scala/org/apache/pekko/stream/impl/io/TLSActor.scala b/stream/src/main/scala/org/apache/pekko/stream/impl/io/TLSActor.scala index 9f220074d8..7871dd9535 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/impl/io/TLSActor.scala +++ b/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( 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..fda0652f4e 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 @@ -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