Return required size if buffer is too small
The documentation for DiceMainFlow() says that the required buffer size
is returned if the provided buffer was too small. This stopped happening
in the CBOR implementation during the migration away from cn-cbor, so
fix that.
As the buffer is used both for the TBS data and the final result, care
is taken to find the size of buffer that is required for the whole
function to complete.
There are a number of APIs that follow a similar argument pattern so
also fix those to return the required size when the buffer is too small.
Change-Id: Ia5b01e36216761a58972a1afc8f747182470a631
Reviewed-on: https://pigweed-review.googlesource.com/c/open-dice/+/67760
Reviewed-by: Darren Krahn <dkrahn@google.com>
Commit-Queue: Andrew Scull <ascull@google.com>
diff --git a/include/dice/android/bcc.h b/include/dice/android/bcc.h
index 06f9fcb..366c51f 100644
--- a/include/dice/android/bcc.h
+++ b/include/dice/android/bcc.h
@@ -40,7 +40,10 @@
   uint64_t component_version;
 } BccConfigValues;
 
-// Formats a configuration descriptor following the BCC's specification.
+// Formats a configuration descriptor following the BCC's specification. On
+// success, |actual_size| is set to the number of bytes used. If
+// kDiceResultBufferTooSmall is returned |actual_size| will be set to the
+// required size of the buffer.
 DiceResult BccFormatConfigDescriptor(const BccConfigValues* input_values,
                                      size_t buffer_size, uint8_t* buffer,
                                      size_t* actual_size);
@@ -52,7 +55,10 @@
 // the BCC handover format, use BccHandoverMainFlow instead.
 //
 // Given a full set of input values along with the current BCC and CDI values,
-// computes the next CDI values and matching updated BCC.
+// computes the next CDI values and matching updated BCC. On success,
+// |actual_size| is set to the number of bytes used. If
+// kDiceResultBufferTooSmall is returned |actual_size| will be set to the
+// required size of the buffer.
 DiceResult BccMainFlow(void* context,
                        const uint8_t current_cdi_attest[DICE_CDI_SIZE],
                        const uint8_t current_cdi_seal[DICE_CDI_SIZE],
@@ -68,7 +74,9 @@
 // combine the BCC and CDIs in a single CBOR object.
 //
 // Given a full set of input values and the current BCC handover data, computes
-// the next BCC handover data.
+// the next BCC handover data. On success, |actual_size| is set to the number
+// of bytes used. If kDiceResultBufferTooSmall is returned |actual_size| will
+// be set to the required size of the buffer.
 //
 // Using a CBOR object to bundle is one option for passing the values passed
 // between boot stages. This function can take the current boot stage's bundle
diff --git a/include/dice/ops/trait/cose.h b/include/dice/ops/trait/cose.h
index 63404a5..78cb838 100644
--- a/include/dice/ops/trait/cose.h
+++ b/include/dice/ops/trait/cose.h
@@ -28,8 +28,10 @@
 extern "C" {
 #endif
 
-// Encodes a public key into |buffer| as a COSE_Key structure, setting
-// |encoded_size| to the number of bytes used.
+// Encodes a public key into |buffer| as a COSE_Key structure. On success,
+// |encoded_size| is set to the number of bytes used. If
+// kDiceResultBufferTooSmall is returned |encoded_size| will be set to the
+// required size of the buffer.
 DiceResult DiceCoseEncodePublicKey(
     void* context, const uint8_t public_key[DICE_PUBLIC_KEY_SIZE],
     size_t buffer_size, uint8_t* buffer, size_t* encoded_size);
@@ -39,8 +41,9 @@
 // the result.
 //
 // |buffer| is used to hold the intermediate To-Be-Signed (TBS) structure and
-// then the final result. |encoded_size| is set to the size of the final result
-// in |buffer|.
+// then the final result. On success, |encoded_size| is set to the size of the
+// final result in |buffer|. If kDiceResultBufferTooSmall is returned,
+// |encoded_size| will be set to the required size of the buffer.
 DiceResult DiceCoseSignAndEncodeSign1(
     void* context, const uint8_t* payload, size_t payload_size,
     const uint8_t* aad, size_t aad_size,
diff --git a/src/android/bcc.c b/src/android/bcc.c
index 5dd8620..262e6d3 100644
--- a/src/android/bcc.c
+++ b/src/android/bcc.c
@@ -18,9 +18,9 @@
 
 #include "dice/cbor_reader.h"
 #include "dice/cbor_writer.h"
-#include "dice/ops/trait/cose.h"
 #include "dice/dice.h"
 #include "dice/ops.h"
+#include "dice/ops/trait/cose.h"
 
 // Completely gratuitous bit twiddling.
 static size_t PopulationCount(uint32_t n) {
@@ -57,10 +57,10 @@
     CborWriteInt(kResettableLabel, &out);
     CborWriteNull(&out);
   }
+  *actual_size = CborOutSize(&out);
   if (CborOutOverflowed(&out)) {
     return kDiceResultBufferTooSmall;
   }
-  *actual_size = CborOutSize(&out);
   return kDiceResultOk;
 }
 
@@ -111,22 +111,20 @@
   CborWriteArray(bcc_item_count + 1, &out);
   if (CborOutOverflowed(&out) ||
       bcc_items_size > buffer_size - CborOutSize(&out)) {
-    return kDiceResultBufferTooSmall;
+    // Continue with an empty buffer to measure the required size.
+    buffer_size = 0;
+  } else {
+    memcpy(buffer + CborOutSize(&out), bcc + bcc_items_offset, bcc_items_size);
+    buffer += CborOutSize(&out) + bcc_items_size;
+    buffer_size -= CborOutSize(&out) + bcc_items_size;
   }
-  memcpy(buffer + CborOutSize(&out), bcc + bcc_items_offset, bcc_items_size);
 
   size_t certificate_size;
-  result =
-      DiceMainFlow(context, current_cdi_attest, current_cdi_seal, input_values,
-                   buffer_size - (CborOutSize(&out) + bcc_items_size),
-                   buffer + CborOutSize(&out) + bcc_items_size,
-                   &certificate_size, next_cdi_attest, next_cdi_seal);
-  if (result != kDiceResultOk) {
-    return result;
-  }
-
+  result = DiceMainFlow(context, current_cdi_attest, current_cdi_seal,
+                        input_values, buffer_size, buffer, &certificate_size,
+                        next_cdi_attest, next_cdi_seal);
   *actual_size = CborOutSize(&out) + bcc_items_size + certificate_size;
-  return kDiceResultOk;
+  return result;
 }
 
 static DiceResult BccMainFlowWithNewBcc(
@@ -157,31 +155,35 @@
   struct CborOut out;
   CborOutInit(buffer, buffer_size, &out);
   CborWriteArray(2, &out);
-  if (CborOutOverflowed(&out)) {
-    result = kDiceResultBufferTooSmall;
-    goto out;
-  }
   size_t encoded_size_used = CborOutSize(&out);
-  buffer += encoded_size_used;
-  buffer_size -= encoded_size_used;
+  if (CborOutOverflowed(&out)) {
+    // Continue with an empty buffer to measure the required size.
+    buffer_size = 0;
+  } else {
+    buffer += encoded_size_used;
+    buffer_size -= encoded_size_used;
+  }
 
   size_t encoded_pub_key_size = 0;
   result = DiceCoseEncodePublicKey(context, attestation_public_key, buffer_size,
                                    buffer, &encoded_pub_key_size);
-  if (result != kDiceResultOk) {
+  if (result == kDiceResultOk) {
+    buffer += encoded_pub_key_size;
+    buffer_size -= encoded_pub_key_size;
+  } else if (result == kDiceResultBufferTooSmall) {
+    // Continue with an empty buffer to measure the required size.
+    buffer_size = 0;
+  } else {
     goto out;
   }
 
-  buffer += encoded_pub_key_size;
-  buffer_size -= encoded_pub_key_size;
-
   result = DiceMainFlow(context, current_cdi_attest, current_cdi_seal,
                         input_values, buffer_size, buffer, bcc_size,
                         next_cdi_attest, next_cdi_seal);
+  *bcc_size += encoded_size_used + encoded_pub_key_size;
   if (result != kDiceResultOk) {
     return result;
   }
-  *bcc_size += encoded_size_used + encoded_pub_key_size;
 
 out:
   DiceClearMemory(context, sizeof(current_cdi_private_key_seed),
@@ -224,30 +226,34 @@
   uint8_t* next_cdi_seal = CborAllocBstr(DICE_CDI_SIZE, &out);
   CborWriteInt(kBccLabel, &out);
 
-  if (CborOutOverflowed(&out) || !next_cdi_attest || !next_cdi_seal) {
-    return kDiceResultBufferTooSmall;
+  uint8_t ignored_cdi_attest[DICE_CDI_SIZE];
+  uint8_t ignored_cdi_seal[DICE_CDI_SIZE];
+  if (CborOutOverflowed(&out)) {
+    // Continue with an empty buffer and placeholders for the output CDIs to
+    // measure the required size.
+    buffer_size = 0;
+    next_cdi_attest = ignored_cdi_attest;
+    next_cdi_seal = ignored_cdi_seal;
+  } else {
+    buffer += CborOutSize(&out);
+    buffer_size -= CborOutSize(&out);
   }
 
   if (bcc_size != 0) {
     // If BCC is present in the bcc_handover, append the next certificate to the
     // existing BCC.
     result = BccMainFlow(context, current_cdi_attest, current_cdi_seal, bcc,
-                         bcc_size, input_values, buffer_size - CborOutSize(&out),
-                         buffer + CborOutSize(&out), &bcc_size, next_cdi_attest,
-                         next_cdi_seal);
+                         bcc_size, input_values, buffer_size, buffer, &bcc_size,
+                         next_cdi_attest, next_cdi_seal);
   } else {
-    // If BCC is not present in the bcc_handover, construct BCC from the public key
-    // derived from the current CDI attest and the next CDI certificate.
+    // If BCC is not present in the bcc_handover, construct BCC from the public
+    // key derived from the current CDI attest and the next CDI certificate.
     result = BccMainFlowWithNewBcc(
         context, current_cdi_attest, current_cdi_seal, input_values,
-        buffer_size - CborOutSize(&out), buffer + CborOutSize(&out), &bcc_size,
-        next_cdi_attest, next_cdi_seal);
-  }
-  if (result != kDiceResultOk) {
-      return result;
+        buffer_size, buffer, &bcc_size, next_cdi_attest, next_cdi_seal);
   }
   *actual_size = CborOutSize(&out) + bcc_size;
-  return kDiceResultOk;
+  return result;
 }
 
 DiceResult BccHandoverParse(const uint8_t* bcc_handover,
diff --git a/src/android/bcc_test.cc b/src/android/bcc_test.cc
index 0c45e06..ba062b9 100644
--- a/src/android/bcc_test.cc
+++ b/src/android/bcc_test.cc
@@ -31,6 +31,15 @@
   EXPECT_EQ(0xa0, buffer[0]);
 }
 
+TEST(BccConfigTest, NoInputsMeasurement) {
+  BccConfigValues input_values = {};
+  size_t buffer_size;
+  DiceResult result =
+      BccFormatConfigDescriptor(&input_values, 0, NULL, &buffer_size);
+  EXPECT_EQ(kDiceResultBufferTooSmall, result);
+  EXPECT_EQ(1u, buffer_size);
+}
+
 TEST(BccConfigTest, AllInputs) {
   BccConfigValues input_values = {
       .inputs = BCC_INPUT_COMPONENT_NAME | BCC_INPUT_COMPONENT_VERSION |
@@ -38,18 +47,21 @@
       .component_name = "Test Component Name",
       .component_version = 0x232a13dec90f42b5,
   };
-  uint8_t buffer[256];
   size_t buffer_size;
-  DiceResult result = BccFormatConfigDescriptor(&input_values, sizeof(buffer),
-                                                buffer, &buffer_size);
-  EXPECT_EQ(kDiceResultOk, result);
+  DiceResult result =
+      BccFormatConfigDescriptor(&input_values, 0, NULL, &buffer_size);
+  EXPECT_EQ(kDiceResultBufferTooSmall, result);
+  std::vector<uint8_t> buffer(buffer_size);
   const uint8_t expected[] = {
       0xa3, 0x3a, 0x00, 0x01, 0x11, 0x71, 0x73, 'T',  'e',  's',  't',  ' ',
       'C',  'o',  'm',  'p',  'o',  'n',  'e',  'n',  't',  ' ',  'N',  'a',
       'm',  'e',  0x3a, 0x00, 0x01, 0x11, 0x72, 0x1b, 0x23, 0x2a, 0x13, 0xde,
       0xc9, 0x0f, 0x42, 0xb5, 0x3a, 0x00, 0x01, 0x11, 0x73, 0xf6};
+  EXPECT_EQ(sizeof(expected), buffer.size());
+  result = BccFormatConfigDescriptor(&input_values, buffer.size(),
+                                     buffer.data(), &buffer_size);
   EXPECT_EQ(sizeof(expected), buffer_size);
-  EXPECT_EQ(0, memcmp(expected, buffer, buffer_size));
+  EXPECT_EQ(0, memcmp(expected, buffer.data(), buffer.size()));
 }
 
 TEST(BccTest, PreservesPreviousEntries) {
@@ -68,19 +80,24 @@
   const uint8_t fake_cdi_attest[DICE_CDI_SIZE] = {};
   const uint8_t fake_cdi_seal[DICE_CDI_SIZE] = {};
   DiceInputValues input_values = {};
-  uint8_t next_bcc[2048] = {};
   size_t next_bcc_size;
   uint8_t next_cdi_attest[DICE_CDI_SIZE];
   uint8_t next_cdi_seal[DICE_CDI_SIZE];
-  DiceResult result =
+  DiceResult result = BccMainFlow(
+      /*context=*/NULL, fake_cdi_attest, fake_cdi_seal, bcc, sizeof(bcc),
+      &input_values, 0, NULL, &next_bcc_size, next_cdi_attest, next_cdi_seal);
+  EXPECT_EQ(kDiceResultBufferTooSmall, result);
+  EXPECT_GT(next_bcc_size, sizeof(bcc));
+  std::vector<uint8_t> next_bcc(next_bcc_size);
+  result =
       BccMainFlow(/*context=*/NULL, fake_cdi_attest, fake_cdi_seal, bcc,
-                  sizeof(bcc), &input_values, sizeof(next_bcc), next_bcc,
+                  sizeof(bcc), &input_values, next_bcc.size(), next_bcc.data(),
                   &next_bcc_size, next_cdi_attest, next_cdi_seal);
   EXPECT_EQ(kDiceResultOk, result);
-  EXPECT_GT(next_bcc_size, sizeof(bcc));
+  EXPECT_EQ(next_bcc_size, next_bcc.size());
   EXPECT_EQ(0x84, next_bcc[0]);
-  EXPECT_NE(0, memcmp(next_bcc + 1, bcc + 1, sizeof(bcc) - 1));
-  EXPECT_EQ(0, memcmp(next_bcc + 1, bcc + 1, sizeof(bcc) - 8 - 1));
+  EXPECT_NE(0, memcmp(next_bcc.data() + 1, bcc + 1, sizeof(bcc) - 1));
+  EXPECT_EQ(0, memcmp(next_bcc.data() + 1, bcc + 1, sizeof(bcc) - 8 - 1));
 }
 
 TEST(BccHandoverTest, PreservesPreviousEntries) {
@@ -100,18 +117,24 @@
       // 8-bytes of trailing data that aren't part of the BCC.
       0x84, 0x41, 0x55, 0xa0, 0x42, 0x11, 0x22, 0x40};
   DiceInputValues input_values = {};
-  uint8_t next_bcc_handover[2048] = {};
   size_t next_bcc_handover_size;
   DiceResult result = BccHandoverMainFlow(
-      /*context=*/NULL, bcc_handover, sizeof(bcc_handover), &input_values,
-      sizeof(next_bcc_handover), next_bcc_handover, &next_bcc_handover_size);
-  EXPECT_EQ(kDiceResultOk, result);
+      /*context=*/NULL, bcc_handover, sizeof(bcc_handover), &input_values, 0,
+      NULL, &next_bcc_handover_size);
+  EXPECT_EQ(kDiceResultBufferTooSmall, result);
   EXPECT_GT(next_bcc_handover_size, sizeof(bcc_handover));
+  std::vector<uint8_t> next_bcc_handover(next_bcc_handover_size);
+  result = BccHandoverMainFlow(
+      /*context=*/NULL, bcc_handover, sizeof(bcc_handover), &input_values,
+      next_bcc_handover.size(), next_bcc_handover.data(),
+      &next_bcc_handover_size);
+  EXPECT_EQ(kDiceResultOk, result);
+  EXPECT_EQ(next_bcc_handover_size, next_bcc_handover.size());
   EXPECT_EQ(0xa3, next_bcc_handover[0]);
   EXPECT_EQ(0x83, next_bcc_handover[72]);
-  EXPECT_NE(0, memcmp(next_bcc_handover + 73, bcc_handover + 73,
+  EXPECT_NE(0, memcmp(next_bcc_handover.data() + 73, bcc_handover + 73,
                       sizeof(bcc_handover) - 73));
-  EXPECT_EQ(0, memcmp(next_bcc_handover + 73, bcc_handover + 73,
+  EXPECT_EQ(0, memcmp(next_bcc_handover.data() + 73, bcc_handover + 73,
                       sizeof(bcc_handover) - 8 - 73));
 }
 
@@ -129,13 +152,19 @@
       // 8-bytes of trailing data that aren't part of the BCC.
       0x00, 0x41, 0x55, 0xa0, 0x42, 0x11, 0x22, 0x40};
   DiceInputValues input_values = {};
-  uint8_t next_bcc_handover[1024] = {};
   size_t next_bcc_handover_size;
   DiceResult result = BccHandoverMainFlow(
-      /*context=*/NULL, bcc_handover, sizeof(bcc_handover), &input_values,
-      sizeof(next_bcc_handover), next_bcc_handover, &next_bcc_handover_size);
-  EXPECT_EQ(kDiceResultOk, result);
+      /*context=*/NULL, bcc_handover, sizeof(bcc_handover), &input_values, 0,
+      NULL, &next_bcc_handover_size);
+  EXPECT_EQ(kDiceResultBufferTooSmall, result);
   EXPECT_GT(next_bcc_handover_size, sizeof(bcc_handover));
+  std::vector<uint8_t> next_bcc_handover(next_bcc_handover_size);
+  result = BccHandoverMainFlow(
+      /*context=*/NULL, bcc_handover, sizeof(bcc_handover), &input_values,
+      next_bcc_handover.size(), next_bcc_handover.data(),
+      &next_bcc_handover_size);
+  EXPECT_EQ(kDiceResultOk, result);
+  EXPECT_EQ(next_bcc_handover_size, next_bcc_handover.size());
   EXPECT_EQ(0xa3, next_bcc_handover[0]);
 }
 
@@ -155,13 +184,19 @@
       // 8-bytes of trailing data that aren't part of the BCC.
       0x00, 0x41, 0x55, 0xa0, 0x42, 0x11, 0x22, 0x40};
   DiceInputValues input_values = {};
-  uint8_t next_bcc_handover[1024] = {};
   size_t next_bcc_handover_size;
   DiceResult result = BccHandoverMainFlow(
-      /*context=*/NULL, bcc_handover, sizeof(bcc_handover), &input_values,
-      sizeof(next_bcc_handover), next_bcc_handover, &next_bcc_handover_size);
-  EXPECT_EQ(kDiceResultOk, result);
+      /*context=*/NULL, bcc_handover, sizeof(bcc_handover), &input_values, 0,
+      NULL, &next_bcc_handover_size);
+  EXPECT_EQ(kDiceResultBufferTooSmall, result);
   EXPECT_GT(next_bcc_handover_size, sizeof(bcc_handover));
+  std::vector<uint8_t> next_bcc_handover(next_bcc_handover_size);
+  result = BccHandoverMainFlow(
+      /*context=*/NULL, bcc_handover, sizeof(bcc_handover), &input_values,
+      next_bcc_handover.size(), next_bcc_handover.data(),
+      &next_bcc_handover_size);
+  EXPECT_EQ(kDiceResultOk, result);
+  EXPECT_EQ(next_bcc_handover_size, next_bcc_handover.size());
   EXPECT_EQ(0xa3, next_bcc_handover[0]);
 }
 
diff --git a/src/cbor_cert_op.c b/src/cbor_cert_op.c
index 5df1583..d0670de 100644
--- a/src/cbor_cert_op.c
+++ b/src/cbor_cert_op.c
@@ -44,10 +44,10 @@
   // Add the algorithm.
   CborWriteInt(kCoseHeaderAlgLabel, &out);
   CborWriteInt(DICE_COSE_KEY_ALG_VALUE, &out);
+  *encoded_size = CborOutSize(&out);
   if (CborOutOverflowed(&out)) {
     return kDiceResultBufferTooSmall;
   }
-  *encoded_size = CborOutSize(&out);
   return kDiceResultOk;
 }
 
@@ -69,10 +69,10 @@
   CborWriteBstr(aad_size, aad, &out);
   // Payload from COSE_Sign1.
   CborWriteBstr(payload_size, payload, &out);
+  *encoded_size = CborOutSize(&out);
   if (CborOutOverflowed(&out)) {
     return kDiceResultBufferTooSmall;
   }
-  *encoded_size = CborOutSize(&out);
   return kDiceResultOk;
 }
 
@@ -94,10 +94,10 @@
   CborWriteBstr(payload_size, payload, &out);
   // Signature.
   CborWriteBstr(/*num_elements=*/DICE_SIGNATURE_SIZE, signature, &out);
+  *encoded_size = CborOutSize(&out);
   if (CborOutOverflowed(&out)) {
     return kDiceResultBufferTooSmall;
   }
-  *encoded_size = CborOutSize(&out);
   return kDiceResultOk;
 }
 
@@ -118,7 +118,7 @@
                                      protected_attributes,
                                      &protected_attributes_size);
   if (result != kDiceResultOk) {
-    return result;
+    return kDiceResultPlatformError;
   }
 
   // Construct a To-Be-Signed (TBS) structure based on the relevant fields of
@@ -127,6 +127,14 @@
                          payload, payload_size, aad, aad_size, buffer_size,
                          buffer, encoded_size);
   if (result != kDiceResultOk) {
+    // Check how big the buffer needs to be in total.
+    size_t final_encoded_size = 0;
+    EncodeCoseSign1(protected_attributes, protected_attributes_size, payload,
+                    payload_size, /*signature=*/NULL, /*buffer_size=*/0,
+                    /*buffer=*/NULL, &final_encoded_size);
+    if (*encoded_size < final_encoded_size) {
+      *encoded_size = final_encoded_size;
+    }
     return result;
   }
 
@@ -239,10 +247,10 @@
   // Add the key usage.
   CborWriteInt(kKeyUsageLabel, &out);
   CborWriteBstr(/*data_size=*/1, &key_usage, &out);
+  *encoded_size = CborOutSize(&out);
   if (CborOutOverflowed(&out)) {
     return kDiceResultBufferTooSmall;
   }
-  *encoded_size = CborOutSize(&out);
   return kDiceResultOk;
 }
 
@@ -308,6 +316,7 @@
       context, subject_public_key, sizeof(encoded_public_key),
       encoded_public_key, &encoded_public_key_size);
   if (result != kDiceResultOk) {
+    result = kDiceResultPlatformError;
     goto out;
   }
 
diff --git a/src/cbor_cert_op_test.cc b/src/cbor_cert_op_test.cc
index cb41c5c..94a437b 100644
--- a/src/cbor_cert_op_test.cc
+++ b/src/cbor_cert_op_test.cc
@@ -53,6 +53,22 @@
                       next_state.certificate, next_state.certificate_size));
 }
 
+TEST(DiceOpsTest, KnownAnswerZeroInputMeasurement) {
+  DiceStateForTest current_state = {};
+  DiceStateForTest next_state = {};
+  DiceInputValues input_values = {};
+  ASSERT_LE(sizeof(dice::test::kExpectedCborEd25519Cert_ZeroInput) / 2,
+            sizeof(next_state.certificate));
+  DiceResult result = DiceMainFlow(
+      NULL, current_state.cdi_attest, current_state.cdi_seal, &input_values,
+      sizeof(dice::test::kExpectedCborEd25519Cert_ZeroInput) / 2,
+      next_state.certificate, &next_state.certificate_size,
+      next_state.cdi_attest, next_state.cdi_seal);
+  EXPECT_EQ(kDiceResultBufferTooSmall, result);
+  EXPECT_EQ(sizeof(dice::test::kExpectedCborEd25519Cert_ZeroInput),
+            next_state.certificate_size);
+}
+
 TEST(DiceOpsTest, KnownAnswerHashOnlyInput) {
   DiceStateForTest current_state = {};
   DeriveFakeInputValue("cdi_attest", DICE_CDI_SIZE, current_state.cdi_attest);
diff --git a/src/cbor_ed25519_cert_op.c b/src/cbor_ed25519_cert_op.c
index a461ae1..a2c5fbf 100644
--- a/src/cbor_ed25519_cert_op.c
+++ b/src/cbor_ed25519_cert_op.c
@@ -61,9 +61,10 @@
   // Add the public key.
   CborWriteInt(kCoseOkpXLabel, &out);
   CborWriteBstr(/*data_size=*/DICE_PUBLIC_KEY_SIZE, public_key, &out);
+
+  *encoded_size = CborOutSize(&out);
   if (CborOutOverflowed(&out)) {
     return kDiceResultBufferTooSmall;
   }
-  *encoded_size = CborOutSize(&out);
   return kDiceResultOk;
 }
diff --git a/src/cbor_p384_cert_op.c b/src/cbor_p384_cert_op.c
index de078b6..a263573 100644
--- a/src/cbor_p384_cert_op.c
+++ b/src/cbor_p384_cert_op.c
@@ -72,9 +72,9 @@
   CborWriteBstr(/*data_size=*/DICE_PUBLIC_KEY_SIZE / 2,
                 &public_key[DICE_PUBLIC_KEY_SIZE / 2], &out);
 
+  *encoded_size = CborOutSize(&out);
   if (CborOutOverflowed(&out)) {
     return kDiceResultBufferTooSmall;
   }
-  *encoded_size = CborOutSize(&out);
   return kDiceResultOk;
 }