#Centipede more protection against blob sequence buffer corruption.

Also, do not check-fail when failed to read a batch result - an adversely corrupted buffer could pass the check anyway.

PiperOrigin-RevId: 734579646

centipede/centipede_callbacks.cc

@@ -215,7 +215,8 @@ int CentipedeCallbacks::ExecuteCentipedeSancovBinaryWithShmem(
 
   // Get results.
   batch_result.exit_code() = retval;
-  CHECK(batch_result.Read(outputs_blobseq_));
+  const bool read_success = batch_result.Read(outputs_blobseq_);
+  LOG_IF(ERROR, !read_success) << "Failed to read batch result!";
   outputs_blobseq_.ReleaseSharedMemory();  // Outputs are already consumed.
 
   // We may have fewer feature blobs than inputs if
@@ -225,7 +226,8 @@ int CentipedeCallbacks::ExecuteCentipedeSancovBinaryWithShmem(
   //   * Will be logged by the caller.
   // * some outputs were not written because the outputs_blobseq_ overflown.
   //   * Logged by the following code.
-  if (retval == 0 && batch_result.num_outputs_read() != num_inputs_written) {
+  if (retval == 0 && read_success &&
+      batch_result.num_outputs_read() != num_inputs_written) {
     LOG(INFO) << "Read " << batch_result.num_outputs_read() << "/"
               << num_inputs_written
               << " outputs; shmem_size_mb might be too small: "

centipede/runner_result.cc

@@ -101,15 +101,18 @@ bool BatchResult::Read(BlobSequence &blobseq) {
       continue;
     }
     if (blob.tag == kTagMetadata) {
+      if (current_execution_result == nullptr) return false;
       current_execution_result->metadata().Read(blob);
       continue;
     }
     if (blob.tag == kTagStats) {
+      if (current_execution_result == nullptr) return false;
       if (blob.size != sizeof(ExecutionResult::Stats)) return false;
       memcpy(&current_execution_result->stats(), blob.data, blob.size);
       continue;
     }
     if (blob.tag == kTagFeatures) {
+      if (current_execution_result == nullptr) return false;
       const size_t features_size = blob.size / sizeof(feature_t);
       FeatureVec &features = current_execution_result->mutable_features();
       features.resize(features_size);

centipede/runner_result_test.cc

@@ -194,5 +194,26 @@ TEST(MutationResult, WriteThenRead) {
       ElementsAre(ByteArray{1, 2, 3}, ByteArray{4, 5, 6}, ByteArray{7, 8, 9}));
 }
 
+TEST(ExecutionResult, ReadResultSucceedsOnlyWithInputBegin) {
+  auto buffer = std::make_unique<uint8_t[]>(1000);
+  BlobSequence blobseq(buffer.get(), 1000);
+  BatchResult batch_result;
+
+  EXPECT_TRUE(BatchResult::WriteInputBegin(blobseq));
+  EXPECT_TRUE(BatchResult::WriteOneFeatureVec({}, 0, blobseq));
+  EXPECT_TRUE(BatchResult::WriteInputEnd(blobseq));
+  blobseq.Reset();
+  batch_result.ClearAndResize(1);
+  EXPECT_TRUE(batch_result.Read(blobseq));
+
+  blobseq.Reset();
+  EXPECT_TRUE(BatchResult::WriteOneFeatureVec({}, 0, blobseq));
+  EXPECT_TRUE(BatchResult::WriteInputEnd(blobseq));
+
+  blobseq.Reset();
+  batch_result.ClearAndResize(1);
+  EXPECT_FALSE(batch_result.Read(blobseq));
+}
+
 }  // namespace
 }  // namespace centipede