Unify logic of adding to CAS in ByteStreamServer and CASServer

...by calling the generalized CASUtils's implementation.

4 files changed, 8 insertions, 123 deletions

diff --git a/src/buildtool/execution_api/execution_service/bytestream_server.cpp b/src/buildtool/execution_api/execution_service/bytestream_server.cpp
index d1e13d53..1740c1f1 100644
--- a/src/buildtool/execution_api/execution_service/bytestream_server.cpp
+++ b/src/buildtool/execution_api/execution_service/bytestream_server.cpp
@@ -150,53 +150,12 @@ auto BytestreamServiceImpl::Write(
         } while (not request.finish_write() and reader->Read(&request));
     }
 
-    // Before storing a tree, we have to verify that its parts are present
-    if (write_digest->IsTree()) {
-        // ... unfortunately, this requires us to read the whole tree object
-        // into memory
-        auto const content = FileSystemManager::ReadFile(tmp);
-        if (not content) {
-            auto const msg =
-                fmt::format("Failed to read temporary file {} for {}",
-                            tmp.string(),
-                            write_digest->hash());
-            logger_.Emit(LogLevel::Error, "{}", msg);
-            return ::grpc::Status{::grpc::StatusCode::INTERNAL, msg};
-        }
-
-        if (auto err = CASUtils::EnsureTreeInvariant(
-                *write_digest, *content, storage_)) {
-            auto const str = fmt::format("Write: {}", *std::move(err));
-            logger_.Emit(LogLevel::Error, "{}", str);
-            return ::grpc::Status{grpc::StatusCode::FAILED_PRECONDITION, str};
-        }
-    }
-
-    // Store blob and verify hash
-    static constexpr bool kOwner = true;
-    auto const stored =
-        write_digest->IsTree()
-            ? storage_.CAS().StoreTree<kOwner>(tmp)
-            : storage_.CAS().StoreBlob<kOwner>(tmp, /*is_executable=*/false);
-    if (not stored) {
-        // This is a serious problem: we have a sequence of bytes, but cannot
-        // write them to CAS.
-        auto const str =
-            fmt::format("Failed to store object {}", write_digest->hash());
+    auto const status = CASUtils::AddFileToCAS(*write_digest, tmp, storage_);
+    if (not status.ok()) {
+        auto const str = fmt::format("Write: {}", status.error_message());
         logger_.Emit(LogLevel::Error, "{}", str);
-        return ::grpc::Status{::grpc::StatusCode::INTERNAL, str};
+        return ::grpc::Status{status.error_code(), str};
     }
-
-    if (*stored != *write_digest) {
-        // User error: did not get a file with the announced hash
-        auto const str =
-            fmt::format("In upload for {} received object with hash {}",
-                        write_digest->hash(),
-                        stored->hash());
-        logger_.Emit(LogLevel::Error, "{}", str);
-        return ::grpc::Status{::grpc::StatusCode::INVALID_ARGUMENT, str};
-    }
-
     response->set_committed_size(
         static_cast<google::protobuf::int64>(std::filesystem::file_size(tmp)));
     return ::grpc::Status::OK;
diff --git a/src/buildtool/execution_api/execution_service/cas_server.cpp b/src/buildtool/execution_api/execution_service/cas_server.cpp
index 87f2e92f..6ca79348 100644
--- a/src/buildtool/execution_api/execution_service/cas_server.cpp
+++ b/src/buildtool/execution_api/execution_service/cas_server.cpp
@@ -47,31 +47,6 @@ namespace {
             return "[Unknown Chunking Algorithm Type]";
     }
 }
-
-[[nodiscard]] auto CheckDigestConsistency(
-    HashFunction::Type hash_type,
-    ArtifactDigest const& ref,
-    ArtifactDigest const& computed) noexcept -> std::optional<std::string> {
-    bool valid = ref.hash() == computed.hash();
-    if (valid) {
-        bool const check_sizes =
-            not ProtocolTraits::IsNative(hash_type) or ref.size() != 0;
-        if (check_sizes) {
-            valid = ref.size() == computed.size();
-        }
-    }
-    if (not valid) {
-        return fmt::format(
-            "Blob {} is corrupted: provided digest {}:{} and digest computed "
-            "from data {}:{} do not correspond.",
-            ref.hash(),
-            ref.hash(),
-            ref.size(),
-            computed.hash(),
-            computed.size());
-    }
-    return std::nullopt;
-}
 }  // namespace
 
 auto CASServiceImpl::FindMissingBlobs(
@@ -138,39 +113,12 @@ auto CASServiceImpl::BatchUpdateBlobs(
         auto* r = response->add_responses();
         r->mutable_digest()->CopyFrom(x.digest());
 
-        if (digest->IsTree()) {
-            // In native mode: for trees, check whether the tree invariant holds
-            // before storing the actual tree object.
-            if (auto err = CASUtils::EnsureTreeInvariant(
-                    *digest, x.data(), storage_)) {
-                auto const str =
-                    fmt::format("BatchUpdateBlobs: {}", *std::move(err));
-                logger_.Emit(LogLevel::Error, "{}", str);
-                return ::grpc::Status{grpc::StatusCode::FAILED_PRECONDITION,
-                                      str};
-            }
-        }
-
-        auto const cas_digest =
-            digest->IsTree()
-                ? storage_.CAS().StoreTree(x.data())
-                : storage_.CAS().StoreBlob(x.data(), /*is_executable=*/false);
-
-        if (not cas_digest) {
+        auto const status = CASUtils::AddDataToCAS(*digest, x.data(), storage_);
+        if (not status.ok()) {
             auto const str =
-                fmt::format("BatchUpdateBlobs: could not upload {} {}",
-                            digest->IsTree() ? "tree" : "blob",
-                            digest->hash());
+                fmt::format("BatchUpdateBlobs: {}", status.error_message());
             logger_.Emit(LogLevel::Error, "{}", str);
-            return ::grpc::Status{grpc::StatusCode::INTERNAL, str};
-        }
-
-        if (auto err =
-                CheckDigestConsistency(hash_type, *digest, *cas_digest)) {
-            auto const str =
-                fmt::format("BatchUpdateBlobs: {}", *std::move(err));
-            logger_.Emit(LogLevel::Error, "{}", str);
-            return ::grpc::Status{grpc::StatusCode::INVALID_ARGUMENT, str};
+            return ::grpc::Status{status.error_code(), str};
         }
     }
     return ::grpc::Status::OK;
@@ -369,12 +317,6 @@ auto CASServiceImpl::SpliceBlob(::grpc::ServerContext* /*context*/,
         logger_.Emit(LogLevel::Error, "{}", str);
         return ::grpc::Status{status.error_code(), str};
     }
-    if (auto err =
-            CheckDigestConsistency(hash_type, *blob_digest, *splice_result)) {
-        auto const str = fmt::format("SpliceBlob: {}", *std::move(err));
-        logger_.Emit(LogLevel::Error, "{}", str);
-        return ::grpc::Status{grpc::StatusCode::INVALID_ARGUMENT, str};
-    }
 
     (*response->mutable_blob_digest()) =
         ArtifactDigestFactory::ToBazel(*splice_result);
diff --git a/src/buildtool/execution_api/execution_service/cas_utils.cpp b/src/buildtool/execution_api/execution_service/cas_utils.cpp
index 42d7d074..ecd75c4f 100644
--- a/src/buildtool/execution_api/execution_service/cas_utils.cpp
+++ b/src/buildtool/execution_api/execution_service/cas_utils.cpp
@@ -124,17 +124,6 @@ auto CASUtils::AddFileToCAS(ArtifactDigest const& digest,
     return CASContentValidator{&storage, is_owner}.Add(digest, file);
 }
 
-auto CASUtils::EnsureTreeInvariant(ArtifactDigest const& digest,
-                                   std::string const& tree_data,
-                                   Storage const& storage) noexcept
-    -> std::optional<std::string> {
-    auto error = storage.CAS().CheckTreeInvariant(digest, tree_data);
-    if (error) {
-        return std::move(*error).Message();
-    }
-    return std::nullopt;
-}
-
 auto CASUtils::SplitBlobIdentity(ArtifactDigest const& blob_digest,
                                  Storage const& storage) noexcept
     -> expected<std::vector<ArtifactDigest>, grpc::Status> {
diff --git a/src/buildtool/execution_api/execution_service/cas_utils.hpp b/src/buildtool/execution_api/execution_service/cas_utils.hpp
index 8de56bdd..b98e3223 100644
--- a/src/buildtool/execution_api/execution_service/cas_utils.hpp
+++ b/src/buildtool/execution_api/execution_service/cas_utils.hpp
@@ -27,11 +27,6 @@
 
 class CASUtils {
   public:
-    [[nodiscard]] static auto EnsureTreeInvariant(
-        ArtifactDigest const& digest,
-        std::string const& tree_data,
-        Storage const& storage) noexcept -> std::optional<std::string>;
-
     [[nodiscard]] static auto AddDataToCAS(ArtifactDigest const& digest,
                                            std::string const& content,
                                            Storage const& storage) noexcept