Wipe packet buffers that held serialized WipedString (apple#10018)

* Extend WipedString guarantees to serialized packets

* Apply review suggestions

fdbserver/workloads/UnitTests.actor.cpp

@@ -49,6 +49,7 @@ void forceLinkArenaStringTests();
 void forceLinkActorCollectionTests();
 void forceLinkDDSketchTests();
 void forceLinkCommitProxyTests();
+void forceLinkWipedStringTests();
 
 struct UnitTestWorkload : TestWorkload {
 	static constexpr auto NAME = "UnitTests";
@@ -113,6 +114,7 @@ struct UnitTestWorkload : TestWorkload {
 		forceLinkArenaStringTests();
 		forceLinkActorCollectionTests();
 		forceLinkDDSketchTests();
+		forceLinkWipedStringTests();
 	}
 
 	Future<Void> setup(Database const& cx) override {

flow/Arena.cpp

@@ -344,22 +344,6 @@ void* ArenaBlock::allocate(Reference<ArenaBlock>& self, int bytes, IsSecureMem i
 	return result;
 }
 
-force_inline ArenaBlock* newBigArenaBlock(int size) {
-	if (keepalive_allocator::isActive()) [[unlikely]] {
-		return (ArenaBlock*)keepalive_allocator::allocate(size);
-	} else {
-		return (ArenaBlock*)new uint8_t[size];
-	}
-}
-
-force_inline void deleteBigArenaBlock(ArenaBlock* block) {
-	if (keepalive_allocator::isActive()) [[unlikely]] {
-		return keepalive_allocator::invalidate(block);
-	} else {
-		return delete[] reinterpret_cast<uint8_t*>(block);
-	}
-}
-
 // Return an appropriately-sized ArenaBlock to store the given data
 ArenaBlock* ArenaBlock::create(int dataSize, Reference<ArenaBlock>& next) {
 	ArenaBlock* b;
@@ -402,23 +386,23 @@ ArenaBlock* ArenaBlock::create(int dataSize, Reference<ArenaBlock>& next) {
 				b->bigSize = 256;
 				INSTRUMENT_ALLOCATE("Arena256");
 			} else if (reqSize <= 512) {
-				b = newBigArenaBlock(512);
+				b = (ArenaBlock*)allocateAndMaybeKeepalive(512);
 				b->bigSize = 512;
 				INSTRUMENT_ALLOCATE("Arena512");
 			} else if (reqSize <= 1024) {
-				b = newBigArenaBlock(1024);
+				b = (ArenaBlock*)allocateAndMaybeKeepalive(1024);
 				b->bigSize = 1024;
 				INSTRUMENT_ALLOCATE("Arena1024");
 			} else if (reqSize <= 2048) {
-				b = newBigArenaBlock(2048);
+				b = (ArenaBlock*)allocateAndMaybeKeepalive(2048);
 				b->bigSize = 2048;
 				INSTRUMENT_ALLOCATE("Arena2048");
 			} else if (reqSize <= 4096) {
-				b = newBigArenaBlock(4096);
+				b = (ArenaBlock*)allocateAndMaybeKeepalive(4096);
 				b->bigSize = 4096;
 				INSTRUMENT_ALLOCATE("Arena4096");
 			} else {
-				b = newBigArenaBlock(8192);
+				b = (ArenaBlock*)allocateAndMaybeKeepalive(8192);
 				b->bigSize = 8192;
 				INSTRUMENT_ALLOCATE("Arena8192");
 			}
@@ -430,7 +414,7 @@ ArenaBlock* ArenaBlock::create(int dataSize, Reference<ArenaBlock>& next) {
 #ifdef ALLOC_INSTRUMENTATION
 			allocInstr["ArenaHugeKB"].alloc((reqSize + 1023) >> 10);
 #endif
-			b = newBigArenaBlock(reqSize);
+			b = (ArenaBlock*)allocateAndMaybeKeepalive(reqSize);
 			b->tinySize = b->tinyUsed = NOT_TINY;
 			b->bigSize = reqSize;
 			b->totalSizeEstimate = b->bigSize;
@@ -522,26 +506,26 @@ void ArenaBlock::destroyLeaf() {
 			FastAllocator<256>::release(this);
 			INSTRUMENT_RELEASE("Arena256");
 		} else if (bigSize <= 512) {
-			deleteBigArenaBlock(this);
+			freeOrMaybeKeepalive(this);
 			INSTRUMENT_RELEASE("Arena512");
 		} else if (bigSize <= 1024) {
-			deleteBigArenaBlock(this);
+			freeOrMaybeKeepalive(this);
 			INSTRUMENT_RELEASE("Arena1024");
 		} else if (bigSize <= 2048) {
-			deleteBigArenaBlock(this);
+			freeOrMaybeKeepalive(this);
 			INSTRUMENT_RELEASE("Arena2048");
 		} else if (bigSize <= 4096) {
-			deleteBigArenaBlock(this);
+			freeOrMaybeKeepalive(this);
 			INSTRUMENT_RELEASE("Arena4096");
 		} else if (bigSize <= 8192) {
-			deleteBigArenaBlock(this);
+			freeOrMaybeKeepalive(this);
 			INSTRUMENT_RELEASE("Arena8192");
 		} else {
 #ifdef ALLOC_INSTRUMENTATION
 			allocInstr["ArenaHugeKB"].dealloc((bigSize + 1023) >> 10);
 #endif
 			g_hugeArenaMemory.fetch_sub(bigSize);
-			deleteBigArenaBlock(this);
+			freeOrMaybeKeepalive(this);
 		}
 	}
 }

flow/FastAlloc.cpp

@@ -310,6 +310,7 @@ namespace detail {
 
 std::set<void*> g_allocatedSet;
 std::set<void*> g_freedSet;
+std::vector<std::pair<const uint8_t*, int>> g_wipedSet;
 bool g_active = false;
 
 } // namespace detail
@@ -319,33 +320,52 @@ ActiveScope::ActiveScope() {
 	ASSERT(!detail::g_active);
 	ASSERT(detail::g_allocatedSet.empty());
 	ASSERT(detail::g_freedSet.empty());
+	ASSERT(detail::g_wipedSet.empty());
 	detail::g_active = true;
+	// As of writing, TraceEvent uses eventname-based throttling keyed by Standalone<StringRef>,
+	// which uses Arena and stays allocated after scope.
+	// Therefore, we disable allocation tracing (e.g. hugeArenaSample()) while this scope is active.
+	g_allocation_tracing_disabled++;
 }
 
 ActiveScope::~ActiveScope() {
 	ASSERT_ABORT(detail::g_active);
 	ASSERT_ABORT(detail::g_allocatedSet == detail::g_freedSet);
+	g_allocation_tracing_disabled--;
 	for (auto memory : detail::g_allocatedSet) {
-		delete[] reinterpret_cast<uint8_t*>(memory);
+		delete[] static_cast<uint8_t*>(memory);
 	}
 	detail::g_allocatedSet.clear();
 	detail::g_freedSet.clear();
+	detail::g_wipedSet.clear();
 	detail::g_active = false;
 }
 
 void* allocate(size_t size) {
+	ASSERT_ABORT(detail::g_active);
 	auto ptr = new uint8_t[size];
 	auto [_, inserted] = detail::g_allocatedSet.insert(ptr);
-	ASSERT(inserted); // no duplicates
+	ASSERT_ABORT(inserted); // no duplicates
 	return ptr;
 }
 
 void invalidate(void* ptr) {
-	ASSERT(detail::g_allocatedSet.contains(ptr));
-	ASSERT(!detail::g_freedSet.contains(ptr));
+	ASSERT_ABORT(detail::g_active);
+	ASSERT_ABORT(detail::g_allocatedSet.contains(ptr));
+	ASSERT_ABORT(!detail::g_freedSet.contains(ptr));
 	detail::g_freedSet.insert(ptr);
 }
 
+void trackWipedArea(const uint8_t* begin, int size) {
+	ASSERT_ABORT(detail::g_active);
+	detail::g_wipedSet.emplace_back(begin, size);
+}
+
+std::vector<std::pair<const uint8_t*, int>> const& getWipedAreaSet() {
+	ASSERT_ABORT(detail::g_active);
+	return detail::g_wipedSet;
+}
+
 } // namespace keepalive_allocator
 
 template <int Size>

flow/Knobs.cpp

@@ -145,6 +145,7 @@ void FlowKnobs::initialize(Randomize randomize, IsSimulated isSimulated) {
 	init( PUBLIC_KEY_FILE_REFRESH_INTERVAL_SECONDS,            300 );
 	init( AUDIT_TIME_WINDOW,                                   5.0 );
 	init( TOKEN_CACHE_SIZE,                                   2000 );
+	init( WIPE_SENSITIVE_DATA_FROM_PACKET_BUFFER,             true );
 
 	//AsyncFileCached
 	init( PAGE_CACHE_4K,                                   2LL<<30 );