Add QuicPskCache to FizzClientQuicHandshakeContext

facebook · Feb 26, 2020 · bde83b2 · bde83b2
1 parent 37dfc0f
commit bde83b2
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 4 deletions.
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_;
   };
 };