From 4b5d78f29e3a249992903836b0e68354e9b27279 Mon Sep 17 00:00:00 2001
From: sahnjeok <71247142+sahnjeok@users.noreply.github.com>
Date: Fri, 17 Jun 2022 12:42:21 +0900
Subject: [PATCH] Update SSLSocketChannel2.java

'handshakeStartTime' long variable is added and isHandShakeComplete()  function is updated for #896.

If wss handshake is not completed in 10s, close this channel to prevent cpu overload or unexpected channel error. See #896.
---
 .../org/java_websocket/SSLSocketChannel2.java | 25 ++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/main/java/org/java_websocket/SSLSocketChannel2.java b/src/main/java/org/java_websocket/SSLSocketChannel2.java
index d9b9a4c3..a1e8a7c9 100644
--- a/src/main/java/org/java_websocket/SSLSocketChannel2.java
+++ b/src/main/java/org/java_websocket/SSLSocketChannel2.java
@@ -105,6 +105,12 @@ public class SSLSocketChannel2 implements ByteChannel, WrappedByteChannel, ISSLC
    **/
   protected int bufferallocations = 0;
 
+  /**
+   * 2022-06-17 Handshake start time in WSS for the underlying channel. 
+   *            If wss handshake is not completed in 10s, close this channel to prevent cpu overload or unexpected channel error. see #896.
+   */
+  protected long	  handshakeStartTime = System.currentTimeMillis() ;
+  
   public SSLSocketChannel2(SocketChannel channel, SSLEngine sslEngine, ExecutorService exec,
       SelectionKey key) throws IOException {
     if (channel == null || sslEngine == null || exec == null) {
@@ -393,8 +399,21 @@ public void close() throws IOException {
 
   private boolean isHandShakeComplete() {
     HandshakeStatus status = sslEngine.getHandshakeStatus();
-    return status == SSLEngineResult.HandshakeStatus.FINISHED
-        || status == SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING;
+    
+    // handshake status
+    boolean ret = status == SSLEngineResult.HandshakeStatus.FINISHED
+               || status == SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING;
+    
+    if ( ret == false )
+    {
+      // 2022-06-17 If wss handshake is not completed in 10s, close this channel to prevent cpu overload or unexpected channel error. see #896.
+      if ( handshakeStartTime > 0 && ( System.currentTimeMillis() - handshakeStartTime ) > 10000 )
+      {
+        try{close() ;}catch(Exception ex){} ;
+      }
+    }
+
+    return ret;
   }
 
   public SelectableChannel configureBlocking(boolean b) throws IOException {
@@ -495,4 +514,4 @@ private void tryRestoreCryptedData() {
       saveCryptData = null;
     }
   }
-}
\ No newline at end of file
+}