Add hooks allowing a MessageBuilder to be initialized from existing memory.

c++/src/capnp/arena.c++

@@ -122,6 +122,33 @@ kj::Maybe<kj::Own<ClientHook>> ReaderArena::extractCap(uint index) {
 
 BuilderArena::BuilderArena(MessageBuilder* message)
     : message(message), segment0(nullptr, SegmentId(0), nullptr, nullptr) {}
+
+BuilderArena::BuilderArena(MessageBuilder* message,
+                           kj::ArrayPtr<MessageBuilder::SegmentInit> segments)
+    : message(message),
+      segment0(this, SegmentId(0), segments[0].space, &this->dummyLimiter, segments[0].wordsUsed) {
+  if (segments.size() > 1) {
+    kj::Vector<kj::Own<SegmentBuilder>> builders(segments.size() - 1);
+
+    uint i = 1;
+    for (auto& segment: segments.slice(1, segments.size())) {
+      builders.add(kj::heap<SegmentBuilder>(
+          this, SegmentId(i++), segment.space, &this->dummyLimiter, segment.wordsUsed));
+    }
+
+    kj::Vector<kj::ArrayPtr<const word>> forOutput;
+    forOutput.resize(segments.size());
+
+    segmentWithSpace = builders.back();
+
+    this->moreSegments = kj::heap<MultiSegmentState>(
+        MultiSegmentState { kj::mv(builders), kj::mv(forOutput) });
+
+  } else {
+    segmentWithSpace = &segment0;
+  }
+}
+
 BuilderArena::~BuilderArena() noexcept(false) {}
 
 SegmentBuilder* BuilderArena::getSegment(SegmentId id) {

c++/src/capnp/arena.h

@@ -143,7 +143,7 @@ private:
 class SegmentBuilder: public SegmentReader {
 public:
   inline SegmentBuilder(BuilderArena* arena, SegmentId id, kj::ArrayPtr<word> ptr,
-                        ReadLimiter* readLimiter);
+                        ReadLimiter* readLimiter, size_t wordsUsed = 0);
   inline SegmentBuilder(BuilderArena* arena, SegmentId id, kj::ArrayPtr<const word> ptr,
                         ReadLimiter* readLimiter);
   inline SegmentBuilder(BuilderArena* arena, SegmentId id, decltype(nullptr),
@@ -247,7 +247,8 @@ class BuilderArena final: public Arena {
   // A BuilderArena that does not allow the injection of capabilities.
 
 public:
-  BuilderArena(MessageBuilder* message);
+  explicit BuilderArena(MessageBuilder* message);
+  BuilderArena(MessageBuilder* message, kj::ArrayPtr<MessageBuilder::SegmentInit> segments);
   ~BuilderArena() noexcept(false);
   KJ_DISALLOW_COPY(BuilderArena);
 
@@ -379,8 +380,9 @@ inline void SegmentReader::unread(WordCount64 amount) { readLimiter->unread(amou
 // -------------------------------------------------------------------
 
 inline SegmentBuilder::SegmentBuilder(
-    BuilderArena* arena, SegmentId id, kj::ArrayPtr<word> ptr, ReadLimiter* readLimiter)
-    : SegmentReader(arena, id, ptr, readLimiter), pos(ptr.begin()), readOnly(false) {}
+    BuilderArena* arena, SegmentId id, kj::ArrayPtr<word> ptr, ReadLimiter* readLimiter,
+    size_t wordsUsed)
+    : SegmentReader(arena, id, ptr, readLimiter), pos(ptr.begin() + wordsUsed), readOnly(false) {}
 inline SegmentBuilder::SegmentBuilder(
     BuilderArena* arena, SegmentId id, kj::ArrayPtr<const word> ptr, ReadLimiter* readLimiter)
     : SegmentReader(arena, id, ptr, readLimiter),

c++/src/capnp/message-test.c++

@@ -20,6 +20,10 @@
 // THE SOFTWARE.
 
 #include "message.h"
+#include "test-util.h"
+#include <kj/array.h>
+#include <kj/vector.h>
+#include <kj/debug.h>
 #include <gtest/gtest.h>
 
 namespace capnp {
@@ -44,6 +48,110 @@ TEST(Message, MallocBuilderWithFirstSegment) {
   EXPECT_EQ(16u, segment.size());
 }
 
+class TestInitMessageBuilder: public MessageBuilder {
+public:
+  TestInitMessageBuilder(kj::ArrayPtr<SegmentInit> segments): MessageBuilder(segments) {}
+
+  kj::ArrayPtr<word> allocateSegment(uint minimumSize) override {
+    auto array = kj::heapArray<word>(minimumSize);
+    memset(array.begin(), 0, array.asBytes().size());
+    allocations.add(kj::mv(array));
+    return allocations.back();
+  }
+
+  kj::Vector<kj::Array<word>> allocations;
+};
+
+TEST(Message, MessageBuilderInit) {
+  MallocMessageBuilder builder(2048);
+  initTestMessage(builder.getRoot<TestAllTypes>());
+
+  // Pull the segments out and make a segment init table out of them.
+  //
+  // We const_cast for simplicity of implementing the test, but you shouldn't do that at home. :)
+  auto segs = builder.getSegmentsForOutput();
+  ASSERT_EQ(1, segs.size());
+
+  auto segInits = KJ_MAP(seg, segs) -> MessageBuilder::SegmentInit {
+    return { kj::arrayPtr(const_cast<word*>(seg.begin()), seg.size()), seg.size() };
+  };
+
+  // Init a new builder from the old segments.
+  TestInitMessageBuilder builder2(segInits);
+  checkTestMessage(builder2.getRoot<TestAllTypes>());
+
+  // Verify that they're really using the same underlying memory.
+  builder2.getRoot<TestAllTypes>().setInt64Field(123321);
+  EXPECT_EQ(123321, builder.getRoot<TestAllTypes>().getInt64Field());
+
+  // Force builder2 to allocate new space.
+  EXPECT_EQ(0, builder2.allocations.size());
+  builder2.getRoot<TestAllTypes>().setTextField("foobarbaz");
+  EXPECT_EQ(1, builder2.allocations.size());
+}
+
+TEST(Message, MessageBuilderInitMultiSegment) {
+  // Same as previous test, but with a message containing many segments.
+
+  MallocMessageBuilder builder(1, AllocationStrategy::FIXED_SIZE);
+  initTestMessage(builder.getRoot<TestAllTypes>());
+
+  // Pull the segments out and make a segment init table out of them.
+  //
+  // We const_cast for simplicity of implementing the test, but you shouldn't do that at home. :)
+  auto segs = builder.getSegmentsForOutput();
+  ASSERT_NE(1, segs.size());
+
+  auto segInits = KJ_MAP(seg, segs) -> MessageBuilder::SegmentInit {
+    return { kj::arrayPtr(const_cast<word*>(seg.begin()), seg.size()), seg.size() };
+  };
+
+  // Init a new builder from the old segments.
+  TestInitMessageBuilder builder2(segInits);
+  checkTestMessage(builder2.getRoot<TestAllTypes>());
+
+  // Verify that they're really using the same underlying memory.
+  builder2.getRoot<TestAllTypes>().setInt64Field(123321);
+  EXPECT_EQ(123321, builder.getRoot<TestAllTypes>().getInt64Field());
+
+  // Force builder2 to allocate new space.
+  EXPECT_EQ(0, builder2.allocations.size());
+  builder2.getRoot<TestAllTypes>().setTextField("foobarbaz");
+  EXPECT_EQ(1, builder2.allocations.size());
+}
+
+TEST(Message, MessageBuilderInitSpaceAvailable) {
+  word buffer[2048];
+  memset(buffer, 0, sizeof(buffer));
+  MallocMessageBuilder builder(buffer);
+  initTestMessage(builder.getRoot<TestAllTypes>());
+
+  // Find out how much space in `buffer` was used in order to use in initializing the new message.
+  auto segs = builder.getSegmentsForOutput();
+  ASSERT_EQ(1, segs.size());
+  KJ_ASSERT(segs[0].begin() == buffer);
+
+  MessageBuilder::SegmentInit init = { buffer, segs[0].size() };
+  KJ_DBG(init.space.size(), init.wordsUsed);
+
+  // Init a new builder from the old segments.
+  TestInitMessageBuilder builder2(kj::arrayPtr(&init, 1));
+  checkTestMessage(builder2.getRoot<TestAllTypes>());
+
+  // Verify that they're really using the same underlying memory.
+  builder2.getRoot<TestAllTypes>().setInt64Field(123321);
+  EXPECT_EQ(123321, builder.getRoot<TestAllTypes>().getInt64Field());
+
+  // Ask builder2 to allocate new space. It should go into the free space at the end of the
+  // segment.
+  EXPECT_EQ(0, builder2.allocations.size());
+  builder2.getRoot<TestAllTypes>().setTextField("foobarbaz");
+  EXPECT_EQ(0, builder2.allocations.size());
+
+  EXPECT_EQ(kj::implicitCast<void*>(buffer + segs[0].size()),
+            kj::implicitCast<void*>(builder2.getRoot<TestAllTypes>().getTextField().begin()));
+}
+
 // TODO(test):  More tests.
 
 }  // namespace

c++/src/capnp/message.c++

@@ -62,12 +62,19 @@ AnyPointer::Reader MessageReader::getRootInternal() {
 // -------------------------------------------------------------------
 
 MessageBuilder::MessageBuilder(): allocatedArena(false) {}
+
 MessageBuilder::~MessageBuilder() noexcept(false) {
   if (allocatedArena) {
     kj::dtor(*arena());
   }
 }
 
+MessageBuilder::MessageBuilder(kj::ArrayPtr<SegmentInit> segments)
+    : allocatedArena(false) {
+  kj::ctor(*arena(), this, segments);
+  allocatedArena = true;
+}
+
 _::SegmentBuilder* MessageBuilder::getRootSegment() {
   if (allocatedArena) {
     return arena()->getSegment(_::SegmentId(0));

c++/src/capnp/message.h

@@ -163,6 +163,35 @@ public:
   virtual ~MessageBuilder() noexcept(false);
   KJ_DISALLOW_COPY(MessageBuilder);
 
+  struct SegmentInit {
+    kj::ArrayPtr<word> space;
+
+    size_t wordsUsed;
+    // Number of words in `space` which are used; the rest are free space in which additional
+    // objects may be allocated.
+  };
+
+  explicit MessageBuilder(kj::ArrayPtr<SegmentInit> segments);
+  // Create a MessageBuilder backed by existing memory. This is an advanced interface that most
+  // people should not use. THIS METHOD IS INSECURE; see below.
+  //
+  // This allows a MessageBuilder to be constructed to modify an in-memory message without first
+  // making a copy of the content. This is especially useful in conjunction with mmap().
+  //
+  // The contents of each segment must outlive the MessageBuilder, but the SegmentInit array itself
+  // only need outlive the constructor.
+  //
+  // SECURITY: Do not use this in conjunction with untrusted data. This constructor assumes that
+  //   the input message is valid. This constructor is designed to be used with data you control,
+  //   e.g. an mmap'd file which is owned and accessed by only one program. When reading data you
+  //   do not trust, you *must* load it into a Reader and then copy into a Builder as a means of
+  //   validating the content.
+  //
+  // WARNING: It is NOT safe to initialize a MessageBuilder in this way from memory that is
+  //   currently in use by another MessageBuilder or MessageReader. Other readers/builders will
+  //   not observe changes to the segment sizes nor newly-allocated segments caused by allocating
+  //   new objects in this message.
+
   virtual kj::ArrayPtr<word> allocateSegment(uint minimumSize) = 0;
   // Allocates an array of at least the given number of words, throwing an exception or crashing if
   // this is not possible.  It is expected that this method will usually return more space than