From bdaec845ffba5b8292bc0fd490ccd990dace0c01 Mon Sep 17 00:00:00 2001
From: Andrew Huang <huangandrew@fb.com>
Date: Sun, 7 Nov 2021 17:24:51 -0800
Subject: [PATCH] Add support for remove session callbacks in SSLContext;
 remove from SSLSessionCallbacks

Reviewed By: mingtaoy

Differential Revision: D32218077

fbshipit-source-id: 5a966ced07387340d02ad2a85c095939c9e7f530
---
 folly/io/async/SSLContext.cpp | 10 +++++++
 folly/io/async/SSLContext.h   | 50 +++++++++++++++++++++++++++++++++--
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/folly/io/async/SSLContext.cpp b/folly/io/async/SSLContext.cpp
index 5f48e1a6e23..ba9ba33b0f3 100644
--- a/folly/io/async/SSLContext.cpp
+++ b/folly/io/async/SSLContext.cpp
@@ -838,6 +838,7 @@ void SSLContext::setupCtx(SSL_CTX* ctx) {
 
   SSL_CTX_set_ex_data(ctx, getExDataIndex(), this);
   SSL_CTX_sess_set_new_cb(ctx, SSLContext::newSessionCallback);
+  SSL_CTX_sess_set_remove_cb(ctx, SSLContext::removeSessionCallback);
 }
 
 SSLContext* SSLContext::getFromSSLCtx(const SSL_CTX* ctx) {
@@ -866,6 +867,15 @@ int SSLContext::newSessionCallback(SSL* ssl, SSL_SESSION* session) {
   return 1;
 }
 
+void SSLContext::removeSessionCallback(SSL_CTX* ctx, SSL_SESSION* session) {
+  SSLContext* context = getFromSSLCtx(ctx);
+
+  auto& cb = context->sessionLifecycleCallbacks_;
+  if (cb) {
+    cb->onRemoveSession(ctx, session);
+  }
+}
+
 void SSLContext::setSessionLifecycleCallbacks(
     std::unique_ptr<SessionLifecycleCallbacks> cb) {
   sessionLifecycleCallbacks_ = std::move(cb);
diff --git a/folly/io/async/SSLContext.h b/folly/io/async/SSLContext.h
index ed66f4d2f29..926fa6913cc 100644
--- a/folly/io/async/SSLContext.h
+++ b/folly/io/async/SSLContext.h
@@ -86,11 +86,55 @@ class SSLAcceptRunner {
  */
 class SSLContext {
  public:
+  /**
+   * SessionLifecycleCallbacks can be used to receive notifications about
+   * `SSL_SESSION`s that are constructed by OpenSSL after establishing a TLS
+   * connection.
+   *
+   * SSL_SESSIONs contain properties of the TLS connection, such as the traffic
+   * keys negotiated as part of the handshake, the certificate of the peer, etc.
+   * This information can be stored in a cache, so that it can later be used for
+   * TLS session resumption (see AsyncSSLSocket::setSSLSession)
+   *
+   * SessionLifecycleCallbacks is intended to allow an implementation of a SSL
+   * session cache.
+   */
   struct SessionLifecycleCallbacks {
-    virtual void onNewSession(SSL*, folly::ssl::SSLSessionUniquePtr) = 0;
+    /**
+     * SessionLifecycleCallbacks::onNewSession is invoked when a new session has
+     * been created by OpenSSL which can be stored in a session cache.
+     *
+     * Multiple `onNewSession` invocations can occur for a given connection.
+     * Implementations must be prepared to handle this.
+     *
+     * @param ssl      The `ssl` object corresponding to the connection that
+     * established the session
+     * @param session  The SSL_SESSION object that should be stored.
+     */
+    virtual void onNewSession(
+        SSL* /*ssl */, folly::ssl::SSLSessionUniquePtr /* session */) = 0;
+
+    /**
+     * SessionLifecycleCallbacks::onRemoveSession is invoked when OpenSSL
+     * considers a session expired for any reason. (For example, OpenSSL may
+     * want to remove a session after it was used for a resumed connection). The
+     * session should be considered "invalid".
+     *
+     * It's important to note that for TLS 1.3 connections, OpenSSL will invoke
+     * this after the handshake to discourage session reuse.
+     *
+     * The interface is asymmetric w.r.t `onNewSession` intentionally; OpenSSL's
+     * underlying functions require this signature.
+     *
+     * @param ctx     The SSL_CTX of the SSL that established the original
+     * session.
+     * @param session A *non-owning* pointer to the SSL_SESSION that should be
+     * removed. Do not attempt to SSL_SESSION_free this.
+     */
+    virtual void onRemoveSession(
+        SSL_CTX* /* ctx */, SSL_SESSION* /* session */) = 0;
     virtual ~SessionLifecycleCallbacks() = default;
   };
-
   enum SSLVersion {
     SSLv2,
     SSLv3,
@@ -714,6 +758,8 @@ class SSLContext {
       nullptr};
 
   static int newSessionCallback(SSL* ssl, SSL_SESSION* session);
+
+  static void removeSessionCallback(SSL_CTX* ctx, SSL_SESSION* session);
 };
 
 typedef std::shared_ptr<SSLContext> SSLContextPtr;