Protected by FLAGS_enable_quic_stateless_reject_support. Add support

net/tools/quic/quic_dispatcher_test.cc

Index: net/tools/quic/quic_dispatcher_test.cc
diff --git a/net/tools/quic/quic_dispatcher_test.cc b/net/tools/quic/quic_dispatcher_test.cc
index 129f7402b600c827290f7ff127b84b7f40cb88a4..95b0030fb350d3fb8f3353908c7e301817077bf1 100644
--- a/net/tools/quic/quic_dispatcher_test.cc
+++ b/net/tools/quic/quic_dispatcher_test.cc
@@ -4,6 +4,7 @@
 
 #include "net/tools/quic/quic_dispatcher.h"
 
+#include <ostream>
 #include <string>
 
 #include "base/strings/string_piece.h"
@@ -31,6 +32,7 @@ using net::test::MockConnection;
 using net::test::MockSession;
 using net::test::ValueRestore;
 using std::string;
+using std::vector;
 using testing::DoAll;
 using testing::InSequence;
 using testing::Invoke;
@@ -52,21 +54,28 @@ class TestServerSession : public QuicServerSession {
   MOCK_METHOD1(CreateIncomingDataStream, QuicDataStream*(QuicStreamId id));
   MOCK_METHOD0(CreateOutgoingDataStream, QuicDataStream*());
 
+  void SetCryptoStream(QuicCryptoServerStream* crypto_stream) {
+    crypto_stream_ = crypto_stream;
+  }
+
+  QuicCryptoServerStream* GetCryptoStream() override { return crypto_stream_; }
+
  private:
+  QuicCryptoServerStream* crypto_stream_;
+
   DISALLOW_COPY_AND_ASSIGN(TestServerSession);
 };
 
 class TestDispatcher : public QuicDispatcher {
  public:
-  explicit TestDispatcher(const QuicConfig& config,
-                          const QuicCryptoServerConfig* crypto_config,
-                          EpollServer* eps)
+  TestDispatcher(const QuicConfig& config,
+                 const QuicCryptoServerConfig* crypto_config,
+                 EpollServer* eps)
       : QuicDispatcher(config,
                        crypto_config,
                        QuicSupportedVersions(),
                        new QuicDispatcher::DefaultPacketWriterFactory(),
-                       new QuicEpollConnectionHelper(eps)) {
-  }
+                       new QuicEpollConnectionHelper(eps)) {}
 
   MOCK_METHOD3(CreateQuicSession,
                QuicServerSession*(QuicConnectionId connection_id,
@@ -77,8 +86,8 @@ class TestDispatcher : public QuicDispatcher {
   using QuicDispatcher::current_client_address;
 };
 
-// A Connection class which unregisters the session from the dispatcher
-// when sending connection close.
+// A Connection class which unregisters the session from the dispatcher when
+// sending connection close.
 // It'd be slightly more realistic to do this from the Session but it would
 // involve a lot more mocking.
 class MockServerConnection : public MockConnection {
@@ -289,6 +298,125 @@ TEST_F(QuicDispatcherTest, NoVersionPacketToTimeWaitListManager) {
   ProcessPacket(client_address, connection_id, false, "data");
 }
 
+// Enables mocking of the handshake-confirmation for stateless rejects.
+class MockQuicCryptoServerStream : public QuicCryptoServerStream {
+ public:
+  MockQuicCryptoServerStream(const QuicCryptoServerConfig& crypto_config,
+                             QuicSession* session)
+      : QuicCryptoServerStream(&crypto_config, session) {}
+  void set_handshake_confirmed_for_testing(bool handshake_confirmed) {
+    handshake_confirmed_ = handshake_confirmed;
+  }
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(MockQuicCryptoServerStream);
+};
+
+struct StatelessRejectTestParams {
+  StatelessRejectTestParams(bool enable_stateless_rejects_via_flag,
+                            bool use_stateless_rejects_if_peer_supported,
+                            bool client_supports_statelesss_rejects,
+                            bool crypto_handshake_successful)
+      : enable_stateless_rejects_via_flag(enable_stateless_rejects_via_flag),
+        use_stateless_rejects_if_peer_supported(
+            use_stateless_rejects_if_peer_supported),
+        client_supports_statelesss_rejects(client_supports_statelesss_rejects),
+        crypto_handshake_successful(crypto_handshake_successful) {}
+
+  friend std::ostream& operator<<(std::ostream& os,
+                                  const StatelessRejectTestParams& p) {
+    os << "  enable_stateless_rejects_via_flag: "
+       << p.enable_stateless_rejects_via_flag << std::endl;
+    os << "{ use_stateless_rejects_if_peer_supported: "
+       << p.use_stateless_rejects_if_peer_supported << std::endl;
+    os << "{ client_supports_statelesss_rejects: "
+       << p.client_supports_statelesss_rejects << std::endl;
+    os << "  crypto_handshake_successful: " << p.crypto_handshake_successful
+       << " }";
+    return os;
+  }
+
+  // This only enables the stateless reject feature via the feature-flag.
+  // It does not force the crypto server to emit stateless rejects.
+  bool enable_stateless_rejects_via_flag;
+  // If true, this forces the server to send a stateless reject when rejecting
+  // messages.  This should be a no-op if enable_stateless_rejects_via_flag is
+  // false or the peer does not support them.
+  bool use_stateless_rejects_if_peer_supported;
+  // Whether or not the client supports stateless rejects.
+  bool client_supports_statelesss_rejects;
+  // Should the initial crypto handshake succeed or not.
+  bool crypto_handshake_successful;
+};
+
+// Constructs various test permutations for stateless rejects.
+vector<StatelessRejectTestParams> GetStatelessRejectTestParams() {
+  vector<StatelessRejectTestParams> params;
+  for (bool enable_stateless_rejects_via_flag : {true, false}) {
+    for (bool use_stateless_rejects_if_peer_supported : {true, false}) {
+      for (bool client_supports_statelesss_rejects : {true, false}) {
+        for (bool crypto_handshake_successful : {true, false}) {
+          params.push_back(StatelessRejectTestParams(
+              enable_stateless_rejects_via_flag,
+              use_stateless_rejects_if_peer_supported,
+              client_supports_statelesss_rejects, crypto_handshake_successful));
+        }
+      }
+    }
+  }
+  return params;
+}
+
+class QuicDispatcherStatelessRejectTest
+    : public QuicDispatcherTest,
+      public ::testing::WithParamInterface<StatelessRejectTestParams> {
+ public:
+  QuicDispatcherStatelessRejectTest() : crypto_stream1_(nullptr) {}
+
+  ~QuicDispatcherStatelessRejectTest() override {
+    if (crypto_stream1_) {
+      delete crypto_stream1_;
+    }
+  }
+
+  // This test setup assumes that all testing will be done using
+  // crypto_stream1_.
+  void SetUp() override {
+    FLAGS_enable_quic_stateless_reject_support =
+        GetParam().enable_stateless_rejects_via_flag;
+  }
+
+  // Returns true or false, depending on whether the server will emit
+  // a stateless reject, depending upon the parameters of the test.
+  bool ExpectStatelessReject() {
+    return GetParam().enable_stateless_rejects_via_flag &&
+           GetParam().use_stateless_rejects_if_peer_supported &&
+           !GetParam().crypto_handshake_successful &&
+           GetParam().client_supports_statelesss_rejects;
+  }
+
+  // Sets up dispatcher_, sesession1_, and crypto_stream1_ based on
+  // the test parameters.
+  QuicServerSession* CreateSessionBasedOnTestParams(
+      QuicConnectionId connection_id,
+      const IPEndPoint& client_address) {
+    CreateSession(&dispatcher_, config_, connection_id, client_address,
+                  &session1_);
+
+    crypto_stream1_ = new MockQuicCryptoServerStream(crypto_config_, session1_);
+    session1_->SetCryptoStream(crypto_stream1_);
+    crypto_stream1_->set_use_stateless_rejects_if_peer_supported(
+        GetParam().use_stateless_rejects_if_peer_supported);
+    crypto_stream1_->set_handshake_confirmed_for_testing(
+        GetParam().crypto_handshake_successful);
+    crypto_stream1_->set_peer_supports_stateless_rejects(
+        GetParam().client_supports_statelesss_rejects);
+    return session1_;
+  }
+
+  MockQuicCryptoServerStream* crypto_stream1_;
+};
+
 TEST_F(QuicDispatcherTest, ProcessPacketWithZeroPort) {
   CreateTimeWaitListManager();
 
@@ -339,6 +467,52 @@ TEST_F(QuicDispatcherTest, TooBigSeqNoPacketToTimeWaitListManager) {
                 QuicDispatcher::kMaxReasonableInitialSequenceNumber + 1);
 }
 
+INSTANTIATE_TEST_CASE_P(QuicDispatcherStatelessRejectTests,
+                        QuicDispatcherStatelessRejectTest,
+                        ::testing::ValuesIn(GetStatelessRejectTestParams()));
+
+// Parameterized test for stateless rejects.  Should test all
+// combinations of enabling/disabling, reject/no-reject for stateless
+// rejects.
+TEST_P(QuicDispatcherStatelessRejectTest, ParameterizedBasicTest) {
+  CreateTimeWaitListManager();
+
+  IPEndPoint client_address(net::test::Loopback4(), 1);
+  QuicConnectionId connection_id = 1;
+  EXPECT_CALL(dispatcher_, CreateQuicSession(connection_id, _, client_address))
+      .WillOnce(testing::Return(
+          CreateSessionBasedOnTestParams(connection_id, client_address)));
+
+  // Process the first packet for the connection.
+  if (ExpectStatelessReject()) {
+    // If this is a stateless reject, we expect the connection to close.
+    EXPECT_CALL(*session1_, OnConnectionClosed(_, _))
+        .Times(1)
+        .WillOnce(WithoutArgs(Invoke(
+            reinterpret_cast<MockServerConnection*>(session1_->connection()),
+            &MockServerConnection::UnregisterOnConnectionClosed)));
+  }
+  ProcessPacket(client_address, connection_id, true, "foo");
+
+  // Send a second packet and check the results.  If this is a stateless reject,
+  // the existing connection_id will go on the time-wait list.
+  EXPECT_EQ(ExpectStatelessReject(),
+            time_wait_list_manager_->IsConnectionIdInTimeWait(connection_id));
+  if (ExpectStatelessReject()) {
+    // The second packet will be processed on the time-wait list.
+    EXPECT_CALL(*time_wait_list_manager_,
+                ProcessPacket(_, _, connection_id, _, _)).Times(1);
+  } else {
+    // The second packet will trigger a packet-validation
+    EXPECT_CALL(*reinterpret_cast<MockConnection*>(session1_->connection()),
+                ProcessUdpPacket(_, _, _))
+        .Times(1)
+        .WillOnce(testing::WithArgs<2>(
+            Invoke(this, &QuicDispatcherTest::ValidatePacket)));
+  }
+  ProcessPacket(client_address, connection_id, true, "foo");
+}
+
 // Verify the stopgap test: Packets with truncated connection IDs should be
 // dropped.
 class QuicDispatcherTestStrayPacketConnectionId