[llvm][cas] Improve UnifiedOnDiskActionCache validation to check cas refs (llvm#171732)

Check that action cache references point to valid CAS objects by
ensuring they are contained within the corresponding CAS and also that
the offsets match. This prevents accidentally referencing "dead" index
records that were not properly flushed to disk, which can lead to the
action cache pointing to the wrong data or to garbage data.

rdar://126642956
(cherry picked from commit a451ff0)

Author: benlangmuir

llvm/include/llvm/CAS/OnDiskGraphDB.h

@@ -278,16 +278,23 @@ class OnDiskGraphDB {
   ///
   /// Returns \p nullopt if the object is not stored in this CAS.
   LLVM_ABI_FOR_TEST std::optional<ObjectID>
-  getExistingReference(ArrayRef<uint8_t> Digest);
+  getExistingReference(ArrayRef<uint8_t> Digest, bool CheckUpstream = true);
 
   /// Check whether the object associated with \p Ref is stored in the CAS.
   /// Note that this function will fault-in according to the policy.
   Expected<bool> isMaterialized(ObjectID Ref);
 
   /// Check whether the object associated with \p Ref is stored in the CAS.
   /// Note that this function does not fault-in.
-  bool containsObject(ObjectID Ref) const {
-    return containsObject(Ref, /*CheckUpstream=*/true);
+  bool containsObject(ObjectID Ref, bool CheckUpstream = true) const {
+    switch (getObjectPresence(Ref, CheckUpstream)) {
+    case ObjectPresence::Missing:
+      return false;
+    case ObjectPresence::InPrimaryDB:
+      return true;
+    case ObjectPresence::OnlyInUpstreamDB:
+      return true;
+    }
   }
 
   /// \returns the data part of the provided object handle.
@@ -360,17 +367,6 @@ class OnDiskGraphDB {
   LLVM_ABI_FOR_TEST ObjectPresence
   getObjectPresence(ObjectID Ref, bool CheckUpstream) const;
 
-  bool containsObject(ObjectID Ref, bool CheckUpstream) const {
-    switch (getObjectPresence(Ref, CheckUpstream)) {
-    case ObjectPresence::Missing:
-      return false;
-    case ObjectPresence::InPrimaryDB:
-      return true;
-    case ObjectPresence::OnlyInUpstreamDB:
-      return true;
-    }
-  }
-
   /// When \p load is called for a node that doesn't exist, this function tries
   /// to load it from the upstream store and copy it to the primary one.
   Expected<std::optional<ObjectHandle>> faultInFromUpstream(ObjectID PrimaryID);
@@ -404,6 +400,10 @@ class OnDiskGraphDB {
 
   IndexProxy getIndexProxyFromRef(InternalRef Ref) const;
 
+  // FIXME: on newer branches we have refactored getIndexProxyFromRef to return
+  // Expected<IndexProxy>. As a stop gap, provide a checked API.
+  Expected<IndexProxy> getIndexProxyFromRefChecked(InternalRef Ref) const;
+
   static InternalRef makeInternalRef(FileOffset IndexOffset);
 
   IndexProxy

llvm/lib/CAS/OnDiskGraphDB.cpp

@@ -1102,10 +1102,11 @@ ObjectID OnDiskGraphDB::getExternalReference(const IndexProxy &I) {
 }
 
 std::optional<ObjectID>
-OnDiskGraphDB::getExistingReference(ArrayRef<uint8_t> Digest) {
+OnDiskGraphDB::getExistingReference(ArrayRef<uint8_t> Digest,
+                                    bool CheckUpstream) {
   auto tryUpstream =
       [&](std::optional<IndexProxy> I) -> std::optional<ObjectID> {
-    if (!UpstreamDB)
+    if (!CheckUpstream || !UpstreamDB)
       return std::nullopt;
     std::optional<ObjectID> UpstreamID =
         UpstreamDB->getExistingReference(Digest);
@@ -1138,6 +1139,15 @@ OnDiskGraphDB::getIndexProxyFromRef(InternalRef Ref) const {
   return getIndexProxyFromPointer(P);
 }
 
+Expected<OnDiskGraphDB::IndexProxy>
+OnDiskGraphDB::getIndexProxyFromRefChecked(InternalRef Ref) const {
+  OnDiskHashMappedTrie::const_pointer P =
+      Index.recoverFromFileOffset(Ref.getFileOffset());
+  if (LLVM_UNLIKELY(!P))
+    return createStringError(make_error_code(std::errc::protocol_error), "corrupt internal reference");
+  return getIndexProxyFromPointer(P);
+}
+
 ArrayRef<uint8_t> OnDiskGraphDB::getDigest(InternalRef Ref) const {
   IndexProxy I = getIndexProxyFromRef(Ref);
   return I.Hash;
@@ -1232,14 +1242,19 @@ OnDiskGraphDB::ObjectPresence
 OnDiskGraphDB::getObjectPresence(ObjectID ExternalRef,
                                  bool CheckUpstream) const {
   InternalRef Ref = getInternalRef(ExternalRef);
-  IndexProxy I = getIndexProxyFromRef(Ref);
-  TrieRecord::Data Object = I.Ref.load();
+  Expected<IndexProxy> I = getIndexProxyFromRefChecked(Ref);
+  if (!I) {
+    // FIXME: this decision should be migrated to callers.
+    consumeError(I.takeError());
+    return ObjectPresence::Missing;
+  }
+  TrieRecord::Data Object = I->Ref.load();
   if (Object.SK != TrieRecord::StorageKind::Unknown)
     return ObjectPresence::InPrimaryDB;
   if (!CheckUpstream || !UpstreamDB)
     return ObjectPresence::Missing;
   std::optional<ObjectID> UpstreamID =
-      UpstreamDB->getExistingReference(getDigest(I));
+      UpstreamDB->getExistingReference(getDigest(*I));
   return UpstreamID.has_value() ? ObjectPresence::OnlyInUpstreamDB
                                 : ObjectPresence::Missing;
 }

llvm/lib/CAS/UnifiedOnDiskCache.cpp

@@ -165,11 +165,24 @@ Error UnifiedOnDiskCache::validateActionCache() {
       return createStringError(
           llvm::errc::illegal_byte_sequence,
           "bad record at 0x" +
-              utohexstr((unsigned)Offset.get(), /*LowerCase=*/true) + ": " +
-              Msg.str());
+              utohexstr((unsigned)Offset.get(), /*LowerCase=*/true) +
+              " ref=0x" + utohexstr(ID.getOpaqueData(), /*LowerCase=*/true) +
+              ": " + Msg.str());
     };
     if (ID.getOpaqueData() == 0)
       return formatError("zero is not a valid ref");
+    // Check containsObject first, because other API assumes a valid ObjectID.
+    if (!getGraphDB().containsObject(ID, /*CheckUpstream=*/false))
+      return formatError("ref is not in cas index");
+    auto Hash = getGraphDB().getDigest(ID);
+    auto Ref = getGraphDB().getExistingReference(Hash, /*CheckUpstream=*/false);
+    assert(Ref && "missing object passed containsObject check?");
+    if (!Ref)
+      return formatError("ref is not in cas index after contains");
+    if (*Ref != ID)
+      return formatError("ref does not match indexed offset " +
+                         utohexstr(Ref->getOpaqueData(), /*LowerCase=*/true) +
+                         " for hash " + toHex(Hash));
     return Error::success();
   };
   if (Error E = PrimaryKVDB->validate(ValidateRef))

llvm/test/tools/llvm-cas/validation.test

@@ -26,6 +26,14 @@ RUN:   --data - >%t/abc.casid
 
 RUN: llvm-cas --cas %t/ac --put-cache-key @%t/abc.casid @%t/empty.casid
 RUN: llvm-cas --cas %t/ac --validate
+
+# Check that validation fails if the objects referenced are missing.
+RUN: mv %t/ac/v1.1/v9.index %t/tmp.v9.index
+RUN: not llvm-cas --cas %t/ac --validate
+
+RUN: mv %t/tmp.v9.index %t/ac/v1.1/v9.index
+RUN: llvm-cas --cas %t/ac --validate
+
 # Note: records are 40 bytes (32 hash bytes + 8 byte value), so trim the last
 # allocated record, leaving it invalid.
 RUN: truncate -s -40 %t/ac/v1.1/v4.actions