Add more record header sanity checks

Summary: When reading a file from disk, the data could be corrupt. This diff adds defensive checks to prevent crashes when the data is boggus. Reviewed By: finik Differential Revision: D66267999 fbshipit-source-id: 43ed9e3309bdf4a2243325582f3c07a819eb27d5
facebookresearch · Nov 21, 2024 · 1474631 · 1474631
1 parent 12bf698
commit 1474631
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 0 deletions.
diff --git a/vrs/FileFormat.cpp b/vrs/FileFormat.cpp
@@ -187,6 +187,28 @@ void RecordHeader::initDescriptionHeader(
   this->timestamp.set(Record::kMaxTimestamp);
 }
 
+bool RecordHeader::isSanityCheckOk() const {
+  if (!XR_VERIFY(recordSize.get() >= sizeof(RecordHeader)) ||
+      !XR_VERIFY(
+          previousRecordSize.get() == 0 || previousRecordSize.get() >= sizeof(RecordHeader))) {
+    return false;
+  }
+  if (!XR_VERIFY(recordType.get() > static_cast<uint8_t>(Record::Type::UNDEFINED)) ||
+      !XR_VERIFY(recordType.get() < static_cast<uint8_t>(Record::Type::COUNT))) {
+    return false;
+  }
+  if (uncompressedSize.get() > 0) {
+    if (!XR_VERIFY(uncompressedSize.get() >= recordSize.get() / 2)) {
+      return false;
+    }
+    if (!XR_VERIFY(compressionType.get() != static_cast<uint8_t>(CompressionType::None)) ||
+        !XR_VERIFY(compressionType.get() < static_cast<uint8_t>(CompressionType::COUNT))) {
+      return false;
+    }
+  }
+  return true;
+}
+
 RecordableTypeId readRecordableTypeId(const FileFormat::LittleEndian<int32_t>& recordableTypeId) {
   int32_t rawTypeId = recordableTypeId.get();
   // reinterpret ids for test & sample devices in their legacy space...

diff --git a/vrs/FileFormat.h b/vrs/FileFormat.h
@@ -232,6 +232,10 @@ struct RecordHeader {
       uint32_t formatVersion,
       uint32_t descriptionRecordSize,
       uint32_t previousRecordSize);
+
+  /// Do some surperficial sanity checks on the record header.
+  /// @return True if the header looks valid, false otherwise.
+  bool isSanityCheckOk() const;
 };
 
 #pragma pack(pop)

diff --git a/vrs/IndexRecord.cpp b/vrs/IndexRecord.cpp
@@ -301,6 +301,10 @@ int IndexRecord::Reader::readRecord(
     XR_LOGE("Record size too small. Corrupt?");
     return INDEX_RECORD_ERROR;
   }
+  if (!recordHeader->isSanityCheckOk()) {
+    XR_LOGE("Record header sanity check failed. Corrupt?");
+    return INDEX_RECORD_ERROR;
+  }
   size_t indexByteSize = recordHeader->recordSize.get() - static_cast<uint32_t>(recordHeaderSize);
   if (recordHeader->formatVersion.get() == kClassicIndexFormatVersion) {
     return readClassicIndexRecord(
@@ -672,6 +676,11 @@ int IndexRecord::Reader::rebuildIndex(bool writeFixedIndex) {
       error = REINDEXING_ERROR;
       break;
     }
+    if (!recordHeader->isSanityCheckOk()) {
+      XR_LOGW("Reindexing: record #{} header sanity check failed.", index_.size());
+      error = REINDEXING_ERROR;
+      break;
+    }
     RecordableTypeId recordableTypeId = recordHeader->getRecordableTypeId();
     uint32_t dataSize = recordSize - static_cast<uint32_t>(recordHeaderSize);
     if (recordableTypeId == RecordableTypeId::VRSIndex &&

diff --git a/vrs/RecordFileReader.cpp b/vrs/RecordFileReader.cpp
@@ -1052,6 +1052,10 @@ bool RecordFileReader::isRecordAvailableOrPrefetch(const IndexRecord::RecordInfo
         errorCodeToMessageWithCode(error));
     return false;
   }
+  if (!recordHeader.isSanityCheckOk()) {
+    XR_LOGE("Record #{} Record header doesn't look right.", getRecordIndex(&recordInfo));
+    return false;
+  }
   uint32_t recordSize = recordHeader.recordSize.get();
   return file_->isAvailableOrPrefetch(recordSize - sizeof(recordHeader));
 }
@@ -1111,6 +1115,10 @@ int RecordFileReader::readRecord(
   }
 
   bool integrityCheck = true;
+  if (!recordHeader.isSanityCheckOk()) {
+    integrityCheck = false;
+    XR_LOGE("Record #{} sanity check failed.", static_cast<int>(&recordInfo - &getIndex()[0]));
+  }
   if (recordInfo.timestamp != recordHeader.timestamp.get()) {
     integrityCheck = false;
     XR_LOGE(