From f1c93b8e7e8e92b4a071ec2a51303ad7d1ec04dd Mon Sep 17 00:00:00 2001
From: Erik van Zijst <erik.van.zijst@gmail.com>
Date: Wed, 1 Sep 2021 10:51:43 +0200
Subject: [PATCH] Close connection on CancelledKeyException

---
 build.sbt                                          |  2 +-
 .../scalatlsproxy/TlsProxyHandler.scala            | 14 ++++++--------
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/build.sbt b/build.sbt
index c6f2846..da8e25b 100644
--- a/build.sbt
+++ b/build.sbt
@@ -1,7 +1,7 @@
 name := "scala-tlsproxy"
 organization := "io.github.erikvanzijst"
 
-version := "0.4.1"
+version := "0.4.2"
 
 scalaVersion := "2.12.14"
 
diff --git a/src/main/scala/io/github/erikvanzijst/scalatlsproxy/TlsProxyHandler.scala b/src/main/scala/io/github/erikvanzijst/scalatlsproxy/TlsProxyHandler.scala
index 9d4df04..5e9ac9e 100644
--- a/src/main/scala/io/github/erikvanzijst/scalatlsproxy/TlsProxyHandler.scala
+++ b/src/main/scala/io/github/erikvanzijst/scalatlsproxy/TlsProxyHandler.scala
@@ -3,9 +3,8 @@ package io.github.erikvanzijst.scalatlsproxy
 import java.io.IOException
 import java.net.InetSocketAddress
 import java.nio.ByteBuffer
-import java.nio.channels.{SelectionKey, Selector, SocketChannel, UnresolvedAddressException}
+import java.nio.channels.{CancelledKeyException, SelectionKey, Selector, SocketChannel, UnresolvedAddressException}
 import java.nio.charset.StandardCharsets
-
 import com.typesafe.scalalogging.StrictLogging
 
 import scala.util.Try
@@ -152,7 +151,7 @@ class TlsProxyHandler(selector: Selector, clientChannel: SocketChannel, config:
           }.get
 
       if (phase == Response)
-        if (clientKey.isValid && clientKey.isWritable) {
+        if (clientKey.isWritable) {
           serverBuffer.flip()
           clientChannel.write(serverBuffer)
           serverBuffer.compact()
@@ -180,7 +179,7 @@ class TlsProxyHandler(selector: Selector, clientChannel: SocketChannel, config:
       }
 
       if (phase == Error)
-        if (clientKey.isValid && clientKey.isWritable) {
+        if (clientKey.isWritable) {
           serverBuffer.flip()
           clientChannel.write(serverBuffer)
           serverBuffer.compact()
@@ -190,11 +189,10 @@ class TlsProxyHandler(selector: Selector, clientChannel: SocketChannel, config:
         }
 
     } catch {
-      case e: IOException =>
-        val msg = s"$clientAddress -> $getServerAddress" +
+      case e @ (_: IOException | _: CancelledKeyException) =>
+        logger.warn(s"$clientAddress -> $getServerAddress" +
           (if (phase == Established) s" (up: ${upstreamPipe.bytes} down: ${downstreamPipe.bytes})" else "") +
-          s" connection failed: ${e.getClass.getSimpleName}: ${e.getMessage}"
-        logger.warn(msg)
+          s" connection failed: ${e.getClass.getSimpleName}: ${e.getMessage}")
         close()
     }