1 files changed, 71 insertions, 27 deletions

diff --git a/security/nss/gtests/ssl_gtest/tls_protect.cc b/security/nss/gtests/ssl_gtest/tls_protect.cc
index 6c87d0a05..6187660a5 100644
--- a/security/nss/gtests/ssl_gtest/tls_protect.cc
+++ b/security/nss/gtests/ssl_gtest/tls_protect.cc
@@ -1,4 +1,5 @@
 /* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=2 et sw=2 tw=80: */
 /* This Source Code Form is subject to the terms of the Mozilla Public
  * License, v. 2.0. If a copy of the MPL was not distributed with this file,
  * You can obtain one at http://mozilla.org/MPL/2.0/. */
@@ -24,39 +25,68 @@ TlsCipherSpec::TlsCipherSpec(bool dtls, uint16_t epoc)
 
 bool TlsCipherSpec::SetKeys(SSLCipherSuiteInfo* cipherinfo,
                             PK11SymKey* secret) {
-  SSLAeadContext* ctx;
-  SECStatus rv = SSL_MakeAead(SSL_LIBRARY_VERSION_TLS_1_3,
-                              cipherinfo->cipherSuite, secret, "",
-                              0,  // Use the default labels.
-                              &ctx);
+  SSLAeadContext* aead_ctx;
+  SSLProtocolVariant variant =
+      dtls_ ? ssl_variant_datagram : ssl_variant_stream;
+  SECStatus rv =
+      SSL_MakeVariantAead(SSL_LIBRARY_VERSION_TLS_1_3, cipherinfo->cipherSuite,
+                          variant, secret, "", 0,  // Use the default labels.
+                          &aead_ctx);
+  if (rv != SECSuccess) {
+    return false;
+  }
+  aead_.reset(aead_ctx);
+
+  SSLMaskingContext* mask_ctx;
+  const char kHkdfPurposeSn[] = "sn";
+  rv = SSL_CreateVariantMaskingContext(
+      SSL_LIBRARY_VERSION_TLS_1_3, cipherinfo->cipherSuite, variant, secret,
+      kHkdfPurposeSn, strlen(kHkdfPurposeSn), &mask_ctx);
   if (rv != SECSuccess) {
     return false;
   }
-  aead_.reset(ctx);
+  mask_.reset(mask_ctx);
   return true;
 }
 
 bool TlsCipherSpec::Unprotect(const TlsRecordHeader& header,
                               const DataBuffer& ciphertext,
-                              DataBuffer* plaintext) {
-  if (aead_ == nullptr) {
+                              DataBuffer* plaintext,
+                              TlsRecordHeader* out_header) {
+  if (!aead_ || !out_header) {
     return false;
   }
+  *out_header = header;
+
   // Make space.
   plaintext->Allocate(ciphertext.len());
 
-  auto header_bytes = header.header();
   unsigned int len;
-  uint64_t seqno;
-  if (dtls_) {
-    seqno = header.sequence_number();
-  } else {
-    seqno = in_seqno_;
+  uint64_t seqno = dtls_ ? header.sequence_number() : in_seqno_;
+  SECStatus rv;
+
+  if (header.is_dtls13_ciphertext()) {
+    if (!mask_ || !out_header) {
+      return false;
+    }
+    PORT_Assert(ciphertext.len() >= 16);
+    DataBuffer mask(2);
+    rv = SSL_CreateMask(mask_.get(), ciphertext.data(), ciphertext.len(),
+                        mask.data(), mask.len());
+    if (rv != SECSuccess) {
+      return false;
+    }
+
+    if (!out_header->MaskSequenceNumber(mask)) {
+      return false;
+    }
+    seqno = out_header->sequence_number();
   }
-  SECStatus rv =
-      SSL_AeadDecrypt(aead_.get(), seqno, header_bytes.data(),
-                      header_bytes.len(), ciphertext.data(), ciphertext.len(),
-                      plaintext->data(), &len, plaintext->len());
+
+  auto header_bytes = out_header->header();
+  rv = SSL_AeadDecrypt(aead_.get(), seqno, header_bytes.data(),
+                       header_bytes.len(), ciphertext.data(), ciphertext.len(),
+                       plaintext->data(), &len, plaintext->len());
   if (rv != SECSuccess) {
     return false;
   }
@@ -68,11 +98,14 @@ bool TlsCipherSpec::Unprotect(const TlsRecordHeader& header,
 }
 
 bool TlsCipherSpec::Protect(const TlsRecordHeader& header,
-                            const DataBuffer& plaintext,
-                            DataBuffer* ciphertext) {
-  if (aead_ == nullptr) {
+                            const DataBuffer& plaintext, DataBuffer* ciphertext,
+                            TlsRecordHeader* out_header) {
+  if (!aead_ || !out_header) {
     return false;
   }
+
+  *out_header = header;
+
   // Make a padded buffer.
   ciphertext->Allocate(plaintext.len() +
                        32);  // Room for any plausible auth tag
@@ -80,12 +113,7 @@ bool TlsCipherSpec::Protect(const TlsRecordHeader& header,
 
   DataBuffer header_bytes;
   (void)header.WriteHeader(&header_bytes, 0, plaintext.len() + 16);
-  uint64_t seqno;
-  if (dtls_) {
-    seqno = header.sequence_number();
-  } else {
-    seqno = out_seqno_;
-  }
+  uint64_t seqno = dtls_ ? header.sequence_number() : out_seqno_;
 
   SECStatus rv =
       SSL_AeadEncrypt(aead_.get(), seqno, header_bytes.data(),
@@ -95,6 +123,22 @@ bool TlsCipherSpec::Protect(const TlsRecordHeader& header,
     return false;
   }
 
+  if (header.is_dtls13_ciphertext()) {
+    if (!mask_ || !out_header) {
+      return false;
+    }
+    PORT_Assert(ciphertext->len() >= 16);
+    DataBuffer mask(2);
+    rv = SSL_CreateMask(mask_.get(), ciphertext->data(), ciphertext->len(),
+                        mask.data(), mask.len());
+    if (rv != SECSuccess) {
+      return false;
+    }
+    if (!out_header->MaskSequenceNumber(mask)) {
+      return false;
+    }
+  }
+
   RecordProtected();
   ciphertext->Truncate(len);