Validate all blobs received over the network.

3 files changed, 27 insertions, 16 deletions

diff --git a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
index 70f82c2d..4f9214d8 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
@@ -340,21 +340,7 @@ auto BazelCasClient::ReadSingleBlob(
     std::string const& instance_name,
     bazel_re::Digest const& digest) const noexcept -> std::optional<BazelBlob> {
     if (auto data = stream_->Read(ToResourceName(instance_name, digest))) {
-        // Recompute the digest from the received content to cross-check a
-        // correct transmission.
-        auto real_digest = static_cast<bazel_re::Digest>(
-            NativeSupport::IsTree(digest.hash())
-                ? ArtifactDigest::Create<ObjectType::Tree>(*data)
-                : ArtifactDigest::Create<ObjectType::File>(*data));
-        if (digest.hash() != real_digest.hash()) {
-            logger_.Emit(LogLevel::Warning,
-                         "Requested {}, but received {}",
-                         digest.hash(),
-                         real_digest.hash());
-            return std::nullopt;
-        }
-        return BazelBlob{
-            std::move(real_digest), std::move(*data), /*is_exec=*/false};
+        return BazelBlob{digest, std::move(*data), /*is_exec=*/false};
     }
     return std::nullopt;
 }
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp
index f4ffcb87..1bd4a90a 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp
@@ -140,7 +140,8 @@ auto BazelNetworkReader::MakeAuxiliaryMap(
 
 auto BazelNetworkReader::ReadSingleBlob(ArtifactDigest const& digest)
     const noexcept -> std::optional<ArtifactBlob> {
-    if (auto blob = cas_.ReadSingleBlob(instance_name_, digest)) {
+    auto blob = cas_.ReadSingleBlob(instance_name_, digest);
+    if (blob and BazelNetworkReader::Validate(*blob)) {
         return ArtifactBlob{
             ArtifactDigest{blob->digest}, blob->data, blob->is_exec};
     }
@@ -153,6 +154,12 @@ auto BazelNetworkReader::BatchReadBlobs(
     std::vector<BazelBlob> result =
         cas_.BatchReadBlobs(instance_name_, blobs.begin(), blobs.end());
 
+    auto it =
+        std::remove_if(result.begin(), result.end(), [](BazelBlob const& blob) {
+            return not BazelNetworkReader::Validate(blob);
+        });
+    result.erase(it, result.end());
+
     std::vector<ArtifactBlob> artifacts;
     artifacts.reserve(result.size());
     std::transform(result.begin(),
@@ -165,3 +172,18 @@ auto BazelNetworkReader::BatchReadBlobs(
                    });
     return artifacts;
 }
+
+auto BazelNetworkReader::Validate(BazelBlob const& blob) noexcept -> bool {
+    ArtifactDigest const rehashed_digest =
+        NativeSupport::IsTree(blob.digest.hash())
+            ? ArtifactDigest::Create<ObjectType::Tree>(*blob.data)
+            : ArtifactDigest::Create<ObjectType::File>(*blob.data);
+    if (rehashed_digest == ArtifactDigest{blob.digest}) {
+        return true;
+    }
+    Logger::Log(LogLevel::Warning,
+                "Requested {}, but received {}",
+                ArtifactDigest{blob.digest}.hash(),
+                rehashed_digest.hash());
+    return false;
+}
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp
index fb807a79..66be173b 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp
@@ -25,6 +25,7 @@
 #include "src/buildtool/common/artifact.hpp"
 #include "src/buildtool/common/artifact_digest.hpp"
 #include "src/buildtool/common/bazel_types.hpp"
+#include "src/buildtool/execution_api/bazel_msg/bazel_blob_container.hpp"
 #include "src/buildtool/execution_api/common/artifact_blob_container.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp"
 #include "src/buildtool/file_system/git_repo.hpp"
@@ -72,6 +73,8 @@ class BazelNetworkReader final {
     [[nodiscard]] auto BatchReadBlobs(
         std::vector<bazel_re::Digest> const& blobs) const noexcept
         -> std::vector<ArtifactBlob>;
+
+    [[nodiscard]] static auto Validate(BazelBlob const& blob) noexcept -> bool;
 };
 
 #endif  // INCLUDED_SRC_BUILDTOOL_EXECUTION_API_REMOTE_BAZEL_BAZEL_TREE_READER_HPP