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

Validate the ActionCache hash-mapped trie structure and sanity check the
resulting values. Unlike the CAS itself there is no direct way to check
the values are "correct", but at least we can check for invalid zero
offsets, which is what we would get if we dropped page writes or
truncated the file.

(cherry picked from commit 2966de4)

Author: benlangmuir

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