From d22e9b86123bbd4c2cac5179dbfa2561cb4f220d Mon Sep 17 00:00:00 2001
From: Eden Zik <ezik@meta.com>
Date: Fri, 13 Oct 2023 09:14:15 -0700
Subject: [PATCH] stopTLS leverages ExtractedIdentityCert to keep credentials
 after dropping certificate

Summary:
When dropping down to plain, we currently keep the X509 payload in order to eventually extract identities.

This was attempted to be fixed in D49045451 - however that caused the identities to not be kept and authentication failure ensued.

This diff attempts to:
- Extract identities when dropping down to plain
- Drop the SSL context as soon as we do so

This ensures that we have identities downstream, without needing the X509 payload.

Reviewed By: frqiu

Differential Revision: D49611022

fbshipit-source-id: 7f4246c938c613dd13c973cb07300b53769f36df
---
 mcrouter/lib/network/McSSLUtil.cpp       | 16 ++++++++++------
 mcrouter/lib/network/McSSLUtil.h         |  8 ++++++--
 mcrouter/lib/network/McServerSession.cpp |  4 +++-
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/mcrouter/lib/network/McSSLUtil.cpp b/mcrouter/lib/network/McSSLUtil.cpp
index 952de50a2..2df19c824 100644
--- a/mcrouter/lib/network/McSSLUtil.cpp
+++ b/mcrouter/lib/network/McSSLUtil.cpp
@@ -174,10 +174,9 @@ folly::AsyncTransportWrapper::UniquePtr McSSLUtil::moveToPlaintext(
   SSL_shutdown(ssl);
 
   // fallback to plaintext
-  auto selfCert =
-      folly::ssl::BasicTransportCertificate::create(sock.getSelfCertificate());
-  auto peerCert =
-      folly::ssl::BasicTransportCertificate::create(sock.getPeerCertificate());
+  // Clear X509 Payload, get certificates out
+  auto [selfCert, peerCert] = dropCertificateX509Payload(sock);
+
   auto evb = sock.getEventBase();
   auto zcId = sock.getZeroCopyBufId();
   auto fd = sock.detachNetworkSocket();
@@ -200,13 +199,18 @@ folly::AsyncTransportWrapper::UniquePtr McSSLUtil::moveToKtls(
   return nullptr;
 }
 
-void McSSLUtil::dropCertificateX509Payload(
+McSSLUtil::CertPair McSSLUtil::dropCertificateX509Payload(
     folly::AsyncSSLSocket& sock) noexcept {
   folly::SharedMutex::ReadHolder rh(getMutex());
   auto& func = getDropCertificateX509PayloadFuncRef();
   if (func) {
-    func(sock);
+    return func(sock);
   }
+  /* By default, initialize a BasicTransportCertificate which is OSS compatible
+   */
+  return std::make_tuple(
+      folly::ssl::BasicTransportCertificate::create(sock.getSelfCertificate()),
+      folly::ssl::BasicTransportCertificate::create(sock.getPeerCertificate()));
 }
 
 folly::Optional<SecurityTransportStats> McSSLUtil::getKtlsStats(
diff --git a/mcrouter/lib/network/McSSLUtil.h b/mcrouter/lib/network/McSSLUtil.h
index 4cacdacfa..15440c54c 100644
--- a/mcrouter/lib/network/McSSLUtil.h
+++ b/mcrouter/lib/network/McSSLUtil.h
@@ -28,8 +28,11 @@ class McSSLUtil {
       folly::Function<bool(folly::AsyncSSLSocket*, bool, X509_STORE_CTX*)
                           const noexcept>;
 
+  using CertPair = std::tuple<
+      std::unique_ptr<folly::AsyncTransportCertificate>,
+      std::unique_ptr<folly::AsyncTransportCertificate>>;
   using DropCertificateX509PayloadFunction =
-      folly::Function<bool(folly::AsyncSSLSocket&) const noexcept>;
+      folly::Function<CertPair(folly::AsyncSSLSocket&) const noexcept>;
 
   static const std::string kTlsToPlainProtocolName;
 
@@ -105,7 +108,8 @@ class McSSLUtil {
   /**
    * Drops certificate x509 payload after connection is established
    */
-  static void dropCertificateX509Payload(folly::AsyncSSLSocket& sock) noexcept;
+  static CertPair dropCertificateX509Payload(
+      folly::AsyncSSLSocket& sock) noexcept;
 
   /**
    * Move the existing transport to kTLS if possible.  Return value of
diff --git a/mcrouter/lib/network/McServerSession.cpp b/mcrouter/lib/network/McServerSession.cpp
index 1c8cc4741..43df7da96 100644
--- a/mcrouter/lib/network/McServerSession.cpp
+++ b/mcrouter/lib/network/McServerSession.cpp
@@ -764,7 +764,9 @@ void McServerSession::onAccepted() {
       std::make_unique<const McServerThriftRequestContext>(transport_.get());
   /* Trims the certificate memory */
   if (auto sock = transport_->getUnderlyingTransport<folly::AsyncSSLSocket>()) {
-    McSSLUtil::dropCertificateX509Payload(*sock);
+    auto [selfCert, peerCert] = McSSLUtil::dropCertificateX509Payload(*sock);
+    sock->setSelfCertificate(std::move(selfCert));
+    sock->setPeerCertificate(std::move(peerCert));
   }
   debugFifo_ = getDebugFifo(
       options_.debugFifoPath, transport_.get(), onRequest_->name());