From 8901350a83780edd463ddd64d5196fdc46f7e8e8 Mon Sep 17 00:00:00 2001
From: Scott Bertin <sbertin@telular.com>
Date: Mon, 23 Apr 2018 12:00:35 -0400
Subject: [PATCH] Determine buffer size needed without encoding.

If a NULL buffer is passed to cn_cbor_encoder_write, it will determine
the necessary size of the buffer without doing the encoding. The
buf_size parameter can still set an upper limit on the size.
---
 include/cn-cbor/cn-cbor.h |  7 +++++--
 src/cn-encoder.c          | 19 +++++++++++++++----
 test/cbor_test.c          | 13 +++++++++++++
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/include/cn-cbor/cn-cbor.h b/include/cn-cbor/cn-cbor.h
index 187a55c..362e223 100644
--- a/include/cn-cbor/cn-cbor.h
+++ b/include/cn-cbor/cn-cbor.h
@@ -265,10 +265,13 @@ void cn_cbor_free(cn_cbor* cb CBOR_CONTEXT);
 /**
  * Write a CBOR value and all of the child values.
  *
- * @param[in]  buf        The buffer into which to write
+ * @param[in]  buf        The buffer into which to write. May be NULL to
+ *                        determine the necessary size.
  * @param[in]  buf_offset The offset (in bytes) from the beginning of the buffer
  *                        to start writing at
- * @param[in]  buf_size   The total length (in bytes) of the buffer
+ * @param[in]  buf_size   The total length (in bytes) of the buffer. If buf is
+ *                        NULL, this is an upper limit and may be 0 to specify
+ *                        no limit.
  * @param[in]  cb         [description]
  * @return                -1 on fail, or number of bytes written
  */
diff --git a/src/cn-encoder.c b/src/cn-encoder.c
index d8a4d49..1f9aa8d 100644
--- a/src/cn-encoder.c
+++ b/src/cn-encoder.c
@@ -41,12 +41,20 @@ typedef struct _write_state
 }
 
 #define write_byte_and_data(b, data, sz) \
-ws->buf[ws->offset++] = (b); \
-memcpy(ws->buf+ws->offset, (data), (sz)); \
+if(ws->buf) { \
+  ws->buf[ws->offset++] = (b); \
+  memcpy(ws->buf+ws->offset, (data), (sz)); \
+} else { \
+  ws->offset++; \
+} \
 ws->offset += sz;
 
 #define write_byte(b) \
-ws->buf[ws->offset++] = (b); \
+if(ws->buf) { \
+  ws->buf[ws->offset++] = (b); \
+} else { \
+  ws->offset++; \
+}
 
 #define write_byte_ensured(b) \
 ensure_writable(1); \
@@ -249,7 +257,9 @@ void _encoder_visitor(const cn_cbor *cb, int depth, void *context)
   case CN_CBOR_BYTES:
     CHECK(_write_positive(ws, cb->type, cb->length));
     ensure_writable(cb->length);
-    memcpy(ws->buf+ws->offset, cb->v.str, cb->length);
+    if (ws->buf) {
+      memcpy(ws->buf+ws->offset, cb->v.str, cb->length);
+    }
     ws->offset += cb->length;
     break;
 
@@ -302,6 +312,7 @@ ssize_t cn_cbor_encoder_write(uint8_t *buf,
 			      const cn_cbor *cb)
 {
   cn_write_state ws = { buf, buf_offset, buf_size };
+  if (!ws.buf && ws.size <= 0) { ws.size = (ssize_t)(((size_t)-1) / 2); }
   _visit(cb, _encoder_visitor, _encoder_breaker, &ws);
   if (ws.offset < 0) { return -1; }
   return ws.offset - buf_offset;
diff --git a/test/cbor_test.c b/test/cbor_test.c
index eafea5d..8ea549d 100644
--- a/test/cbor_test.c
+++ b/test/cbor_test.c
@@ -116,6 +116,7 @@ CTEST(cbor, parse)
     size_t i;
     unsigned char encoded[1024];
     ssize_t enc_sz;
+    ssize_t enc_sz2;
 
     for (i=0; i<sizeof(tests)/sizeof(char*); i++) {
         ASSERT_TRUE(parse_hex(tests[i], &b));
@@ -125,7 +126,9 @@ CTEST(cbor, parse)
         ASSERT_EQUAL(err.err, CN_CBOR_NO_ERROR);
         ASSERT_NOT_NULL(cb);
 
+        enc_sz2 = cn_cbor_encoder_write(NULL, 0, sizeof(encoded), cb);
         enc_sz = cn_cbor_encoder_write(encoded, 0, sizeof(encoded), cb);
+        ASSERT_EQUAL(enc_sz, enc_sz2);
         ASSERT_DATA(b.ptr, b.sz, encoded, enc_sz);
         free(b.ptr);
         cn_cbor_free(cb CONTEXT_NULL);
@@ -164,6 +167,7 @@ CTEST(cbor, parse_normalize)
     size_t i;
     unsigned char encoded[1024];
     ssize_t enc_sz;
+    ssize_t enc_sz2;
 
     for (i=0; i<sizeof(basic_tests)/sizeof(char*); i+=2) {
         ASSERT_TRUE(parse_hex(basic_tests[i], &b));
@@ -174,7 +178,9 @@ CTEST(cbor, parse_normalize)
         ASSERT_EQUAL(err.err, CN_CBOR_NO_ERROR);
         ASSERT_NOT_NULL(cb);
 
+        enc_sz2 = cn_cbor_encoder_write(NULL, 0, sizeof(encoded), cb);
         enc_sz = cn_cbor_encoder_write(encoded, 0, sizeof(encoded), cb);
+        ASSERT_EQUAL(enc_sz, enc_sz2);
         ASSERT_DATA(b2.ptr, b2.sz, encoded, enc_sz);
         free(b.ptr);
         free(b2.ptr);
@@ -195,7 +201,9 @@ CTEST(cbor, parse_normalize)
         ASSERT_NULL(cb);
 #endif /* CBOR_NO_FLOAT */
 
+        /* enc_sz2 = cn_cbor_encoder_write(NULL, 0, sizeof(encoded), cb); */
         /* enc_sz = cn_cbor_encoder_write(encoded, 0, sizeof(encoded), cb); */
+        /* ASSERT_EQUAL(enc_sz, enc_sz2); */
         /* ASSERT_DATA(b2.ptr, b2.sz, encoded, enc_sz); */
         free(b.ptr);
         free(b2.ptr);
@@ -258,13 +266,16 @@ CTEST(cbor, float)
     size_t i;
     unsigned char encoded[1024];
     ssize_t enc_sz;
+    ssize_t enc_sz2;
 
     for (i=0; i<sizeof(tests)/sizeof(char*); i++) {
         ASSERT_TRUE(parse_hex(tests[i], &b));
         cb = cn_cbor_decode(b.ptr, b.sz CONTEXT_NULL, &err);
         ASSERT_NOT_NULL(cb);
 
+        enc_sz2 = cn_cbor_encoder_write(NULL, 0, sizeof(encoded), cb);
         enc_sz = cn_cbor_encoder_write(encoded, 0, sizeof(encoded), cb);
+        ASSERT_EQUAL(enc_sz, enc_sz2);
         ASSERT_DATA(b.ptr, b.sz, encoded, enc_sz);
 
         free(b.ptr);
@@ -446,6 +457,8 @@ CTEST(cbor, create_encode)
   ASSERT_NOT_NULL(cdata);
 
   ASSERT_TRUE(cn_cbor_mapput_int(map, 0, cdata, CONTEXT_NULL_COMMA NULL));
+  enc_sz = cn_cbor_encoder_write(NULL, 0, sizeof(encoded), map);
+  ASSERT_EQUAL(7, enc_sz);
   enc_sz = cn_cbor_encoder_write(encoded, 0, sizeof(encoded), map);
   ASSERT_EQUAL(7, enc_sz);
 }