From 20882507f4837080c6a2486b389ea848e79a23f1 Mon Sep 17 00:00:00 2001
From: Dan Lambright <dlambrig@gmail.com>
Date: Tue, 2 Jan 2024 12:09:32 -0500
Subject: [PATCH] sanity checks, fix knob

---
 fdbclient/ServerKnobs.cpp           |  2 +-
 fdbserver/ApplyMetadataMutation.cpp |  9 ++++--
 fdbserver/storageserver.actor.cpp   | 43 ++++++++++++++---------------
 3 files changed, 28 insertions(+), 26 deletions(-)

diff --git a/fdbclient/ServerKnobs.cpp b/fdbclient/ServerKnobs.cpp
index 948585d27a..68df893dc8 100644
--- a/fdbclient/ServerKnobs.cpp
+++ b/fdbclient/ServerKnobs.cpp
@@ -892,7 +892,7 @@ void ServerKnobs::initialize(Randomize randomize, ClientKnobs* clientKnobs, IsSi
 	init( HOT_SHARD_THROTTLING_TRACKED,                             1 );
 	init( HOT_SHARD_MONITOR_FREQUENCY,                            5.0 );
 
-	init( GENERATE_DATA_ENABLED,                                 true );
+	init( GENERATE_DATA_ENABLED,                                false );
 
 	//Storage Metrics
 	init( STORAGE_METRICS_AVERAGE_INTERVAL,                    120.0 );
diff --git a/fdbserver/ApplyMetadataMutation.cpp b/fdbserver/ApplyMetadataMutation.cpp
index 4755803a8f..8800b4b752 100644
--- a/fdbserver/ApplyMetadataMutation.cpp
+++ b/fdbserver/ApplyMetadataMutation.cpp
@@ -593,7 +593,10 @@ private:
 			uint64_t valSize, keyCount, seed;
 			Standalone<StringRef> prefix;
 			std::tie(prefix, valSize, keyCount, seed) = decodeConstructKeys(m.param2);
-			uint8_t keyBuf[prefix.size() + 4];
+			if (keyCount >= UINT16_MAX || valSize >= CLIENT_KNOBS->VALUE_SIZE_LIMIT) {
+				return;
+			}
+			uint8_t keyBuf[prefix.size() + sizeof(uint16_t)];
 			uint8_t* keyPos = prefix.copyTo(keyBuf);
 			*keyPos = '\xff';
 			StringRef keyEnd(keyBuf, keyPos - keyBuf + 1);
@@ -610,8 +613,8 @@ private:
 			toCommit->addTags(allTags);
 			TraceEvent(SevInfo, "ConstructDataRequest")
 			    .detail("Prefix", prefix)
-			    .detail("ValSize", valSize)
-			    .detail("KeyCount", keyCount);
+			    .detail("KeyCount", keyCount)
+			    .detail("ValSize", valSize);
 		}
 	}
 
diff --git a/fdbserver/storageserver.actor.cpp b/fdbserver/storageserver.actor.cpp
index 0b25b3cf10..33a13aa9f7 100644
--- a/fdbserver/storageserver.actor.cpp
+++ b/fdbserver/storageserver.actor.cpp
@@ -10860,12 +10860,12 @@ private:
 			uint64_t valSize, keyCount, seed;
 			Standalone<StringRef> prefix;
 			std::tie(prefix, valSize, keyCount, seed) = decodeConstructKeys(m.param2);
-			// ASSERT
+			ASSERT(keyCount < UINT16_MAX && valSize >= CLIENT_KNOBS->VALUE_SIZE_LIMIT);
 			for (auto& r : data->shards.ranges()) {
 				KeyRange keyRange = KeyRange(r.range());
 				if (keyRange.contains(prefix) && r.value() &&
 				    (r.value()->adding || r.value()->moveInShard || r.value()->readWrite)) {
-					uint8_t keyBuf[prefix.size() + 4];
+					uint8_t keyBuf[prefix.size() + sizeof(uint16_t)];
 					uint8_t* keyPos = prefix.copyTo(keyBuf);
 					uint8_t valBuf[valSize];
 					setThreadLocalDeterministicRandomSeed(seed);
@@ -10878,13 +10878,12 @@ private:
 						data->constructedData.emplace_back(
 						    Standalone<StringRef>(StringRef(keyBuf, keyPos - keyBuf + 1)),
 						    Standalone<StringRef>(StringRef(valBuf, valSize)));
-						TraceEvent(SevDebug, "ConstructDataBuilder")
-						    .detail("Prefix", prefix)
-						    .detail("ValSize", valSize)
-						    .detail("KeyCount", keyCount)
-						    .detail("Seed", seed)
-						    .detail("Key", StringRef(keyBuf, keyPos - keyBuf + 1));
 					}
+					TraceEvent(SevDebug, "ConstructDataBuilder")
+					    .detail("Prefix", prefix)
+					    .detail("KeyCount", keyCount)
+					    .detail("ValSize", valSize)
+					    .detail("Seed", seed);
 					break;
 				}
 			}
@@ -11402,20 +11401,6 @@ ACTOR Future<Void> update(StorageServer* data, bool* pReceivedUpdate) {
 						++data->counters.atomicMutations;
 						break;
 					}
-					while (SERVER_KNOBS->GENERATE_DATA_ENABLED && data->constructedData.size()) {
-						// TraceEvent(SevDebug, "ConstructDataCommit").detail("Key", constructedMutation.param1);
-						MutationRef constructedMutation(MutationRef::SetValue,
-						                                data->constructedData.front().first,
-						                                data->constructedData.front().second);
-						MutationRefAndCipherKeys encryptedMutation;
-						updater.applyMutation(data, constructedMutation, encryptedMutation, ver, false);
-						data->constructedData.pop_front();
-						mutationBytes += constructedMutation.totalSize();
-						data->counters.mutationBytes += constructedMutation.totalSize();
-						data->counters.logicalBytesInput += constructedMutation.expectedSize();
-						++data->counters.mutations;
-						++data->counters.setMutations;
-					}
 				} else
 					TraceEvent(SevError, "DiscardingPeekedData", data->thisServerID)
 					    .detail("Mutation", msg)
@@ -11423,6 +11408,20 @@ ACTOR Future<Void> update(StorageServer* data, bool* pReceivedUpdate) {
 			}
 		}
 
+		while (SERVER_KNOBS->GENERATE_DATA_ENABLED && data->constructedData.size()) {
+			MutationRef constructedMutation(
+			    MutationRef::SetValue, data->constructedData.front().first, data->constructedData.front().second);
+			// TraceEvent(SevDebug, "ConstructDataCommit").detail("Key", constructedMutation.param1);
+			MutationRefAndCipherKeys encryptedMutation;
+			updater.applyMutation(data, constructedMutation, encryptedMutation, ver, false);
+			data->constructedData.pop_front();
+			mutationBytes += constructedMutation.totalSize();
+			data->counters.mutationBytes += constructedMutation.totalSize();
+			data->counters.logicalBytesInput += constructedMutation.expectedSize();
+			++data->counters.mutations;
+			++data->counters.setMutations;
+		}
+
 		data->tLogMsgsPTreeUpdatesLatencyHistogram->sampleSeconds(now() - beforeTLogMsgsUpdates);
 		if (data->currentChangeFeeds.size()) {
 			data->changeFeedVersions.emplace_back(