1 files changed, 118 insertions, 1 deletions

diff --git a/src/buildtool/execution_api/execution_service/cas_utils.cpp b/src/buildtool/execution_api/execution_service/cas_utils.cpp
index a406bdec..42d7d074 100644
--- a/src/buildtool/execution_api/execution_service/cas_utils.cpp
+++ b/src/buildtool/execution_api/execution_service/cas_utils.cpp
@@ -14,14 +14,18 @@
 
 #include "src/buildtool/execution_api/execution_service/cas_utils.hpp"
 
+#include <type_traits>
 #include <utility>
 
 #include "fmt/core.h"
+#include "gsl/gsl"
+#include "src/buildtool/common/protocol_traits.hpp"
 #include "src/buildtool/file_system/file_system_manager.hpp"
 #include "src/buildtool/storage/large_object_cas.hpp"
 #include "src/buildtool/storage/local_cas.hpp"
 
-static auto ToGrpc(LargeObjectError&& error) noexcept -> grpc::Status {
+namespace {
+[[nodiscard]] auto ToGrpc(LargeObjectError&& error) noexcept -> grpc::Status {
     switch (error.Code()) {
         case LargeObjectErrorCode::Internal:
             return grpc::Status{grpc::StatusCode::INTERNAL,
@@ -37,6 +41,89 @@ static auto ToGrpc(LargeObjectError&& error) noexcept -> grpc::Status {
     return grpc::Status{grpc::StatusCode::INTERNAL, "an unknown error"};
 }
 
+class CASContentValidator final {
+  public:
+    explicit CASContentValidator(gsl::not_null<Storage const*> const& storage,
+                                 bool is_owner = true) noexcept;
+
+    template <typename TData>
+    [[nodiscard]] auto Add(ArtifactDigest const& digest,
+                           TData const& data) const noexcept -> grpc::Status {
+        if (digest.IsTree()) {
+            // For trees, check whether the tree invariant holds before storing
+            // the actual tree object.
+            if (auto err = storage_.CAS().CheckTreeInvariant(digest, data)) {
+                return ToGrpc(std::move(*err));
+            }
+        }
+
+        auto const cas_digest =
+            digest.IsTree() ? StoreTree(data) : StoreBlob(data);
+        if (not cas_digest) {
+            // This is a serious problem: we have a sequence of bytes, but
+            // cannot write them to CAS.
+            return ::grpc::Status{grpc::StatusCode::INTERNAL,
+                                  fmt::format("Could not upload {} {}",
+                                              digest.IsTree() ? "tree" : "blob",
+                                              digest.hash())};
+        }
+
+        if (auto err = CheckDigestConsistency(digest, *cas_digest)) {
+            // User error: did not get a file with the announced hash
+            return ::grpc::Status{grpc::StatusCode::INVALID_ARGUMENT,
+                                  *std::move(err)};
+        }
+        return ::grpc::Status::OK;
+    }
+
+  private:
+    Storage const& storage_;
+    bool const is_owner_;
+
+    template <typename TData>
+    [[nodiscard]] auto StoreTree(TData const& data) const noexcept
+        -> std::optional<ArtifactDigest> {
+        if constexpr (std::is_same_v<TData, std::string>) {
+            return storage_.CAS().StoreTree(data);
+        }
+        else {
+            return is_owner_ ? storage_.CAS().StoreTree<true>(data)
+                             : storage_.CAS().StoreTree<false>(data);
+        }
+    }
+
+    template <typename TData>
+    [[nodiscard]] auto StoreBlob(TData const& data) const noexcept
+        -> std::optional<ArtifactDigest> {
+        static constexpr bool kIsExec = false;
+        if constexpr (std::is_same_v<TData, std::string>) {
+            return storage_.CAS().StoreBlob(data, kIsExec);
+        }
+        else {
+            return is_owner_ ? storage_.CAS().StoreBlob<true>(data, kIsExec)
+                             : storage_.CAS().StoreBlob<false>(data, kIsExec);
+        }
+    }
+
+    [[nodiscard]] auto CheckDigestConsistency(ArtifactDigest const& ref,
+                                              ArtifactDigest const& computed)
+        const noexcept -> std::optional<std::string>;
+};
+}  // namespace
+
+auto CASUtils::AddDataToCAS(ArtifactDigest const& digest,
+                            std::string const& content,
+                            Storage const& storage) noexcept -> grpc::Status {
+    return CASContentValidator{&storage}.Add(digest, content);
+}
+
+auto CASUtils::AddFileToCAS(ArtifactDigest const& digest,
+                            std::filesystem::path const& file,
+                            Storage const& storage,
+                            bool is_owner) noexcept -> grpc::Status {
+    return CASContentValidator{&storage, is_owner}.Add(digest, file);
+}
+
 auto CASUtils::EnsureTreeInvariant(ArtifactDigest const& digest,
                                    std::string const& tree_data,
                                    Storage const& storage) noexcept
@@ -115,3 +202,33 @@ auto CASUtils::SpliceBlob(ArtifactDigest const& blob_digest,
     }
     return *std::move(splice);
 }
+
+namespace {
+CASContentValidator::CASContentValidator(
+    gsl::not_null<Storage const*> const& storage,
+    bool is_owner) noexcept
+    : storage_{*storage}, is_owner_{is_owner} {}
+
+auto CASContentValidator::CheckDigestConsistency(ArtifactDigest const& ref,
+                                                 ArtifactDigest const& computed)
+    const noexcept -> std::optional<std::string> {
+    bool valid = ref == computed;
+    if (valid) {
+        bool const check_sizes = not ProtocolTraits::IsNative(
+                                     storage_.GetHashFunction().GetType()) or
+                                 ref.size() != 0;
+        if (check_sizes) {
+            valid = ref.size() == computed.size();
+        }
+    }
+    if (not valid) {
+        return fmt::format(
+            "Expected digest {}:{} and computed digest {}:{} do not match.",
+            ref.hash(),
+            ref.size(),
+            computed.hash(),
+            computed.size());
+    }
+    return std::nullopt;
+}
+}  // namespace