Abort TLS transport on session verification failure

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

@@ -21,12 +21,12 @@ import javax.net.ssl._
 
 import scala.collection.immutable
 import scala.concurrent.Await
-import scala.concurrent.Future
+import scala.concurrent.{ Future, Promise }
 import scala.concurrent.duration._
-import scala.util.{ Random, Success }
+import scala.util.{ Failure, Random, Success }
 
 import org.apache.pekko
-import pekko.NotUsed
+import pekko.{ Done, NotUsed }
 import pekko.pattern.{ after => later }
 import pekko.stream._
 import pekko.stream.TLSProtocol._
@@ -503,6 +503,30 @@ abstract class AbstractTlsSpec(useLegacyActor: Boolean)
         clientErrText should include("unable to find valid certification path to requested target")
       }
 
+      "abort the transport when session verification fails" in {
+        val verificationFailure = new SSLException("session verification failed")
+        val rejectingClientTls = tlsBidi(
+          () => createSSLEngine(sslContext, Client),
+          verifySession = _ => Failure(verificationFailure),
+          closing = IgnoreBoth)
+        val clientToServerDone = Promise[Done]()
+        val clientToServer = Flow[ByteString].watchTermination { (_, done) =>
+          done.onComplete(clientToServerDone.tryComplete)
+          NotUsed
+        }
+        val transport = BidiFlow.fromFlows(clientToServer, Flow[ByteString])
+        val serverApplication: Flow[SslTlsInbound, SslTlsOutbound, NotUsed] =
+          Flow.fromSinkAndSource(Sink.ignore, Source.maybe[SslTlsOutbound])
+
+        Source
+          .maybe[SslTlsOutbound]
+          .via(rejectingClientTls.atop(transport).atop(serverTls(IgnoreBoth).reversed).join(serverApplication))
+          .runWith(Sink.ignore)
+
+        val failure = intercept[SSLException](Await.result(clientToServerDone.future, 3.seconds.dilated))
+        (failure should be).theSameInstanceAs(verificationFailure)
+      }
+
       "reliably cancel subscriptions when TransportIn fails early" in {
         val ex = new Exception("hello")
         val (sub, out1, out2) =

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

@@ -380,26 +380,29 @@ import pekko.util.ByteString
         val result = wrapIntoTransportBuffer()
         lastHandshakeStatus = result.getHandshakeStatus
 
-        if (lastHandshakeStatus == FINISHED) handshakeFinished()
-        runDelegatedTasks()
+        if (lastHandshakeStatus == FINISHED && !handshakeFinished()) {
+          // stop processing after failed session verification
+        } else {
+          runDelegatedTasks()
 
-        result.getStatus match {
-          case OK =>
-            if (transportOutBuffer.position() == 0 && lastHandshakeStatus == NEED_WRAP)
-              throw new IllegalStateException("SSLEngine trying to loop NEED_WRAP without producing output")
+          result.getStatus match {
+            case OK =>
+              if (transportOutBuffer.position() == 0 && lastHandshakeStatus == NEED_WRAP)
+                throw new IllegalStateException("SSLEngine trying to loop NEED_WRAP without producing output")
 
-            flushToTransportIfNeeded(result)
+              flushToTransportIfNeeded(result)
 
-          case CLOSED =>
-            flushToTransport()
-            if (engine.isInboundDone) nextPhase(CompletedPhase)
-            else nextPhase(AwaitingClosePhase)
+            case CLOSED =>
+              flushToTransport()
+              if (engine.isInboundDone) nextPhase(CompletedPhase)
+              else nextPhase(AwaitingClosePhase)
 
-          case BUFFER_OVERFLOW =>
-            growTransportOutBuffer()
+            case BUFFER_OVERFLOW =>
+              growTransportOutBuffer()
 
-          case status =>
-            failTls(new IllegalStateException(s"unexpected status $status in doWrap()"))
+            case status =>
+              failTls(new IllegalStateException(s"unexpected status $status in doWrap()"))
+          }
         }
       }
 
@@ -450,8 +453,8 @@ import pekko.util.ByteString
 
               case FINISHED =>
                 flushToUser()
-                handshakeFinished()
-                transportInput.putBackUnreadBuffer(transportInBuffer)
+                if (handshakeFinished())
+                  transportInput.putBackUnreadBuffer(transportInBuffer)
 
               case NEED_UNWRAP
                   if transportInBuffer.hasRemaining &&
@@ -484,16 +487,18 @@ import pekko.util.ByteString
         if (TlsEngineHelpers.runDelegatedTasks(engine) > 0 || lastHandshakeStatus == NEED_TASK)
           lastHandshakeStatus = engine.getHandshakeStatus
 
-      private def handshakeFinished(): Unit = {
+      private def handshakeFinished(): Boolean = {
         val session = engine.getSession
         verifySession(session) match {
           case Success(()) =>
             currentSession = session
             stateBits |= PlainDataAllowedFlag
             flushToUser()
+            true
 
           case Failure(ex) =>
-            throw ex
+            failTls(ex)
+            false
         }
       }