Add security context to target

9 files changed, 112 insertions, 3 deletions

diff --git a/jrt/src/com/yahoo/jrt/Connection.java b/jrt/src/com/yahoo/jrt/Connection.java
index 6521c7cb8a5..f6d9989324d 100644
--- a/jrt/src/com/yahoo/jrt/Connection.java
+++ b/jrt/src/com/yahoo/jrt/Connection.java
@@ -9,6 +9,7 @@ import java.nio.channels.SocketChannel;
 import java.util.HashMap;
 import java.util.IdentityHashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -439,6 +440,12 @@ class Connection extends Target {
         return lostReason;
     }
 
+    @Override
+    public Optional<SecurityContext> getSecurityContext() {
+        return Optional.ofNullable(socket)
+                .flatMap(CryptoSocket::securityContext);
+    }
+
     public boolean isClient() {
         return !server;
     }
diff --git a/jrt/src/com/yahoo/jrt/CryptoSocket.java b/jrt/src/com/yahoo/jrt/CryptoSocket.java
index 9ae714edbd9..59f4d0e6650 100644
--- a/jrt/src/com/yahoo/jrt/CryptoSocket.java
+++ b/jrt/src/com/yahoo/jrt/CryptoSocket.java
@@ -5,6 +5,7 @@ package com.yahoo.jrt;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.channels.SocketChannel;
+import java.util.Optional;
 
 
 /**
@@ -93,4 +94,10 @@ public interface CryptoSocket {
      * io event has triggered.
      **/
     public FlushResult flush() throws IOException;
+
+    /**
+     * Returns the security context for the current connection (given handshake completed),
+     * or empty if the current connection is not secure.
+     */
+    public Optional<SecurityContext> securityContext();
 }
diff --git a/jrt/src/com/yahoo/jrt/MaybeTlsCryptoSocket.java b/jrt/src/com/yahoo/jrt/MaybeTlsCryptoSocket.java
index 2090f69c80f..99d47399436 100644
--- a/jrt/src/com/yahoo/jrt/MaybeTlsCryptoSocket.java
+++ b/jrt/src/com/yahoo/jrt/MaybeTlsCryptoSocket.java
@@ -4,6 +4,7 @@ package com.yahoo.jrt;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.channels.SocketChannel;
+import java.util.Optional;
 
 /**
  * A crypto socket for the server side of a connection that
@@ -128,4 +129,5 @@ public class MaybeTlsCryptoSocket implements CryptoSocket {
     @Override public int drain(ByteBuffer dst) throws IOException { return socket.drain(dst); }
     @Override public int write(ByteBuffer src) throws IOException { return socket.write(src); }
     @Override public FlushResult flush() throws IOException { return socket.flush(); }
+    @Override public Optional<SecurityContext> securityContext() { return Optional.ofNullable(socket).flatMap(CryptoSocket::securityContext); }
 }
diff --git a/jrt/src/com/yahoo/jrt/NullCryptoSocket.java b/jrt/src/com/yahoo/jrt/NullCryptoSocket.java
index 83359bb65a5..6124d4918a3 100644
--- a/jrt/src/com/yahoo/jrt/NullCryptoSocket.java
+++ b/jrt/src/com/yahoo/jrt/NullCryptoSocket.java
@@ -5,6 +5,7 @@ package com.yahoo.jrt;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.channels.SocketChannel;
+import java.util.Optional;
 
 
 /**
@@ -30,4 +31,5 @@ public class NullCryptoSocket implements CryptoSocket {
     @Override public int drain(ByteBuffer dst) throws IOException { return 0; }
     @Override public int write(ByteBuffer src) throws IOException { return channel.write(src); }
     @Override public FlushResult flush() throws IOException { return FlushResult.DONE; }
+    @Override public Optional<SecurityContext> securityContext() { return Optional.empty(); }
 }
diff --git a/jrt/src/com/yahoo/jrt/SecurityContext.java b/jrt/src/com/yahoo/jrt/SecurityContext.java
new file mode 100644
index 00000000000..0ac49013f2d
--- /dev/null
+++ b/jrt/src/com/yahoo/jrt/SecurityContext.java
@@ -0,0 +1,24 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.jrt;
+
+import java.security.cert.X509Certificate;
+import java.util.List;
+
+/**
+ * @author bjorncs
+ */
+public class SecurityContext {
+
+    private final List<X509Certificate> peerCertificateChain;
+
+    public SecurityContext(List<X509Certificate> peerCertificateChain) {
+        this.peerCertificateChain = peerCertificateChain;
+    }
+
+    /**
+     * @return the peer certificate chain if the peer was authenticated, empty list if not.
+     */
+    public List<X509Certificate> peerCertificateChain() {
+        return peerCertificateChain;
+    }
+}
diff --git a/jrt/src/com/yahoo/jrt/Target.java b/jrt/src/com/yahoo/jrt/Target.java
index 1bfae3b712d..e18f0bba7b6 100644
--- a/jrt/src/com/yahoo/jrt/Target.java
+++ b/jrt/src/com/yahoo/jrt/Target.java
@@ -2,6 +2,8 @@
 package com.yahoo.jrt;
 
 
+import java.util.Optional;
+
 /**
  * A Target represents a connection endpoint with RPC
  * capabilities. Each such connection has a client and a server
@@ -68,6 +70,11 @@ public abstract class Target {
     public Exception getConnectionLostReason() { return null; }
 
     /**
+     * @return the security context associated with this target, or empty if no connection or is insecure.
+     */
+    public abstract Optional<SecurityContext> getSecurityContext();
+
+    /**
      * Check if this target represents the client side of a
      * connection.
      *
diff --git a/jrt/src/com/yahoo/jrt/TlsCryptoSocket.java b/jrt/src/com/yahoo/jrt/TlsCryptoSocket.java
index 31d76bc7362..e9f72ee12e0 100644
--- a/jrt/src/com/yahoo/jrt/TlsCryptoSocket.java
+++ b/jrt/src/com/yahoo/jrt/TlsCryptoSocket.java
@@ -8,13 +8,19 @@ import javax.net.ssl.SSLEngine;
 import javax.net.ssl.SSLEngineResult;
 import javax.net.ssl.SSLException;
 import javax.net.ssl.SSLHandshakeException;
+import javax.net.ssl.SSLPeerUnverifiedException;
 import javax.net.ssl.SSLSession;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.channels.ClosedChannelException;
 import java.nio.channels.SocketChannel;
+import java.security.cert.X509Certificate;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Optional;
 import java.util.logging.Logger;
 
+import static java.util.stream.Collectors.toList;
 import static javax.net.ssl.SSLEngineResult.HandshakeStatus;
 import static javax.net.ssl.SSLEngineResult.Status;
 
@@ -211,6 +217,22 @@ public class TlsCryptoSocket implements CryptoSocket {
         return wrapBuffer.bytes() > 0 ? FlushResult.NEED_WRITE : FlushResult.DONE;
     }
 
+    @Override
+    public Optional<SecurityContext> securityContext() {
+        try {
+            if (handshakeState != HandshakeState.COMPLETED) {
+                return Optional.empty();
+            }
+            List<X509Certificate> peerCertificateChain =
+                    Arrays.stream(sslEngine.getSession().getPeerCertificates())
+                            .map(X509Certificate.class::cast)
+                            .collect(toList());
+            return Optional.of(new SecurityContext(peerCertificateChain));
+        } catch (SSLPeerUnverifiedException e) { // unverified peer: non-certificate based ciphers or peer did not provide a certificate
+            return Optional.of(new SecurityContext(List.of())); // secure connection, but peer does not have a certificate chain.
+        }
+    }
+
     private boolean handshakeWrap() throws IOException {
         SSLEngineResult result = sslEngineWrap(NULL_BUFFER);
         switch (result.getStatus()) {
diff --git a/jrt/src/com/yahoo/jrt/XorCryptoSocket.java b/jrt/src/com/yahoo/jrt/XorCryptoSocket.java
index db76ac62f3c..fecb6ac54a3 100644
--- a/jrt/src/com/yahoo/jrt/XorCryptoSocket.java
+++ b/jrt/src/com/yahoo/jrt/XorCryptoSocket.java
@@ -6,6 +6,7 @@ import java.nio.ByteBuffer;
 import java.nio.channels.SocketChannel;
 import java.security.SecureRandom;
 import java.util.ArrayDeque;
+import java.util.Optional;
 import java.util.Queue;
 
 /**
@@ -119,4 +120,8 @@ public class XorCryptoSocket implements CryptoSocket {
             return FlushResult.DONE;
         }
     }
+
+    @Override public Optional<SecurityContext> securityContext() {
+        return Optional.empty();
+    }
 }
diff --git a/jrt/tests/com/yahoo/jrt/EchoTest.java b/jrt/tests/com/yahoo/jrt/EchoTest.java
index 8fe98ff3510..16f18afb58c 100644
--- a/jrt/tests/com/yahoo/jrt/EchoTest.java
+++ b/jrt/tests/com/yahoo/jrt/EchoTest.java
@@ -9,8 +9,13 @@ import org.junit.runners.Parameterized;
 import org.junit.runners.Parameterized.Parameter;
 import org.junit.runners.Parameterized.Parameters;
 
+import java.security.cert.X509Certificate;
+import java.util.List;
+
 import static com.yahoo.jrt.CryptoUtils.createTestTlsContext;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
 @RunWith(Parameterized.class)
@@ -23,13 +28,19 @@ public class EchoTest {
     Supervisor client;
     Target     target;
     Values     refValues;
+    SecurityContext securityContext;
 
     private interface MetricsAssertions {
         void assertMetrics(TransportMetrics.Snapshot snapshot) throws AssertionError;
     }
 
+    private interface SecurityContextAssertion {
+        void assertSecurityContext(SecurityContext securityContext) throws AssertionError;
+    }
+
     @Parameter(value = 0) public CryptoEngine crypto;
     @Parameter(value = 1) public MetricsAssertions metricsAssertions;
+    @Parameter(value = 2) public SecurityContextAssertion securityContextAssertion;
 
 
     @Parameters(name = "{0}") public static Object[] engines() {
@@ -39,25 +50,40 @@ public class EchoTest {
                         (MetricsAssertions) metrics -> {
                             assertEquals(1, metrics.serverUnencryptedConnectionsEstablished());
                             assertEquals(1, metrics.clientUnencryptedConnectionsEstablished());
-                        }},
-                {new XorCryptoEngine(), null},
+                        },
+                        null},
+                {
+                        new XorCryptoEngine(),
+                        null,
+                        null},
                 {
                         new TlsCryptoEngine(createTestTlsContext()),
                         (MetricsAssertions) metrics -> {
                             assertEquals(1, metrics.serverTlsConnectionsEstablished());
                             assertEquals(1, metrics.clientTlsConnectionsEstablished());
+                        },
+                        (SecurityContextAssertion) context -> {
+                            List<X509Certificate> chain = context.peerCertificateChain();
+                            assertEquals(1, chain.size());
+                            assertEquals(CryptoUtils.certificate, chain.get(0));
                         }},
                 {
                         new MaybeTlsCryptoEngine(new TlsCryptoEngine(createTestTlsContext()), false),
                         (MetricsAssertions) metrics -> {
                             assertEquals(1, metrics.serverUnencryptedConnectionsEstablished());
                             assertEquals(1, metrics.clientUnencryptedConnectionsEstablished());
-                        }},
+                        },
+                        null},
                 {
                         new MaybeTlsCryptoEngine(new TlsCryptoEngine(createTestTlsContext()), true),
                         (MetricsAssertions) metrics -> {
                              assertEquals(1, metrics.serverTlsConnectionsEstablished());
                              assertEquals(1, metrics.clientTlsConnectionsEstablished());
+                        },
+                        (SecurityContextAssertion) context -> {
+                            List<X509Certificate> chain = context.peerCertificateChain();
+                            assertEquals(1, chain.size());
+                            assertEquals(CryptoUtils.certificate, chain.get(0));
                         }}};
     }
 
@@ -120,6 +146,7 @@ public class EchoTest {
         for (int i = 0; i < p.size(); i++) {
             r.add(p.get(i));
         }
+        securityContext = req.target().getSecurityContext().orElse(null);
     }
 
     @org.junit.Test
@@ -137,5 +164,11 @@ public class EchoTest {
         if (metricsAssertions != null) {
             metricsAssertions.assertMetrics(metrics.snapshot().changesSince(startSnapshot));
         }
+        if (securityContextAssertion != null) {
+            assertNotNull(securityContext);
+            securityContextAssertion.assertSecurityContext(securityContext);
+        } else {
+            assertNull(securityContext);
+        }
     }
 }