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

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