Store PSK info in handshake info

sunnysingh85 · Jun 14, 2024 · 99d644a · 99d644a
1 parent 0ef6c0a
commit 99d644a
Show file tree

Hide file tree

Showing 5 changed files with 50 additions and 10 deletions.
diff --git a/zuul-core/src/main/java/com/netflix/netty/common/ssl/SslHandshakeInfo.java b/zuul-core/src/main/java/com/netflix/netty/common/ssl/SslHandshakeInfo.java
@@ -16,6 +16,7 @@
 
 package com.netflix.netty.common.ssl;
 
+import com.netflix.zuul.netty.server.psk.ClientPSKIdentityInfo;
 import io.netty.handler.ssl.ClientAuth;
 import java.security.cert.Certificate;
 import java.security.cert.X509Certificate;
@@ -31,20 +32,26 @@ public class SslHandshakeInfo {
     private final Certificate serverCertificate;
     private final X509Certificate clientCertificate;
     private final boolean isOfIntermediary;
+    private final boolean usingExternalPSK;
+    private final ClientPSKIdentityInfo clientPSKIdentityInfo;
 
     public SslHandshakeInfo(
             boolean isOfIntermediary,
             String protocol,
             String cipherSuite,
             ClientAuth clientAuthRequirement,
             Certificate serverCertificate,
-            X509Certificate clientCertificate) {
+            X509Certificate clientCertificate,
+            boolean usingExternalPSK,
+            ClientPSKIdentityInfo clientPSKIdentityInfo) {
         this.protocol = protocol;
         this.cipherSuite = cipherSuite;
         this.clientAuthRequirement = clientAuthRequirement;
         this.serverCertificate = serverCertificate;
         this.clientCertificate = clientCertificate;
         this.isOfIntermediary = isOfIntermediary;
+        this.usingExternalPSK = usingExternalPSK;
+        this.clientPSKIdentityInfo = clientPSKIdentityInfo;
     }
 
     public boolean isOfIntermediary() {
@@ -71,6 +78,14 @@ public X509Certificate getClientCertificate() {
         return clientCertificate;
     }
 
+    public boolean usingExternalPSK() {
+        return usingExternalPSK;
+    }
+
+    public ClientPSKIdentityInfo geClientPSKIdentityInfo() {
+        return clientPSKIdentityInfo;
+    }
+
     @Override
     public String toString() {
         return "SslHandshakeInfo{" + "protocol='"

diff --git a/zuul-core/src/main/java/com/netflix/zuul/netty/server/psk/ClientPSKIdentityInfo.java b/zuul-core/src/main/java/com/netflix/zuul/netty/server/psk/ClientPSKIdentityInfo.java
@@ -0,0 +1,3 @@
+package com.netflix.zuul.netty.server.psk;
+
+public record ClientPSKIdentityInfo(byte[] clientPSKIdentity) {}
diff --git a/zuul-core/src/main/java/com/netflix/zuul/netty/server/psk/ExternalTlsPskProvider.java b/zuul-core/src/main/java/com/netflix/zuul/netty/server/psk/ExternalTlsPskProvider.java
@@ -16,6 +16,7 @@
 
 package com.netflix.zuul.netty.server.psk;
 
+
 public interface ExternalTlsPskProvider {
     byte[] provide(byte[] clientPskIdentity, byte[] clientRandom) throws PskCreationFailureException;
 }
diff --git a/zuul-core/src/main/java/com/netflix/zuul/netty/server/psk/TlsPskHandler.java b/zuul-core/src/main/java/com/netflix/zuul/netty/server/psk/TlsPskHandler.java
@@ -27,6 +27,7 @@
 import io.netty.channel.ChannelPromise;
 import io.netty.handler.codec.ByteToMessageDecoder;
 import io.netty.handler.ssl.SslHandshakeCompletionEvent;
+import io.netty.util.AttributeKey;
 import io.netty.util.ReferenceCountUtil;
 import lombok.SneakyThrows;
 import org.bouncycastle.tls.AbstractTlsServer;
@@ -74,6 +75,12 @@ public class TlsPskHandler extends ByteToMessageDecoder
             CipherSuite.TLS_AES_256_GCM_SHA384,
             "TLS_AES_256_GCM_SHA384");
 
+    public static final AttributeKey<ClientPSKIdentityInfo> CLIENT_PSK_IDENTITY_ATTRIBUTE_KEY =
+            AttributeKey.newInstance("_client_psk_identity_info");
+
+    public static final AttributeKey<Boolean> TLS_HANDSHAKE_USING_EXTERNAL_PSK =
+            AttributeKey.newInstance("_tls_handshake_using_external_psk");
+
     private final Registry registry;
     private final ExternalTlsPskProvider externalTlsPskProvider;
 
@@ -264,15 +271,17 @@ protected Vector getProtocolNames() {
 
         @Override
         public void notifyHandshakeBeginning() throws IOException {
+            pskTimings.recordHandshakeStarting();
+            this.ctx.channel().attr(TLS_HANDSHAKE_USING_EXTERNAL_PSK).set(false);
             // TODO: sunnys - handshake timeouts
             super.notifyHandshakeBeginning();
-            pskTimings.recordHandshakeStarting();
         }
 
         @Override
         public void notifyHandshakeComplete() throws IOException {
-            super.notifyHandshakeComplete();
             pskTimings.recordHandshakeComplete();
+            this.ctx.channel().attr(TLS_HANDSHAKE_USING_EXTERNAL_PSK).set(true);
+            super.notifyHandshakeComplete();
             ctx.fireUserEventTriggered(SslHandshakeCompletionEvent.SUCCESS);
         }
 
@@ -301,6 +310,7 @@ public TlsPSKExternal getExternalPSK(Vector clientPskIdentities) {
             byte[] clientPskIdentity = ((PskIdentity)clientPskIdentities.get(0)).getIdentity();
             byte[] psk;
             try{
+                this.ctx.channel().attr(CLIENT_PSK_IDENTITY_ATTRIBUTE_KEY).set(new ClientPSKIdentityInfo(clientPskIdentity));
                 psk = externalTlsPskProvider.provide(clientPskIdentity, this.context.getSecurityParametersHandshake().getClientRandom());
             }catch (PskCreationFailureException e) {
                 throw switch (e.getTlsAlertMessage()) {

diff --git a/zuul-core/src/main/java/com/netflix/zuul/netty/server/ssl/SslHandshakeInfoHandler.java b/zuul-core/src/main/java/com/netflix/zuul/netty/server/ssl/SslHandshakeInfoHandler.java
@@ -23,6 +23,7 @@
 import com.netflix.spectator.api.NoopRegistry;
 import com.netflix.spectator.api.Registry;
 import com.netflix.zuul.netty.ChannelUtils;
+import com.netflix.zuul.netty.server.psk.ClientPSKIdentityInfo;
 import com.netflix.zuul.netty.server.psk.TlsPskHandler;
 import com.netflix.zuul.passport.CurrentPassport;
 import com.netflix.zuul.passport.PassportState;
@@ -34,16 +35,15 @@
 import io.netty.handler.ssl.SslHandler;
 import io.netty.handler.ssl.SslHandshakeCompletionEvent;
 import io.netty.util.AttributeKey;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import javax.net.ssl.SSLException;
-import javax.net.ssl.SSLSession;
 import java.nio.channels.ClosedChannelException;
 import java.security.cert.Certificate;
 import java.security.cert.X509Certificate;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import javax.net.ssl.SSLException;
+import javax.net.ssl.SSLSession;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * Stores info about the client and server's SSL certificates in the context, after a successful handshake.
@@ -103,13 +103,23 @@ public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exc
                         serverCert = session.getLocalCertificates()[0];
                     }
 
+                    Boolean tlsHandshakeUsingExternalPSK = ctx.channel()
+                            .attr(TlsPskHandler.TLS_HANDSHAKE_USING_EXTERNAL_PSK)
+                            .get();
+
+                    ClientPSKIdentityInfo clientPSKIdentityInfo = ctx.channel()
+                            .attr(TlsPskHandler.CLIENT_PSK_IDENTITY_ATTRIBUTE_KEY)
+                            .get();
+
                     SslHandshakeInfo info = new SslHandshakeInfo(
                             isSSlFromIntermediary,
                             session.getProtocol(),
                             session.getCipherSuite(),
                             clientAuth,
                             serverCert,
-                            peerCert);
+                            peerCert,
+                            tlsHandshakeUsingExternalPSK,
+                            clientPSKIdentityInfo);
                     ctx.channel().attr(ATTR_SSL_INFO).set(info);
 
                     // Metrics.
@@ -134,7 +144,8 @@ public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exc
                         // without sending anything.
                         // So don't treat these as SSL handshake failures.
                         logger.debug(
-                                "Client closed connection or it idle timed-out without doing an ssl handshake. , client_ip = {}, channel_info = {}",
+                                "Client closed connection or it idle timed-out without doing an ssl handshake. ,"
+                                        + " client_ip = {}, channel_info = {}",
                                 clientIP,
                                 ChannelUtils.channelInfoForLogging(ctx.channel()));
                     } else if (cause instanceof SSLException