Bug 1558681 - Anti-replay contexts, r=jcj,kjacobs

Stop using a global anti-replay context and enable creating a context directly.
This increases the overhead of managing anti-replay for applications marginally,
but allows much greater flexibility in use of anti-replay mechanisms.  In
particular, it enables the testing of 0-RTT in a threaded environment.

The comments in sslexp should be clear enough in explaining how this works.
Basically, this is a new reference-counted object that can be created and
tracked by applications.

The only thing that I can see might be a problem with the API is that I haven't
exposed a function to add a reference for use by applications.  My thinking is
that reference counting is an internal thing; it seems like applications won't
need to worry about that.

selfserv is updated to create a context and attach it to sockets.  This shows
that the management overhead is minor.

The gtests have been tweaked to create a context during setup. The context is
owned by the overall test framework and is passed to server instances after the
sockets are initialized.

Bonus changes:

* ESNI keys are copied from the model socket when calling SSL_ReConfigFD().
* Some better tracing in the anti-replay functions.

Neither of these seemed worth the overhead of a bug to fix.

Differential Revision: https://phabricator.services.mozilla.com/D34660

12 files changed, 264 insertions, 109 deletions

diff --git a/cmd/selfserv/selfserv.c b/cmd/selfserv/selfserv.c
index 4b1adb028..533b3f053 100644
--- a/cmd/selfserv/selfserv.c
+++ b/cmd/selfserv/selfserv.c
@@ -805,6 +805,7 @@ PRBool enableSessionTickets = PR_FALSE;
 PRBool failedToNegotiateName = PR_FALSE;
 PRBool enableExtendedMasterSecret = PR_FALSE;
 PRBool zeroRTT = PR_FALSE;
+SSLAntiReplayContext *antiReplay = NULL;
 PRBool enableALPN = PR_FALSE;
 PRBool enablePostHandshakeAuth = PR_FALSE;
 SSLNamedGroup *enabledGroups = NULL;
@@ -1954,7 +1955,7 @@ server_main(
         if (enabledVersions.max < SSL_LIBRARY_VERSION_TLS_1_3) {
             errExit("You tried enabling 0RTT without enabling TLS 1.3!");
         }
-        rv = SSL_InitAntiReplay(PR_Now(), 10L * PR_USEC_PER_SEC, 7, 14);
+        rv = SSL_SetAntiReplayContext(model_sock, antiReplay);
         if (rv != SECSuccess) {
             errExit("error configuring anti-replay ");
         }
@@ -2469,6 +2470,12 @@ main(int argc, char **argv)
 
             case 'Z':
                 zeroRTT = PR_TRUE;
+                rv = SSL_CreateAntiReplayContext(PR_Now(), 10L * PR_USEC_PER_SEC, 7, 14, &antiReplay);
+                if (rv != SECSuccess) {
+                    PL_DestroyOptState(optstate);
+                    fprintf(stderr, "Unable to create anti-replay context for 0-RTT.\n");
+                    exit(1);
+                }
                 break;
 
             case 'Q':
@@ -2798,6 +2805,9 @@ cleanup:
     if (enabledGroups) {
         PORT_Free(enabledGroups);
     }
+    if (antiReplay) {
+        SSL_ReleaseAntiReplayContext(antiReplay);
+    }
     if (NSS_Shutdown() != SECSuccess) {
         SECU_PrintError(progName, "NSS_Shutdown");
         if (loggerThread) {
diff --git a/cpputil/scoped_ptrs_ssl.h b/cpputil/scoped_ptrs_ssl.h
index 581fe423a..474187540 100644
--- a/cpputil/scoped_ptrs_ssl.h
+++ b/cpputil/scoped_ptrs_ssl.h
@@ -11,10 +11,13 @@
 #include "sslexp.h"
 
 struct ScopedDeleteSSL {
+  void operator()(SSLAeadContext* ctx) { SSL_DestroyAead(ctx); }
+  void operator()(SSLAntiReplayContext* ctx) {
+    SSL_ReleaseAntiReplayContext(ctx);
+  }
   void operator()(SSLResumptionTokenInfo* token) {
     SSL_DestroyResumptionTokenInfo(token);
   }
-  void operator()(SSLAeadContext* ctx) { SSL_DestroyAead(ctx); }
 };
 
 template <class T>
@@ -29,8 +32,9 @@ struct ScopedMaybeDeleteSSL {
 
 #define SCOPED(x) typedef std::unique_ptr<x, ScopedMaybeDeleteSSL<x> > Scoped##x
 
-SCOPED(SSLResumptionTokenInfo);
 SCOPED(SSLAeadContext);
+SCOPED(SSLAntiReplayContext);
+SCOPED(SSLResumptionTokenInfo);
 
 #undef SCOPED
 
diff --git a/gtests/ssl_gtest/ssl_0rtt_unittest.cc b/gtests/ssl_gtest/ssl_0rtt_unittest.cc
index 979351518..c9036291c 100644
--- a/gtests/ssl_gtest/ssl_0rtt_unittest.cc
+++ b/gtests/ssl_gtest/ssl_0rtt_unittest.cc
@@ -74,10 +74,11 @@ TEST_P(TlsConnectTls13, ZeroRttApplicationReject) {
 }
 
 TEST_P(TlsConnectTls13, ZeroRttApparentReplayAfterRestart) {
-  // The test fixtures call SSL_InitAntiReplay() in SetUp().  This results in
-  // 0-RTT being rejected until at least one window passes.  SetupFor0Rtt()
-  // forces a rollover of the anti-replay filters, which clears this state.
-  // Here, we do the setup manually here without that forced rollover.
+  // The test fixtures enable anti-replay in SetUp().  This results in 0-RTT
+  // being rejected until at least one window passes.  SetupFor0Rtt() forces a
+  // rollover of the anti-replay filters, which clears that state and allows
+  // 0-RTT to work.  Make the first connection manually to avoid that rollover
+  // and cause 0-RTT to be rejected.
 
   ConfigureSessionCache(RESUME_BOTH, RESUME_TICKET);
   ConfigureVersion(SSL_LIBRARY_VERSION_TLS_1_3);
@@ -217,7 +218,7 @@ TEST_P(TlsConnectTls13, ZeroRttServerOnly) {
 // delay exceeds half the anti-replay window.
 TEST_P(TlsConnectTls13, ZeroRttRejectOldTicket) {
   static const PRTime kWindow = 10 * PR_USEC_PER_SEC;
-  EXPECT_EQ(SECSuccess, SSL_InitAntiReplay(now(), kWindow, 1, 3));
+  ResetAntiReplay(kWindow);
   SetupForZeroRtt();
 
   Reset();
@@ -242,7 +243,7 @@ TEST_P(TlsConnectTls13, ZeroRttRejectOldTicket) {
 // arrive prematurely, causing the server to reject early data.
 TEST_P(TlsConnectTls13, ZeroRttRejectPrematureTicket) {
   static const PRTime kWindow = 10 * PR_USEC_PER_SEC;
-  EXPECT_EQ(SECSuccess, SSL_InitAntiReplay(now(), kWindow, 1, 3));
+  ResetAntiReplay(kWindow);
   ConfigureSessionCache(RESUME_BOTH, RESUME_TICKET);
   ConfigureVersion(SSL_LIBRARY_VERSION_TLS_1_3);
   server_->Set0RttEnabled(true);
@@ -904,20 +905,21 @@ TEST_F(TlsConnectDatagram13, ZeroRttShortReadDtls) {
 // that a small sleep results in rejection of early data. 0-RTT has a
 // configurable timer, which makes it ideal for this.
 TEST_F(TlsConnectStreamTls13, TimePassesByDefault) {
-  // Set a tiny anti-replay window.  This has to be at least 2 milliseconds to
-  // have any chance of being relevant as that is the smallest window that we
-  // can detect.  Anything smaller rounds to zero.
-  static const unsigned int kTinyWindowMs = 5;
-  EXPECT_EQ(SECSuccess, SSL_InitAntiReplay(
-                            PR_Now(), kTinyWindowMs * PR_USEC_PER_MSEC, 1, 5));
-
   // Calling EnsureTlsSetup() replaces the time function on client and server,
-  // which we don't want, so initialize each directly.
+  // and sets up anti-replay, which we don't want, so initialize each directly.
   client_->EnsureTlsSetup();
   server_->EnsureTlsSetup();
-  client_->StartConnect();  // Also avoid StartConnect().
+  // StartConnect() calls EnsureTlsSetup(), so avoid that too.
+  client_->StartConnect();
   server_->StartConnect();
 
+  // Set a tiny anti-replay window.  This has to be at least 2 milliseconds to
+  // have any chance of being relevant as that is the smallest window that we
+  // can detect.  Anything smaller rounds to zero.
+  static const unsigned int kTinyWindowMs = 5;
+  ResetAntiReplay(static_cast<PRTime>(kTinyWindowMs * PR_USEC_PER_MSEC));
+  server_->SetAntiReplayContext(anti_replay_);
+
   ConfigureSessionCache(RESUME_BOTH, RESUME_TICKET);
   ConfigureVersion(SSL_LIBRARY_VERSION_TLS_1_3);
   server_->Set0RttEnabled(true);
@@ -951,6 +953,35 @@ TEST_F(TlsConnectStreamTls13, TimePassesByDefault) {
   CheckConnected();
 }
 
+// Test that SSL_CreateAntiReplayContext doesn't pass bad inputs.
+TEST_F(TlsConnectStreamTls13, BadAntiReplayArgs) {
+  SSLAntiReplayContext* p;
+  // Zero or negative window.
+  EXPECT_EQ(SECFailure, SSL_CreateAntiReplayContext(0, -1, 1, 1, &p));
+  EXPECT_EQ(SEC_ERROR_INVALID_ARGS, PORT_GetError());
+  EXPECT_EQ(SECFailure, SSL_CreateAntiReplayContext(0, 0, 1, 1, &p));
+  EXPECT_EQ(SEC_ERROR_INVALID_ARGS, PORT_GetError());
+  // Zero k.
+  EXPECT_EQ(SECFailure, SSL_CreateAntiReplayContext(0, 1, 0, 1, &p));
+  EXPECT_EQ(SEC_ERROR_INVALID_ARGS, PORT_GetError());
+  // Zero bits.
+  EXPECT_EQ(SECFailure, SSL_CreateAntiReplayContext(0, 1, 1, 0, &p));
+  EXPECT_EQ(SEC_ERROR_INVALID_ARGS, PORT_GetError());
+  EXPECT_EQ(SECFailure, SSL_CreateAntiReplayContext(0, 1, 1, 1, nullptr));
+  EXPECT_EQ(SEC_ERROR_INVALID_ARGS, PORT_GetError());
+
+  // Prove that these parameters do work, even if they are useless..
+  EXPECT_EQ(SECSuccess, SSL_CreateAntiReplayContext(0, 1, 1, 1, &p));
+  ASSERT_NE(nullptr, p);
+  ScopedSSLAntiReplayContext ctx(p);
+
+  // The socket isn't a client or server until later, so configuring a client
+  // should work OK.
+  client_->EnsureTlsSetup();
+  EXPECT_EQ(SECSuccess, SSL_SetAntiReplayContext(client_->ssl_fd(), ctx.get()));
+  EXPECT_EQ(SECSuccess, SSL_SetAntiReplayContext(client_->ssl_fd(), nullptr));
+}
+
 #ifndef NSS_DISABLE_TLS_1_3
 INSTANTIATE_TEST_CASE_P(Tls13ZeroRttReplayTest, TlsZeroRttReplayTest,
                         TlsConnectTestBase::kTlsVariantsAll);
diff --git a/gtests/ssl_gtest/tls_agent.cc b/gtests/ssl_gtest/tls_agent.cc
index e726c6a61..697da6df2 100644
--- a/gtests/ssl_gtest/tls_agent.cc
+++ b/gtests/ssl_gtest/tls_agent.cc
@@ -251,6 +251,10 @@ bool TlsAgent::MaybeSetResumptionToken() {
   return true;
 }
 
+void TlsAgent::SetAntiReplayContext(ScopedSSLAntiReplayContext& ctx) {
+  EXPECT_EQ(SECSuccess, SSL_SetAntiReplayContext(ssl_fd_.get(), ctx.get()));
+}
+
 void TlsAgent::SetupClientAuth() {
   EXPECT_TRUE(EnsureTlsSetup());
   ASSERT_EQ(CLIENT, role_);
diff --git a/gtests/ssl_gtest/tls_agent.h b/gtests/ssl_gtest/tls_agent.h
index a2ffe7711..0e1b3df6d 100644
--- a/gtests/ssl_gtest/tls_agent.h
+++ b/gtests/ssl_gtest/tls_agent.h
@@ -108,6 +108,8 @@ class TlsAgent : public PollTarget {
   void PrepareForRenegotiate();
   // Prepares for renegotiation, then actually triggers it.
   void StartRenegotiate();
+  void SetAntiReplayContext(ScopedSSLAntiReplayContext& ctx);
+
   static bool LoadCertificate(const std::string& name,
                               ScopedCERTCertificate* cert,
                               ScopedSECKEYPrivateKey* priv);
diff --git a/gtests/ssl_gtest/tls_connect.cc b/gtests/ssl_gtest/tls_connect.cc
index fa3376937..520950829 100644
--- a/gtests/ssl_gtest/tls_connect.cc
+++ b/gtests/ssl_gtest/tls_connect.cc
@@ -215,7 +215,7 @@ void TlsConnectTestBase::SetUp() {
   SSL_ConfigServerSessionIDCache(1024, 0, 0, g_working_dir_path.c_str());
   SSLInt_ClearSelfEncryptKey();
   now_ = PR_Now();
-  SSL_InitAntiReplay(now_, kAntiReplayWindow, 1, 3);
+  ResetAntiReplay(kAntiReplayWindow);
   ClearStats();
   SaveAlgorithmPolicy();
   Init();
@@ -240,6 +240,14 @@ void TlsConnectTestBase::Init() {
   }
 }
 
+void TlsConnectTestBase::ResetAntiReplay(PRTime window) {
+  SSLAntiReplayContext* p_anti_replay;
+  EXPECT_EQ(SECSuccess,
+            SSL_CreateAntiReplayContext(now_, window, 1, 3, &p_anti_replay));
+  EXPECT_NE(nullptr, p_anti_replay);
+  anti_replay_.reset(p_anti_replay);
+}
+
 void TlsConnectTestBase::Reset() {
   // Take a copy of the names because they are about to disappear.
   std::string server_name = server_->name();
@@ -258,7 +266,8 @@ void TlsConnectTestBase::Reset(const std::string& server_name,
     server_->SkipVersionChecks();
   }
 
-  std::cerr << "Reset" << std::endl;
+  std::cerr << "Reset server:" << server_name << ", client:" << client_name
+            << std::endl;
   Init();
 }
 
@@ -290,6 +299,7 @@ void TlsConnectTestBase::EnsureTlsSetup() {
                                                     : nullptr));
   EXPECT_TRUE(client_->EnsureTlsSetup(client_model_ ? client_model_->ssl_fd()
                                                     : nullptr));
+  server_->SetAntiReplayContext(anti_replay_);
   EXPECT_EQ(SECSuccess, SSL_SetTimeFunc(client_->ssl_fd(),
                                         TlsConnectTestBase::TimeFunc, &now_));
   EXPECT_EQ(SECSuccess, SSL_SetTimeFunc(server_->ssl_fd(),
diff --git a/gtests/ssl_gtest/tls_connect.h b/gtests/ssl_gtest/tls_connect.h
index 6d74df972..4321828f0 100644
--- a/gtests/ssl_gtest/tls_connect.h
+++ b/gtests/ssl_gtest/tls_connect.h
@@ -134,6 +134,8 @@ class TlsConnectTestBase : public ::testing::Test {
   // Move the DTLS timers for both endpoints to pop the next timer.
   void ShiftDtlsTimers();
   void AdvanceTime(PRTime time_shift);
+
+  void ResetAntiReplay(PRTime window);
   void RolloverAntiReplay();
 
   void SaveAlgorithmPolicy();
@@ -149,6 +151,7 @@ class TlsConnectTestBase : public ::testing::Test {
   SessionResumptionMode expected_resumption_mode_;
   uint8_t expected_resumptions_;
   std::vector<std::vector<uint8_t>> session_ids_;
+  ScopedSSLAntiReplayContext anti_replay_;
 
   // A simple value of "a", "b".  Note that the preferred value of "a" is placed
   // at the end, because the NSS API follows the now defunct NPN specification,
diff --git a/lib/ssl/sslexp.h b/lib/ssl/sslexp.h
index 7eb53cfee..96159f56b 100644
--- a/lib/ssl/sslexp.h
+++ b/lib/ssl/sslexp.h
@@ -159,11 +159,16 @@ typedef SECStatus(PR_CALLBACK *SSLExtensionHandler)(
                           handler, handlerArg))
 
 /*
- * Initialize the anti-replay buffer for supporting 0-RTT in TLS 1.3 on servers.
+ * Create an anti-replay context for supporting 0-RTT in TLS 1.3 on servers.
  *
- * To use 0-RTT on a server, you must call this function.  Failing to call this
- * function will result in all 0-RTT being rejected.  Connections will complete,
- * but early data will be rejected.
+ * To use 0-RTT on a server, you must create an anti-replay context using
+ * SSL_CreateAntiReplayContext and set that on the socket with
+ * SSL_SetAntiReplayContext.  Failing to set a context on the server will result
+ * in all 0-RTT being rejected.  Connections will complete, but early data will
+ * be rejected.
+ *
+ * Anti-replay contexts are reference counted and are released with
+ * SSL_ReleaseAntiReplayContext.
  *
  * NSS uses a Bloom filter to track the ClientHello messages that it receives
  * (specifically, it uses the PSK binder).  This function initializes a pair of
@@ -219,11 +224,23 @@ typedef SECStatus(PR_CALLBACK *SSLExtensionHandler)(
  * Early data can be replayed at least once with every server instance that will
  * accept tickets that are encrypted with the same key.
  */
-#define SSL_InitAntiReplay(now, window, k, bits)                \
-    SSL_EXPERIMENTAL_API("SSL_InitAntiReplay",                  \
-                         (PRTime _now, PRTime _window,          \
-                          unsigned int _k, unsigned int _bits), \
-                         (now, window, k, bits))
+typedef struct SSLAntiReplayContextStr SSLAntiReplayContext;
+#define SSL_CreateAntiReplayContext(now, window, k, bits, ctx) \
+    SSL_EXPERIMENTAL_API("SSL_CreateAntiReplayContext",        \
+                         (PRTime _now, PRTime _window,         \
+                          unsigned int _k, unsigned int _bits, \
+                          SSLAntiReplayContext **_ctx),        \
+                         (now, window, k, bits, ctx))
+
+#define SSL_SetAntiReplayContext(fd, ctx)                                 \
+    SSL_EXPERIMENTAL_API("SSL_SetAntiReplayContext",                      \
+                         (PRFileDesc * _fd, SSLAntiReplayContext * _ctx), \
+                         (fd, ctx))
+
+#define SSL_ReleaseAntiReplayContext(ctx)                \
+    SSL_EXPERIMENTAL_API("SSL_ReleaseAntiReplayContext", \
+                         (SSLAntiReplayContext * _ctx),  \
+                         (ctx))
 
 /*
  * This function allows a server application to generate a session ticket that
@@ -743,6 +760,7 @@ typedef PRTime(PR_CALLBACK *SSLTimeFunc)(void *arg);
 /* Deprecated experimental APIs */
 #define SSL_UseAltServerHelloType(fd, enable) SSL_DEPRECATED_EXPERIMENTAL_API
 #define SSL_SetupAntiReplay(a, b, c) SSL_DEPRECATED_EXPERIMENTAL_API
+#define SSL_InitAntiReplay(a, b, c) SSL_DEPRECATED_EXPERIMENTAL_API
 
 SEC_END_PROTOS
 
diff --git a/lib/ssl/sslimpl.h b/lib/ssl/sslimpl.h
index 8c472580b..aa6f804ff 100644
--- a/lib/ssl/sslimpl.h
+++ b/lib/ssl/sslimpl.h
@@ -1094,6 +1094,9 @@ struct sslSocketStr {
     /* The information from the ESNI keys record
      * (also the private key for the server). */
     sslEsniKeys *esniKeys;
+
+    /* Anti-replay for TLS 1.3 0-RTT. */
+    SSLAntiReplayContext *antiReplay;
 };
 
 struct sslSelfEncryptKeysStr {
diff --git a/lib/ssl/sslsock.c b/lib/ssl/sslsock.c
index 802ef6714..b53075be0 100644
--- a/lib/ssl/sslsock.c
+++ b/lib/ssl/sslsock.c
@@ -372,6 +372,13 @@ ssl_DupSocket(sslSocket *os)
                 goto loser;
             }
         }
+        if (os->antiReplay) {
+            ss->antiReplay = tls13_RefAntiReplayContext(os->antiReplay);
+            PORT_Assert(ss->antiReplay); /* Can't fail. */
+            if (!ss->antiReplay) {
+                goto loser;
+            }
+        }
 
         /* Create security data */
         rv = ssl_CopySecurityInfo(ss, os);
@@ -460,6 +467,7 @@ ssl_DestroySocketContents(sslSocket *ss)
     ssl_ClearPRCList(&ss->ssl3.hs.dtlsRcvdHandshake, NULL);
 
     tls13_DestroyESNIKeys(ss->esniKeys);
+    tls13_ReleaseAntiReplayContext(ss->antiReplay);
 }
 
 /*
@@ -2292,6 +2300,29 @@ SSL_ReconfigFD(PRFileDesc *model, PRFileDesc *fd)
         }
     }
 
+    /* Copy ESNI. */
+    tls13_DestroyESNIKeys(ss->esniKeys);
+    ss->esniKeys = NULL;
+    if (sm->esniKeys) {
+        ss->esniKeys = tls13_CopyESNIKeys(sm->esniKeys);
+        if (!ss->esniKeys) {
+            return NULL;
+        }
+    }
+
+    /* Copy anti-replay context. */
+    if (ss->antiReplay) {
+        tls13_ReleaseAntiReplayContext(ss->antiReplay);
+        ss->antiReplay = NULL;
+    }
+    if (sm->antiReplay) {
+        ss->antiReplay = tls13_RefAntiReplayContext(sm->antiReplay);
+        PORT_Assert(ss->antiReplay);
+        if (!ss->antiReplay) {
+            return NULL;
+        }
+    }
+
     if (sm->authCertificate)
         ss->authCertificate = sm->authCertificate;
     if (sm->authCertificateArg)
@@ -3988,6 +4019,7 @@ ssl_NewSocket(PRBool makeLocks, SSLProtocolVariant protocolVariant)
     dtls_InitTimers(ss);
 
     ss->esniKeys = NULL;
+    ss->antiReplay = NULL;
 
     if (makeLocks) {
         rv = ssl_MakeLocks(ss);
@@ -4056,6 +4088,7 @@ struct {
 #ifndef SSL_DISABLE_EXPERIMENTAL_API
     EXP(AeadDecrypt),
     EXP(AeadEncrypt),
+    EXP(CreateAntiReplayContext),
     EXP(DestroyAead),
     EXP(DestroyResumptionTokenInfo),
     EXP(EnableESNI),
@@ -4064,7 +4097,6 @@ struct {
     EXP(GetExtensionSupport),
     EXP(GetResumptionTokenInfo),
     EXP(HelloRetryRequestCallback),
-    EXP(InitAntiReplay),
     EXP(InstallExtensionHooks),
     EXP(HkdfExtract),
     EXP(HkdfExpandLabel),
@@ -4073,9 +4105,11 @@ struct {
     EXP(MakeAead),
     EXP(RecordLayerData),
     EXP(RecordLayerWriteCallback),
+    EXP(ReleaseAntiReplayContext),
     EXP(SecretCallback),
     EXP(SendCertificateRequest),
     EXP(SendSessionTicket),
+    EXP(SetAntiReplayContext),
     EXP(SetESNIKeyPair),
     EXP(SetMaxEarlyDataSize),
     EXP(SetResumptionTokenCallback),
diff --git a/lib/ssl/tls13con.h b/lib/ssl/tls13con.h
index d634f4b84..bd309419f 100644
--- a/lib/ssl/tls13con.h
+++ b/lib/ssl/tls13con.h
@@ -115,10 +115,16 @@ SECStatus tls13_NegotiateVersion(sslSocket *ss,
 PRBool tls13_ShouldRequestClientAuth(sslSocket *ss);
 
 PRBool tls13_IsReplay(const sslSocket *ss, const sslSessionID *sid);
-void tls13_AntiReplayRollover(PRTime now);
-
-SECStatus SSLExp_InitAntiReplay(PRTime now, PRTime window, unsigned int k,
-                                unsigned int bits);
+void tls13_AntiReplayRollover(SSLAntiReplayContext *ctx, PRTime now);
+SSLAntiReplayContext *tls13_RefAntiReplayContext(SSLAntiReplayContext *ctx);
+void tls13_ReleaseAntiReplayContext(SSLAntiReplayContext *ctx);
+
+SECStatus SSLExp_CreateAntiReplayContext(
+    PRTime now, PRTime window, unsigned int k, unsigned int bits,
+    SSLAntiReplayContext **ctx);
+SECStatus SSLExp_SetAntiReplayContext(PRFileDesc *fd,
+                                      SSLAntiReplayContext *ctx);
+SECStatus SSLExp_ReleaseAntiReplayContext(SSLAntiReplayContext *ctx);
 
 SECStatus SSLExp_HelloRetryRequestCallback(PRFileDesc *fd,
                                            SSLHelloRetryRequestCallback cb,
diff --git a/lib/ssl/tls13replay.c b/lib/ssl/tls13replay.c
index 2770be761..ac44811f0 100644
--- a/lib/ssl/tls13replay.c
+++ b/lib/ssl/tls13replay.c
@@ -9,7 +9,6 @@
 #include "nss.h"      /* for NSS_RegisterShutdown */
 #include "nssilock.h" /* for PZMonitor */
 #include "pk11pub.h"
-#include "prinit.h" /* for PR_CallOnce */
 #include "prmon.h"
 #include "prtime.h"
 #include "secerr.h"
@@ -18,10 +17,10 @@
 #include "sslimpl.h"
 #include "tls13hkdf.h"
 
-static struct {
-    /* Used to ensure that we only initialize the cleanup function once. */
-    PRCallOnceType init;
-    /* Used to serialize access to the filters. */
+struct SSLAntiReplayContextStr {
+    /* The number of outstanding references to this context. */
+    PRInt32 refCount;
+    /* Used to serialize access. */
     PZMonitor *lock;
     /* The filters, use of which alternates. */
     sslBloomFilter filters[2];
@@ -33,44 +32,55 @@ static struct {
     PRTime window;
     /* This key ensures that the bloom filter index is unpredictable. */
     PK11SymKey *key;
-} ssl_anti_replay;
+};
 
-/* Clear the current state and free any resources we allocated. The signature
- * here is odd to allow this to be called during shutdown. */
-static SECStatus
-tls13_AntiReplayReset(void *appData, void *nssData)
+void
+tls13_ReleaseAntiReplayContext(SSLAntiReplayContext *ctx)
 {
-    if (ssl_anti_replay.key) {
-        PK11_FreeSymKey(ssl_anti_replay.key);
-        ssl_anti_replay.key = NULL;
+    if (!ctx) {
+        return;
     }
-    if (ssl_anti_replay.lock) {
-        PZ_DestroyMonitor(ssl_anti_replay.lock);
-        ssl_anti_replay.lock = NULL;
+    if (PR_ATOMIC_DECREMENT(&ctx->refCount) >= 1) {
+        return;
     }
-    sslBloom_Destroy(&ssl_anti_replay.filters[0]);
-    sslBloom_Destroy(&ssl_anti_replay.filters[1]);
+
+    if (ctx->lock) {
+        PZ_DestroyMonitor(ctx->lock);
+        ctx->lock = NULL;
+    }
+    PK11_FreeSymKey(ctx->key);
+    ctx->key = NULL;
+    sslBloom_Destroy(&ctx->filters[0]);
+    sslBloom_Destroy(&ctx->filters[1]);
+}
+
+/* Clear the current state and free any resources we allocated. The signature
+ * here is odd to allow this to be called during shutdown. */
+SECStatus
+SSLExp_ReleaseAntiReplayContext(SSLAntiReplayContext *ctx)
+{
+    tls13_ReleaseAntiReplayContext(ctx);
     return SECSuccess;
 }
 
-static PRStatus
-tls13_AntiReplayInit(void)
+SSLAntiReplayContext *
+tls13_RefAntiReplayContext(SSLAntiReplayContext *ctx)
 {
-    SECStatus rv = NSS_RegisterShutdown(tls13_AntiReplayReset, NULL);
-    if (rv != SECSuccess) {
-        return PR_FAILURE;
-    }
-    return PR_SUCCESS;
+    PORT_Assert(ctx);
+    PR_ATOMIC_INCREMENT(&ctx->refCount);
+    return ctx;
 }
 
 static SECStatus
-tls13_AntiReplayKeyGen()
+tls13_AntiReplayKeyGen(SSLAntiReplayContext *ctx)
 {
     PRUint8 buf[32];
     SECItem keyItem = { siBuffer, buf, sizeof(buf) };
     PK11SlotInfo *slot;
     SECStatus rv;
 
+    PORT_Assert(ctx);
+
     slot = PK11_GetInternalSlot();
     if (!slot) {
         PORT_SetError(SEC_ERROR_LIBRARY_FAILURE);
@@ -81,10 +91,10 @@ tls13_AntiReplayKeyGen()
         goto loser;
     }
 
-    ssl_anti_replay.key = PK11_ImportSymKey(slot, CKM_NSS_HKDF_SHA256,
-                                            PK11_OriginUnwrap, CKA_DERIVE,
-                                            &keyItem, NULL);
-    if (!ssl_anti_replay.key) {
+    ctx->key = PK11_ImportSymKey(slot, CKM_NSS_HKDF_SHA256,
+                                 PK11_OriginUnwrap, CKA_DERIVE,
+                                 &keyItem, NULL);
+    if (!ctx->key) {
         goto loser;
     }
 
@@ -100,20 +110,18 @@ loser:
 #define SSL_MAX_BLOOM_FILTER_SIZE 64
 
 /*
- * The structures created by this function can be called concurrently on
- * multiple threads if the server is multi-threaded.  A monitor is used to
- * ensure that only one thread can access the structures that change over time,
- * but no such guarantee is provided for configuration data.
- *
- * Functions that read from static configuration data depend on there being a
- * memory barrier between the setup and use of this function.
+ * The context created by this function can be called concurrently on multiple
+ * threads if the server is multi-threaded.  A monitor is used to ensure that
+ * only one thread can access the structures that change over time, but no such
+ * guarantee is provided for configuration data.
  */
 SECStatus
-SSLExp_InitAntiReplay(PRTime now, PRTime window, unsigned int k, unsigned int bits)
+SSLExp_CreateAntiReplayContext(PRTime now, PRTime window, unsigned int k,
+                               unsigned int bits, SSLAntiReplayContext **pctx)
 {
     SECStatus rv;
 
-    if (k == 0 || bits == 0) {
+    if (window <= 0 || k == 0 || bits == 0 || pctx == NULL) {
         PORT_SetError(SEC_ERROR_INVALID_ARGS);
         return SECFailure;
     }
@@ -122,54 +130,69 @@ SSLExp_InitAntiReplay(PRTime now, PRTime window, unsigned int k, unsigned int bi
         return SECFailure;
     }
 
-    if (PR_SUCCESS != PR_CallOnce(&ssl_anti_replay.init,
-                                  tls13_AntiReplayInit)) {
-        PORT_SetError(SEC_ERROR_LIBRARY_FAILURE);
-        return SECFailure;
+    SSLAntiReplayContext *ctx = PORT_ZNew(SSLAntiReplayContext);
+    if (!ctx) {
+        return SECFailure; /* Code already set. */
     }
 
-    (void)tls13_AntiReplayReset(NULL, NULL);
-
-    ssl_anti_replay.lock = PZ_NewMonitor(nssILockSSL);
-    if (!ssl_anti_replay.lock) {
+    ctx->refCount = 1;
+    ctx->lock = PZ_NewMonitor(nssILockSSL);
+    if (!ctx->lock) {
         goto loser; /* Code already set. */
     }
 
-    rv = tls13_AntiReplayKeyGen();
+    rv = tls13_AntiReplayKeyGen(ctx);
     if (rv != SECSuccess) {
         goto loser; /* Code already set. */
     }
 
-    rv = sslBloom_Init(&ssl_anti_replay.filters[0], k, bits);
+    rv = sslBloom_Init(&ctx->filters[0], k, bits);
     if (rv != SECSuccess) {
         goto loser; /* Code already set. */
     }
-    rv = sslBloom_Init(&ssl_anti_replay.filters[1], k, bits);
+    rv = sslBloom_Init(&ctx->filters[1], k, bits);
     if (rv != SECSuccess) {
         goto loser; /* Code already set. */
     }
     /* When starting out, ensure that 0-RTT is not accepted until the window is
      * updated.  A ClientHello might have been accepted prior to a restart. */
-    sslBloom_Fill(&ssl_anti_replay.filters[1]);
+    sslBloom_Fill(&ctx->filters[1]);
 
-    ssl_anti_replay.current = 0;
-    ssl_anti_replay.nextUpdate = now + window;
-    ssl_anti_replay.window = window;
+    ctx->current = 0;
+    ctx->nextUpdate = now + window;
+    ctx->window = window;
+    *pctx = ctx;
     return SECSuccess;
 
 loser:
-    (void)tls13_AntiReplayReset(NULL, NULL);
+    tls13_ReleaseAntiReplayContext(ctx);
     return SECFailure;
 }
 
+SECStatus
+SSLExp_SetAntiReplayContext(PRFileDesc *fd, SSLAntiReplayContext *ctx)
+{
+    sslSocket *ss = ssl_FindSocket(fd);
+    if (!ss) {
+        return SECFailure; /* Code already set. */
+    }
+    tls13_ReleaseAntiReplayContext(ss->antiReplay);
+    if (ctx != NULL) {
+        ss->antiReplay = tls13_RefAntiReplayContext(ctx);
+    } else {
+        ss->antiReplay = NULL;
+    }
+    return SECSuccess;
+}
+
 static void
-tls13_AntiReplayUpdate(PRTime now)
+tls13_AntiReplayUpdate(SSLAntiReplayContext *ctx, PRTime now)
 {
-    PR_ASSERT_CURRENT_THREAD_IN_MONITOR(ssl_anti_replay.lock);
-    if (now >= ssl_anti_replay.nextUpdate) {
-        ssl_anti_replay.current ^= 1;
-        ssl_anti_replay.nextUpdate = now + ssl_anti_replay.window;
-        sslBloom_Zero(ssl_anti_replay.filters + ssl_anti_replay.current);
+    PR_ASSERT_CURRENT_THREAD_IN_MONITOR(ctx->lock);
+    if (now >= ctx->nextUpdate) {
+        ctx->current ^= 1;
+        ctx->nextUpdate = now + ctx->window;
+        sslBloom_Zero(ctx->filters + ctx->current);
     }
 }
 
@@ -207,7 +230,10 @@ tls13_InWindow(const sslSocket *ss, const sslSessionID *sid)
      * prevent the same 0-RTT attempt from being accepted during window 1 and
      * later window 3.
      */
-    return PR_ABS(timeDelta) < (ssl_anti_replay.window / (PR_USEC_PER_MSEC * 2));
+    PRInt32 allowance = ss->antiReplay->window / (PR_USEC_PER_MSEC * 2);
+    SSL_TRC(10, ("%d: TLS13[%d]: replay check time delta=%d, allow=%d",
+                 SSL_GETPID(), ss->fd, timeDelta, allowance));
+    return PR_ABS(timeDelta) < allowance;
 }
 
 /* Checks for a duplicate in the two filters we have.  Performs maintenance on
@@ -222,12 +248,13 @@ tls13_IsReplay(const sslSocket *ss, const sslSessionID *sid)
     unsigned int size;
     PRUint8 index;
     SECStatus rv;
-    static const char *label = "tls13 anti-replay";
+    static const char *label = "anti-replay";
     PRUint8 buf[SSL_MAX_BLOOM_FILTER_SIZE];
+    SSLAntiReplayContext *ctx = ss->antiReplay;
 
-    /* If SSL_SetupAntiReplay hasn't been called, then treat all attempts at
-     * 0-RTT as a replay. */
-    if (!ssl_anti_replay.init.initialized) {
+    /* If SSL_SetAntiReplayContext hasn't been called with a valid context, then
+     * treat all attempts at 0-RTT as a replay. */
+    if (ctx == NULL) {
         return PR_TRUE;
     }
 
@@ -235,10 +262,9 @@ tls13_IsReplay(const sslSocket *ss, const sslSessionID *sid)
         return PR_TRUE;
     }
 
-    size = ssl_anti_replay.filters[0].k *
-           (ssl_anti_replay.filters[0].bits + 7) / 8;
+    size = ctx->filters[0].k * (ctx->filters[0].bits + 7) / 8;
     PORT_Assert(size <= SSL_MAX_BLOOM_FILTER_SIZE);
-    rv = tls13_HkdfExpandLabelRaw(ssl_anti_replay.key, ssl_hash_sha256,
+    rv = tls13_HkdfExpandLabelRaw(ctx->key, ssl_hash_sha256,
                                   ss->xtnData.pskBinder.data,
                                   ss->xtnData.pskBinder.len,
                                   label, strlen(label),
@@ -247,15 +273,19 @@ tls13_IsReplay(const sslSocket *ss, const sslSessionID *sid)
         return PR_TRUE;
     }
 
-    PZ_EnterMonitor(ssl_anti_replay.lock);
-    tls13_AntiReplayUpdate(ssl_Time(ss));
+    PZ_EnterMonitor(ctx->lock);
+    tls13_AntiReplayUpdate(ctx, ssl_Time(ss));
 
-    index = ssl_anti_replay.current;
-    replay = sslBloom_Add(&ssl_anti_replay.filters[index], buf);
+    index = ctx->current;
+    replay = sslBloom_Add(&ctx->filters[index], buf);
+    SSL_TRC(10, ("%d: TLS13[%d]: replay check current window: %s",
+                 SSL_GETPID(), ss->fd, replay ? "replay" : "ok"));
     if (!replay) {
-        replay = sslBloom_Check(&ssl_anti_replay.filters[index ^ 1], buf);
+        replay = sslBloom_Check(&ctx->filters[index ^ 1], buf);
+        SSL_TRC(10, ("%d: TLS13[%d]: replay check previous window: %s",
+                     SSL_GETPID(), ss->fd, replay ? "replay" : "ok"));
     }
 
-    PZ_ExitMonitor(ssl_anti_replay.lock);
+    PZ_ExitMonitor(ctx->lock);
     return replay;
 }