Support uninitialised fields for skipping pre-zeroing

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

@@ -158,22 +158,24 @@ void ReaderArena::reportReadLimitReached() {
 // =======================================================================================
 
 BuilderArena::BuilderArena(MessageBuilder* message)
-    : message(message), segment0(nullptr, SegmentId(0), nullptr, nullptr) {}
+    : message(message), segment0(nullptr, SegmentId(0), nullptr, nullptr),
+      needLazyZero(message->needLazyZero()) {}
 
 BuilderArena::BuilderArena(MessageBuilder* message,
                            kj::ArrayPtr<MessageBuilder::SegmentInit> segments)
     : message(message),
       segment0(this, SegmentId(0), segments[0].space.begin(),
                verifySegment(segments[0].space),
-               &this->dummyLimiter, verifySegmentSize(segments[0].wordsUsed)) {
+               &this->dummyLimiter, verifySegmentSize(segments[0].wordsUsed), segments[0].needLazyZero),
+      needLazyZero(message->needLazyZero()){
   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.begin(), verifySegment(segment.space),
-          &this->dummyLimiter, verifySegmentSize(segment.wordsUsed)));
+          &this->dummyLimiter, verifySegmentSize(segment.wordsUsed), segment.needLazyZero));
     }
 
     kj::Vector<kj::ArrayPtr<const word>> forOutput;
@@ -224,6 +226,9 @@ SegmentBuilder* BuilderArena::getSegment(SegmentId id) {
 }
 
 BuilderArena::AllocateResult BuilderArena::allocate(SegmentWordCount amount) {
+  // Check if memory needs lazy zero
+  bool needLazyZero = this->needLazyZero;
+
   if (segment0.getArena() == nullptr) {
     // We're allocating the first segment.
     kj::ArrayPtr<word> ptr = message->allocateSegment(unbound(amount / WORDS));
@@ -232,10 +237,13 @@ BuilderArena::AllocateResult BuilderArena::allocate(SegmentWordCount amount) {
     // Re-allocate segment0 in-place.  This is a bit of a hack, but we have not returned any
     // pointers to this segment yet, so it should be fine.
     kj::dtor(segment0);
-    kj::ctor(segment0, this, SegmentId(0), ptr.begin(), actualSize, &this->dummyLimiter);
+    kj::ctor(segment0, this, SegmentId(0), ptr.begin(), actualSize, &this->dummyLimiter, ZERO * WORDS, needLazyZero);
 
     segmentWithSpace = &segment0;
-    return AllocateResult { &segment0, segment0.allocate(amount) };
+    word* resultPtr = segment0.allocate(amount);
+    // Zero the root pointer field if the memory is not pre-zeroed.
+    if (needLazyZero) memset(resultPtr, 0, static_cast<size_t>(POINTER_SIZE_IN_WORDS) * sizeof(word));
+    return AllocateResult { &segment0, resultPtr };
   } else {
     if (segmentWithSpace != nullptr) {
       // Check if there is space in an existing segment.
@@ -252,13 +260,17 @@ BuilderArena::AllocateResult BuilderArena::allocate(SegmentWordCount amount) {
     }
 
     // Need to allocate a new segment.
-    SegmentBuilder* result = addSegmentInternal(message->allocateSegment(unbound(amount / WORDS)));
+    SegmentBuilder* result = addSegmentInternal(message->allocateSegment(unbound(amount / WORDS)),
+      needLazyZero);
 
     // Check this new segment first the next time we need to allocate.
     segmentWithSpace = result;
 
     // Allocating from the new segment is guaranteed to succeed since we made it big enough.
-    return AllocateResult { result, result->allocate(amount) };
+    word* resultPtr = result->allocate(amount);
+    // Zero the root pointer field if the memory is not pre-zeroed.
+    if (needLazyZero) memset(resultPtr, 0, static_cast<size_t>(POINTER_SIZE_IN_WORDS) * sizeof(word));
+    return AllocateResult { result, resultPtr };
   }
 }
 
@@ -267,7 +279,7 @@ SegmentBuilder* BuilderArena::addExternalSegment(kj::ArrayPtr<const word> conten
 }
 
 template <typename T>
-SegmentBuilder* BuilderArena::addSegmentInternal(kj::ArrayPtr<T> content) {
+SegmentBuilder* BuilderArena::addSegmentInternal(kj::ArrayPtr<T> content, bool notZeroed) {
   // This check should never fail in practice, since you can't get an Orphanage without allocating
   // the root segment.
   KJ_REQUIRE(segment0.getArena() != nullptr,
@@ -286,7 +298,7 @@ SegmentBuilder* BuilderArena::addSegmentInternal(kj::ArrayPtr<T> content) {
 
   kj::Own<SegmentBuilder> newBuilder = kj::heap<SegmentBuilder>(
       this, SegmentId(segmentState->builders.size() + 1),
-      content.begin(), contentSize, &this->dummyLimiter);
+      content.begin(), contentSize, &this->dummyLimiter, ZERO * WORDS, notZeroed);
   SegmentBuilder* result = newBuilder.get();
   segmentState->builders.add(kj::mv(newBuilder));
 

c++/src/capnp/arena.h

@@ -185,9 +185,11 @@ class SegmentReader {
 class SegmentBuilder: public SegmentReader {
 public:
   inline SegmentBuilder(BuilderArena* arena, SegmentId id, word* ptr, SegmentWordCount size,
-                        ReadLimiter* readLimiter, SegmentWordCount wordsUsed = ZERO * WORDS);
+                        ReadLimiter* readLimiter, SegmentWordCount wordsUsed = ZERO * WORDS,
+                        bool needLazyZero = false);
   inline SegmentBuilder(BuilderArena* arena, SegmentId id, const word* ptr, SegmentWordCount size,
-                        ReadLimiter* readLimiter);
+                        ReadLimiter* readLimiter, SegmentWordCount wordsUsed = ZERO * WORDS,
+                        bool needLazyZero = false);
   inline SegmentBuilder(BuilderArena* arena, SegmentId id, decltype(nullptr),
                         ReadLimiter* readLimiter);
 
@@ -208,6 +210,8 @@ class SegmentBuilder: public SegmentReader {
 
   inline bool isWritable() { return !readOnly; }
 
+  inline bool needLazyZero() { return needLazyZero_; }
+
   inline void tryTruncate(word* from, word* to);
   // If `from` points just past the current end of the segment, then move the end back to `to`.
   // Otherwise, do nothing.
@@ -224,6 +228,8 @@ class SegmentBuilder: public SegmentReader {
 
   bool readOnly;
 
+  bool needLazyZero_;
+
   [[noreturn]] void throwNotWritable();
 
   KJ_DISALLOW_COPY_AND_MOVE(SegmentBuilder);
@@ -376,7 +382,9 @@ class BuilderArena final: public Arena {
   // segment that is already-full, in which case we don't update this pointer.
 
   template <typename T>  // Can be `word` or `const word`.
-  SegmentBuilder* addSegmentInternal(kj::ArrayPtr<T> content);
+  SegmentBuilder* addSegmentInternal(kj::ArrayPtr<T> content, bool needLazyZero = false);
+
+  bool needLazyZero = false;
 };
 
 // =======================================================================================
@@ -452,20 +460,20 @@ inline void SegmentReader::unread(WordCount64 amount) { readLimiter->unread(amou
 
 inline SegmentBuilder::SegmentBuilder(
     BuilderArena* arena, SegmentId id, word* ptr, SegmentWordCount size,
-    ReadLimiter* readLimiter, SegmentWordCount wordsUsed)
+    ReadLimiter* readLimiter, SegmentWordCount wordsUsed, bool needLazyZero)
     : SegmentReader(arena, id, ptr, size, readLimiter),
-      pos(ptr + wordsUsed), readOnly(false) {}
+      pos(ptr + wordsUsed), readOnly(false), needLazyZero_(needLazyZero) {}
 inline SegmentBuilder::SegmentBuilder(
     BuilderArena* arena, SegmentId id, const word* ptr, SegmentWordCount size,
-    ReadLimiter* readLimiter)
+    ReadLimiter* readLimiter, SegmentWordCount /*wordsUsed*/, bool needLazyZero)
     : SegmentReader(arena, id, ptr, size, readLimiter),
       // const_cast is safe here because the member won't ever be dereferenced because it appears
       // to point to the end of the segment anyway.
-      pos(const_cast<word*>(ptr + size)), readOnly(true) {}
+      pos(const_cast<word*>(ptr + size)), readOnly(true), needLazyZero_(needLazyZero) {}
 inline SegmentBuilder::SegmentBuilder(BuilderArena* arena, SegmentId id, decltype(nullptr),
                                       ReadLimiter* readLimiter)
     : SegmentReader(arena, id, nullptr, ZERO * WORDS, readLimiter),
-      pos(nullptr), readOnly(false) {}
+      pos(nullptr), readOnly(false), needLazyZero_(false) {}
 
 inline word* SegmentBuilder::allocate(SegmentWordCount amount) {
   if (intervalLength(pos, ptr.end(), MAX_SEGMENT_WORDS) < amount) {

c++/src/capnp/compiler/capnpc-c++.c++

@@ -1756,6 +1756,8 @@ private:
                 "  inline ::capnp::BuilderFor<T_> init", titleCase, "As(unsigned int size);\n"),
               COND(!shouldTemplatizeInit,
                 "  ", maybeInline, builderType, " init", titleCase, "(unsigned int size);\n")),
+            COND(shouldIncludeSizedInit && !shouldTemplatizeInit && typeSchema.which() == schema::Type::DATA,
+              "  ", maybeInline, builderType, " uninitialized", titleCase, "(unsigned int size);\n"),
             "  ", maybeInline, "void adopt", titleCase, "(::capnp::Orphan<", type, ">&& value);\n"
             "  ", maybeInline, "::capnp::Orphan<", type, "> disown", titleCase, "();\n",
             COND(shouldExcludeInLiteMode, "#endif  // !CAPNP_LITE\n"),
@@ -1826,6 +1828,13 @@ private:
                 "  return ::capnp::_::PointerHelpers<", type, ">::init(_builder.getPointerField(\n"
                 "      ::capnp::bounded<", offset, ">() * ::capnp::POINTERS), size);\n"
                 "}\n"),
+            COND(shouldIncludeSizedInit && !shouldTemplatizeInit && typeSchema.which() == schema::Type::DATA,
+                templateContext.allDecls(),
+                maybeInline, builderType, " ", scope, "Builder::uninitialized", titleCase, "(unsigned int size) {\n",
+                unionDiscrim.set,
+                "  return _builder.getPointerField(\n"
+                "      ::capnp::bounded<", offset, ">() * ::capnp::POINTERS).uninitializedBlob<", type, ">(size);\n"
+                "}\n"),
             templateContext.allDecls(),
             maybeInline, "void ", scope, "Builder::adopt", titleCase, "(\n"
             "    ::capnp::Orphan<", type, ">&& value) {\n",

c++/src/capnp/layout.c++

@@ -458,7 +458,7 @@ struct WireHelpers {
 
   static KJ_ALWAYS_INLINE(word* allocate(
       WirePointer*& ref, SegmentBuilder*& segment, CapTableBuilder* capTable,
-      SegmentWordCount amount, WirePointer::Kind kind, BuilderArena* orphanArena)) {
+      SegmentWordCount amount, WirePointer::Kind kind, BuilderArena* orphanArena, bool requireCleanMemory = true)) {
     // Allocate space in the message for a new object, creating far pointers if necessary. The
     // space is guaranteed to be zero'd (because MessageBuilder implementations are required to
     // return zero'd memory).
@@ -515,10 +515,22 @@ struct WireHelpers {
         ref = reinterpret_cast<WirePointer*>(ptr);
         ref->setKindAndTarget(kind, ptr + POINTER_SIZE_IN_WORDS, segment);
 
+        // If clean memory is wanted and the segment is not pre-zeroed, zero the data portion
+        // after the landing pad.
+        if (requireCleanMemory && segment->needLazyZero()) {
+          WireHelpers::zeroMemory(ptr + POINTER_SIZE_IN_WORDS, amount);
+        }
+
         // Allocated space follows new pointer.
         return ptr + POINTER_SIZE_IN_WORDS;
       } else {
         ref->setKindAndTarget(kind, ptr, segment);
+
+        // If clean memory is wanted and the segment is not pre-zeroed, zero the data portion
+        // after the landing pad.
+        if (requireCleanMemory && segment->needLazyZero()) {
+          WireHelpers::zeroMemory(ptr, amount);
+        }
         return ptr;
       }
     } else {
@@ -527,6 +539,10 @@ struct WireHelpers {
       auto allocation = orphanArena->allocate(amount);
       segment = allocation.segment;
       ref->setKindForOrphan(kind);
+      // Check for the need to lazy zero the segment allocated by OrphanArena.
+      if (requireCleanMemory && segment->needLazyZero()) {
+        WireHelpers::zeroMemory(allocation.words, amount);
+      }
       return allocation.words;
     }
   }
@@ -1717,6 +1733,32 @@ struct WireHelpers {
     return { segment, Data::Builder(reinterpret_cast<byte*>(ptr), unbound(size / BYTES)) };
   }
 
+  static KJ_ALWAYS_INLINE(SegmentAnd<Data::Builder> uninitializedDataPointer(
+      WirePointer* ref, SegmentBuilder* segment, CapTableBuilder* capTable, BlobSize size,
+      BuilderArena* orphanArena = nullptr)) {
+    // Allocate the space with requireCleanMemory = false
+    word* ptr = allocate(ref, segment, capTable, roundBytesUpToWords(size),
+                         WirePointer::LIST, orphanArena, false);
+
+    // Initialize the pointer.
+    ref->listRef.set(ElementSize::BYTE, size * (ONE * ELEMENTS / BYTES));
+
+    // Zero-out padding bytes if memory is dirty for security reason.
+    if (segment->needLazyZero()) {
+      size_t byteSize = unbound(size / BYTES);
+      // Calculate the actual allocated size in bytes (aligned to Word boundary).
+      size_t allocatedSize = unbound(roundBytesUpToWords(size)) * sizeof(word);
+
+      if (byteSize < allocatedSize) {
+        // Zero-out the trailing padding bytes to prevent information leakage.
+        memset(reinterpret_cast<byte*>(ptr) + byteSize, 0, allocatedSize - byteSize);
+      }
+    }
+
+    // Build the Data::Builder.
+    return { segment, Data::Builder(reinterpret_cast<byte*>(ptr), unbound(size / BYTES)) };
+  }
+
   static KJ_ALWAYS_INLINE(SegmentAnd<Data::Builder> setDataPointer(
       WirePointer* ref, SegmentBuilder* segment, CapTableBuilder* capTable, Data::Reader value,
       BuilderArena* orphanArena = nullptr)) {
@@ -2595,6 +2637,11 @@ Data::Builder PointerBuilder::initBlob<Data>(ByteCount size) {
       assertMaxBits<BLOB_SIZE_BITS>(size, ThrowOverflow())).value;
 }
 template <>
+Data::Builder PointerBuilder::uninitializedBlob<Data>(ByteCount size) {
+  return WireHelpers::uninitializedDataPointer(pointer, segment, capTable,
+      assertMaxBits<BLOB_SIZE_BITS>(size, ThrowOverflow())).value;
+}
+template <>
 void PointerBuilder::setBlob<Data>(Data::Reader value) {
   WireHelpers::setDataPointer(pointer, segment, capTable, value);
 }

c++/src/capnp/layout.h

@@ -358,6 +358,7 @@ class PointerBuilder: public kj::DisallowConstCopy {
   ListBuilder initList(ElementSize elementSize, ElementCount elementCount);
   ListBuilder initStructList(ElementCount elementCount, StructSize size);
   template <typename T> typename T::Builder initBlob(ByteCount size);
+  template <typename T> typename T::Builder uninitializedBlob(ByteCount size);
   // Init methods:  Initialize the pointer to a newly-allocated object, discarding the existing
   // object.
 
@@ -967,6 +968,7 @@ template <> typename Text::Reader PointerReader::getBlob<Text>(
     const void* defaultValue, ByteCount defaultSize) const;
 
 template <> typename Data::Builder PointerBuilder::initBlob<Data>(ByteCount size);
+template <> typename Data::Builder PointerBuilder::uninitializedBlob<Data>(ByteCount size);
 template <> void PointerBuilder::setBlob<Data>(typename Data::Reader value);
 template <> typename Data::Builder PointerBuilder::getBlob<Data>(
     const void* defaultValue, ByteCount defaultSize);