Do heavy computation work as an explicit handshake step

Move execution of SSLEngine's delegated tasks to a separate method doHandshakeWork().
This is a preparation for performing computation heavy handshake work in a
separate thread.

6 files changed, 43 insertions, 16 deletions

diff --git a/jrt/src/com/yahoo/jrt/Connection.java b/jrt/src/com/yahoo/jrt/Connection.java
index 76ef831850e..df7acf881df 100644
--- a/jrt/src/com/yahoo/jrt/Connection.java
+++ b/jrt/src/com/yahoo/jrt/Connection.java
@@ -1,6 +1,8 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.jrt;
 
+import com.yahoo.jrt.CryptoSocket.HandshakeResult;
+
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.channels.SelectionKey;
@@ -215,7 +217,11 @@ class Connection extends Target {
     }
 
     private void handshake() throws IOException {
-        switch (socket.handshake()) {
+        HandshakeResult result;
+        while ((result = socket.handshake()) == HandshakeResult.NEED_WORK) {
+            socket.doHandshakeWork();
+        }
+        switch (result) {
         case DONE:
             if (socket.getMinimumReadBufferSize() > readSize) {
                 readSize = socket.getMinimumReadBufferSize();
diff --git a/jrt/src/com/yahoo/jrt/CryptoSocket.java b/jrt/src/com/yahoo/jrt/CryptoSocket.java
index 626eed47468..9ae714edbd9 100644
--- a/jrt/src/com/yahoo/jrt/CryptoSocket.java
+++ b/jrt/src/com/yahoo/jrt/CryptoSocket.java
@@ -25,7 +25,7 @@ public interface CryptoSocket {
      **/
     public SocketChannel channel();
 
-    public enum HandshakeResult { DONE, NEED_READ, NEED_WRITE }
+    public enum HandshakeResult { DONE, NEED_READ, NEED_WRITE, NEED_WORK }
 
     /**
      * Try to progress the initial connection handshake. Handshaking
@@ -35,10 +35,18 @@ public interface CryptoSocket {
      * the status is either DONE or an IOException is thrown. When
      * NEED_READ or NEED_WRITE is returned, the handshake function
      * will be called again when the appropriate io event has
-     * triggered.
+     * triggered. When NEED_WORK is returned, the {@link #doHandshakeWork()}
+     * will be called (possibly in another thread) before this function is called again.
      **/
     public HandshakeResult handshake() throws IOException;
 
+
+    /**
+     * Called when {@link #handshake()} returns {@link HandshakeResult#NEED_WORK} to perform compute-heavy tasks.
+     * This method may be called from another thread to avoid blocking the transport thread.
+     */
+    public void doHandshakeWork();
+
     /**
      * This function should be called after handshaking has completed
      * before calling the read function. It dictates the minimum size
diff --git a/jrt/src/com/yahoo/jrt/MaybeTlsCryptoSocket.java b/jrt/src/com/yahoo/jrt/MaybeTlsCryptoSocket.java
index 2e0d41b28d1..2090f69c80f 100644
--- a/jrt/src/com/yahoo/jrt/MaybeTlsCryptoSocket.java
+++ b/jrt/src/com/yahoo/jrt/MaybeTlsCryptoSocket.java
@@ -122,6 +122,7 @@ public class MaybeTlsCryptoSocket implements CryptoSocket {
 
     @Override public SocketChannel channel() { return socket.channel(); }
     @Override public HandshakeResult handshake() throws IOException { return socket.handshake(); }
+    @Override public void doHandshakeWork() { socket.doHandshakeWork(); }
     @Override public int getMinimumReadBufferSize() { return socket.getMinimumReadBufferSize(); }
     @Override public int read(ByteBuffer dst) throws IOException { return socket.read(dst); }
     @Override public int drain(ByteBuffer dst) throws IOException { return socket.drain(dst); }
diff --git a/jrt/src/com/yahoo/jrt/NullCryptoSocket.java b/jrt/src/com/yahoo/jrt/NullCryptoSocket.java
index 0d7b83f1c7d..83359bb65a5 100644
--- a/jrt/src/com/yahoo/jrt/NullCryptoSocket.java
+++ b/jrt/src/com/yahoo/jrt/NullCryptoSocket.java
@@ -24,6 +24,7 @@ public class NullCryptoSocket implements CryptoSocket {
         }
         return HandshakeResult.DONE;
     }
+    @Override public void doHandshakeWork() {}
     @Override public int getMinimumReadBufferSize() { return 1; }
     @Override public int read(ByteBuffer dst) throws IOException { return channel.read(dst); }
     @Override public int drain(ByteBuffer dst) throws IOException { return 0; }
diff --git a/jrt/src/com/yahoo/jrt/TlsCryptoSocket.java b/jrt/src/com/yahoo/jrt/TlsCryptoSocket.java
index f25a45169a8..c12c8ac88a3 100644
--- a/jrt/src/com/yahoo/jrt/TlsCryptoSocket.java
+++ b/jrt/src/com/yahoo/jrt/TlsCryptoSocket.java
@@ -29,7 +29,7 @@ public class TlsCryptoSocket implements CryptoSocket {
 
     private static final Logger log = Logger.getLogger(TlsCryptoSocket.class.getName());
 
-    private enum HandshakeState { NOT_STARTED, NEED_READ, NEED_WRITE, COMPLETED }
+    private enum HandshakeState { NOT_STARTED, NEED_READ, NEED_WRITE, NEED_WORK, COMPLETED }
 
     private final TransportMetrics metrics = TransportMetrics.getInstance();
     private final SocketChannel channel;
@@ -72,6 +72,14 @@ public class TlsCryptoSocket implements CryptoSocket {
         return toHandshakeResult(newHandshakeState);
     }
 
+    @Override
+    public void doHandshakeWork() {
+        Runnable task;
+        while ((task = sslEngine.getDelegatedTask()) != null) {
+            task.run();
+        }
+    }
+
     private HandshakeState processHandshakeState(HandshakeState state) throws IOException {
         try {
             switch (state) {
@@ -85,6 +93,17 @@ public class TlsCryptoSocket implements CryptoSocket {
                 case NEED_READ:
                     channelRead();
                     break;
+                case NEED_WORK:
+                    if (authorizationResult != null) {
+                        PeerAuthorizerTrustManager.getAuthorizationResult(sslEngine) // only available during handshake
+                                .ifPresent(result ->  {
+                                    if (!result.succeeded()) {
+                                        metrics.incrementPeerAuthorizationFailures();
+                                    }
+                                    authorizationResult = result;
+                                });
+                    }
+                    break;
                 case COMPLETED:
                     return HandshakeState.COMPLETED;
                 default:
@@ -108,17 +127,7 @@ public class TlsCryptoSocket implements CryptoSocket {
                         }
                         return HandshakeState.COMPLETED;
                     case NEED_TASK:
-                        sslEngine.getDelegatedTask().run();
-                        if (authorizationResult != null) {
-                            PeerAuthorizerTrustManager.getAuthorizationResult(sslEngine) // only available during handshake
-                                    .ifPresent(result ->  {
-                                        if (!result.succeeded()) {
-                                            metrics.incrementPeerAuthorizationFailures();
-                                        }
-                                        authorizationResult = result;
-                                    });
-                        }
-                        break;
+                        return HandshakeState.NEED_WORK;
                     case NEED_UNWRAP:
                         if (wrapBuffer.bytes() > 0) return HandshakeState.NEED_WRITE;
                         if (!handshakeUnwrap()) return HandshakeState.NEED_READ;
@@ -145,6 +154,8 @@ public class TlsCryptoSocket implements CryptoSocket {
                 return HandshakeResult.NEED_READ;
             case NEED_WRITE:
                 return HandshakeResult.NEED_WRITE;
+            case NEED_WORK:
+                return HandshakeResult.NEED_WORK;
             case COMPLETED:
                 return HandshakeResult.DONE;
             default:
diff --git a/jrt/src/com/yahoo/jrt/XorCryptoSocket.java b/jrt/src/com/yahoo/jrt/XorCryptoSocket.java
index fa784f62b12..db76ac62f3c 100644
--- a/jrt/src/com/yahoo/jrt/XorCryptoSocket.java
+++ b/jrt/src/com/yahoo/jrt/XorCryptoSocket.java
@@ -7,7 +7,6 @@ import java.nio.channels.SocketChannel;
 import java.security.SecureRandom;
 import java.util.ArrayDeque;
 import java.util.Queue;
-import java.util.Random;
 
 /**
  * A very simple CryptoSocket that performs connection handshaking and
@@ -81,6 +80,7 @@ public class XorCryptoSocket implements CryptoSocket {
         }
         return HandshakeResult.DONE;
     }
+    @Override public void doHandshakeWork() {}
     @Override public int getMinimumReadBufferSize() { return 1; }
     @Override public int read(ByteBuffer dst) throws IOException {
         if (input.bytes() == 0) {