Return the record number out of the DTLS record layer

For now this is ignored by the caller, but ACK management will require
this value. To do this, I've introduced a DTLSRecordNumber type. This
isn't quite used in all the places that we could use, so I've left a
TODO about this.

As part of this, also fix some of the weirdness where some code would
mix up seq and epoch+seq across 1.2 and 1.3.

Fixing that up also revealed that actually we'd accidentally changed
SSL_get_read_sequence slightly in the epoch state refactor. This CL ends
up tweaking it slightly more in the DTLS 1.3 case (which doesn't work
anyway). I've updated the TODO comment accordingly.

Bug: 42290594
Change-Id: Ie7e7da02c432a137ae2314fb3557d1c3717952f3
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/72271
Reviewed-by: Nick Harper <nharper@chromium.org>
Commit-Queue: David Benjamin <davidben@google.com>
diff --git a/ssl/d1_both.cc b/ssl/d1_both.cc
index dd7d146..4bd30b6 100644
--- a/ssl/d1_both.cc
+++ b/ssl/d1_both.cc
@@ -412,8 +412,10 @@
 ssl_open_record_t dtls1_open_handshake(SSL *ssl, size_t *out_consumed,
                                        uint8_t *out_alert, Span<uint8_t> in) {
   uint8_t type;
+  DTLSRecordNumber record_number;
   Span<uint8_t> record;
-  auto ret = dtls_open_record(ssl, &type, &record, out_consumed, out_alert, in);
+  auto ret = dtls_open_record(ssl, &type, &record_number, &record, out_consumed,
+                              out_alert, in);
   if (ret != ssl_open_record_success) {
     return ret;
   }
@@ -709,7 +711,8 @@
       return seal_no_progress;
     }
 
-    if (!dtls_seal_record(ssl, out, out_len, max_out,
+    DTLSRecordNumber record_number;
+    if (!dtls_seal_record(ssl, &record_number, out, out_len, max_out,
                           SSL3_RT_CHANGE_CIPHER_SPEC, kChangeCipherSpec,
                           sizeof(kChangeCipherSpec), msg->epoch)) {
       return seal_error;
@@ -763,8 +766,10 @@
   ssl_do_msg_callback(ssl, 1 /* write */, SSL3_RT_HANDSHAKE,
                       MakeSpan(frag, frag_len));
 
-  if (!dtls_seal_record(ssl, out, out_len, max_out, SSL3_RT_HANDSHAKE,
-                        out + prefix, frag_len, msg->epoch)) {
+  DTLSRecordNumber record_number;
+  if (!dtls_seal_record(ssl, &record_number, out, out_len, max_out,
+                        SSL3_RT_HANDSHAKE, out + prefix, frag_len,
+                        msg->epoch)) {
     return seal_error;
   }
 
diff --git a/ssl/d1_pkt.cc b/ssl/d1_pkt.cc
index e768e03..18cd5ee 100644
--- a/ssl/d1_pkt.cc
+++ b/ssl/d1_pkt.cc
@@ -149,8 +149,10 @@
   assert(!SSL_in_init(ssl));
 
   uint8_t type;
+  DTLSRecordNumber record_number;
   Span<uint8_t> record;
-  auto ret = dtls_open_record(ssl, &type, &record, out_consumed, out_alert, in);
+  auto ret = dtls_open_record(ssl, &type, &record_number, &record, out_consumed,
+                              out_alert, in);
   if (ret != ssl_open_record_success) {
     return ret;
   }
@@ -259,12 +261,13 @@
     return -1;
   }
 
+  DTLSRecordNumber record_number;
   size_t ciphertext_len;
   if (!buf->EnsureCap(dtls_seal_prefix_len(ssl, epoch),
                       in.size() + SSL_max_seal_overhead(ssl)) ||
-      !dtls_seal_record(ssl, buf->remaining().data(), &ciphertext_len,
-                        buf->remaining().size(), type, in.data(), in.size(),
-                        epoch)) {
+      !dtls_seal_record(ssl, &record_number, buf->remaining().data(),
+                        &ciphertext_len, buf->remaining().size(), type,
+                        in.data(), in.size(), epoch)) {
     buf->Clear();
     return -1;
   }
diff --git a/ssl/dtls_record.cc b/ssl/dtls_record.cc
index f769696..5ab89bb 100644
--- a/ssl/dtls_record.cc
+++ b/ssl/dtls_record.cc
@@ -123,8 +123,6 @@
 
 BSSL_NAMESPACE_BEGIN
 
-static constexpr uint64_t kMaxSequenceNumber = (uint64_t{1} << 48) - 1;
-
 bool DTLSReplayBitmap::ShouldDiscard(uint64_t seq_num) const {
   const size_t kWindowSize = map_.size();
 
@@ -167,6 +165,15 @@
                                                      : ssl->s3->version;
 }
 
+static uint64_t dtls_aead_sequence(const SSL *ssl, DTLSRecordNumber num) {
+  // DTLS 1.3 uses the sequence number with the AEAD, while DTLS 1.2 uses the
+  // combined value. If the version is not known, the epoch is unencrypted and
+  // the value is ignored.
+  return (ssl->s3->version != 0 && ssl_protocol_version(ssl) >= TLS1_3_VERSION)
+             ? num.sequence()
+             : num.combined();
+}
+
 // reconstruct_epoch finds the largest epoch that ends with the epoch bits from
 // |wire_epoch| that is less than or equal to |current_epoch|, to match the
 // epoch reconstruction algorithm described in RFC 9147 section 4.2.2.
@@ -186,7 +193,7 @@
   // bytes), no payload, and 16 byte AEAD overhead, sending 2^48 records would
   // require 5 petabytes. This allows us to continue to pack a DTLS record
   // number into an 8-byte structure.
-  assert(max_valid_seqnum <= kMaxSequenceNumber);
+  assert(max_valid_seqnum <= DTLSRecordNumber::kMaxSequence);
   assert(seq_mask == 0xff || seq_mask == 0xffff);
 
   uint64_t max_seqnum_plus_one = max_valid_seqnum + 1;
@@ -195,7 +202,7 @@
   // This addition cannot overflow. It is at most 2^48 + seq_mask. It, however,
   // may exceed 2^48-1.
   uint64_t seqnum = max_seqnum_plus_one + diff;
-  bool too_large = seqnum > kMaxSequenceNumber;
+  bool too_large = seqnum > DTLSRecordNumber::kMaxSequence;
   // If the diff is larger than half the step size, then the closest seqnum
   // to max_seqnum_plus_one (in Z_{2^64}) is seqnum minus step instead of
   // seqnum.
@@ -205,7 +212,7 @@
   if (too_large || (closer_is_less && !would_underflow)) {
     seqnum -= step;
   }
-  assert(seqnum <= kMaxSequenceNumber);
+  assert(seqnum <= DTLSRecordNumber::kMaxSequence);
   return seqnum;
 }
 
@@ -215,11 +222,9 @@
 
 struct ParsedDTLSRecord {
   // read_epoch will be null if the record is for an unrecognized epoch. In that
-  // case, sequence may be zero if we are unable to decrypt the sequence number.
+  // case, |number| may be unset.
   DTLSReadEpoch *read_epoch = nullptr;
-  // sequence includes the epoch for DTLS 1.2 records and does not include it
-  // for DTLS 1.3.
-  uint64_t sequence = 0;
+  DTLSRecordNumber number;
   CBS header, body;
   uint8_t type = 0;
   uint16_t version = 0;
@@ -282,21 +287,24 @@
       writable_seq[i] ^= mask[i];
       seq = (seq << 8) | writable_seq[i];
     }
-    out->sequence = reconstruct_seqnum(seq, (1 << (seq_len * 8)) - 1,
-                                       out->read_epoch->bitmap.max_seq_num());
+    uint64_t full_seq = reconstruct_seqnum(
+        seq, (1 << (seq_len * 8)) - 1, out->read_epoch->bitmap.max_seq_num());
+    out->number = DTLSRecordNumber(epoch, full_seq);
   }
 
   return true;
 }
 
 static bool parse_dtls12_record(SSL *ssl, CBS *in, ParsedDTLSRecord *out) {
+  uint64_t epoch_and_seq;
   if (!CBS_get_u16(in, &out->version) ||  //
-      !CBS_get_u64(in, &out->sequence) ||
+      !CBS_get_u64(in, &epoch_and_seq) ||
       !CBS_get_u16_length_prefixed(in, &out->body)) {
     return false;
   }
+  out->number = DTLSRecordNumber::FromCombined(epoch_and_seq);
 
-  uint16_t epoch = static_cast<uint16_t>(out->sequence >> 48);
+  uint16_t epoch = out->number.epoch();
   bool version_ok;
   if (epoch == 0) {
     // Only check the first byte. Enforcing beyond that can prevent decoding
@@ -345,6 +353,7 @@
 }
 
 enum ssl_open_record_t dtls_open_record(SSL *ssl, uint8_t *out_type,
+                                        DTLSRecordNumber *out_number,
                                         Span<uint8_t> *out,
                                         size_t *out_consumed,
                                         uint8_t *out_alert, Span<uint8_t> in) {
@@ -368,7 +377,7 @@
   ssl_do_msg_callback(ssl, 0 /* read */, SSL3_RT_HEADER, record.header);
 
   if (record.read_epoch == nullptr ||
-      record.read_epoch->bitmap.ShouldDiscard(record.sequence)) {
+      record.read_epoch->bitmap.ShouldDiscard(record.number.sequence())) {
     // Drop this record. It's from an unknown epoch or is a replay. Note that if
     // the record is from next epoch, it could be buffered for later. For
     // simplicity, drop it and expect retransmit to handle it later; DTLS must
@@ -377,9 +386,10 @@
     return ssl_open_record_discard;
   }
 
-  // discard the body in-place.
+  // Decrypt the body in-place.
   if (!record.read_epoch->aead->Open(out, record.type, record.version,
-                                     record.sequence, record.header,
+                                     dtls_aead_sequence(ssl, record.number),
+                                     record.header,
                                      cbs_to_writable_bytes(record.body))) {
     // Bad packets are silently dropped in DTLS. See section 4.2.1 of RFC 6347.
     // Clear the error queue of any errors decryption may have added. Drop the
@@ -416,7 +426,7 @@
     } while (record.type == 0);
   }
 
-  record.read_epoch->bitmap.Record(record.sequence);
+  record.read_epoch->bitmap.Record(record.number.sequence());
 
   // TODO(davidben): Limit the number of empty records as in TLS? This is only
   // useful if we also limit discarded packets.
@@ -428,6 +438,7 @@
   ssl->s3->warning_alert_count = 0;
 
   *out_type = record.type;
+  *out_number = record.number;
   return ssl_open_record_success;
 }
 
@@ -479,9 +490,9 @@
          write_epoch->aead->ExplicitNonceLen();
 }
 
-bool dtls_seal_record(SSL *ssl, uint8_t *out, size_t *out_len, size_t max_out,
-                      uint8_t type, const uint8_t *in, size_t in_len,
-                      uint16_t epoch) {
+bool dtls_seal_record(SSL *ssl, DTLSRecordNumber *out_number, uint8_t *out,
+                      size_t *out_len, size_t max_out, uint8_t type,
+                      const uint8_t *in, size_t in_len, uint16_t epoch) {
   const size_t prefix = dtls_seal_prefix_len(ssl, epoch);
   if (buffers_alias(in, in_len, out, max_out) &&
       (max_out < prefix || out + prefix != in)) {
@@ -499,7 +510,7 @@
   const size_t record_header_len = dtls_record_header_write_len(ssl, epoch);
 
   // Ensure the sequence number update does not overflow.
-  if (write_epoch->next_seq + 1 > kMaxSequenceNumber) {
+  if (write_epoch->next_seq + 1 > DTLSRecordNumber::kMaxSequence) {
     OPENSSL_PUT_ERROR(SSL, ERR_R_OVERFLOW);
     return false;
   }
@@ -524,7 +535,7 @@
   }
 
   uint16_t record_version = dtls_record_version(ssl);
-  uint64_t aead_seq;
+  DTLSRecordNumber record_number(epoch, write_epoch->next_seq);
   if (dtls13_header) {
     // The first byte of the DTLS 1.3 record header has the following format:
     // 0 1 2 3 4 5 6 7
@@ -549,25 +560,21 @@
     // omit the length.
     out[3] = ciphertext_len >> 8;
     out[4] = ciphertext_len & 0xff;
-    // DTLS 1.3 uses the sequence number without the epoch for the AEAD.
-    aead_seq = write_epoch->next_seq;
   } else {
     out[0] = type;
     out[1] = record_version >> 8;
     out[2] = record_version & 0xff;
-    // DTLS 1.2 uses the sequence number with the epoch for the AEAD.
-    aead_seq = (uint64_t{epoch} << 48) | write_epoch->next_seq;
-    CRYPTO_store_u64_be(&out[3], aead_seq);
+    CRYPTO_store_u64_be(&out[3], record_number.combined());
     out[11] = ciphertext_len >> 8;
     out[12] = ciphertext_len & 0xff;
   }
   Span<const uint8_t> header = MakeConstSpan(out, record_header_len);
 
 
-  if (!write_epoch->aead->SealScatter(out + record_header_len, out + prefix,
-                                      out + prefix + in_len, type,
-                                      record_version, aead_seq, header, in,
-                                      in_len, extra_in, extra_in_len)) {
+  if (!write_epoch->aead->SealScatter(
+          out + record_header_len, out + prefix, out + prefix + in_len, type,
+          record_version, dtls_aead_sequence(ssl, record_number), header, in,
+          in_len, extra_in, extra_in_len)) {
     return false;
   }
 
@@ -587,6 +594,7 @@
     out[2] ^= mask[1];
   }
 
+  *out_number = record_number;
   write_epoch->next_seq++;
   *out_len = record_header_len + ciphertext_len;
   ssl_do_msg_callback(ssl, 1 /* write */, SSL3_RT_HEADER, header);
diff --git a/ssl/internal.h b/ssl/internal.h
index 0f3bd82..d9b6519 100644
--- a/ssl/internal.h
+++ b/ssl/internal.h
@@ -1231,6 +1231,31 @@
 
 // Record layer.
 
+// TODO(davidben): Use this type more extensively in the epoch state.
+class DTLSRecordNumber {
+ public:
+  static constexpr uint64_t kMaxSequence = (uint64_t{1} << 48) - 1;
+
+  DTLSRecordNumber() = default;
+  DTLSRecordNumber(uint16_t epoch, uint64_t sequence) {
+    BSSL_CHECK(sequence <= kMaxSequence);
+    combined_ = (uint64_t{epoch} << 48) | sequence;
+  }
+
+  static DTLSRecordNumber FromCombined(uint64_t combined) {
+    return DTLSRecordNumber(combined);
+  }
+
+  uint64_t combined() const { return combined_; }
+  uint16_t epoch() const { return combined_ >> 48; }
+  uint64_t sequence() const { return combined_ & kMaxSequence; }
+
+ private:
+  explicit DTLSRecordNumber(uint64_t combined) : combined_(combined) {}
+
+  uint64_t combined_ = 0;
+};
+
 class RecordNumberEncrypter {
  public:
   static constexpr bool kAllowUniquePtr = true;
@@ -1309,8 +1334,10 @@
 
 // dtls_open_record implements |tls_open_record| for DTLS. It only returns
 // |ssl_open_record_partial| if |in| was empty and sets |*out_consumed| to
-// zero. The caller should read one packet and try again.
+// zero. The caller should read one packet and try again. On success,
+// |*out_number| is set to the record number of the record.
 enum ssl_open_record_t dtls_open_record(SSL *ssl, uint8_t *out_type,
+                                        DTLSRecordNumber *out_number,
                                         Span<uint8_t> *out,
                                         size_t *out_consumed,
                                         uint8_t *out_alert, Span<uint8_t> in);
@@ -1348,10 +1375,11 @@
 // dtls_seal_record implements |tls_seal_record| for DTLS. |epoch| selects which
 // epoch's cipher state to use. Unlike |tls_seal_record|, |in| and |out| may
 // alias but, if they do, |in| must be exactly |dtls_seal_prefix_len| bytes
-// ahead of |out|.
-bool dtls_seal_record(SSL *ssl, uint8_t *out, size_t *out_len, size_t max_out,
-                      uint8_t type, const uint8_t *in, size_t in_len,
-                      uint16_t epoch);
+// ahead of |out|. On success, |*out_number| is set to the record number of the
+// record.
+bool dtls_seal_record(SSL *ssl, DTLSRecordNumber *out_number, uint8_t *out,
+                      size_t *out_len, size_t max_out, uint8_t type,
+                      const uint8_t *in, size_t in_len, uint16_t epoch);
 
 // ssl_process_alert processes |in| as an alert and updates |ssl|'s shutdown
 // state. It returns one of |ssl_open_record_discard|, |ssl_open_record_error|,
diff --git a/ssl/ssl_lib.cc b/ssl/ssl_lib.cc
index ef4b7d5..a5023c8 100644
--- a/ssl/ssl_lib.cc
+++ b/ssl/ssl_lib.cc
@@ -2952,20 +2952,24 @@
 
 uint64_t SSL_get_read_sequence(const SSL *ssl) {
   if (SSL_is_dtls(ssl)) {
-    // TODO(crbug.com/42290608): The API for read sequences in DTLS 1.3 needs to
-    // reworked. In DTLS 1.3, the read epoch is updated once new keys are
-    // derived (before we receive a message encrypted with those keys), which
-    // results in the read epoch being ahead of the highest record received.
-    // Additionally, when we process a KeyUpdate, we will install new read keys
-    // for the new epoch, but we may receive messages from the old epoch for
-    // some time if the ACK gets lost or there is reordering.
-
-    // max_seq_num already includes the epoch. However, the current epoch may
-    // be one ahead of the highest record received, immediately after a key
-    // change.
+    // TODO(crbug.com/42290608): This API needs to reworked. Right at an epoch
+    // transition, it is possible that |read_epoch| has not received any
+    // records. We will then return that sequence 0 is the highest received, but
+    // this is not quite right.
+    //
+    // This is mostly moot in DTLS 1.2 because, after the handshake, we will
+    // never be in this state. In DTLS 1.3, there is a key transition
+    // immediately after the handshake, and in the steady state with KeyUpdate.
+    // While not yet implemented, DTLS 1.3 will handle key changes by having two
+    // epochs live at once (current and optional next), only cycling forward
+    // when we receive a record at the new epoch.
+    //
+    // When we implement this, this sequence 0 edge case will be gone, but
+    // replaced with a different issue: our record layer APIs have no way to
+    // report transition state. We'll likely need a new API for DTLS offload.
     const DTLSReadEpoch *read_epoch = &ssl->d1->read_epoch;
-    assert(read_epoch->epoch >= read_epoch->bitmap.max_seq_num() >> 48);
-    return read_epoch->bitmap.max_seq_num();
+    return DTLSRecordNumber(read_epoch->epoch, read_epoch->bitmap.max_seq_num())
+        .combined();
   }
   return ssl->s3->read_sequence;
 }
@@ -2973,12 +2977,8 @@
 uint64_t SSL_get_write_sequence(const SSL *ssl) {
   if (SSL_is_dtls(ssl)) {
     const DTLSWriteEpoch *write_epoch = &ssl->d1->write_epoch;
-    uint64_t ret = write_epoch->next_seq;
-    if (SSL_is_dtls(ssl)) {
-      assert((ret >> 48) == 0);
-      ret |= uint64_t{write_epoch->epoch} << 48;
-    }
-    return ret;
+    return DTLSRecordNumber(write_epoch->epoch, write_epoch->next_seq)
+        .combined();
   }
 
   return ssl->s3->write_sequence;
diff --git a/ssl/ssl_test.cc b/ssl/ssl_test.cc
index eff0712..06f90f4 100644
--- a/ssl/ssl_test.cc
+++ b/ssl/ssl_test.cc
@@ -2809,22 +2809,28 @@
 
   if (is_dtls()) {
     if (version() == DTLS1_3_EXPERIMENTAL_VERSION) {
-      // Client and server write epochs should be at 3 (application data).
+      // Both client and server must be at epoch 3 (application data).
+      EXPECT_EQ(EpochFromSequence(client_read_seq), 3);
       EXPECT_EQ(EpochFromSequence(client_write_seq), 3);
+      EXPECT_EQ(EpochFromSequence(server_read_seq), 3);
       EXPECT_EQ(EpochFromSequence(server_write_seq), 3);
-      // TODO(crbug.com/42290608): The read sequences aren't checked because the
-      // SSL_get_read_sequence API needs to be reworked for DTLS 1.3.
+
+      // TODO(crbug.com/42290608): The next record to be written should exceed
+      // the largest received, but they'll actually be equal because the
+      // |SSL_get_read_sequence| API cannot represent DTLS key transitions.
+      EXPECT_GE(client_write_seq, server_read_seq);
+      EXPECT_GE(server_write_seq, client_read_seq);
     } else {
       // Both client and server must be at epoch 1.
       EXPECT_EQ(EpochFromSequence(client_read_seq), 1);
       EXPECT_EQ(EpochFromSequence(client_write_seq), 1);
       EXPECT_EQ(EpochFromSequence(server_read_seq), 1);
       EXPECT_EQ(EpochFromSequence(server_write_seq), 1);
-    }
 
-    // The next record to be written should exceed the largest received.
-    EXPECT_GT(client_write_seq, server_read_seq);
-    EXPECT_GT(server_write_seq, client_read_seq);
+      // The next record to be written should exceed the largest received.
+      EXPECT_GT(client_write_seq, server_read_seq);
+      EXPECT_GT(server_write_seq, client_read_seq);
+    }
   } else {
     // The next record to be written should equal the next to be received.
     EXPECT_EQ(client_write_seq, server_read_seq);