1 files changed, 63 insertions, 29 deletions

diff --git a/src/crypto/bytestring/cbb.c b/src/crypto/bytestring/cbb.c
index 434ec13..8fc5187 100644
--- a/src/crypto/bytestring/cbb.c
+++ b/src/crypto/bytestring/cbb.c
@@ -25,6 +25,7 @@ void CBB_zero(CBB *cbb) {
 }
 
 static int cbb_init(CBB *cbb, uint8_t *buf, size_t cap) {
+  /* This assumes that |cbb| has already been zeroed. */
   struct cbb_buffer_st *base;
 
   base = OPENSSL_malloc(sizeof(struct cbb_buffer_st));
@@ -37,16 +38,15 @@ static int cbb_init(CBB *cbb, uint8_t *buf, size_t cap) {
   base->cap = cap;
   base->can_resize = 1;
 
-  memset(cbb, 0, sizeof(CBB));
   cbb->base = base;
   cbb->is_top_level = 1;
   return 1;
 }
 
 int CBB_init(CBB *cbb, size_t initial_capacity) {
-  uint8_t *buf;
+  CBB_zero(cbb);
 
-  buf = OPENSSL_malloc(initial_capacity);
+  uint8_t *buf = OPENSSL_malloc(initial_capacity);
   if (initial_capacity > 0 && buf == NULL) {
     return 0;
   }
@@ -60,6 +60,8 @@ int CBB_init(CBB *cbb, size_t initial_capacity) {
 }
 
 int CBB_init_fixed(CBB *cbb, uint8_t *buf, size_t len) {
+  CBB_zero(cbb);
+
   if (!cbb_init(cbb, buf, len)) {
     return 0;
   }
@@ -82,8 +84,8 @@ void CBB_cleanup(CBB *cbb) {
   cbb->base = NULL;
 }
 
-static int cbb_buffer_add(struct cbb_buffer_st *base, uint8_t **out,
-                          size_t len) {
+static int cbb_buffer_reserve(struct cbb_buffer_st *base, uint8_t **out,
+                              size_t len) {
   size_t newlen;
 
   if (base == NULL) {
@@ -119,7 +121,17 @@ static int cbb_buffer_add(struct cbb_buffer_st *base, uint8_t **out,
   if (out) {
     *out = base->buf + base->len;
   }
-  base->len = newlen;
+
+  return 1;
+}
+
+static int cbb_buffer_add(struct cbb_buffer_st *base, uint8_t **out,
+                          size_t len) {
+  if (!cbb_buffer_reserve(base, out, len)) {
+    return 0;
+  }
+  /* This will not overflow or |cbb_buffer_reserve| would have failed. */
+  base->len += len;
   return 1;
 }
 
@@ -177,28 +189,28 @@ int CBB_flush(CBB *cbb) {
     return 0;
   }
 
-  if (cbb->child == NULL || cbb->pending_len_len == 0) {
+  if (cbb->child == NULL || cbb->child->pending_len_len == 0) {
     return 1;
   }
 
-  child_start = cbb->offset + cbb->pending_len_len;
+  child_start = cbb->child->offset + cbb->child->pending_len_len;
 
   if (!CBB_flush(cbb->child) ||
-      child_start < cbb->offset ||
+      child_start < cbb->child->offset ||
       cbb->base->len < child_start) {
     return 0;
   }
 
   len = cbb->base->len - child_start;
 
-  if (cbb->pending_is_asn1) {
+  if (cbb->child->pending_is_asn1) {
     /* For ASN.1 we assume that we'll only need a single byte for the length.
      * If that turned out to be incorrect, we have to move the contents along
      * in order to make space. */
     size_t len_len;
     uint8_t initial_length_byte;
 
-    assert (cbb->pending_len_len == 1);
+    assert (cbb->child->pending_len_len == 1);
 
     if (len > 0xfffffffe) {
       /* Too large. */
@@ -230,12 +242,13 @@ int CBB_flush(CBB *cbb) {
       memmove(cbb->base->buf + child_start + extra_bytes,
               cbb->base->buf + child_start, len);
     }
-    cbb->base->buf[cbb->offset++] = initial_length_byte;
-    cbb->pending_len_len = len_len - 1;
+    cbb->base->buf[cbb->child->offset++] = initial_length_byte;
+    cbb->child->pending_len_len = len_len - 1;
   }
 
-  for (i = cbb->pending_len_len - 1; i < cbb->pending_len_len; i--) {
-    cbb->base->buf[cbb->offset + i] = len;
+  for (i = cbb->child->pending_len_len - 1; i < cbb->child->pending_len_len;
+       i--) {
+    cbb->base->buf[cbb->child->offset + i] = len;
     len >>= 8;
   }
   if (len != 0) {
@@ -244,17 +257,20 @@ int CBB_flush(CBB *cbb) {
 
   cbb->child->base = NULL;
   cbb->child = NULL;
-  cbb->pending_len_len = 0;
-  cbb->pending_is_asn1 = 0;
-  cbb->offset = 0;
 
   return 1;
 }
 
+const uint8_t *CBB_data(const CBB *cbb) {
+  assert(cbb->child == NULL);
+  return cbb->base->buf + cbb->offset + cbb->pending_len_len;
+}
+
 size_t CBB_len(const CBB *cbb) {
   assert(cbb->child == NULL);
+  assert(cbb->offset + cbb->pending_len_len <= cbb->base->len);
 
-  return cbb->base->len;
+  return cbb->base->len - cbb->offset - cbb->pending_len_len;
 }
 
 static int cbb_add_length_prefixed(CBB *cbb, CBB *out_contents,
@@ -265,7 +281,7 @@ static int cbb_add_length_prefixed(CBB *cbb, CBB *out_contents,
     return 0;
   }
 
-  cbb->offset = cbb->base->len;
+  size_t offset = cbb->base->len;
   if (!cbb_buffer_add(cbb->base, &prefix_bytes, len_len)) {
     return 0;
   }
@@ -274,8 +290,9 @@ static int cbb_add_length_prefixed(CBB *cbb, CBB *out_contents,
   memset(out_contents, 0, sizeof(CBB));
   out_contents->base = cbb->base;
   cbb->child = out_contents;
-  cbb->pending_len_len = len_len;
-  cbb->pending_is_asn1 = 0;
+  cbb->child->offset = offset;
+  cbb->child->pending_len_len = len_len;
+  cbb->child->pending_is_asn1 = 0;
 
   return 1;
 }
@@ -303,7 +320,7 @@ int CBB_add_asn1(CBB *cbb, CBB *out_contents, uint8_t tag) {
     return 0;
   }
 
-  cbb->offset = cbb->base->len;
+  size_t offset = cbb->base->len;
   if (!CBB_add_u8(cbb, 0)) {
     return 0;
   }
@@ -311,8 +328,9 @@ int CBB_add_asn1(CBB *cbb, CBB *out_contents, uint8_t tag) {
   memset(out_contents, 0, sizeof(CBB));
   out_contents->base = cbb->base;
   cbb->child = out_contents;
-  cbb->pending_len_len = 1;
-  cbb->pending_is_asn1 = 1;
+  cbb->child->offset = offset;
+  cbb->child->pending_len_len = 1;
+  cbb->child->pending_is_asn1 = 1;
 
   return 1;
 }
@@ -336,6 +354,25 @@ int CBB_add_space(CBB *cbb, uint8_t **out_data, size_t len) {
   return 1;
 }
 
+int CBB_reserve(CBB *cbb, uint8_t **out_data, size_t len) {
+  if (!CBB_flush(cbb) ||
+      !cbb_buffer_reserve(cbb->base, out_data, len)) {
+    return 0;
+  }
+  return 1;
+}
+
+int CBB_did_write(CBB *cbb, size_t len) {
+  size_t newlen = cbb->base->len + len;
+  if (cbb->child != NULL ||
+      newlen < cbb->base->len ||
+      newlen > cbb->base->cap) {
+    return 0;
+  }
+  cbb->base->len = newlen;
+  return 1;
+}
+
 int CBB_add_u8(CBB *cbb, uint8_t value) {
   if (!CBB_flush(cbb)) {
     return 0;
@@ -365,13 +402,10 @@ void CBB_discard_child(CBB *cbb) {
     return;
   }
 
-  cbb->base->len = cbb->offset;
+  cbb->base->len = cbb->child->offset;
 
   cbb->child->base = NULL;
   cbb->child = NULL;
-  cbb->pending_len_len = 0;
-  cbb->pending_is_asn1 = 0;
-  cbb->offset = 0;
 }
 
 int CBB_add_asn1_uint64(CBB *cbb, uint64_t value) {