Bug 1623374 Need to support the new PKCS #11 Message interface for AES GCM and ChaCha Poly r=mt

Update ssl to use the new PK11_AEADOp() interface.
 1. We restore the use of PK11Context_Create() for AEAD operations.
 2. AES GCM and CHACHA/Poly specific functions are no longer needed as
  PK11_AEADOp() handles all the mechanism specific processing.
 3. TLS semantic differences between the two algorithms is handled by their
  parameters:
       1. Nonce length is the length of the nonce counter. If it's zero, then
       XOR_Counter is used (and the nonce length is the
       sizeof(sslSequenceNumber)).
       2. IV length is the full IV length - nonce length.
       3. TLS 1.3 always uses XOR_Counter.
 4. The IV is returned from the token in the encrypt case. Only in the explict
  nonce case is it examined. (The code depends on the fact that the count in
  the token will match sslSequenceNumber). I did have assert code to verify
  this was happening for testing, but it's removed from this patch it can be
  added back.
 5. All the decrypt instances of XOR_Counter IV creation have been colapsed
  into tls13_WriteNonce().
 6. Even tough PK11_AEADOp returns and accepts the tag separately (for encrypt
  and decrypt respectively). The SSL code still returns the values as
  buffer||tag.
 7. tls13_AEAD() has been enhanced so all uses of AEAD outside of the TLS
  stream can use it instead of their own wrapped version. It can handle streams
  (CreateContext() tls13_AEAD() tls13_AEAD() DestroyContext()) or single shot
  tls13_AEAD(context=NULL). In the later case, the keys for the single shot
  operation should not be resued.
 8. libssl_internals.c in the gtests directory has been updated to handle
  advancing the internal iv counter when we artifically advance the seqNum.
  Since we don't have access to any token iv counter (including softoken),
  The code switches to simulated message mode, and updates the simulated state
  as appropriate. (obviously this is for testing only code as it reaches into
  normally private data structures).

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

10 files changed, 454 insertions, 405 deletions

diff --git a/gtests/ssl_gtest/libssl_internals.c b/gtests/ssl_gtest/libssl_internals.c
index 6479dbb66..fff310b97 100644
--- a/gtests/ssl_gtest/libssl_internals.c
+++ b/gtests/ssl_gtest/libssl_internals.c
@@ -9,8 +9,11 @@
 
 #include "nss.h"
 #include "pk11pub.h"
+#include "pk11priv.h"
 #include "seccomon.h"
 #include "selfencrypt.h"
+#include "secmodti.h"
+#include "sslproto.h"
 
 SECStatus SSLInt_SetDCAdvertisedSigSchemes(PRFileDesc *fd,
                                            const SSLSignatureScheme *schemes,
@@ -331,6 +334,9 @@ SECStatus SSLInt_AdvanceReadSeqNum(PRFileDesc *fd, PRUint64 to) {
 
 SECStatus SSLInt_AdvanceWriteSeqNum(PRFileDesc *fd, PRUint64 to) {
   sslSocket *ss;
+  ssl3CipherSpec *spec;
+  PK11Context *pk11ctxt;
+  const ssl3BulkCipherDef *cipher_def;
 
   ss = ssl_FindSocket(fd);
   if (!ss) {
@@ -341,7 +347,43 @@ SECStatus SSLInt_AdvanceWriteSeqNum(PRFileDesc *fd, PRUint64 to) {
     return SECFailure;
   }
   ssl_GetSpecWriteLock(ss);
-  ss->ssl3.cwSpec->nextSeqNum = to;
+  spec = ss->ssl3.cwSpec;
+  cipher_def = spec->cipherDef;
+  spec->nextSeqNum = to;
+  if (cipher_def->type != type_aead) {
+    ssl_ReleaseSpecWriteLock(ss);
+    return SECSuccess;
+  }
+  /* If we are using aead, we need to advance the counter in the
+   * internal IV generator as well.
+   * This could be in the token or software. */
+  pk11ctxt = spec->cipherContext;
+  /* If counter is in the token, we need to switch it to software,
+   * since we don't have access to the internal state of the token. We do
+   * that by turning on the simulated message interface, then setting up the
+   * software IV generator */
+  if (pk11ctxt->ivCounter == 0) {
+    _PK11_ContextSetAEADSimulation(pk11ctxt);
+    pk11ctxt->ivLen = cipher_def->iv_size + cipher_def->explicit_nonce_size;
+    pk11ctxt->ivMaxCount = PR_UINT64(0xffffffffffffffff);
+    if ((cipher_def->explicit_nonce_size == 0) ||
+        (spec->version >= SSL_LIBRARY_VERSION_TLS_1_3)) {
+      pk11ctxt->ivFixedBits =
+          (pk11ctxt->ivLen - sizeof(sslSequenceNumber)) * BPB;
+      pk11ctxt->ivGen = CKG_GENERATE_COUNTER_XOR;
+    } else {
+      pk11ctxt->ivFixedBits = cipher_def->iv_size * BPB;
+      pk11ctxt->ivGen = CKG_GENERATE_COUNTER;
+    }
+    /* DTLS included the epoch in the fixed portion of the IV */
+    if (IS_DTLS(ss)) {
+      pk11ctxt->ivFixedBits += 2 * BPB;
+    }
+  }
+  /* now we can update the internal counter (either we are already using
+   * the software IV generator, or we just switched to it above */
+  pk11ctxt->ivCounter = to;
+
   ssl_ReleaseSpecWriteLock(ss);
   return SECSuccess;
 }
diff --git a/lib/pk11wrap/exports.gyp b/lib/pk11wrap/exports.gyp
index b3d4bf468..365006080 100644
--- a/lib/pk11wrap/exports.gyp
+++ b/lib/pk11wrap/exports.gyp
@@ -26,7 +26,8 @@
         {
           'files': [
             'dev3hack.h',
-            'secmodi.h'
+            'secmodi.h',
+            'secmodti.h'
           ],
           'destination': '<(nss_private_dist_dir)/<(module)'
         }
diff --git a/lib/pk11wrap/manifest.mn b/lib/pk11wrap/manifest.mn
index 1a2d2446d..1df7e09bf 100644
--- a/lib/pk11wrap/manifest.mn
+++ b/lib/pk11wrap/manifest.mn
@@ -17,6 +17,7 @@ EXPORTS = \
 
 PRIVATE_EXPORTS = \
 	secmodi.h \
+	secmodti.h \
 	dev3hack.h \
 	$(NULL)
 
diff --git a/lib/ssl/ssl3con.c b/lib/ssl/ssl3con.c
index 8beda41a7..073263e73 100644
--- a/lib/ssl/ssl3con.c
+++ b/lib/ssl/ssl3con.c
@@ -392,7 +392,7 @@ static const SSLCipher2Mech alg2Mech[] = {
     { ssl_calg_camellia, CKM_CAMELLIA_CBC },
     { ssl_calg_seed, CKM_SEED_CBC },
     { ssl_calg_aes_gcm, CKM_AES_GCM },
-    { ssl_calg_chacha20, CKM_NSS_CHACHA20_POLY1305 },
+    { ssl_calg_chacha20, CKM_CHACHA20_POLY1305 },
 };
 
 const PRUint8 tls12_downgrade_random[] = { 0x44, 0x4F, 0x57, 0x4E,
@@ -1740,117 +1740,6 @@ ssl3_BuildRecordPseudoHeader(DTLSEpoch epoch,
     return SECSuccess;
 }
 
-static SECStatus
-ssl3_AESGCM(const ssl3KeyMaterial *keys,
-            PRBool doDecrypt,
-            unsigned char *out,
-            unsigned int *outlen,
-            unsigned int maxout,
-            const unsigned char *in,
-            unsigned int inlen,
-            const unsigned char *additionalData,
-            unsigned int additionalDataLen)
-{
-    SECItem param;
-    SECStatus rv = SECFailure;
-    unsigned char nonce[12];
-    unsigned int uOutLen;
-    CK_NSS_GCM_PARAMS gcmParams; /* future use CK_GCM_PARAMS_V3 with fallback */
-
-    const int tagSize = 16;
-    const int explicitNonceLen = 8;
-
-    /* See https://tools.ietf.org/html/rfc5288#section-3 for details of how the
-     * nonce is formed. */
-    memcpy(nonce, keys->iv, 4);
-    if (doDecrypt) {
-        memcpy(nonce + 4, in, explicitNonceLen);
-        in += explicitNonceLen;
-        inlen -= explicitNonceLen;
-        *outlen = 0;
-    } else {
-        if (maxout < explicitNonceLen) {
-            PORT_SetError(SEC_ERROR_INPUT_LEN);
-            return SECFailure;
-        }
-        /* Use the 64-bit sequence number as the explicit nonce. */
-        memcpy(nonce + 4, additionalData, explicitNonceLen);
-        memcpy(out, additionalData, explicitNonceLen);
-        out += explicitNonceLen;
-        maxout -= explicitNonceLen;
-        *outlen = explicitNonceLen;
-    }
-
-    param.type = siBuffer;
-    param.data = (unsigned char *)&gcmParams;
-    param.len = sizeof(gcmParams);
-    gcmParams.pIv = nonce;
-    gcmParams.ulIvLen = sizeof(nonce);
-    gcmParams.pAAD = (unsigned char *)additionalData; /* const cast */
-    gcmParams.ulAADLen = additionalDataLen;
-    gcmParams.ulTagBits = tagSize * 8;
-
-    if (doDecrypt) {
-        rv = PK11_Decrypt(keys->key, CKM_AES_GCM, &param, out, &uOutLen,
-                          maxout, in, inlen);
-    } else {
-        rv = PK11_Encrypt(keys->key, CKM_AES_GCM, &param, out, &uOutLen,
-                          maxout, in, inlen);
-    }
-    *outlen += (int)uOutLen;
-
-    return rv;
-}
-
-static SECStatus
-ssl3_ChaCha20Poly1305(const ssl3KeyMaterial *keys, PRBool doDecrypt,
-                      unsigned char *out, unsigned int *outlen, unsigned int maxout,
-                      const unsigned char *in, unsigned int inlen,
-                      const unsigned char *additionalData,
-                      unsigned int additionalDataLen)
-{
-    size_t i;
-    SECItem param;
-    SECStatus rv = SECFailure;
-    unsigned int uOutLen;
-    unsigned char nonce[12];
-    CK_NSS_AEAD_PARAMS aeadParams;
-
-    const int tagSize = 16;
-
-    /* See
-     * https://tools.ietf.org/html/draft-ietf-tls-chacha20-poly1305-04#section-2
-     * for details of how the nonce is formed. */
-    PORT_Memcpy(nonce, keys->iv, 12);
-
-    /* XOR the last 8 bytes of the IV with the sequence number. */
-    PORT_Assert(additionalDataLen >= 8);
-    for (i = 0; i < 8; ++i) {
-        nonce[4 + i] ^= additionalData[i];
-    }
-
-    param.type = siBuffer;
-    param.len = sizeof(aeadParams);
-    param.data = (unsigned char *)&aeadParams;
-    memset(&aeadParams, 0, sizeof(aeadParams));
-    aeadParams.pNonce = nonce;
-    aeadParams.ulNonceLen = sizeof(nonce);
-    aeadParams.pAAD = (unsigned char *)additionalData;
-    aeadParams.ulAADLen = additionalDataLen;
-    aeadParams.ulTagLen = tagSize;
-
-    if (doDecrypt) {
-        rv = PK11_Decrypt(keys->key, CKM_NSS_CHACHA20_POLY1305, &param,
-                          out, &uOutLen, maxout, in, inlen);
-    } else {
-        rv = PK11_Encrypt(keys->key, CKM_NSS_CHACHA20_POLY1305, &param,
-                          out, &uOutLen, maxout, in, inlen);
-    }
-    *outlen = (int)uOutLen;
-
-    return rv;
-}
-
 /* Initialize encryption and MAC contexts for pending spec.
  * Master Secret already is derived.
  * Caller holds Spec write lock.
@@ -1868,40 +1757,26 @@ ssl3_InitPendingContexts(sslSocket *ss, ssl3CipherSpec *spec)
     PORT_Assert(ss->opt.noLocks || ssl_HaveSSL3HandshakeLock(ss));
     PORT_Assert(ss->opt.noLocks || ssl_HaveSpecWriteLock(ss));
 
-    macLength = spec->macDef->mac_size;
     calg = spec->cipherDef->calg;
     PORT_Assert(alg2Mech[calg].calg == calg);
 
-    if (spec->cipherDef->type == type_aead) {
-        spec->cipher = NULL;
-        spec->cipherContext = NULL;
-        switch (calg) {
-            case ssl_calg_aes_gcm:
-                spec->aead = ssl3_AESGCM;
-                break;
-            case ssl_calg_chacha20:
-                spec->aead = ssl3_ChaCha20Poly1305;
-                break;
-            default:
-                PORT_SetError(SEC_ERROR_LIBRARY_FAILURE);
-                return SECFailure;
-        }
-        return SECSuccess;
-    }
+    if (spec->cipherDef->type != type_aead) {
+        macLength = spec->macDef->mac_size;
 
-    /*
-    ** Now setup the MAC contexts,
-    **   crypto contexts are setup below.
-    */
-    macParam.data = (unsigned char *)&macLength;
-    macParam.len = sizeof(macLength);
-    macParam.type = siBuffer;
-
-    spec->keyMaterial.macContext = PK11_CreateContextBySymKey(
-        spec->macDef->mmech, CKA_SIGN, spec->keyMaterial.macKey, &macParam);
-    if (!spec->keyMaterial.macContext) {
-        ssl_MapLowLevelError(SSL_ERROR_SYM_KEY_CONTEXT_FAILURE);
-        return SECFailure;
+        /*
+        ** Now setup the MAC contexts,
+        **   crypto contexts are setup below.
+        */
+        macParam.data = (unsigned char *)&macLength;
+        macParam.len = sizeof(macLength);
+        macParam.type = siBuffer;
+
+        spec->keyMaterial.macContext = PK11_CreateContextBySymKey(
+            spec->macDef->mmech, CKA_SIGN, spec->keyMaterial.macKey, &macParam);
+        if (!spec->keyMaterial.macContext) {
+            ssl_MapLowLevelError(SSL_ERROR_SYM_KEY_CONTEXT_FAILURE);
+            return SECFailure;
+        }
     }
 
     /*
@@ -1912,15 +1787,21 @@ ssl3_InitPendingContexts(sslSocket *ss, ssl3CipherSpec *spec)
         return SECSuccess;
     }
 
-    spec->cipher = (SSLCipher)PK11_CipherOp;
     encMechanism = ssl3_Alg2Mech(calg);
     encMode = (spec->direction == ssl_secret_write) ? CKA_ENCRYPT : CKA_DECRYPT;
+    if (spec->cipherDef->type == type_aead) {
+        encMode |= CKA_NSS_MESSAGE;
+        iv.data = NULL;
+        iv.len = 0;
+    } else {
+        spec->cipher = (SSLCipher)PK11_CipherOp;
+        iv.data = spec->keyMaterial.iv;
+        iv.len = spec->cipherDef->iv_size;
+    }
 
     /*
      * build the context
      */
-    iv.data = spec->keyMaterial.iv;
-    iv.len = spec->cipherDef->iv_size;
     spec->cipherContext = PK11_CreateContextBySymKey(encMechanism, encMode,
                                                      spec->keyMaterial.key,
                                                      &iv);
@@ -2239,26 +2120,55 @@ ssl3_MACEncryptRecord(ssl3CipherSpec *cwSpec,
         isDTLS, contentLen, &pseudoHeader);
     PORT_Assert(rv == SECSuccess);
     if (cwSpec->cipherDef->type == type_aead) {
-        const int nonceLen = cwSpec->cipherDef->explicit_nonce_size;
-        const int tagLen = cwSpec->cipherDef->tag_size;
+        const unsigned int nonceLen = cwSpec->cipherDef->explicit_nonce_size;
+        const unsigned int tagLen = cwSpec->cipherDef->tag_size;
+        unsigned int ivOffset = 0;
+        CK_GENERATOR_FUNCTION gen;
+        /* ivOut includes the iv and the nonce and is the internal iv/nonce
+         * for the AEAD function. On Encrypt, this is an in/out parameter */
+        unsigned char ivOut[MAX_IV_LENGTH];
+        ivLen = cwSpec->cipherDef->iv_size;
+
+        PORT_Assert((ivLen + nonceLen) <= MAX_IV_LENGTH);
+        PORT_Assert((ivLen + nonceLen) >= sizeof(sslSequenceNumber));
 
         if (nonceLen + contentLen + tagLen > SSL_BUFFER_SPACE(wrBuf)) {
             PORT_SetError(SEC_ERROR_LIBRARY_FAILURE);
             return SECFailure;
         }
 
-        rv = cwSpec->aead(
-            &cwSpec->keyMaterial,
-            PR_FALSE,                /* do encrypt */
-            SSL_BUFFER_NEXT(wrBuf),  /* output  */
-            &len,                    /* out len */
-            SSL_BUFFER_SPACE(wrBuf), /* max out */
-            pIn, contentLen,         /* input   */
-            SSL_BUFFER_BASE(&pseudoHeader), SSL_BUFFER_LEN(&pseudoHeader));
+        if (nonceLen == 0) {
+            ivOffset = ivLen - sizeof(sslSequenceNumber);
+            gen = CKG_GENERATE_COUNTER_XOR;
+        } else {
+            ivOffset = ivLen;
+            gen = CKG_GENERATE_COUNTER;
+        }
+        ivOffset = tls13_SetupAeadIv(isDTLS, ivOut, cwSpec->keyMaterial.iv,
+                                     ivOffset, ivLen, cwSpec->epoch);
+        rv = tls13_AEAD(cwSpec->cipherContext,
+                        PR_FALSE,
+                        gen, ivOffset * BPB,                /* iv generator params */
+                        ivOut,                              /* iv in  */
+                        ivOut,                              /* iv out */
+                        ivLen + nonceLen,                   /* full iv length */
+                        NULL, 0,                            /* nonce is generated*/
+                        SSL_BUFFER_BASE(&pseudoHeader),     /* aad */
+                        SSL_BUFFER_LEN(&pseudoHeader),      /* aadlen */
+                        SSL_BUFFER_NEXT(wrBuf) + nonceLen,  /* output  */
+                        &len,                               /* out len */
+                        SSL_BUFFER_SPACE(wrBuf) - nonceLen, /* max out */
+                        tagLen,
+                        pIn, contentLen); /* input   */
         if (rv != SECSuccess) {
             PORT_SetError(SSL_ERROR_ENCRYPTION_FAILURE);
             return SECFailure;
         }
+        len += nonceLen; /* include the nonce at the beginning */
+        /* copy out the generated iv if we are using explict nonces */
+        if (nonceLen) {
+            PORT_Memcpy(SSL_BUFFER_NEXT(wrBuf), ivOut + ivLen, nonceLen);
+        }
 
         rv = sslBuffer_Skip(wrBuf, len, NULL);
         PORT_Assert(rv == SECSuccess); /* Can't fail. */
@@ -12707,21 +12617,50 @@ ssl3_UnprotectRecord(sslSocket *ss,
          * ciphertext by a fixed byte count, but it is not true in general.
          * Each AEAD cipher should provide a function that returns the
          * plaintext length for a given ciphertext. */
-        unsigned int decryptedLen =
-            cText->buf->len - cipher_def->explicit_nonce_size -
-            cipher_def->tag_size;
+        const unsigned int explicitNonceLen = cipher_def->explicit_nonce_size;
+        const unsigned int tagLen = cipher_def->tag_size;
+        unsigned int nonceLen = explicitNonceLen;
+        unsigned int decryptedLen = cText->buf->len - nonceLen - tagLen;
+        /* even though read doesn't return and IV, we still need a space to put
+         * the combined iv/nonce n the gcm 1.2 case*/
+        unsigned char ivOut[MAX_IV_LENGTH];
+        unsigned char *iv = NULL;
+        unsigned char *nonce = NULL;
+
+        ivLen = cipher_def->iv_size;
+
         rv = ssl3_BuildRecordPseudoHeader(
             spec->epoch, cText->seqNum,
             rType, isTLS, rVersion, IS_DTLS(ss), decryptedLen, &header);
         PORT_Assert(rv == SECSuccess);
-        rv = spec->aead(&spec->keyMaterial,
-                        PR_TRUE,          /* do decrypt */
-                        plaintext->buf,   /* out */
-                        &plaintext->len,  /* outlen */
-                        plaintext->space, /* maxout */
-                        cText->buf->buf,  /* in */
-                        cText->buf->len,  /* inlen */
-                        SSL_BUFFER_BASE(&header), SSL_BUFFER_LEN(&header));
+
+        /* build the iv */
+        if (explicitNonceLen == 0) {
+            nonceLen = sizeof(cText->seqNum);
+            iv = spec->keyMaterial.iv;
+            nonce = SSL_BUFFER_BASE(&header);
+        } else {
+            PORT_Memcpy(ivOut, spec->keyMaterial.iv, ivLen);
+            PORT_Memset(ivOut + ivLen, 0, explicitNonceLen);
+            iv = ivOut;
+            nonce = cText->buf->buf;
+            nonceLen = explicitNonceLen;
+        }
+        rv = tls13_AEAD(spec->cipherContext, PR_TRUE,
+                        CKG_NO_GENERATE, 0,       /* iv generator params
+                                                        * (not used in decrypt)*/
+                        iv,                       /* iv in */
+                        NULL,                     /* iv out */
+                        ivLen + explicitNonceLen, /* full iv length */
+                        nonce, nonceLen,          /* nonce in */
+                        SSL_BUFFER_BASE(&header), /* aad */
+                        SSL_BUFFER_LEN(&header),  /* aadlen */
+                        plaintext->buf,           /* output  */
+                        &plaintext->len,          /* out len */
+                        plaintext->space,         /* max out */
+                        tagLen,
+                        cText->buf->buf + explicitNonceLen,  /* input */
+                        cText->buf->len - explicitNonceLen); /* input len */
         if (rv != SECSuccess) {
             good = 0;
         }
diff --git a/lib/ssl/sslprimitive.c b/lib/ssl/sslprimitive.c
index fcba5de88..2afecfb16 100644
--- a/lib/ssl/sslprimitive.c
+++ b/lib/ssl/sslprimitive.c
@@ -20,8 +20,14 @@
 #include "tls13hkdf.h"
 
 struct SSLAeadContextStr {
-    CK_MECHANISM_TYPE mech;
-    ssl3KeyMaterial keys;
+    /* sigh, the API creates a single context, but then uses either encrypt
+     * and decrypt on that context. We should take an encrypt/decrypt
+     * variable here, but for now create two contexts. */
+    PK11Context *encryptContext;
+    PK11Context *decryptContext;
+    int tagLen;
+    int ivLen;
+    unsigned char iv[MAX_IV_LENGTH];
 };
 
 SECStatus
@@ -33,6 +39,9 @@ SSLExp_MakeVariantAead(PRUint16 version, PRUint16 cipherSuite, SSLProtocolVarian
     char label[255]; // Maximum length label.
     static const char *const keySuffix = "key";
     static const char *const ivSuffix = "iv";
+    CK_MECHANISM_TYPE mech;
+    SECItem nullParams = { siBuffer, NULL, 0 };
+    PK11SymKey *key = NULL;
 
     PORT_Assert(strlen(keySuffix) >= strlen(ivSuffix));
     if (secret == NULL || ctx == NULL ||
@@ -54,7 +63,9 @@ SSLExp_MakeVariantAead(PRUint16 version, PRUint16 cipherSuite, SSLProtocolVarian
     if (out == NULL) {
         goto loser;
     }
-    out->mech = ssl3_Alg2Mech(cipher->calg);
+    mech = ssl3_Alg2Mech(cipher->calg);
+    out->ivLen = cipher->iv_size + cipher->explicit_nonce_size;
+    out->tagLen = cipher->tag_size;
 
     memcpy(label, labelPrefix, labelPrefixLen);
     memcpy(label + labelPrefixLen, ivSuffix, strlen(ivSuffix));
@@ -63,7 +74,7 @@ SSLExp_MakeVariantAead(PRUint16 version, PRUint16 cipherSuite, SSLProtocolVarian
     rv = tls13_HkdfExpandLabelRaw(secret, hash,
                                   NULL, 0, // Handshake hash.
                                   label, labelLen, variant,
-                                  out->keys.iv, ivLen);
+                                  out->iv, ivLen);
     if (rv != SECSuccess) {
         goto loser;
     }
@@ -72,16 +83,36 @@ SSLExp_MakeVariantAead(PRUint16 version, PRUint16 cipherSuite, SSLProtocolVarian
     labelLen = labelPrefixLen + strlen(keySuffix);
     rv = tls13_HkdfExpandLabel(secret, hash,
                                NULL, 0, // Handshake hash.
-                               label, labelLen, out->mech, cipher->key_size,
-                               variant, &out->keys.key);
+                               label, labelLen, mech, cipher->key_size,
+                               variant, &key);
     if (rv != SECSuccess) {
         goto loser;
     }
 
+    /* We really need to change the API to Create a context for each
+     * encrypt and decrypt rather than a single call that does both. it's
+     * almost certain that the underlying application tries to use the same
+     * context for both. */
+    out->encryptContext = PK11_CreateContextBySymKey(mech,
+                                                     CKA_NSS_MESSAGE | CKA_ENCRYPT,
+                                                     key, &nullParams);
+    if (out->encryptContext == NULL) {
+        goto loser;
+    }
+
+    out->decryptContext = PK11_CreateContextBySymKey(mech,
+                                                     CKA_NSS_MESSAGE | CKA_DECRYPT,
+                                                     key, &nullParams);
+    if (out->decryptContext == NULL) {
+        goto loser;
+    }
+
+    PK11_FreeSymKey(key);
     *ctx = out;
     return SECSuccess;
 
 loser:
+    PK11_FreeSymKey(key);
     SSLExp_DestroyAead(out);
     return SECFailure;
 }
@@ -100,71 +131,48 @@ SSLExp_DestroyAead(SSLAeadContext *ctx)
     if (!ctx) {
         return SECSuccess;
     }
+    if (ctx->encryptContext) {
+        PK11_DestroyContext(ctx->encryptContext, PR_TRUE);
+    }
+    if (ctx->decryptContext) {
+        PK11_DestroyContext(ctx->decryptContext, PR_TRUE);
+    }
 
-    PK11_FreeSymKey(ctx->keys.key);
     PORT_ZFree(ctx, sizeof(*ctx));
     return SECSuccess;
 }
 
 /* Bug 1529440 exists to refactor this and the other AEAD uses. */
 static SECStatus
-ssl_AeadInner(const SSLAeadContext *ctx, PRBool decrypt, PRUint64 counter,
+ssl_AeadInner(const SSLAeadContext *ctx, PK11Context *context,
+              PRBool decrypt, PRUint64 counter,
               const PRUint8 *aad, unsigned int aadLen,
-              const PRUint8 *plaintext, unsigned int plaintextLen,
+              const PRUint8 *in, unsigned int inLen,
               PRUint8 *out, unsigned int *outLen, unsigned int maxOut)
 {
-    if (ctx == NULL || (aad == NULL && aadLen > 0) || plaintext == NULL ||
+    if (ctx == NULL || (aad == NULL && aadLen > 0) || in == NULL ||
         out == NULL || outLen == NULL) {
         PORT_SetError(SEC_ERROR_INVALID_ARGS);
         return SECFailure;
     }
 
     // Setup the nonce.
-    PRUint8 nonce[12] = { 0 };
-    sslBuffer nonceBuf = SSL_BUFFER_FIXED(nonce + sizeof(nonce) - sizeof(counter),
-                                          sizeof(counter));
+    PRUint8 nonce[sizeof(counter)] = { 0 };
+    sslBuffer nonceBuf = SSL_BUFFER_FIXED(nonce, sizeof(counter));
     SECStatus rv = sslBuffer_AppendNumber(&nonceBuf, counter, sizeof(counter));
     if (rv != SECSuccess) {
         PORT_Assert(0);
         return SECFailure;
     }
-    for (int i = 0; i < sizeof(nonce); ++i) {
-        nonce[i] ^= ctx->keys.iv[i];
-    }
-
-    // Build AEAD parameters.
-    CK_NSS_GCM_PARAMS gcmParams = { 0 }; /* future, use CK_GCM_PARAMS_V3 with fallback */
-    CK_NSS_AEAD_PARAMS aeadParams = { 0 };
-    unsigned char *params;
-    unsigned int paramsLen;
-    switch (ctx->mech) {
-        case CKM_AES_GCM:
-            gcmParams.pIv = nonce;
-            gcmParams.ulIvLen = sizeof(nonce);
-            gcmParams.pAAD = (unsigned char *)aad; // const cast :(
-            gcmParams.ulAADLen = aadLen;
-            gcmParams.ulTagBits = 128; // GCM measures in bits.
-            params = (unsigned char *)&gcmParams;
-            paramsLen = sizeof(gcmParams);
-            break;
-
-        case CKM_NSS_CHACHA20_POLY1305:
-            aeadParams.pNonce = nonce;
-            aeadParams.ulNonceLen = sizeof(nonce);
-            aeadParams.pAAD = (unsigned char *)aad; // const cast :(
-            aeadParams.ulAADLen = aadLen;
-            aeadParams.ulTagLen = 16; // AEAD measures in octets.
-            params = (unsigned char *)&aeadParams;
-            paramsLen = sizeof(aeadParams);
-            break;
-
-        default:
-            PORT_SetError(SEC_ERROR_LIBRARY_FAILURE);
-            return SECFailure;
-    }
-
-    return tls13_AEAD(&ctx->keys, decrypt, out, outLen, maxOut,
-                      plaintext, plaintextLen, ctx->mech, params, paramsLen);
+    /* at least on encrypt, we should not be using CKG_NO_GENERATE, but
+     * the current experimental API has the application tracking the counter
+     * rather than token. We should look at the QUIC code and see if the
+     * counter can be moved internally where it belongs. That would
+     * also get rid of the  formatting code above and have the API
+     * call tls13_AEAD directly in SSLExp_Aead* */
+    return tls13_AEAD(context, decrypt, CKG_NO_GENERATE, 0, ctx->iv, NULL,
+                      ctx->ivLen, nonce, sizeof(counter), aad, aadLen,
+                      out, outLen, maxOut, ctx->tagLen, in, inLen);
 }
 
 SECStatus
@@ -174,19 +182,21 @@ SSLExp_AeadEncrypt(const SSLAeadContext *ctx, PRUint64 counter,
                    PRUint8 *out, unsigned int *outLen, unsigned int maxOut)
 {
     // false == encrypt
-    return ssl_AeadInner(ctx, PR_FALSE, counter, aad, aadLen,
-                         plaintext, plaintextLen, out, outLen, maxOut);
+    return ssl_AeadInner(ctx, ctx->encryptContext, PR_FALSE, counter,
+                         aad, aadLen, plaintext, plaintextLen,
+                         out, outLen, maxOut);
 }
 
 SECStatus
 SSLExp_AeadDecrypt(const SSLAeadContext *ctx, PRUint64 counter,
                    const PRUint8 *aad, unsigned int aadLen,
-                   const PRUint8 *plaintext, unsigned int plaintextLen,
+                   const PRUint8 *ciphertext, unsigned int ciphertextLen,
                    PRUint8 *out, unsigned int *outLen, unsigned int maxOut)
 {
     // true == decrypt
-    return ssl_AeadInner(ctx, PR_TRUE, counter, aad, aadLen,
-                         plaintext, plaintextLen, out, outLen, maxOut);
+    return ssl_AeadInner(ctx, ctx->decryptContext, PR_TRUE, counter,
+                         aad, aadLen, ciphertext, ciphertextLen,
+                         out, outLen, maxOut);
 }
 
 SECStatus
@@ -346,6 +356,7 @@ ssl_CreateMaskInner(SSLMaskingContext *ctx, const PRUint8 *sample,
 
     SECStatus rv = SECFailure;
     unsigned int outMaskLen = 0;
+    int paramLen = 0;
 
     /* Internal output len/buf, for use if the caller allocated and requested
      * less than one block of output. |oneBlock| should have size equal to the
@@ -375,14 +386,18 @@ ssl_CreateMaskInner(SSLMaskingContext *ctx, const PRUint8 *sample,
             }
             break;
         case CKM_NSS_CHACHA20_CTR:
-            if (sampleLen < 16) {
+            paramLen = 16;
+        /* fall through */
+        case CKM_CHACHA20:
+            paramLen = (paramLen) ? paramLen : sizeof(CK_CHACHA20_PARAMS);
+            if (sampleLen < paramLen) {
                 PORT_SetError(SEC_ERROR_INVALID_ARGS);
                 return SECFailure;
             }
 
             SECItem param;
             param.type = siBuffer;
-            param.len = 16;
+            param.len = paramLen;
             param.data = (PRUint8 *)sample; // const-cast :(
             unsigned char zeros[128] = { 0 };
 
diff --git a/lib/ssl/sslspec.h b/lib/ssl/sslspec.h
index d00b20d76..ad4176bb5 100644
--- a/lib/ssl/sslspec.h
+++ b/lib/ssl/sslspec.h
@@ -105,16 +105,16 @@ typedef SECStatus (*SSLCipher)(void *context,
                                unsigned int maxout,
                                const unsigned char *in,
                                unsigned int inlen);
-typedef SECStatus (*SSLAEADCipher)(
-    const ssl3KeyMaterial *keys,
-    PRBool doDecrypt,
-    unsigned char *out,
-    unsigned int *outlen,
-    unsigned int maxout,
-    const unsigned char *in,
-    unsigned int inlen,
-    const unsigned char *additionalData,
-    unsigned int additionalDataLen);
+typedef SECStatus (*SSLAEADCipher)(PK11Context *context,
+                                   CK_GENERATOR_FUNCTION ivGen,
+                                   unsigned int fixedbits,
+                                   unsigned char *iv, unsigned int ivlen,
+                                   const unsigned char *aad,
+                                   unsigned int aadlen,
+                                   unsigned char *out, unsigned int *outlen,
+                                   unsigned int maxout, unsigned char *tag,
+                                   unsigned int taglen,
+                                   const unsigned char *in, unsigned int inlen);
 
 /* The DTLS anti-replay window in number of packets. Defined here because we
  * need it in the cipher spec. Note that this is a ring buffer but left and
@@ -149,7 +149,6 @@ struct ssl3CipherSpecStr {
     const ssl3MACDef *macDef;
 
     SSLCipher cipher;
-    SSLAEADCipher aead;
     void *cipherContext;
 
     PK11SymKey *masterSecret;
diff --git a/lib/ssl/tls13con.c b/lib/ssl/tls13con.c
index 680358f6f..b8f31c3f0 100644
--- a/lib/ssl/tls13con.c
+++ b/lib/ssl/tls13con.c
@@ -29,24 +29,6 @@
 static SECStatus tls13_SetCipherSpec(sslSocket *ss, PRUint16 epoch,
                                      SSLSecretDirection install,
                                      PRBool deleteSecret);
-static SECStatus tls13_AESGCM(const ssl3KeyMaterial *keys,
-                              PRBool doDecrypt,
-                              unsigned char *out,
-                              unsigned int *outlen,
-                              unsigned int maxout,
-                              const unsigned char *in,
-                              unsigned int inlen,
-                              const unsigned char *additionalData,
-                              unsigned int additionalDataLen);
-static SECStatus tls13_ChaCha20Poly1305(const ssl3KeyMaterial *keys,
-                                        PRBool doDecrypt,
-                                        unsigned char *out,
-                                        unsigned int *outlen,
-                                        unsigned int maxout,
-                                        const unsigned char *in,
-                                        unsigned int inlen,
-                                        const unsigned char *additionalData,
-                                        unsigned int additionalDataLen);
 static SECStatus tls13_SendServerHelloSequence(sslSocket *ss);
 static SECStatus tls13_SendEncryptedExtensions(sslSocket *ss);
 static void tls13_SetKeyExchangeType(sslSocket *ss, const sslNamedGroupDef *group);
@@ -3531,21 +3513,6 @@ tls13_SetSpecRecordVersion(sslSocket *ss, ssl3CipherSpec *spec)
                  SSL_GETPID(), ss->fd, spec, spec->recordVersion));
 }
 
-SSLAEADCipher
-tls13_GetAead(const ssl3BulkCipherDef *cipherDef)
-{
-    switch (cipherDef->calg) {
-        case ssl_calg_aes_gcm:
-            return tls13_AESGCM;
-        case ssl_calg_chacha20:
-            return tls13_ChaCha20Poly1305;
-        default:
-            PORT_Assert(PR_FALSE);
-            PORT_SetError(SEC_ERROR_LIBRARY_FAILURE);
-            return NULL;
-    }
-}
-
 static SECStatus
 tls13_SetupPendingCipherSpec(sslSocket *ss, ssl3CipherSpec *spec)
 {
@@ -3569,10 +3536,6 @@ tls13_SetupPendingCipherSpec(sslSocket *ss, ssl3CipherSpec *spec)
                 SSL_GETPID(), ss->fd, suite));
 
     spec->cipherDef = ssl_GetBulkCipherDef(ssl_LookupCipherSuiteDef(suite));
-    spec->aead = tls13_GetAead(spec->cipherDef);
-    if (!spec->aead) {
-        return SECFailure;
-    }
 
     if (spec->epoch == TrafficKeyEarlyApplicationData) {
         spec->earlyDataRemaining =
@@ -3595,6 +3558,38 @@ tls13_SetupPendingCipherSpec(sslSocket *ss, ssl3CipherSpec *spec)
 }
 
 /*
+ * Initialize the cipher context. All TLS 1.3 operations are AEAD,
+ * so they are all message contexts.
+ */
+static SECStatus
+tls13_InitPendingContext(sslSocket *ss, ssl3CipherSpec *spec)
+{
+    CK_MECHANISM_TYPE encMechanism;
+    CK_ATTRIBUTE_TYPE encMode;
+    SECItem iv;
+    SSLCipherAlgorithm calg;
+
+    calg = spec->cipherDef->calg;
+
+    encMechanism = ssl3_Alg2Mech(calg);
+    encMode = CKA_NSS_MESSAGE | ((spec->direction == ssl_secret_write) ? CKA_ENCRYPT : CKA_DECRYPT);
+    iv.data = NULL;
+    iv.len = 0;
+
+    /*
+     * build the context
+     */
+    spec->cipherContext = PK11_CreateContextBySymKey(encMechanism, encMode,
+                                                     spec->keyMaterial.key,
+                                                     &iv);
+    if (!spec->cipherContext) {
+        ssl_MapLowLevelError(SSL_ERROR_SYM_KEY_CONTEXT_FAILURE);
+        return SECFailure;
+    }
+    return SECSuccess;
+}
+
+/*
  * Called before sending alerts to set up the right key on the client.
  * We might encounter errors during the handshake where the current
  * key is ClearText or EarlyApplicationData. This
@@ -3673,6 +3668,11 @@ tls13_SetCipherSpec(sslSocket *ss, PRUint16 epoch,
         goto loser;
     }
 
+    rv = tls13_InitPendingContext(ss, spec);
+    if (rv != SECSuccess) {
+        goto loser;
+    }
+
     /* Now that we've set almost everything up, finally cut over. */
     specp = (direction == ssl_secret_read) ? &ss->ssl3.crSpec : &ss->ssl3.cwSpec;
     ssl_GetSpecWriteLock(ss);
@@ -3835,100 +3835,113 @@ tls13_DestroyEarlyData(PRCList *list)
  * See RFC 5288 and https://tools.ietf.org/html/draft-ietf-tls-chacha20-poly1305-04#section-2
  */
 static void
-tls13_WriteNonce(const ssl3KeyMaterial *keys,
-                 const unsigned char *seqNumBuf, unsigned int seqNumLen,
-                 unsigned char *nonce, unsigned int nonceLen)
+tls13_WriteNonce(const unsigned char *ivIn, unsigned int ivInLen,
+                 const unsigned char *nonce, unsigned int nonceLen,
+                 unsigned char *ivOut, unsigned int ivOutLen)
 {
     size_t i;
+    unsigned int offset = ivOutLen - nonceLen;
 
-    PORT_Assert(nonceLen == 12);
-    memcpy(nonce, keys->iv, 12);
+    PORT_Assert(ivInLen <= ivOutLen);
+    PORT_Assert(nonceLen <= ivOutLen);
+    PORT_Memset(ivOut, 0, ivOutLen);
+    PORT_Memcpy(ivOut, ivIn, ivInLen);
 
-    /* XOR the last 8 bytes of the IV with the sequence number. */
-    PORT_Assert(seqNumLen == 8);
-    for (i = 0; i < 8; ++i) {
-        nonce[4 + i] ^= seqNumBuf[i];
+    /* XOR the last n bytes of the IV with the nonce (should be a counter). */
+    for (i = 0; i < nonceLen; ++i) {
+        ivOut[offset + i] ^= nonce[i];
     }
-    PRINT_BUF(50, (NULL, "Nonce", nonce, nonceLen));
+    PRINT_BUF(50, (NULL, "Nonce", ivOut, ivOutLen));
 }
 
-/* Implement the SSLAEADCipher interface defined in sslimpl.h.
- *
- * That interface takes the additional data (see below) and reinterprets that as
- * a sequence number. In TLS 1.3 there is no additional data so this value is
- * just the encoded sequence number.
+/* Setup the IV for AEAD encrypt. The PKCS #11 module will add the
+ * counter, but it doesn't know about the DTLS epic, so we add it here.
  */
-SECStatus
-tls13_AEAD(const ssl3KeyMaterial *keys, PRBool doDecrypt,
-           unsigned char *out, unsigned int *outlen, unsigned int maxout,
-           const unsigned char *in, unsigned int inlen,
-           CK_MECHANISM_TYPE mechanism,
-           unsigned char *aeadParams, unsigned int aeadParamLength)
+unsigned int
+tls13_SetupAeadIv(PRBool isDTLS, unsigned char *ivOut, unsigned char *ivIn,
+                  unsigned int offset, unsigned int ivLen, DTLSEpoch epoch)
 {
-    SECItem param = {
-        siBuffer, aeadParams, aeadParamLength
-    };
-
-    if (doDecrypt) {
-        return PK11_Decrypt(keys->key, mechanism, &param,
-                            out, outlen, maxout, in, inlen);
+    PORT_Memcpy(ivOut, ivIn, ivLen);
+    if (isDTLS) {
+        /* handle the tls 1.2 counter mode case, the epoc is copied
+         * instead of xored. We accomplish this by clearing ivOut
+         * before running xor. */
+        if (offset >= ivLen) {
+            ivOut[offset] = ivOut[offset + 1] = 0;
+        }
+        ivOut[offset] ^= (unsigned char)(epoch >> BPB) & 0xff;
+        ivOut[offset + 1] ^= (unsigned char)(epoch)&0xff;
+        offset += 2;
     }
-    return PK11_Encrypt(keys->key, mechanism, &param,
-                        out, outlen, maxout, in, inlen);
+    return offset;
 }
 
-static SECStatus
-tls13_AESGCM(const ssl3KeyMaterial *keys,
-             PRBool doDecrypt,
-             unsigned char *out,
-             unsigned int *outlen,
-             unsigned int maxout,
-             const unsigned char *in,
-             unsigned int inlen,
-             const unsigned char *additionalData,
-             unsigned int additionalDataLen)
+/*
+ * Do a single AEAD for TLS. This differs from PK11_AEADOp in the following
+ * ways.
+ *   1) If context is not supplied, it treats the operation as a single shot
+ *   and creates a context from symKey and mech.
+ *   2) It always assumes the tag will be at the end of the buffer
+ *   (in on decrypt, out on encrypt) just like the old single shot.
+ *   3) If we aren't generating an IV, it uses tls13_WriteNonce to create the
+ *   nonce.
+ * NOTE is context is supplied, symKey and mech are ignored
+ */
+SECStatus
+tls13_AEAD(PK11Context *context, PRBool decrypt,
+           CK_GENERATOR_FUNCTION ivGen, unsigned int fixedbits,
+           const unsigned char *ivIn, unsigned char *ivOut, unsigned int ivLen,
+           const unsigned char *nonceIn, unsigned int nonceLen,
+           const unsigned char *aad, unsigned int aadLen,
+           unsigned char *out, unsigned int *outLen, unsigned int maxout,
+           unsigned int tagLen, const unsigned char *in, unsigned int inLen)
 {
-    CK_NSS_GCM_PARAMS gcmParams; /* future use CK_GCM_PARAMS_V3 with fallback */
-    unsigned char nonce[12];
-
-    PORT_Assert(additionalDataLen >= 8);
-    memset(&gcmParams, 0, sizeof(gcmParams));
-    gcmParams.pIv = nonce;
-    gcmParams.ulIvLen = sizeof(nonce);
-    gcmParams.pAAD = (PRUint8 *)(additionalData + 8);
-    gcmParams.ulAADLen = additionalDataLen - 8;
-    gcmParams.ulTagBits = 128; /* GCM measures tag length in bits. */
-
-    tls13_WriteNonce(keys, additionalData, 8,
-                     nonce, sizeof(nonce));
-    return tls13_AEAD(keys, doDecrypt, out, outlen, maxout, in, inlen,
-                      CKM_AES_GCM,
-                      (unsigned char *)&gcmParams, sizeof(gcmParams));
-}
+    unsigned char *tag;
+    unsigned char iv[MAX_IV_LENGTH];
+    unsigned char tagbuf[HASH_LENGTH_MAX];
+    SECStatus rv;
 
-static SECStatus
-tls13_ChaCha20Poly1305(const ssl3KeyMaterial *keys, PRBool doDecrypt,
-                       unsigned char *out, unsigned int *outlen, unsigned int maxout,
-                       const unsigned char *in, unsigned int inlen,
-                       const unsigned char *additionalData,
-                       unsigned int additionalDataLen)
-{
-    CK_NSS_AEAD_PARAMS aeadParams;
-    unsigned char nonce[12];
-
-    PORT_Assert(additionalDataLen > 8);
-    memset(&aeadParams, 0, sizeof(aeadParams));
-    aeadParams.pNonce = nonce;
-    aeadParams.ulNonceLen = sizeof(nonce);
-    aeadParams.pAAD = (PRUint8 *)(additionalData + 8);
-    aeadParams.ulAADLen = additionalDataLen - 8;
-    aeadParams.ulTagLen = 16; /* The Poly1305 tag is 16 octets. */
-
-    tls13_WriteNonce(keys, additionalData, 8,
-                     nonce, sizeof(nonce));
-    return tls13_AEAD(keys, doDecrypt, out, outlen, maxout, in, inlen,
-                      CKM_NSS_CHACHA20_POLY1305,
-                      (unsigned char *)&aeadParams, sizeof(aeadParams));
+    /* must have either context or the symKey set */
+    if (!context) {
+        PORT_SetError(SEC_ERROR_INVALID_ARGS);
+        return SECFailure;
+    }
+
+    PORT_Assert(ivLen <= MAX_IV_LENGTH);
+    PORT_Assert(tagLen <= HASH_LENGTH_MAX);
+    if (!ivOut) {
+        ivOut = iv; /* caller doesn't need a returned, iv */
+    }
+
+    if (ivGen == CKG_NO_GENERATE) {
+        tls13_WriteNonce(ivIn, ivLen, nonceIn, nonceLen, ivOut, ivLen);
+    } else if (ivIn != ivOut) {
+        PORT_Memcpy(ivOut, ivIn, ivLen);
+    }
+    if (decrypt) {
+        inLen = inLen - tagLen;
+        tag = (unsigned char *)in + inLen;
+        /* tag is const on decrypt, but returned on encrypt */
+    } else {
+        /* tag is written to a separate buffer, then added to the end
+         * of the actual output buffer. This allows output buffer to be larger
+         * than the input buffer and everything still work */
+        tag = tagbuf;
+    }
+    rv = PK11_AEADOp(context, ivGen, fixedbits, ivOut, ivLen, aad, aadLen,
+                     out, (int *)outLen, maxout, tag, tagLen, in, inLen);
+    /* on encrypt SSL always puts the tag at the end of the buffer */
+    if ((rv == SECSuccess) && !(decrypt)) {
+        unsigned int len = *outLen;
+        /* make sure there is still space */
+        if (len + tagLen > maxout) {
+            PORT_SetError(SEC_ERROR_OUTPUT_LEN);
+            return SECFailure;
+        }
+        PORT_Memcpy(out + len, tag, tagLen);
+        *outLen += tagLen;
+    }
+    return rv;
 }
 
 static SECStatus
@@ -5375,6 +5388,10 @@ tls13_ProtectRecord(sslSocket *ss,
         sslBuffer buf = SSL_BUFFER_FIXED(hdr, sizeof(hdr));
         PRBool needsLength;
         PRUint8 aad[21];
+        const int ivLen = cipher_def->iv_size + cipher_def->explicit_nonce_size;
+        unsigned int ivOffset = ivLen - sizeof(sslSequenceNumber);
+        unsigned char ivOut[MAX_IV_LENGTH];
+
         unsigned int aadLen;
         unsigned int len;
 
@@ -5406,14 +5423,22 @@ tls13_ProtectRecord(sslSocket *ss,
         if (rv != SECSuccess) {
             return SECFailure;
         }
-        rv = cwSpec->aead(&cwSpec->keyMaterial,
-                          PR_FALSE,                /* do encrypt */
-                          SSL_BUFFER_NEXT(wrBuf),  /* output  */
-                          &len,                    /* out len */
-                          SSL_BUFFER_SPACE(wrBuf), /* max out */
-                          SSL_BUFFER_NEXT(wrBuf),  /* input */
-                          contentLen + 1,          /* input len */
-                          aad, aadLen);
+        /* set up initial IV value */
+        ivOffset = tls13_SetupAeadIv(IS_DTLS(ss), ivOut, cwSpec->keyMaterial.iv,
+                                     ivOffset, ivLen, cwSpec->epoch);
+
+        rv = tls13_AEAD(cwSpec->cipherContext, PR_FALSE,
+                        CKG_GENERATE_COUNTER_XOR, ivOffset * BPB,
+                        ivOut, ivOut, ivLen,             /* iv */
+                        NULL, 0,                         /* nonce */
+                        aad + sizeof(sslSequenceNumber), /* aad */
+                        aadLen - sizeof(sslSequenceNumber),
+                        SSL_BUFFER_NEXT(wrBuf),  /* output  */
+                        &len,                    /* out len */
+                        SSL_BUFFER_SPACE(wrBuf), /* max out */
+                        tagLen,
+                        SSL_BUFFER_NEXT(wrBuf), /* input */
+                        contentLen + 1);        /* input len */
         if (rv != SECSuccess) {
             PORT_SetError(SSL_ERROR_ENCRYPTION_FAILURE);
             return SECFailure;
@@ -5443,6 +5468,9 @@ tls13_UnprotectRecord(sslSocket *ss,
                       SSL3AlertDescription *alert)
 {
     const ssl3BulkCipherDef *cipher_def = spec->cipherDef;
+    const int ivLen = cipher_def->iv_size + cipher_def->explicit_nonce_size;
+    const int tagLen = cipher_def->tag_size;
+
     PRUint8 aad[21];
     unsigned int aadLen;
     SECStatus rv;
@@ -5470,7 +5498,7 @@ tls13_UnprotectRecord(sslSocket *ss,
 
     /* We can perform this test in variable time because the record's total
      * length and the ciphersuite are both public knowledge. */
-    if (cText->buf->len < cipher_def->tag_size) {
+    if (cText->buf->len < tagLen) {
         SSL_TRC(3,
                 ("%d: TLS13[%d]: record too short to contain valid AEAD data",
                  SSL_GETPID(), ss->fd));
@@ -5497,16 +5525,21 @@ tls13_UnprotectRecord(sslSocket *ss,
                                     spec->epoch, cText->seqNum,
                                     aad, &aadLen, sizeof(aad));
     if (rv != SECSuccess) {
+
         return SECFailure;
     }
-    rv = spec->aead(&spec->keyMaterial,
-                    PR_TRUE,          /* do decrypt */
-                    plaintext->buf,   /* out */
+    rv = tls13_AEAD(spec->cipherContext, PR_TRUE,
+                    CKG_NO_GENERATE, 0,                /* ignored for decrypt */
+                    spec->keyMaterial.iv, NULL, ivLen, /* iv */
+                    aad, sizeof(sslSequenceNumber),    /* nonce */
+                    aad + sizeof(sslSequenceNumber),   /* aad */
+                    aadLen - sizeof(sslSequenceNumber),
+                    plaintext->buf,   /* output  */
                     &plaintext->len,  /* outlen */
                     plaintext->space, /* maxout */
+                    tagLen,
                     cText->buf->buf,  /* in */
-                    cText->buf->len,  /* inlen */
-                    aad, aadLen);
+                    cText->buf->len); /* inlen */
     if (rv != SECSuccess) {
         SSL_TRC(3,
                 ("%d: TLS13[%d]: record has bogus MAC",
diff --git a/lib/ssl/tls13con.h b/lib/ssl/tls13con.h
index 83ecb3756..dd693b377 100644
--- a/lib/ssl/tls13con.h
+++ b/lib/ssl/tls13con.h
@@ -133,12 +133,18 @@ SECStatus tls13_SendKeyUpdate(sslSocket *ss, tls13KeyUpdateRequest request,
                               PRBool buffer);
 SECStatus SSLExp_KeyUpdate(PRFileDesc *fd, PRBool requestUpdate);
 PRBool tls13_MaybeTls13(sslSocket *ss);
-SSLAEADCipher tls13_GetAead(const ssl3BulkCipherDef *cipherDef);
-SECStatus tls13_AEAD(const ssl3KeyMaterial *keys, PRBool doDecrypt,
-                     unsigned char *out, unsigned int *outlen, unsigned int maxout,
-                     const unsigned char *in, unsigned int inlen,
-                     CK_MECHANISM_TYPE mechanism,
-                     unsigned char *aeadParams, unsigned int aeadParamLength);
+unsigned int tls13_SetupAeadIv(PRBool isDTLS, unsigned char *ivOut,
+                               unsigned char *ivIn, unsigned int offset,
+                               unsigned int ivLen, DTLSEpoch epoch);
+SECStatus tls13_AEAD(PK11Context *context, PRBool decrypt,
+                     CK_GENERATOR_FUNCTION ivGen, unsigned int fixedbits,
+                     const unsigned char *ivIn, unsigned char *ivOut,
+                     unsigned int ivLen,
+                     const unsigned char *nonceIn, unsigned int nonceLen,
+                     const unsigned char *aad, unsigned int aadLen,
+                     unsigned char *out, unsigned int *outLen,
+                     unsigned int maxout, unsigned int tagLen,
+                     const unsigned char *in, unsigned int inLen);
 void tls13_SetSpecRecordVersion(sslSocket *ss, ssl3CipherSpec *spec);
 SECStatus SSLExp_SendCertificateRequest(PRFileDesc *fd);
 
diff --git a/lib/ssl/tls13esni.c b/lib/ssl/tls13esni.c
index 4562b86a1..7182f22af 100644
--- a/lib/ssl/tls13esni.c
+++ b/lib/ssl/tls13esni.c
@@ -662,12 +662,6 @@ tls13_FormatEsniAADInput(sslBuffer *aadInput,
 {
     SECStatus rv;
 
-    /* 8 bytes of 0 for the sequence number. */
-    rv = sslBuffer_AppendNumber(aadInput, 0, 8);
-    if (rv != SECSuccess) {
-        return SECFailure;
-    }
-
     /* Key share. */
     PORT_Assert(keyShareLen > 0);
     rv = sslBuffer_Append(aadInput, keyShare, keyShareLen);
@@ -680,12 +674,10 @@ tls13_FormatEsniAADInput(sslBuffer *aadInput,
 
 static SECStatus
 tls13_ServerGetEsniAEAD(const sslSocket *ss, PRUint64 suite,
-                        const ssl3CipherSuiteDef **suiteDefp,
-                        SSLAEADCipher *aeadp)
+                        const ssl3CipherSuiteDef **suiteDefp)
 {
     SECStatus rv;
     const ssl3CipherSuiteDef *suiteDef;
-    SSLAEADCipher aead;
 
     /* Check against the suite list for ESNI */
     PRBool csMatch = PR_FALSE;
@@ -712,13 +704,8 @@ tls13_ServerGetEsniAEAD(const sslSocket *ss, PRUint64 suite,
     if (!suiteDef) {
         return SECFailure;
     }
-    aead = tls13_GetAead(ssl_GetBulkCipherDef(suiteDef));
-    if (!aead) {
-        return SECFailure;
-    }
 
     *suiteDefp = suiteDef;
-    *aeadp = aead;
     return SECSuccess;
 }
 
@@ -729,7 +716,6 @@ tls13_ServerDecryptEsniXtn(const sslSocket *ss, const PRUint8 *in, unsigned int
     sslReader rdr = SSL_READER(in, inLen);
     PRUint64 suite;
     const ssl3CipherSuiteDef *suiteDef = NULL;
-    SSLAEADCipher aead = NULL;
     TLSExtension *keyShareExtension;
     TLS13KeyShareEntry *entry = NULL;
     ssl3KeyMaterial keyMat = { NULL };
@@ -748,7 +734,7 @@ tls13_ServerDecryptEsniXtn(const sslSocket *ss, const PRUint8 *in, unsigned int
     }
 
     /* Find the AEAD */
-    rv = tls13_ServerGetEsniAEAD(ss, suite, &suiteDef, &aead);
+    rv = tls13_ServerGetEsniAEAD(ss, suite, &suiteDef);
     if (rv != SECSuccess) {
         goto loser;
     }
@@ -822,11 +808,25 @@ tls13_ServerDecryptEsniXtn(const sslSocket *ss, const PRUint8 *in, unsigned int
         goto loser;
     }
 
-    rv = aead(&keyMat, PR_TRUE /* Decrypt */,
-              out, outLen, maxLen,
-              buf.buf, buf.len,
-              SSL_BUFFER_BASE(&aadInput),
-              SSL_BUFFER_LEN(&aadInput));
+    const ssl3BulkCipherDef *cipher_def = ssl_GetBulkCipherDef(suiteDef);
+    unsigned char *aad = SSL_BUFFER_BASE(&aadInput);
+    int aadLen = SSL_BUFFER_LEN(&aadInput);
+    int ivLen = cipher_def->iv_size + cipher_def->explicit_nonce_size;
+    SSLCipherAlgorithm calg = cipher_def->calg;
+    unsigned char zero[sizeof(sslSequenceNumber)] = { 0 };
+    SECItem null_params = { siBuffer, NULL, 0 };
+    PK11Context *ctxt = PK11_CreateContextBySymKey(ssl3_Alg2Mech(calg),
+                                                   CKA_NSS_MESSAGE | CKA_DECRYPT,
+                                                   keyMat.key, &null_params);
+    if (!ctxt) {
+        sslBuffer_Clear(&aadInput);
+        goto loser;
+    }
+
+    rv = tls13_AEAD(ctxt, PR_TRUE /* Decrypt */, CKG_NO_GENERATE, 0,
+                    keyMat.iv, NULL, ivLen, zero, sizeof(zero), aad, aadLen,
+                    out, outLen, maxLen, cipher_def->tag_size, buf.buf, buf.len);
+    PK11_DestroyContext(ctxt, PR_TRUE);
     sslBuffer_Clear(&aadInput);
     if (rv != SECSuccess) {
         goto loser;
diff --git a/lib/ssl/tls13exthandle.c b/lib/ssl/tls13exthandle.c
index fadb8ceb9..5768fbce5 100644
--- a/lib/ssl/tls13exthandle.c
+++ b/lib/ssl/tls13exthandle.c
@@ -1154,7 +1154,6 @@ tls13_ClientSendEsniXtn(const sslSocket *ss, TLSExtensionData *xtnData,
     sslBuffer sni = SSL_BUFFER(sniBuf);
     const ssl3CipherSuiteDef *suiteDef;
     ssl3KeyMaterial keyMat;
-    SSLAEADCipher aead;
     PRUint8 outBuf[1024];
     unsigned int outLen;
     unsigned int sniStart;
@@ -1205,10 +1204,6 @@ tls13_ClientSendEsniXtn(const sslSocket *ss, TLSExtensionData *xtnData,
         PORT_SetError(SEC_ERROR_LIBRARY_FAILURE);
         return SECFailure;
     }
-    aead = tls13_GetAead(ssl_GetBulkCipherDef(suiteDef));
-    if (!aead) {
-        return SECFailure;
-    }
 
     /* Format the first part of the extension so we have the
      * encoded KeyShareEntry. */
@@ -1266,15 +1261,33 @@ tls13_ClientSendEsniXtn(const sslSocket *ss, TLSExtensionData *xtnData,
         ssl_DestroyKeyMaterial(&keyMat);
         return SECFailure;
     }
+
     /* Now encrypt. */
-    rv = aead(&keyMat, PR_FALSE /* Encrypt */,
-              outBuf, &outLen, sizeof(outBuf),
-              SSL_BUFFER_BASE(&sni),
-              SSL_BUFFER_LEN(&sni),
-              SSL_BUFFER_BASE(&aadInput),
-              SSL_BUFFER_LEN(&aadInput));
+    unsigned char *aad = SSL_BUFFER_BASE(&aadInput);
+    int aadLen = SSL_BUFFER_LEN(&aadInput);
+    const ssl3BulkCipherDef *cipher_def = ssl_GetBulkCipherDef(suiteDef);
+    int ivLen = cipher_def->iv_size + cipher_def->explicit_nonce_size;
+    unsigned char zero[sizeof(sslSequenceNumber)] = { 0 };
+    SSLCipherAlgorithm calg = cipher_def->calg;
+    SECItem null_params = { siBuffer, NULL, 0 };
+    PK11Context *ctxt = PK11_CreateContextBySymKey(ssl3_Alg2Mech(calg),
+                                                   CKA_NSS_MESSAGE | CKA_ENCRYPT,
+                                                   keyMat.key, &null_params);
+    if (!ctxt) {
+        ssl_DestroyKeyMaterial(&keyMat);
+        sslBuffer_Clear(&aadInput);
+        return SECFailure;
+    }
+
+    /* This function is a single shot, with fresh/unique keys, no need to
+     * generate the IV internally */
+    rv = tls13_AEAD(ctxt, PR_FALSE /* Encrypt */, CKG_NO_GENERATE, 0,
+                    keyMat.iv, NULL, ivLen, zero, sizeof(zero), aad, aadLen,
+                    outBuf, &outLen, sizeof(outBuf), cipher_def->tag_size,
+                    SSL_BUFFER_BASE(&sni), SSL_BUFFER_LEN(&sni));
     ssl_DestroyKeyMaterial(&keyMat);
     sslBuffer_Clear(&aadInput);
+    PK11_DestroyContext(ctxt, PR_TRUE);
     if (rv != SECSuccess) {
         return SECFailure;
     }