[llvm][cas] Extend on-disk CAS validation to ActionCache

llvm/include/llvm/CAS/ActionCache.h

@@ -114,6 +114,9 @@ class ActionCache {
                         Globally, std::move(Callback), CancelObj);
   }
 
+  /// Validate the ActionCache contents.
+  virtual Error validate() const = 0;
+
   virtual ~ActionCache() = default;
 
 protected:

llvm/include/llvm/CAS/OnDiskKeyValueDB.h

@@ -60,6 +60,9 @@ class OnDiskKeyValueDB {
        StringRef ValueName, size_t ValueSize,
        std::shared_ptr<OnDiskCASLogger> Logger = nullptr);
 
+  using CheckValueT = function_ref<Error(FileOffset Offset, ArrayRef<char>)>;
+  Error validate(CheckValueT CheckValue) const;
+
 private:
   OnDiskKeyValueDB(size_t ValueSize, OnDiskHashMappedTrie Cache)
       : ValueSize(ValueSize), Cache(std::move(Cache)) {}

llvm/include/llvm/CAS/UnifiedOnDiskCache.h

@@ -124,6 +124,8 @@ class UnifiedOnDiskCache {
 
   ~UnifiedOnDiskCache();
 
+  Error validateActionCache();
+
 private:
   UnifiedOnDiskCache();
 

llvm/lib/CAS/ActionCaches.cpp

@@ -52,6 +52,10 @@ class InMemoryActionCache final : public ActionCache {
   Expected<std::optional<CASID>> getImpl(ArrayRef<uint8_t> ActionKey,
                                          bool Globally) const final;
 
+  Error validate() const final {
+    return createStringError("InMemoryActionCache doesn't support validate()");
+  }
+
 private:
   using DataT = CacheEntry<sizeof(HashType)>;
   using InMemoryCacheT = ThreadSafeHashMappedTrie<DataT, sizeof(HashType)>;
@@ -68,6 +72,8 @@ class OnDiskActionCache final : public ActionCache {
 
   static Expected<std::unique_ptr<OnDiskActionCache>> create(StringRef Path);
 
+  Error validate() const final;
+
 private:
   static StringRef getHashName() { return "BLAKE3"; }
 
@@ -86,6 +92,8 @@ class UnifiedOnDiskActionCache final : public ActionCache {
 
   UnifiedOnDiskActionCache(std::shared_ptr<ondisk::UnifiedOnDiskCache> UniDB);
 
+  Error validate() const final;
+
 private:
   std::shared_ptr<ondisk::UnifiedOnDiskCache> UniDB;
 };
@@ -198,6 +206,12 @@ Error OnDiskActionCache::putImpl(ArrayRef<uint8_t> Key, const CASID &Result,
       ArrayRef((const uint8_t *)Observed.data(), Observed.size()));
 }
 
+Error OnDiskActionCache::validate() const {
+  // FIXME: without the matching CAS there is nothing we can check about the
+  // cached values. The hash size is already validated by the DB validator.
+  return DB->validate(nullptr);
+}
+
 UnifiedOnDiskActionCache::UnifiedOnDiskActionCache(
     std::shared_ptr<ondisk::UnifiedOnDiskCache> UniDB)
     : ActionCache(builtin::BuiltinCASContext::getDefaultContext()),
@@ -233,6 +247,10 @@ Error UnifiedOnDiskActionCache::putImpl(ArrayRef<uint8_t> Key,
       UniDB->getGraphDB().getDigest(*Observed));
 }
 
+Error UnifiedOnDiskActionCache::validate() const {
+  return UniDB->validateActionCache();
+}
+
 Expected<std::unique_ptr<ActionCache>>
 cas::createOnDiskActionCache(StringRef Path) {
 #if LLVM_ENABLE_ONDISK_CAS

llvm/lib/CAS/OnDiskKeyValueDB.cpp

@@ -81,3 +81,25 @@ OnDiskKeyValueDB::open(StringRef Path, StringRef HashName, unsigned KeySize,
   return std::unique_ptr<OnDiskKeyValueDB>(
       new OnDiskKeyValueDB(ValueSize, std::move(*ActionCache)));
 }
+
+Error OnDiskKeyValueDB::validate(CheckValueT CheckValue) const {
+  return Cache.validate(
+      [&](FileOffset Offset,
+          OnDiskHashMappedTrie::ConstValueProxy Record) -> Error {
+        auto formatError = [&](Twine Msg) {
+          return createStringError(
+              llvm::errc::illegal_byte_sequence,
+              "bad cache value at 0x" +
+                  utohexstr((unsigned)Offset.get(), /*LowerCase=*/true) + ": " +
+                  Msg.str());
+        };
+
+        if (Record.Data.size() != ValueSize)
+          return formatError("wrong cache value size");
+        if (!isAligned(Align(8), Record.Data.size()))
+          return formatError("wrong cache value alignment");
+        if (CheckValue)
+          return CheckValue(Offset, Record.Data);
+        return Error::success();
+      });
+}

llvm/lib/CAS/PluginAPI.h

@@ -108,6 +108,8 @@ struct llcas_functions_t {
                                            bool globally, void *ctx_cb,
                                            llcas_actioncache_put_cb,
                                            llcas_cancellable_t *);
+
+  bool (*actioncache_validate)(llcas_cas_t, char **error);
 };
 
 #endif // LLVM_LIB_CAS_PLUGINAPI_H

llvm/lib/CAS/PluginAPI_functions.def

@@ -7,6 +7,7 @@ CASPLUGINAPI_FUNCTION(actioncache_get_for_digest, true)
 CASPLUGINAPI_FUNCTION(actioncache_get_for_digest_async, true)
 CASPLUGINAPI_FUNCTION(actioncache_put_for_digest, true)
 CASPLUGINAPI_FUNCTION(actioncache_put_for_digest_async, true)
+CASPLUGINAPI_FUNCTION(actioncache_validate, false)
 CASPLUGINAPI_FUNCTION(cancellable_cancel, false)
 CASPLUGINAPI_FUNCTION(cancellable_dispose, false)
 CASPLUGINAPI_FUNCTION(cas_contains_object, true)

llvm/lib/CAS/PluginCAS.cpp

@@ -461,6 +461,8 @@ class PluginActionCache : public ActionCache {
 
   PluginActionCache(std::shared_ptr<PluginCASContext>);
 
+  Error validate() const final;
+
 private:
   std::shared_ptr<PluginCASContext> Ctx;
 };
@@ -596,6 +598,16 @@ void PluginActionCache::putImplAsync(ArrayRef<uint8_t> ResolvedKey,
 PluginActionCache::PluginActionCache(std::shared_ptr<PluginCASContext> CASCtx)
     : ActionCache(*CASCtx), Ctx(std::move(CASCtx)) {}
 
+Error PluginActionCache::validate() const {
+  if (Ctx->Functions.actioncache_validate) {
+    char *c_err = nullptr;
+    if (Ctx->Functions.actioncache_validate(Ctx->c_cas, &c_err))
+      return Ctx->errorAndDispose(c_err);
+    return Error::success();
+  }
+  return createStringError("plugin action cache doesn't support validation");
+}
+
 //===----------------------------------------------------------------------===//
 // createPluginCASDatabases API
 //===----------------------------------------------------------------------===//

llvm/lib/CAS/UnifiedOnDiskCache.cpp

@@ -52,6 +52,7 @@
 #include "llvm/CAS/UnifiedOnDiskCache.h"
 #include "OnDiskCommon.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/CAS/OnDiskCASLogger.h"
 #include "llvm/CAS/OnDiskKeyValueDB.h"
 #include "llvm/Support/Compiler.h"
@@ -123,6 +124,30 @@ UnifiedOnDiskCache::faultInFromUpstreamKV(ArrayRef<uint8_t> Key) {
   return KVPut(Key, *PrimaryID);
 }
 
+Error UnifiedOnDiskCache::validateActionCache() {
+  auto ValidateRef = [&](FileOffset Offset, ArrayRef<char> Value) -> Error {
+    assert(Value.size() == sizeof(uint64_t) && "should be validated already");
+    auto ID = ObjectID::fromOpaqueData(support::endian::read64le(Value.data()));
+    auto formatError = [&](Twine Msg) {
+      return createStringError(
+          llvm::errc::illegal_byte_sequence,
+          "bad record at 0x" +
+              utohexstr((unsigned)Offset.get(), /*LowerCase=*/true) + ": " +
+              Msg.str());
+    };
+    if (ID.getOpaqueData() == 0)
+      return formatError("zero is not a valid ref");
+    if (!PrimaryGraphDB->containsObject(ID))
+      return formatError("cas does not contain ref");
+    return Error::success();
+  };
+  if (Error E = PrimaryKVDB->validate(ValidateRef))
+    return E;
+  if (UpstreamKVDB)
+    return UpstreamKVDB->validate(ValidateRef);
+  return Error::success();
+}
+
 /// \returns all the 'v<version>.<x>' names of sub-directories, sorted with
 /// ascending order of the integer after the dot.
 static Error getAllDBDirs(StringRef Path,

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

@@ -9,3 +9,17 @@ RUN: rm %t/cas/v1.1/v8.data
 RUN: not llvm-cas --cas %t/cas --validate
 RUN: not llvm-cas --cas %t/cas --validate --check-hash
 
+RUN: mkdir %t/ac
+
+RUN: llvm-cas --cas %t/ac --make-blob \
+RUN:   --data /dev/null > %t/empty.casid
+RUN: echo "abc" | \
+RUN:   llvm-cas --cas %t/ac  --make-blob \
+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
+# 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/v3.actions 
+RUN: not llvm-cas --cas %t/ac --validate

llvm/tools/llvm-cas/llvm-cas.cpp

@@ -64,7 +64,7 @@ static int putCacheKey(ObjectStore &CAS, ActionCache &AC,
                        ArrayRef<std::string> Objects);
 static int getCacheResult(ObjectStore &CAS, ActionCache &AC, const CASID &ID);
 static int validateObject(ObjectStore &CAS, const CASID &ID);
-static int validate(ObjectStore &CAS, bool CheckHash);
+static int validate(ObjectStore &CAS, ActionCache &AC, bool CheckHash);
 static int ingestCasIDFile(cas::ObjectStore &CAS, ArrayRef<std::string> CASIDs);
 static int checkLockFiles(StringRef CASPath);
 
@@ -184,7 +184,7 @@ int main(int Argc, char **Argv) {
     return dump(*CAS);
 
   if (Command == Validate)
-    return validate(*CAS, CheckHash);
+    return validate(*CAS, *AC, CheckHash);
 
   if (Command == MakeBlob)
     return makeBlob(*CAS, DataPath);
@@ -722,9 +722,10 @@ int validateObject(ObjectStore &CAS, const CASID &ID) {
   return 0;
 }
 
-int validate(ObjectStore &CAS, bool CheckHash) {
+int validate(ObjectStore &CAS, ActionCache &AC, bool CheckHash) {
   ExitOnError ExitOnErr("llvm-cas: validate: ");
   ExitOnErr(CAS.validate(CheckHash));
+  ExitOnErr(AC.validate());
   outs() << "validated successfully\n";
   return 0;
 }