[llvm][cas] Improve UnifiedOnDiskActionCache validation to check cas refs

llvm/include/llvm/CAS/OnDiskGraphDB.h

@@ -279,16 +279,29 @@ 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 {
+    auto Presence = getObjectPresence(Ref, CheckUpstream);
+    if (!Presence) {
+      consumeError(Presence.takeError());
+      return false;
+    }
+    switch (*Presence) {
+    case ObjectPresence::Missing:
+      return false;
+    case ObjectPresence::InPrimaryDB:
+      return true;
+    case ObjectPresence::OnlyInUpstreamDB:
+      return true;
+    }
+    llvm_unreachable("Unknown ObjectPresence enum");
   }
 
   /// \returns the data part of the provided object handle.
@@ -370,24 +383,6 @@ class OnDiskGraphDB {
   LLVM_ABI_FOR_TEST Expected<ObjectPresence>
   getObjectPresence(ObjectID Ref, bool CheckUpstream) const;
 
-  /// \returns true if object can be found in database.
-  bool containsObject(ObjectID Ref, bool CheckUpstream) const {
-    auto Presence = getObjectPresence(Ref, CheckUpstream);
-    if (!Presence) {
-      consumeError(Presence.takeError());
-      return false;
-    }
-    switch (*Presence) {
-    case ObjectPresence::Missing:
-      return false;
-    case ObjectPresence::InPrimaryDB:
-      return true;
-    case ObjectPresence::OnlyInUpstreamDB:
-      return true;
-    }
-    llvm_unreachable("Unknown ObjectPresence enum");
-  }
-
   /// 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);

llvm/lib/CAS/ActionCaches.cpp

@@ -235,17 +235,31 @@ Error UnifiedOnDiskActionCache::putImpl(ArrayRef<uint8_t> Key,
 }
 
 Error UnifiedOnDiskActionCache::validate() const {
-  auto ValidateRef = [](FileOffset Offset, ArrayRef<char> Value) -> Error {
+  auto ValidateRef = [this](FileOffset Offset, ArrayRef<char> Value) -> Error {
     auto ID = ondisk::UnifiedOnDiskCache::getObjectIDFromValue(Value);
     auto formatError = [&](Twine Msg) {
       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 (!UniDB->getGraphDB().containsObject(ID, /*CheckUpstream=*/false))
+      return formatError("ref is not in cas index");
+    auto Hash = UniDB->getGraphDB().getDigest(ID);
+    auto Ref =
+        UniDB->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();
   };
   return UniDB->getKeyValueDB().validate(ValidateRef);

llvm/lib/CAS/OnDiskGraphDB.cpp

@@ -1109,10 +1109,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);

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

@@ -25,6 +25,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/index.v1 %t/tmp.index.v1
+RUN: not llvm-cas --cas %t/ac --validate
+
+RUN: mv %t/tmp.index.v1 %t/ac/v1.1/index.v1
+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/actions.v1