From d29faccc64757169b80dfcb6083fcb4af21511e0 Mon Sep 17 00:00:00 2001
From: Jason Heeris <jason.heeris@gmail.com>
Date: Mon, 28 Dec 2020 21:10:30 +0800
Subject: [PATCH] Implemented capn_size() for dynamically calculating the
 required buffer size for capn_write_mem().

---
 lib/capn-malloc.c          | 26 ++++++++++++++
 lib/capnp_c.h              |  6 ++++
 tests/capn-stream-test.cpp | 72 ++++++++++++++++++++++++++++++++++----
 3 files changed, 97 insertions(+), 7 deletions(-)

diff --git a/lib/capn-malloc.c b/lib/capn-malloc.c
index 7308061..3badcf9 100644
--- a/lib/capn-malloc.c
+++ b/lib/capn-malloc.c
@@ -407,3 +407,29 @@ int capn_write_fd(struct capn *c, ssize_t (*write_fd)(int fd, const void *p, siz
 
 	return datasz;
 }
+
+int capn_size(struct capn *c)
+{
+	uint32_t headerlen;
+	size_t headersz, datasz = 0;
+	struct capn_ptr root;
+	struct capn_segment *seg;
+	int i;
+
+	if (c->segnum == 0)
+		return -1;
+
+	root = capn_root(c);
+	seg = root.seg;
+	header_calc(c, &headerlen, &headersz);
+
+	for (i = 0; i < c->segnum; i++, seg = seg->next) {
+		if (0 == seg)
+			return -1;
+		datasz += seg->len;
+	}
+	if (0 != seg)
+		return -1;
+
+	return (int) headersz+datasz;
+}
diff --git a/lib/capnp_c.h b/lib/capnp_c.h
index 609826a..7308c7e 100644
--- a/lib/capnp_c.h
+++ b/lib/capnp_c.h
@@ -276,6 +276,12 @@ void capn_init_malloc(struct capn *c);
 int capn_init_fp(struct capn *c, FILE *f, int packed);
 int capn_init_mem(struct capn *c, const uint8_t *p, size_t sz, int packed);
 
+/* capn_size calculates the amount of memory required for the buffer passed to
+ * capn_write_mem. It does not calculate the size for packed serialization, but
+ * that will always be less than the unpacked size.
+ */
+int capn_size(struct capn *c);
+
 /* capn_write_(fp|mem) writes segments to the file/memory buffer in
  * serialized form and returns the number of bytes written.
  */
diff --git a/tests/capn-stream-test.cpp b/tests/capn-stream-test.cpp
index ee55f7f..a05fd7b 100644
--- a/tests/capn-stream-test.cpp
+++ b/tests/capn-stream-test.cpp
@@ -80,6 +80,71 @@ TEST(Stream, ReadStream_Even) {
   capn_free(&ctx);
 }
 
+static struct capn_segment *CreateSmallSegment(void *u, uint32_t id, int sz) {
+  struct capn_segment *s = (struct capn_segment*) calloc(1, sizeof(*s));
+  s->data = (char*) calloc(1, sz);
+  s->cap = sz;
+  return s;
+}
+
+TEST(Stream, SizeEmptyStream) {
+  struct capn ctx;
+  capn_init_malloc(&ctx);
+  struct capn_ptr root = capn_root(&ctx);
+  ASSERT_EQ(CAPN_PTR_LIST, root.type);
+  EXPECT_EQ(2*8, capn_size(&ctx));
+
+  capn_free(&ctx);
+}
+
+TEST(Stream, SizeOneSegment) {
+  struct capn ctx;
+  capn_init_malloc(&ctx);
+  struct capn_ptr root = capn_root(&ctx);
+  struct capn_ptr ptr = capn_new_struct(root.seg, 8, 0);
+  EXPECT_EQ(0, capn_setp(root, 0, ptr));
+  EXPECT_EQ(0, capn_write64(ptr, 0, UINT64_C(0x1011121314151617)));
+  EXPECT_EQ(3*8, capn_size(&ctx));
+
+  capn_free(&ctx);
+}
+
+TEST(Stream, SizeTwoSegments) {
+  struct capn ctx;
+  capn_init_malloc(&ctx);
+  ctx.create = &CreateSmallSegment;
+  struct capn_ptr root = capn_root(&ctx);
+  struct capn_ptr ptr1 = capn_new_struct(root.seg, 8, 0);
+  EXPECT_EQ(0, capn_setp(root, 0, ptr1));
+  EXPECT_EQ(0, capn_write64(ptr1, 0, UINT64_C(0xfffefdfcfbfaf9f8)));
+  EXPECT_EQ(2, ctx.segnum);
+
+  /* 2 words: header
+   * 1 word: segment 1
+   * 2 words: segment 2
+   */
+  EXPECT_EQ(5*8, capn_size(&ctx));
+
+  capn_free(&ctx);
+}
+
+TEST(Stream, SizeThreeSegments) {
+  struct capn ctx;
+  capn_init_malloc(&ctx);
+  ctx.create = &CreateSmallSegment;
+  struct capn_ptr root = capn_root(&ctx);
+  struct capn_ptr ptr1 = capn_new_struct(root.seg, 0, 1);
+  EXPECT_EQ(0, capn_setp(root, 0, ptr1));
+  struct capn_ptr ptr2 = capn_new_struct(ptr1.seg, 4, 0);
+  EXPECT_EQ(0, capn_setp(ptr1, 0, ptr2));
+  EXPECT_EQ(0, capn_write32(ptr2, 0, 0x12345678));
+  EXPECT_EQ(3, ctx.segnum);
+
+  EXPECT_EQ(7*8, capn_size(&ctx));
+
+  capn_free(&ctx);
+}
+
 TEST(Stream, WriteEmptyStream) {
   uint8_t buf[2048];
 
@@ -164,13 +229,6 @@ TEST(Stream, WriteOneSegmentPacked) {
   capn_free(&ctx2);
 }
 
-static struct capn_segment *CreateSmallSegment(void *u, uint32_t id, int sz) {
-  struct capn_segment *s = (struct capn_segment*) calloc(1, sizeof(*s));
-  s->data = (char*) calloc(1, sz);
-  s->cap = sz;
-  return s;
-}
-
 TEST(Stream, WriteTwoSegments) {
   struct capn ctx1, ctx2;
   uint8_t buf[5*8];