gfe-relnote: In QUIC, implement GetHandshakeState() in various handshakers. Not used yet, not protected.

GetHandshakeState() will be used to retrieve current handshake state, and handshaker is the only source of this information. But caller can choose to latch handshake state for performance purpose.

PiperOrigin-RevId: 290142490
Change-Id: I51c33c9ebcbb1ab13eddc0d86197991167642461
diff --git a/quic/core/http/quic_spdy_session_test.cc b/quic/core/http/quic_spdy_session_test.cc
index a0de873..6efaf52 100644
--- a/quic/core/http/quic_spdy_session_test.cc
+++ b/quic/core/http/quic_spdy_session_test.cc
@@ -120,6 +120,9 @@
   bool one_rtt_keys_available() const override {
     return one_rtt_keys_available_;
   }
+  HandshakeState GetHandshakeState() const override {
+    return one_rtt_keys_available() ? HANDSHAKE_COMPLETE : HANDSHAKE_START;
+  }
   const QuicCryptoNegotiatedParameters& crypto_negotiated_params()
       const override {
     return *params_;
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index a29825b..ddb81f5 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -2459,8 +2459,7 @@
 
 char* QuicConnection::GetPacketBuffer() {
   if (version().CanSendCoalescedPackets() &&
-      sent_packet_manager_.handshake_state() <
-          QuicSentPacketManager::HANDSHAKE_CONFIRMED) {
+      sent_packet_manager_.handshake_state() < HANDSHAKE_CONFIRMED) {
     // Do not use writer's packet buffer for coalesced packets which may contain
     // multiple QUIC packets.
     return nullptr;
@@ -4005,8 +4004,7 @@
 SerializedPacketFate QuicConnection::DeterminePacketFate(
     bool is_mtu_discovery) {
   if (version().CanSendCoalescedPackets() &&
-      sent_packet_manager_.handshake_state() <
-          QuicSentPacketManager::HANDSHAKE_CONFIRMED &&
+      sent_packet_manager_.handshake_state() < HANDSHAKE_CONFIRMED &&
       !is_mtu_discovery) {
     // Before receiving ACK for any 1-RTT packets, always try to coalesce
     // packet (except MTU discovery packet).
diff --git a/quic/core/quic_connection.h b/quic/core/quic_connection.h
index 0374cc7..b111f3c 100644
--- a/quic/core/quic_connection.h
+++ b/quic/core/quic_connection.h
@@ -862,8 +862,7 @@
 
   // Whether the handshake completes from this connection's perspective.
   bool IsHandshakeComplete() const {
-    return sent_packet_manager_.handshake_state() >=
-           QuicSentPacketManager::HANDSHAKE_COMPLETE;
+    return sent_packet_manager_.handshake_state() >= HANDSHAKE_COMPLETE;
   }
 
   // Returns the largest received packet number sent by peer.
diff --git a/quic/core/quic_crypto_client_handshaker.cc b/quic/core/quic_crypto_client_handshaker.cc
index 9b6b38d..4d8953a 100644
--- a/quic/core/quic_crypto_client_handshaker.cc
+++ b/quic/core/quic_crypto_client_handshaker.cc
@@ -149,6 +149,10 @@
   return QuicCryptoHandshaker::crypto_message_parser();
 }
 
+HandshakeState QuicCryptoClientHandshaker::GetHandshakeState() const {
+  return one_rtt_keys_available() ? HANDSHAKE_COMPLETE : HANDSHAKE_START;
+}
+
 size_t QuicCryptoClientHandshaker::BufferSizeLimitForLevel(
     EncryptionLevel level) const {
   return QuicCryptoHandshaker::BufferSizeLimitForLevel(level);
diff --git a/quic/core/quic_crypto_client_handshaker.h b/quic/core/quic_crypto_client_handshaker.h
index 2c35796..9729b1d 100644
--- a/quic/core/quic_crypto_client_handshaker.h
+++ b/quic/core/quic_crypto_client_handshaker.h
@@ -45,6 +45,7 @@
   const QuicCryptoNegotiatedParameters& crypto_negotiated_params()
       const override;
   CryptoMessageParser* crypto_message_parser() override;
+  HandshakeState GetHandshakeState() const override;
   size_t BufferSizeLimitForLevel(EncryptionLevel level) const override;
 
   // From QuicCryptoHandshaker
diff --git a/quic/core/quic_crypto_client_stream.cc b/quic/core/quic_crypto_client_stream.cc
index 3286b66..d681eb3 100644
--- a/quic/core/quic_crypto_client_stream.cc
+++ b/quic/core/quic_crypto_client_stream.cc
@@ -86,6 +86,10 @@
   return handshaker_->crypto_message_parser();
 }
 
+HandshakeState QuicCryptoClientStream::GetHandshakeState() const {
+  return handshaker_->GetHandshakeState();
+}
+
 size_t QuicCryptoClientStream::BufferSizeLimitForLevel(
     EncryptionLevel level) const {
   return handshaker_->BufferSizeLimitForLevel(level);
diff --git a/quic/core/quic_crypto_client_stream.h b/quic/core/quic_crypto_client_stream.h
index eb4630b..76d0301 100644
--- a/quic/core/quic_crypto_client_stream.h
+++ b/quic/core/quic_crypto_client_stream.h
@@ -117,6 +117,9 @@
     // Used by QuicCryptoStream to know how much unprocessed data can be
     // buffered at each encryption level.
     virtual size_t BufferSizeLimitForLevel(EncryptionLevel level) const = 0;
+
+    // Returns current handshake state.
+    virtual HandshakeState GetHandshakeState() const = 0;
   };
 
   // ProofHandler is an interface that handles callbacks from the crypto
@@ -162,6 +165,7 @@
       const override;
   CryptoMessageParser* crypto_message_parser() override;
   void OnPacketDecrypted(EncryptionLevel /*level*/) override {}
+  HandshakeState GetHandshakeState() const override;
   size_t BufferSizeLimitForLevel(EncryptionLevel level) const override;
 
   std::string chlo_hash() const;
diff --git a/quic/core/quic_crypto_server_handshaker.cc b/quic/core/quic_crypto_server_handshaker.cc
index 5adbe78..9aff0cc 100644
--- a/quic/core/quic_crypto_server_handshaker.cc
+++ b/quic/core/quic_crypto_server_handshaker.cc
@@ -70,6 +70,7 @@
       process_client_hello_cb_(nullptr),
       encryption_established_(false),
       one_rtt_keys_available_(false),
+      one_rtt_packet_decrypted_(false),
       crypto_negotiated_params_(new QuicCryptoNegotiatedParameters) {}
 
 QuicCryptoServerHandshaker::~QuicCryptoServerHandshaker() {
@@ -365,6 +366,7 @@
 
 void QuicCryptoServerHandshaker::OnPacketDecrypted(EncryptionLevel level) {
   if (level == ENCRYPTION_FORWARD_SECURE) {
+    one_rtt_packet_decrypted_ = true;
     delegate_->NeuterHandshakeData();
   }
 }
@@ -407,6 +409,10 @@
   return QuicCryptoHandshaker::crypto_message_parser();
 }
 
+HandshakeState QuicCryptoServerHandshaker::GetHandshakeState() const {
+  return one_rtt_packet_decrypted_ ? HANDSHAKE_COMPLETE : HANDSHAKE_START;
+}
+
 size_t QuicCryptoServerHandshaker::BufferSizeLimitForLevel(
     EncryptionLevel level) const {
   return QuicCryptoHandshaker::BufferSizeLimitForLevel(level);
diff --git a/quic/core/quic_crypto_server_handshaker.h b/quic/core/quic_crypto_server_handshaker.h
index f873c93..e6d5cae 100644
--- a/quic/core/quic_crypto_server_handshaker.h
+++ b/quic/core/quic_crypto_server_handshaker.h
@@ -59,6 +59,7 @@
   const QuicCryptoNegotiatedParameters& crypto_negotiated_params()
       const override;
   CryptoMessageParser* crypto_message_parser() override;
+  HandshakeState GetHandshakeState() const override;
   size_t BufferSizeLimitForLevel(EncryptionLevel level) const override;
 
   // From QuicCryptoHandshaker
@@ -228,6 +229,7 @@
 
   bool encryption_established_;
   bool one_rtt_keys_available_;
+  bool one_rtt_packet_decrypted_;
   QuicReferenceCountedPointer<QuicCryptoNegotiatedParameters>
       crypto_negotiated_params_;
 };
diff --git a/quic/core/quic_crypto_server_stream.cc b/quic/core/quic_crypto_server_stream.cc
index 19a2b34..abc681f 100644
--- a/quic/core/quic_crypto_server_stream.cc
+++ b/quic/core/quic_crypto_server_stream.cc
@@ -154,6 +154,10 @@
   handshaker_->OnPacketDecrypted(level);
 }
 
+HandshakeState QuicCryptoServerStream::GetHandshakeState() const {
+  return handshaker_->GetHandshakeState();
+}
+
 size_t QuicCryptoServerStream::BufferSizeLimitForLevel(
     EncryptionLevel level) const {
   return handshaker_->BufferSizeLimitForLevel(level);
diff --git a/quic/core/quic_crypto_server_stream.h b/quic/core/quic_crypto_server_stream.h
index 09ccacb..db95116 100644
--- a/quic/core/quic_crypto_server_stream.h
+++ b/quic/core/quic_crypto_server_stream.h
@@ -128,6 +128,9 @@
     // Used by QuicCryptoStream to parse data received on this stream.
     virtual CryptoMessageParser* crypto_message_parser() = 0;
 
+    // Get current handshake state.
+    virtual HandshakeState GetHandshakeState() const = 0;
+
     // Used by QuicCryptoStream to know how much unprocessed data can be
     // buffered at each encryption level.
     virtual size_t BufferSizeLimitForLevel(EncryptionLevel level) const = 0;
@@ -172,6 +175,7 @@
       const override;
   CryptoMessageParser* crypto_message_parser() override;
   void OnPacketDecrypted(EncryptionLevel level) override;
+  HandshakeState GetHandshakeState() const override;
   size_t BufferSizeLimitForLevel(EncryptionLevel level) const override;
   void OnSuccessfulVersionNegotiation(
       const ParsedQuicVersion& version) override;
diff --git a/quic/core/quic_crypto_stream.h b/quic/core/quic_crypto_stream.h
index 15f5a3d..7b9414c 100644
--- a/quic/core/quic_crypto_stream.h
+++ b/quic/core/quic_crypto_stream.h
@@ -85,6 +85,9 @@
   // Called when a packet of encryption |level| has been successfully decrypted.
   virtual void OnPacketDecrypted(EncryptionLevel level) = 0;
 
+  // Returns current handshake state.
+  virtual HandshakeState GetHandshakeState() const = 0;
+
   // Returns the maximum number of bytes that can be buffered at a particular
   // encryption level |level|.
   virtual size_t BufferSizeLimitForLevel(EncryptionLevel level) const;
diff --git a/quic/core/quic_crypto_stream_test.cc b/quic/core/quic_crypto_stream_test.cc
index 62b5633..1b391ab 100644
--- a/quic/core/quic_crypto_stream_test.cc
+++ b/quic/core/quic_crypto_stream_test.cc
@@ -57,6 +57,7 @@
     return QuicCryptoHandshaker::crypto_message_parser();
   }
   void OnPacketDecrypted(EncryptionLevel /*level*/) override {}
+  HandshakeState GetHandshakeState() const override { return HANDSHAKE_START; }
 
  private:
   QuicReferenceCountedPointer<QuicCryptoNegotiatedParameters> params_;
diff --git a/quic/core/quic_sent_packet_manager.h b/quic/core/quic_sent_packet_manager.h
index 39a559c..69a4d23 100644
--- a/quic/core/quic_sent_packet_manager.h
+++ b/quic/core/quic_sent_packet_manager.h
@@ -107,25 +107,6 @@
     PTO_MODE,
   };
 
-  // Handshake state of this connection.
-  enum HandshakeState {
-    // Initial state.
-    HANDSHAKE_START,
-    // Only used in IETF QUIC with TLS handshake. State proceeds to
-    // HANDSHAKE_PROCESSED after a packet of HANDSHAKE packet number space
-    // gets successfully processed, and the initial key can be dropped.
-    HANDSHAKE_PROCESSED,
-    // In QUIC crypto, state proceeds to HANDSHAKE_COMPLETE if client receives
-    // SHLO or server successfully processes an ENCRYPTION_FORWARD_SECURE
-    // packet, such that the handshake packets can be neutered. In IETF QUIC
-    // with TLS handshake, state proceeds to HANDSHAKE_COMPLETE once the
-    // endpoint has both 1-RTT send and receive keys.
-    HANDSHAKE_COMPLETE,
-    // Only used in IETF QUIC with TLS handshake. State proceeds to
-    // HANDSHAKE_CONFIRMED if a 1-RTT packet gets acknowledged.
-    HANDSHAKE_CONFIRMED,
-  };
-
   QuicSentPacketManager(Perspective perspective,
                         const QuicClock* clock,
                         QuicRandom* random,
@@ -596,6 +577,9 @@
   PacingSender pacing_sender_;
 
   // Indicates current handshake state.
+  // TODO(fayang): Stop inferring handshake state, instead, retrieve handshake
+  // state from SessionNotifier::GetHandshakeState and use this variable purely
+  // for performance purpose.
   HandshakeState handshake_state_;
 
   // Records bandwidth from server to client in normal operation, over periods
diff --git a/quic/core/quic_sent_packet_manager_test.cc b/quic/core/quic_sent_packet_manager_test.cc
index a36d86a..25102a0 100644
--- a/quic/core/quic_sent_packet_manager_test.cc
+++ b/quic/core/quic_sent_packet_manager_test.cc
@@ -2852,8 +2852,7 @@
 }
 
 TEST_F(QuicSentPacketManagerTest, ForwardSecurePacketAcked) {
-  EXPECT_LT(manager_.handshake_state(),
-            QuicSentPacketManager::HANDSHAKE_CONFIRMED);
+  EXPECT_LT(manager_.handshake_state(), HANDSHAKE_CONFIRMED);
   SendDataPacket(1, ENCRYPTION_INITIAL);
   // Ack packet 1.
   ExpectAck(1);
@@ -2863,8 +2862,7 @@
   EXPECT_EQ(PACKETS_NEWLY_ACKED,
             manager_.OnAckFrameEnd(clock_.Now(), QuicPacketNumber(1),
                                    ENCRYPTION_INITIAL));
-  EXPECT_LT(manager_.handshake_state(),
-            QuicSentPacketManager::HANDSHAKE_CONFIRMED);
+  EXPECT_LT(manager_.handshake_state(), HANDSHAKE_CONFIRMED);
 
   SendDataPacket(2, ENCRYPTION_ZERO_RTT);
   // Ack packet 2.
@@ -2875,8 +2873,7 @@
   EXPECT_EQ(PACKETS_NEWLY_ACKED,
             manager_.OnAckFrameEnd(clock_.Now(), QuicPacketNumber(2),
                                    ENCRYPTION_FORWARD_SECURE));
-  EXPECT_LT(manager_.handshake_state(),
-            QuicSentPacketManager::HANDSHAKE_CONFIRMED);
+  EXPECT_LT(manager_.handshake_state(), HANDSHAKE_CONFIRMED);
 
   SendDataPacket(3, ENCRYPTION_FORWARD_SECURE);
   // Ack packet 3.
@@ -2887,8 +2884,7 @@
   EXPECT_EQ(PACKETS_NEWLY_ACKED,
             manager_.OnAckFrameEnd(clock_.Now(), QuicPacketNumber(3),
                                    ENCRYPTION_FORWARD_SECURE));
-  EXPECT_EQ(manager_.handshake_state(),
-            QuicSentPacketManager::HANDSHAKE_CONFIRMED);
+  EXPECT_EQ(manager_.handshake_state(), HANDSHAKE_CONFIRMED);
 }
 
 TEST_F(QuicSentPacketManagerTest, PtoTimeoutIncludesMaxAckDelay) {
diff --git a/quic/core/quic_session.cc b/quic/core/quic_session.cc
index 41f8f76..8bf1410 100644
--- a/quic/core/quic_session.cc
+++ b/quic/core/quic_session.cc
@@ -1970,6 +1970,10 @@
   return !streams_waiting_for_acks_.empty();
 }
 
+HandshakeState QuicSession::GetHandshakeState() const {
+  return GetCryptoStream()->GetHandshakeState();
+}
+
 WriteStreamDataResult QuicSession::WriteStreamData(QuicStreamId id,
                                                    QuicStreamOffset offset,
                                                    QuicByteCount data_length,
diff --git a/quic/core/quic_session.h b/quic/core/quic_session.h
index 9762dec..0a79842 100644
--- a/quic/core/quic_session.h
+++ b/quic/core/quic_session.h
@@ -155,6 +155,7 @@
   bool IsFrameOutstanding(const QuicFrame& frame) const override;
   bool HasUnackedCryptoData() const override;
   bool HasUnackedStreamData() const override;
+  HandshakeState GetHandshakeState() const override;
 
   // QuicStreamIdManager::DelegateInterface methods:
   void OnError(QuicErrorCode error_code, std::string error_details) override;
diff --git a/quic/core/quic_session_test.cc b/quic/core/quic_session_test.cc
index 042a671..4686124 100644
--- a/quic/core/quic_session_test.cc
+++ b/quic/core/quic_session_test.cc
@@ -115,6 +115,9 @@
     return QuicCryptoHandshaker::crypto_message_parser();
   }
   void OnPacketDecrypted(EncryptionLevel /*level*/) override {}
+  HandshakeState GetHandshakeState() const override {
+    return one_rtt_keys_available() ? HANDSHAKE_COMPLETE : HANDSHAKE_START;
+  }
 
   MOCK_METHOD0(OnCanWrite, void());
   bool HasPendingCryptoRetransmission() const override { return false; }
diff --git a/quic/core/quic_types.h b/quic/core/quic_types.h
index 917e880..d333fe5 100644
--- a/quic/core/quic_types.h
+++ b/quic/core/quic_types.h
@@ -718,6 +718,27 @@
 QUIC_EXPORT_PRIVATE std::string QuicConnectionCloseTypeString(
     QuicConnectionCloseType type);
 
+// Indicate handshake state of a connection.
+enum HandshakeState {
+  // Initial state.
+  HANDSHAKE_START,
+  // Only used in IETF QUIC with TLS handshake. State proceeds to
+  // HANDSHAKE_PROCESSED after a packet of HANDSHAKE packet number space
+  // gets successfully processed, and the initial key can be dropped.
+  HANDSHAKE_PROCESSED,
+  // In QUIC crypto, state proceeds to HANDSHAKE_COMPLETE if client receives
+  // SHLO or server successfully processes an ENCRYPTION_FORWARD_SECURE
+  // packet, such that the handshake packets can be neutered. In IETF QUIC
+  // with TLS handshake, state proceeds to HANDSHAKE_COMPLETE once the
+  // endpoint has both 1-RTT send and receive keys.
+  HANDSHAKE_COMPLETE,
+  // Only used in IETF QUIC with TLS handshake. State proceeds to
+  // HANDSHAKE_CONFIRMED if a 1-RTT packet gets acknowledged.
+  // TODO(fayang): implement HANDSHAKE_DONE frame to drive handshake to
+  // confirmation according to https://github.com/quicwg/base-drafts/pull/3145.
+  HANDSHAKE_CONFIRMED,
+};
+
 }  // namespace quic
 
 #endif  // QUICHE_QUIC_CORE_QUIC_TYPES_H_
diff --git a/quic/core/session_notifier_interface.h b/quic/core/session_notifier_interface.h
index 89b4a9c..4fb8b50 100644
--- a/quic/core/session_notifier_interface.h
+++ b/quic/core/session_notifier_interface.h
@@ -40,6 +40,9 @@
 
   // Returns true if any stream is waiting for acks.
   virtual bool HasUnackedStreamData() const = 0;
+
+  // Returns current handshake state.
+  virtual HandshakeState GetHandshakeState() const = 0;
 };
 
 }  // namespace quic
diff --git a/quic/core/tls_client_handshaker.cc b/quic/core/tls_client_handshaker.cc
index d3cd24f..53ea536 100644
--- a/quic/core/tls_client_handshaker.cc
+++ b/quic/core/tls_client_handshaker.cc
@@ -238,6 +238,16 @@
   return TlsHandshaker::crypto_message_parser();
 }
 
+HandshakeState TlsClientHandshaker::GetHandshakeState() const {
+  if (one_rtt_keys_available_) {
+    return HANDSHAKE_COMPLETE;
+  }
+  if (state_ >= STATE_ENCRYPTION_HANDSHAKE_DATA_SENT) {
+    return HANDSHAKE_PROCESSED;
+  }
+  return HANDSHAKE_START;
+}
+
 size_t TlsClientHandshaker::BufferSizeLimitForLevel(
     EncryptionLevel level) const {
   return TlsHandshaker::BufferSizeLimitForLevel(level);
diff --git a/quic/core/tls_client_handshaker.h b/quic/core/tls_client_handshaker.h
index 14ef352..a2fa96a 100644
--- a/quic/core/tls_client_handshaker.h
+++ b/quic/core/tls_client_handshaker.h
@@ -49,6 +49,7 @@
   const QuicCryptoNegotiatedParameters& crypto_negotiated_params()
       const override;
   CryptoMessageParser* crypto_message_parser() override;
+  HandshakeState GetHandshakeState() const override;
   size_t BufferSizeLimitForLevel(EncryptionLevel level) const override;
 
   // Override to drop initial keys if trying to write ENCRYPTION_HANDSHAKE data.
diff --git a/quic/core/tls_handshaker.h b/quic/core/tls_handshaker.h
index d0d9bf8..fd554aa 100644
--- a/quic/core/tls_handshaker.h
+++ b/quic/core/tls_handshaker.h
@@ -52,6 +52,7 @@
   virtual const QuicCryptoNegotiatedParameters& crypto_negotiated_params()
       const = 0;
   virtual CryptoMessageParser* crypto_message_parser() { return this; }
+  virtual HandshakeState GetHandshakeState() const = 0;
   size_t BufferSizeLimitForLevel(EncryptionLevel level) const;
 
  protected:
diff --git a/quic/core/tls_handshaker_test.cc b/quic/core/tls_handshaker_test.cc
index 9935839..b175d55 100644
--- a/quic/core/tls_handshaker_test.cc
+++ b/quic/core/tls_handshaker_test.cc
@@ -182,6 +182,10 @@
 
   void OnPacketDecrypted(EncryptionLevel /*level*/) override {}
 
+  HandshakeState GetHandshakeState() const override {
+    return handshaker()->GetHandshakeState();
+  }
+
   const std::vector<std::pair<std::string, EncryptionLevel>>& pending_writes() {
     return pending_writes_;
   }
diff --git a/quic/core/tls_server_handshaker.cc b/quic/core/tls_server_handshaker.cc
index 500cecf..88eb326 100644
--- a/quic/core/tls_server_handshaker.cc
+++ b/quic/core/tls_server_handshaker.cc
@@ -141,6 +141,16 @@
   return TlsHandshaker::crypto_message_parser();
 }
 
+HandshakeState TlsServerHandshaker::GetHandshakeState() const {
+  if (one_rtt_keys_available_) {
+    return HANDSHAKE_COMPLETE;
+  }
+  if (state_ >= STATE_ENCRYPTION_HANDSHAKE_DATA_PROCESSED) {
+    return HANDSHAKE_PROCESSED;
+  }
+  return HANDSHAKE_START;
+}
+
 size_t TlsServerHandshaker::BufferSizeLimitForLevel(
     EncryptionLevel level) const {
   return TlsHandshaker::BufferSizeLimitForLevel(level);
diff --git a/quic/core/tls_server_handshaker.h b/quic/core/tls_server_handshaker.h
index 4bf5d22..919af91 100644
--- a/quic/core/tls_server_handshaker.h
+++ b/quic/core/tls_server_handshaker.h
@@ -56,6 +56,7 @@
   const QuicCryptoNegotiatedParameters& crypto_negotiated_params()
       const override;
   CryptoMessageParser* crypto_message_parser() override;
+  HandshakeState GetHandshakeState() const override;
   size_t BufferSizeLimitForLevel(EncryptionLevel level) const override;
 
  protected:
diff --git a/quic/test_tools/quic_test_utils.h b/quic/test_tools/quic_test_utils.h
index 62eeed9..bdc4670 100644
--- a/quic/test_tools/quic_test_utils.h
+++ b/quic/test_tools/quic_test_utils.h
@@ -689,6 +689,7 @@
       const override;
   CryptoMessageParser* crypto_message_parser() override;
   void OnPacketDecrypted(EncryptionLevel /*level*/) override {}
+  HandshakeState GetHandshakeState() const override { return HANDSHAKE_START; }
 
  private:
   QuicReferenceCountedPointer<QuicCryptoNegotiatedParameters> params_;
@@ -1122,6 +1123,7 @@
   MOCK_CONST_METHOD1(IsFrameOutstanding, bool(const QuicFrame&));
   MOCK_CONST_METHOD0(HasUnackedCryptoData, bool());
   MOCK_CONST_METHOD0(HasUnackedStreamData, bool());
+  MOCK_CONST_METHOD0(GetHandshakeState, HandshakeState());
 };
 
 // Creates a client session for testing.
diff --git a/quic/test_tools/simple_session_notifier.cc b/quic/test_tools/simple_session_notifier.cc
index 9e67c89..0b7fb7d 100644
--- a/quic/test_tools/simple_session_notifier.cc
+++ b/quic/test_tools/simple_session_notifier.cc
@@ -453,6 +453,10 @@
   return false;
 }
 
+HandshakeState SimpleSessionNotifier::GetHandshakeState() const {
+  return HANDSHAKE_COMPLETE;
+}
+
 bool SimpleSessionNotifier::OnControlFrameAcked(const QuicFrame& frame) {
   QuicControlFrameId id = GetControlFrameId(frame);
   if (id == kInvalidControlFrameId) {
diff --git a/quic/test_tools/simple_session_notifier.h b/quic/test_tools/simple_session_notifier.h
index b3d5d72..c9b8964 100644
--- a/quic/test_tools/simple_session_notifier.h
+++ b/quic/test_tools/simple_session_notifier.h
@@ -77,6 +77,7 @@
   bool IsFrameOutstanding(const QuicFrame& frame) const override;
   bool HasUnackedCryptoData() const override;
   bool HasUnackedStreamData() const override;
+  HandshakeState GetHandshakeState() const override;
 
  private:
   struct StreamState {
diff --git a/quic/test_tools/simulator/quic_endpoint.cc b/quic/test_tools/simulator/quic_endpoint.cc
index 626cc64..6ec84b0 100644
--- a/quic/test_tools/simulator/quic_endpoint.cc
+++ b/quic/test_tools/simulator/quic_endpoint.cc
@@ -199,6 +199,10 @@
   return false;
 }
 
+HandshakeState QuicEndpoint::GetHandshakeState() const {
+  return HANDSHAKE_COMPLETE;
+}
+
 WriteStreamDataResult QuicEndpoint::DataProducer::WriteStreamData(
     QuicStreamId /*id*/,
     QuicStreamOffset /*offset*/,
diff --git a/quic/test_tools/simulator/quic_endpoint.h b/quic/test_tools/simulator/quic_endpoint.h
index 06c90c4..b5fc8e2 100644
--- a/quic/test_tools/simulator/quic_endpoint.h
+++ b/quic/test_tools/simulator/quic_endpoint.h
@@ -98,6 +98,7 @@
   bool IsFrameOutstanding(const QuicFrame& frame) const override;
   bool HasUnackedCryptoData() const override;
   bool HasUnackedStreamData() const override;
+  HandshakeState GetHandshakeState() const override;
   // End SessionNotifierInterface implementation.
 
  private: