facebook · deadalnix · Feb 26, 2020
diff --git a/quic/client/handshake/FizzClientQuicHandshakeContext.cpp b/quic/client/handshake/FizzClientQuicHandshakeContext.cpp
@@ -14,15 +14,34 @@ namespace quic {
 
 FizzClientQuicHandshakeContext::FizzClientQuicHandshakeContext(
     std::shared_ptr<const fizz::client::FizzClientContext> context,
-    std::shared_ptr<const fizz::CertificateVerifier> verifier)
-    : context_(std::move(context)), verifier_(std::move(verifier)) {}
+    std::shared_ptr<const fizz::CertificateVerifier> verifier,
+    std::shared_ptr<QuicPskCache> pskCache)
+    : context_(std::move(context)),
+      verifier_(std::move(verifier)),
+      pskCache_(std::move(pskCache)) {}
 
 std::unique_ptr<ClientHandshake>
 FizzClientQuicHandshakeContext::makeClientHandshake(
     QuicClientConnectionState* conn) {
   return std::make_unique<FizzClientHandshake>(conn, shared_from_this());
 }
 
+folly::Optional<QuicCachedPsk> FizzClientQuicHandshakeContext::getPsk(
+    const folly::Optional<std::string>& hostname) {
+  if (!hostname || !pskCache_) {
+    return folly::none;
+  }
+
+  return pskCache_->getPsk(*hostname);
+}
+
+void FizzClientQuicHandshakeContext::removePsk(
+    const folly::Optional<std::string>& hostname) {
+  if (hostname && pskCache_) {
+    pskCache_->removePsk(*hostname);
+  }
+}
+
 std::shared_ptr<FizzClientQuicHandshakeContext>
 FizzClientQuicHandshakeContext::Builder::build() {
   if (!context_) {
@@ -35,7 +54,7 @@ FizzClientQuicHandshakeContext::Builder::build() {
 
   return std::shared_ptr<FizzClientQuicHandshakeContext>(
       new FizzClientQuicHandshakeContext(
-          std::move(context_), std::move(verifier_)));
+          std::move(context_), std::move(verifier_), std::move(pskCache_)));
 }
 
 } // namespace quic
diff --git a/quic/client/handshake/FizzClientQuicHandshakeContext.h b/quic/client/handshake/FizzClientQuicHandshakeContext.h
@@ -10,6 +10,8 @@
 
 #include <quic/client/handshake/ClientHandshakeFactory.h>
 
+#include <quic/client/handshake/QuicPskCache.h>
+
 #include <fizz/client/FizzClientContext.h>
 #include <fizz/protocol/DefaultCertificateVerifier.h>
 
@@ -34,6 +36,10 @@ class FizzClientQuicHandshakeContext
     return verifier_;
   }
 
+  folly::Optional<QuicCachedPsk> getPsk(
+      const folly::Optional<std::string>& hostname);
+  void removePsk(const folly::Optional<std::string>& hostname);
+
  private:
   /**
    * We make the constructor private so that users have to use the Builder
@@ -45,10 +51,12 @@ class FizzClientQuicHandshakeContext
    */
   FizzClientQuicHandshakeContext(
       std::shared_ptr<const fizz::client::FizzClientContext> context,
-      std::shared_ptr<const fizz::CertificateVerifier> verifier);
+      std::shared_ptr<const fizz::CertificateVerifier> verifier,
+      std::shared_ptr<QuicPskCache> pskCache);
 
   std::shared_ptr<const fizz::client::FizzClientContext> context_;
   std::shared_ptr<const fizz::CertificateVerifier> verifier_;
+  std::shared_ptr<QuicPskCache> pskCache_;
 
  public:
   class Builder {
@@ -65,11 +73,17 @@ class FizzClientQuicHandshakeContext
       return *this;
     }
 
+    Builder& setPskCache(std::shared_ptr<QuicPskCache> pskCache) {
+      pskCache_ = std::move(pskCache);
+      return *this;
+    }
+
     std::shared_ptr<FizzClientQuicHandshakeContext> build();
 
    private:
     std::shared_ptr<const fizz::client::FizzClientContext> context_;
     std::shared_ptr<const fizz::CertificateVerifier> verifier_;
+    std::shared_ptr<QuicPskCache> pskCache_;
   };
 };
 

diff --git a/quic/client/handshake/test/ClientHandshakeTest.cpp b/quic/client/handshake/test/ClientHandshakeTest.cpp
@@ -88,6 +88,7 @@ class ClientHandshakeTest : public Test, public boost::static_visitor<> {
     auto handshakeFactory = FizzClientQuicHandshakeContext::Builder()
                                 .setFizzClientContext(clientCtx)
                                 .setCertificateVerifier(verifier)
+                                .setPskCache(getPskCache())
                                 .build();
     conn.reset(new QuicClientConnectionState(handshakeFactory));
     conn->readCodec = std::make_unique<QuicReadCodec>(QuicNodeType::Client);
@@ -118,6 +119,10 @@ class ClientHandshakeTest : public Test, public boost::static_visitor<> {
     fizzServer->accept(&evb, serverCtx, serverTransportParameters);
   }
 
+  virtual std::shared_ptr<QuicPskCache> getPskCache() {
+    return nullptr;
+  }
+
   void clientServerRound() {
     auto writableBytes = getHandshakeWriteBytes();
     serverReadBuf.append(std::move(writableBytes));
@@ -461,6 +466,12 @@ class ClientHandshakeZeroRttTest : public ClientHandshakeTest {
     setupZeroRttServer();
   }
 
+  std::shared_ptr<QuicPskCache> getPskCache() override {
+    auto pskCache = std::make_shared<BasicQuicPskCache>();
+    pskCache->putPsk(hostname, psk);
+    return pskCache;
+  }
+
   void connect() override {
     handshake->connect(
         hostname,

diff --git a/quic/client/test/QuicClientTransportTest.cpp b/quic/client/test/QuicClientTransportTest.cpp
@@ -220,6 +220,7 @@ class QuicClientTransportIntegrationTest : public TestWithParam<TestingParams> {
     auto fizzClientContext = FizzClientQuicHandshakeContext::Builder()
                                  .setFizzClientContext(clientCtx)
                                  .setCertificateVerifier(verifier)
+                                 .setPskCache(pskCache_)
                                  .build();
     client = std::make_shared<TestingQuicClientTransport>(
         &eventbase_,
@@ -1268,17 +1269,29 @@ class QuicClientTransportTest : public Test {
   QuicClientTransportTest()
       : eventbase_(std::make_unique<folly::EventBase>()) {}
 
+  std::shared_ptr<FizzClientQuicHandshakeContext> getFizzClientContext() {
+    if (!fizzClientContext) {
+      fizzClientContext =
+          FizzClientQuicHandshakeContext::Builder()
+              .setCertificateVerifier(createTestCertificateVerifier())
+              .setPskCache(getPskCache())
+              .build();
+    }
+
+    return fizzClientContext;
+  }
+
+  virtual std::shared_ptr<QuicPskCache> getPskCache() {
+    return nullptr;
+  }
+
   void SetUp() override final {
     auto socket = std::make_unique<NiceMock<folly::test::MockAsyncUDPSocket>>(
         eventbase_.get());
     sock = socket.get();
 
-    auto fizzClientContext =
-        FizzClientQuicHandshakeContext::Builder()
-            .setCertificateVerifier(createTestCertificateVerifier())
-            .build();
     client = TestingQuicClientTransport::newClient<TestingQuicClientTransport>(
-        eventbase_.get(), std::move(socket), std::move(fizzClientContext));
+        eventbase_.get(), std::move(socket), getFizzClientContext());
     destructionCallback = std::make_shared<DestructionCallback>();
     client->setDestructionCallback(destructionCallback);
     client->setSupportedVersions(
@@ -1324,8 +1337,8 @@ class QuicClientTransportTest : public Test {
 
   virtual void setupCryptoLayer() {
     // Fake that the handshake has already occured and fix the keys.
-    mockClientHandshake =
-        new FakeOneRttHandshakeLayer(&client->getNonConstConn());
+    mockClientHandshake = new FakeOneRttHandshakeLayer(
+        &client->getNonConstConn());
     client->getNonConstConn().clientHandshakeLayer = mockClientHandshake;
     client->getNonConstConn().handshakeLayer.reset(mockClientHandshake);
     setFakeHandshakeCiphers();
@@ -1663,6 +1676,7 @@ class QuicClientTransportTest : public Test {
   SocketAddress serverAddr{"127.0.0.1", 443};
   AsyncUDPSocket::ReadCallback* networkReadCallback{nullptr};
   FakeOneRttHandshakeLayer* mockClientHandshake;
+  std::shared_ptr<FizzClientQuicHandshakeContext> fizzClientContext;
   std::shared_ptr<TestingQuicClientTransport> client;
   PacketNum initialPacketNum{0}, handshakePacketNum{0}, appDataPacketNum{0};
   std::unique_ptr<ConnectionIdAlgo> connIdAlgo_;
@@ -5071,11 +5085,15 @@ class QuicClientTransportPskCacheTest
     : public QuicClientTransportAfterStartTestBase {
  public:
   void SetUpChild() override {
-    mockPskCache_ = std::make_shared<NiceMock<MockQuicPskCache>>();
     client->setPskCache(mockPskCache_);
     QuicClientTransportAfterStartTestBase::SetUpChild();
   }
 
+  std::shared_ptr<QuicPskCache> getPskCache() override {
+    mockPskCache_ = std::make_shared<NiceMock<MockQuicPskCache>>();
+    return mockPskCache_;
+  }
+
  protected:
   std::shared_ptr<MockQuicPskCache> mockPskCache_;
 };
@@ -5164,12 +5182,16 @@ class QuicZeroRttClientTest : public QuicClientTransportAfterStartTestBase {
         test::createNoOpHeaderCipher());
   }
 
+  std::shared_ptr<QuicPskCache> getPskCache() override {
+    mockQuicPskCache_ = std::make_shared<MockQuicPskCache>();
+    return mockQuicPskCache_;
+  }
+
   void start() override {
     TransportSettings clientSettings;
     // Ignore path mtu to test negotiation.
     clientSettings.canIgnorePathMTU = true;
     client->setTransportSettings(clientSettings);
-    mockQuicPskCache_ = std::make_shared<NiceMock<MockQuicPskCache>>();
     client->setPskCache(mockQuicPskCache_);
   }
 
@@ -5433,7 +5455,6 @@ class QuicZeroRttHappyEyeballsClientTransportTest
  public:
   void SetUpChild() override {
     client->setHostname(hostname_);
-    mockQuicPskCache_ = std::make_shared<NiceMock<MockQuicPskCache>>();
     client->setPskCache(mockQuicPskCache_);
 
     auto secondSocket =
@@ -5453,6 +5474,11 @@ class QuicZeroRttHappyEyeballsClientTransportTest
     setupCryptoLayer();
   }
 
+  std::shared_ptr<QuicPskCache> getPskCache() override {
+    mockQuicPskCache_ = std::make_shared<MockQuicPskCache>();
+    return mockQuicPskCache_;
+  }
+
  protected:
   folly::test::MockAsyncUDPSocket* secondSock;
   SocketAddress firstAddress{"::1", 443};