From 123febeb9899ff2d84a26aee8ab89d9f977f8103 Mon Sep 17 00:00:00 2001
From: Xue-Lei Andrew Fan <xuelei@openjdk.org>
Date: Sun, 20 Oct 2019 13:42:44 -0700
Subject: [PATCH] 8229733: TLS message handling improvements

Reviewed-by: jnimeh, rhalade, ahgross
---
 .../sun/security/ssl/HandshakeContext.java    |  3 +-
 .../security/ssl/PostHandshakeContext.java    | 40 +++++++++++++++----
 .../sun/security/ssl/TransportContext.java    |  5 ++-
 3 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/src/java.base/share/classes/sun/security/ssl/HandshakeContext.java b/src/java.base/share/classes/sun/security/ssl/HandshakeContext.java
index 1f597b55a29..f69a9288745 100644
--- a/src/java.base/share/classes/sun/security/ssl/HandshakeContext.java
+++ b/src/java.base/share/classes/sun/security/ssl/HandshakeContext.java
@@ -209,7 +209,7 @@ abstract class HandshakeContext implements ConnectionContext {
     /**
      * Constructor for PostHandshakeContext
      */
-    HandshakeContext(TransportContext conContext) {
+    protected HandshakeContext(TransportContext conContext) {
         this.sslContext = conContext.sslContext;
         this.conContext = conContext;
         this.sslConfig = conContext.sslConfig;
@@ -219,6 +219,7 @@ abstract class HandshakeContext implements ConnectionContext {
         this.handshakeOutput = new HandshakeOutStream(conContext.outputRecord);
         this.delegatedActions = new LinkedList<>();
 
+        this.handshakeConsumers = new LinkedHashMap<>();
         this.handshakeProducers = null;
         this.handshakeHash = null;
         this.activeProtocols = null;
diff --git a/src/java.base/share/classes/sun/security/ssl/PostHandshakeContext.java b/src/java.base/share/classes/sun/security/ssl/PostHandshakeContext.java
index 516b6554e64..e0af48719d2 100644
--- a/src/java.base/share/classes/sun/security/ssl/PostHandshakeContext.java
+++ b/src/java.base/share/classes/sun/security/ssl/PostHandshakeContext.java
@@ -30,17 +30,11 @@ import java.nio.BufferOverflowException;
 import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
-import java.util.LinkedHashMap;
-import java.util.Map;
 
 /**
  * A compact implementation of HandshakeContext for post-handshake messages
  */
 final class PostHandshakeContext extends HandshakeContext {
-    private final static Map<Byte, SSLConsumer> consumers = Map.of(
-        SSLHandshake.KEY_UPDATE.id, SSLHandshake.KEY_UPDATE,
-        SSLHandshake.NEW_SESSION_TICKET.id, SSLHandshake.NEW_SESSION_TICKET);
-
     PostHandshakeContext(TransportContext context) throws IOException {
         super(context);
 
@@ -49,10 +43,23 @@ final class PostHandshakeContext extends HandshakeContext {
                 "Post-handshake not supported in " + negotiatedProtocol.name);
         }
 
-        this.localSupportedSignAlgs = new ArrayList<SignatureScheme>(
+        this.localSupportedSignAlgs = new ArrayList<>(
             context.conSession.getLocalSupportedSignatureSchemes());
 
-        handshakeConsumers = new LinkedHashMap<>(consumers);
+        // Add the potential post-handshake consumers.
+        if (context.sslConfig.isClientMode) {
+            handshakeConsumers.putIfAbsent(
+                    SSLHandshake.KEY_UPDATE.id,
+                    SSLHandshake.KEY_UPDATE);
+            handshakeConsumers.putIfAbsent(
+                    SSLHandshake.NEW_SESSION_TICKET.id,
+                    SSLHandshake.NEW_SESSION_TICKET);
+        } else {
+            handshakeConsumers.putIfAbsent(
+                    SSLHandshake.KEY_UPDATE.id,
+                    SSLHandshake.KEY_UPDATE);
+        }
+
         handshakeFinished = true;
         handshakeSession = context.conSession;
     }
@@ -83,4 +90,21 @@ final class PostHandshakeContext extends HandshakeContext {
                     SSLHandshake.nameOf(handshakeType), be);
         }
     }
+
+    static boolean isConsumable(TransportContext context, byte handshakeType) {
+        if (handshakeType == SSLHandshake.KEY_UPDATE.id) {
+            // The KeyUpdate handshake message does not apply to TLS 1.2 and
+            // previous protocols.
+            return context.protocolVersion.useTLS13PlusSpec();
+        }
+
+        if (handshakeType == SSLHandshake.NEW_SESSION_TICKET.id) {
+            // The new session ticket handshake message could be consumer in
+            // client side only.
+            return context.sslConfig.isClientMode;
+        }
+
+        // No more post-handshake message supported currently.
+        return false;
+    }
 }
diff --git a/src/java.base/share/classes/sun/security/ssl/TransportContext.java b/src/java.base/share/classes/sun/security/ssl/TransportContext.java
index 233e1bc75e9..92d562ef9a6 100644
--- a/src/java.base/share/classes/sun/security/ssl/TransportContext.java
+++ b/src/java.base/share/classes/sun/security/ssl/TransportContext.java
@@ -164,12 +164,13 @@ final class TransportContext implements ConnectionContext {
                                             " message: " +
                                             SSLHandshake.nameOf(type));
                         }
-                        if (type == SSLHandshake.KEY_UPDATE.id &&
-                                !protocolVersion.useTLS13PlusSpec()) {
+
+                        if (!PostHandshakeContext.isConsumable(this, type)) {
                             throw fatal(Alert.UNEXPECTED_MESSAGE,
                                     "Unexpected post-handshake message: " +
                                     SSLHandshake.nameOf(type));
                         }
+
                         handshakeContext = new PostHandshakeContext(this);
                     } else {
                         handshakeContext = sslConfig.isClientMode ?