From 6b41822007345e1e266e1bf3d3e67ad5e101be11 Mon Sep 17 00:00:00 2001
From: Yang Chi <yangchi@fb.com>
Date: Wed, 2 Oct 2019 22:58:30 -0700
Subject: [PATCH] Fix use after free in QuicServerWorker

Summary:
If the server connection is never generated, source address map may
have the only shared_ptr to the transport. Erase from the map then use the
transport is clearly a use-after-free bug.

Reviewed By: udippant, lnicco

Differential Revision: D17733297

fbshipit-source-id: 80d141293458920a0ba6c5eaed14dcbeec17d3ff
---
 quic/server/QuicServerWorker.cpp    |  6 ++---
 quic/server/test/QuicServerTest.cpp | 35 +++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/quic/server/QuicServerWorker.cpp b/quic/server/QuicServerWorker.cpp
index 603ef1677..6a50824d5 100644
--- a/quic/server/QuicServerWorker.cpp
+++ b/quic/server/QuicServerWorker.cpp
@@ -677,13 +677,13 @@ void QuicServerWorker::onConnectionUnbound(
     const QuicServerTransport::SourceIdentity& source,
     folly::Optional<ConnectionId> connectionId) noexcept {
   VLOG(4) << "Removing from sourceAddressMap_ address=" << source.first;
-  // TODO: verify we are removing the right transport
-  sourceAddressMap_.erase(source);
-
   // Ensures we only process `onConnectionUnbound()` once.
   transport->setRoutingCallback(nullptr);
   boundServerTransports_.erase(transport);
 
+  // TODO: verify we are removing the right transport
+  sourceAddressMap_.erase(source);
+
   if (connectionId) {
     VLOG(4) << "Removing from connectionIdMap_ for CID=" << *connectionId
             << ", workerId=" << (uint32_t)workerId_;
diff --git a/quic/server/test/QuicServerTest.cpp b/quic/server/test/QuicServerTest.cpp
index 623ef0338..4a4f45dda 100644
--- a/quic/server/test/QuicServerTest.cpp
+++ b/quic/server/test/QuicServerTest.cpp
@@ -294,6 +294,41 @@ TEST_F(QuicServerWorkerTest, NoConnFoundTestReset) {
       QuicTransportStatsCallback::PacketDropReason::CONNECTION_NOT_FOUND);
 }
 
+TEST_F(QuicServerWorkerTest, QuicServerWorkerUnbindBeforeCidAvailable) {
+  MockConnectionCallback connCb;
+  auto mockSock =
+      std::make_unique<folly::test::MockAsyncUDPSocket>(&eventbase_);
+  EXPECT_CALL(*mockSock, address()).WillRepeatedly(ReturnRef(fakeAddress_));
+  MockQuicTransport::Ptr testTransport = std::make_shared<MockQuicTransport>(
+      worker_->getEventBase(), std::move(mockSock), connCb, nullptr);
+  EXPECT_CALL(*testTransport, getEventBase())
+      .WillRepeatedly(Return(&eventbase_));
+
+  EXPECT_CALL(*testTransport, getOriginalPeerAddress())
+      .WillRepeatedly(ReturnRef(kClientAddr));
+  auto connId = getTestConnectionId(hostId_);
+  createQuicConnection(kClientAddr, connId, testTransport);
+
+  // Otherwise the mock of _make will hold on to a shared_ptr to the transport
+  Mock::VerifyAndClearExpectations(factory_.get());
+
+  auto& srcAddrMap = worker_->getSrcToTransportMap();
+  EXPECT_EQ(1, srcAddrMap.size());
+  EXPECT_EQ(srcAddrMap.begin()->second.get(), testTransport.get());
+  auto& srcIdentity = srcAddrMap.begin()->first;
+  auto& connIdMap = worker_->getConnectionIdMap();
+  EXPECT_EQ(0, connIdMap.size());
+
+  auto* rawTransport = testTransport.get();
+  // This is fine, server worker still has at one shared_ptr in its map.
+  testTransport.reset();
+
+  EXPECT_CALL(*rawTransport, setRoutingCallback(nullptr)).Times(1);
+  // Now remove it from the maps. Nothing should crash.
+  worker_->onConnectionUnbound(rawTransport, srcIdentity, folly::none);
+  EXPECT_EQ(0, srcAddrMap.size());
+}
+
 // TODO (T54143063) Must change use of connectionIdMap_ before
 // can test multiple conn ids routing to the same connection.
 TEST_F(QuicServerWorkerTest, QuicServerMultipleConnIdsRouting) {