Add CBB_discard

Callers might need to unwrite data, after having written it.

Change-Id: I2588f7b8bd282b59e5c8fb3ae9a26ea4ffe946a6
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/79187
Reviewed-by: Adam Langley <agl@google.com>
Auto-Submit: David Benjamin <davidben@google.com>
Commit-Queue: David Benjamin <davidben@google.com>
diff --git a/crypto/bytestring/bytestring_test.cc b/crypto/bytestring/bytestring_test.cc
index c5fff11..0407f3f 100644
--- a/crypto/bytestring/bytestring_test.cc
+++ b/crypto/bytestring/bytestring_test.cc
@@ -517,6 +517,58 @@
   EXPECT_EQ(Bytes(kExpected), Bytes(buf, buf_len));
 }
 
+TEST(CBBTest, Discard) {
+  bssl::ScopedCBB cbb;
+  ASSERT_TRUE(CBB_init(cbb.get(), 0));
+  CBB_discard(cbb.get(), 0);
+  ASSERT_TRUE(CBB_add_u8(cbb.get(), 1));
+  ASSERT_TRUE(CBB_add_u8(cbb.get(), 2));
+  ASSERT_TRUE(CBB_add_u8(cbb.get(), 3));
+  ASSERT_TRUE(CBB_add_u8(cbb.get(), 4));
+  CBB_discard(cbb.get(), 2);
+  ASSERT_TRUE(CBB_add_u8(cbb.get(), 5));
+  ASSERT_TRUE(CBB_add_u8(cbb.get(), 6));
+  const uint8_t kExpected[] = {1, 2, 5, 6};
+  EXPECT_EQ(Bytes(kExpected), Bytes(CBB_data(cbb.get()), CBB_len(cbb.get())));
+  CBB child;
+  ASSERT_TRUE(CBB_add_u8_length_prefixed(cbb.get(), &child));
+  CBB_discard(&child, 0);
+  ASSERT_TRUE(CBB_add_u8(&child, 7));
+  ASSERT_TRUE(CBB_add_u8(&child, 8));
+  CBB_discard(&child, 2);
+  ASSERT_TRUE(CBB_add_u8(&child, 9));
+  ASSERT_TRUE(CBB_add_u8(&child, 10));
+  CBB_discard(&child, 1);
+  ASSERT_TRUE(CBB_flush(cbb.get()));
+  const uint8_t kExpected2[] = {1, 2, 5, 6, 1, 9};
+  EXPECT_EQ(Bytes(kExpected2), Bytes(CBB_data(cbb.get()), CBB_len(cbb.get())));
+  CBB_discard(cbb.get(), 6);
+  EXPECT_EQ(Bytes(""), Bytes(CBB_data(cbb.get()), CBB_len(cbb.get())));
+}
+
+TEST(CBBDeathTest, DiscardMisuse) {
+  bssl::ScopedCBB cbb;
+  ASSERT_TRUE(CBB_init(cbb.get(), 0));
+  // Discard too many bytes.
+  EXPECT_DEATH_IF_SUPPORTED(CBB_discard(cbb.get(), 1), "");
+  ASSERT_TRUE(CBB_add_u8(cbb.get(), 1));
+  ASSERT_TRUE(CBB_add_u8(cbb.get(), 2));
+  ASSERT_TRUE(CBB_add_u8(cbb.get(), 3));
+  ASSERT_TRUE(CBB_add_u8(cbb.get(), 4));
+  // Discard too many bytes.
+  EXPECT_DEATH_IF_SUPPORTED(CBB_discard(cbb.get(), 5), "");
+  CBB child;
+  ASSERT_TRUE(CBB_add_u8_length_prefixed(cbb.get(), &child));
+  // Discard from a |cbb| with an unflushed child.
+  EXPECT_DEATH_IF_SUPPORTED(CBB_discard(cbb.get(), 1), "");
+  EXPECT_DEATH_IF_SUPPORTED(CBB_discard(&child, 1), "");
+  ASSERT_TRUE(CBB_add_u8(&child, 1));
+  ASSERT_TRUE(CBB_add_u8(&child, 2));
+  ASSERT_TRUE(CBB_add_u8(&child, 3));
+  ASSERT_TRUE(CBB_add_u8(&child, 4));
+  EXPECT_DEATH_IF_SUPPORTED(CBB_discard(&child, 5), "");
+}
+
 TEST(CBBTest, Misuse) {
   bssl::ScopedCBB cbb;
   CBB child, contents;
diff --git a/crypto/bytestring/cbb.cc b/crypto/bytestring/cbb.cc
index b847331..86b772b 100644
--- a/crypto/bytestring/cbb.cc
+++ b/crypto/bytestring/cbb.cc
@@ -468,6 +468,13 @@
   return CBB_add_u64(cbb, CRYPTO_bswap8(value));
 }
 
+void CBB_discard(CBB *cbb, size_t len) {
+  BSSL_CHECK(cbb->child == nullptr);
+  BSSL_CHECK(len <= CBB_len(cbb));
+  struct cbb_buffer_st *base = cbb_get_base(cbb);
+  base->len -= len;
+}
+
 void CBB_discard_child(CBB *cbb) {
   if (cbb->child == NULL) {
     return;
diff --git a/include/openssl/bytestring.h b/include/openssl/bytestring.h
index 1738c08..d5df0b2 100644
--- a/include/openssl/bytestring.h
+++ b/include/openssl/bytestring.h
@@ -597,6 +597,10 @@
 // It returns one on success and zero otherwise.
 OPENSSL_EXPORT int CBB_add_u64le(CBB *cbb, uint64_t value);
 
+// CBB_discard discards the last |len| bytes written to |cbb|. The process will
+// abort if |cbb| has an unflushed child, or its length is smaller than |len|.
+OPENSSL_EXPORT void CBB_discard(CBB *cbb, size_t len);
+
 // CBB_discard_child discards the current unflushed child of |cbb|. Neither the
 // child's contents nor the length prefix will be included in the output.
 OPENSSL_EXPORT void CBB_discard_child(CBB *cbb);