ArtifactBlob: Check access to the content

5 files changed, 64 insertions, 24 deletions

diff --git a/src/buildtool/execution_api/local/local_api.cpp b/src/buildtool/execution_api/local/local_api.cpp
index c9352d25..040b4be6 100644
--- a/src/buildtool/execution_api/local/local_api.cpp
+++ b/src/buildtool/execution_api/local/local_api.cpp
@@ -232,16 +232,22 @@ auto LocalApi::RetrieveToMemory(Artifact::ObjectInfo const& artifact_info)
 
 auto LocalApi::Upload(std::unordered_set<ArtifactBlob>&& blobs,
                       bool /*skip_find_missing*/) const noexcept -> bool {
-    return std::all_of(
+    // Blobs could have been received over the network, so a simple failure
+    // could result in lost traffic. Try add all blobs and fail if at least
+    // one is corrupted.
+    std::size_t const valid_count = std::count_if(
         blobs.begin(),
         blobs.end(),
         [&cas = local_context_.storage->CAS()](ArtifactBlob const& blob) {
-            auto const cas_digest =
-                blob.GetDigest().IsTree()
-                    ? cas.StoreTree(*blob.ReadContent())
-                    : cas.StoreBlob(*blob.ReadContent(), blob.IsExecutable());
+            std::optional<ArtifactDigest> cas_digest;
+            if (auto const content = blob.ReadContent()) {
+                cas_digest = blob.GetDigest().IsTree()
+                                 ? cas.StoreTree(*content)
+                                 : cas.StoreBlob(*content, blob.IsExecutable());
+            }
             return cas_digest and *cas_digest == blob.GetDigest();
         });
+    return valid_count == blobs.size();
 }
 
 auto LocalApi::UploadTree(
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
index ea6b71b9..eb140b3c 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
@@ -237,9 +237,15 @@ auto BazelApi::CreateAction(
         for (std::size_t pos = 0; pos < blobs.size(); ++pos) {
             auto gpos = artifact_pos[count + pos];
             auto const& type = artifacts_info[gpos].type;
-            if (not FileSystemManager::WriteFileAs</*kSetEpochTime=*/true,
-                                                   /*kSetWritable=*/true>(
-                    *blobs[pos].ReadContent(), output_paths[gpos], type)) {
+
+            bool written = false;
+            if (auto const content = blobs[pos].ReadContent()) {
+                written = FileSystemManager::WriteFileAs</*kSetEpochTime=*/true,
+                                                         /*kSetWritable=*/true>(
+                    *content, output_paths[gpos], type);
+            }
+
+            if (not written) {
                 Logger::Log(LogLevel::Warning,
                             "staging to output path {} failed.",
                             output_paths[gpos].string());
@@ -486,7 +492,9 @@ auto BazelApi::CreateAction(
     -> std::optional<std::string> {
     auto reader = network_->CreateReader();
     if (auto blob = reader.ReadSingleBlob(artifact_info.digest)) {
-        return *blob->ReadContent();
+        if (auto const content = blob->ReadContent()) {
+            return *content;
+        }
     }
     return std::nullopt;
 }
@@ -520,7 +528,9 @@ auto BazelApi::CreateAction(
             targets->reserve(digests.size());
             for (auto blobs : reader.ReadIncrementally(&digests)) {
                 for (auto const& blob : blobs) {
-                    targets->emplace_back(*blob.ReadContent());
+                    if (auto const content = blob.ReadContent()) {
+                        targets->emplace_back(*content);
+                    }
                 }
             }
         });
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 20b4b750..dd4ec919 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
@@ -594,14 +594,20 @@ auto BazelCasClient::BatchUpdateBlobs(std::string const& instance_name,
 
     auto const max_content_size = GetMaxBatchTransferSize(instance_name);
 
-    auto request_creator = [&instance_name](ArtifactBlob const& blob) {
+    auto request_creator = [&instance_name](ArtifactBlob const& blob)
+        -> std::optional<bazel_re::BatchUpdateBlobsRequest> {
+        auto const content = blob.ReadContent();
+        if (content == nullptr) {
+            return std::nullopt;
+        }
+
         bazel_re::BatchUpdateBlobsRequest request;
         request.set_instance_name(instance_name);
 
         auto& r = *request.add_requests();
         (*r.mutable_digest()) =
             ArtifactDigestFactory::ToBazel(blob.GetDigest());
-        r.set_data(*blob.ReadContent());
+        r.set_data(*content);
         return request;
     };
 
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 7fe90b72..b8cb7eae 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp
@@ -66,8 +66,10 @@ auto BazelNetworkReader::ReadDirectory(ArtifactDigest const& digest)
     }
 
     if (auto blob = ReadSingleBlob(digest)) {
-        return BazelMsgFactory::MessageFromString<bazel_re::Directory>(
-            *blob->ReadContent());
+        if (auto const content = blob->ReadContent()) {
+            return BazelMsgFactory::MessageFromString<bazel_re::Directory>(
+                *content);
+        }
     }
     Logger::Log(
         LogLevel::Debug, "Directory {} not found in CAS", digest.hash());
@@ -83,6 +85,11 @@ auto BazelNetworkReader::ReadGitTree(ArtifactDigest const& digest)
         Logger::Log(LogLevel::Debug, "Tree {} not found in CAS", digest.hash());
         return std::nullopt;
     }
+    auto const content = read_blob->ReadContent();
+    if (content == nullptr) {
+        return std::nullopt;
+    }
+
     auto check_symlinks = [this](std::vector<ArtifactDigest> const& ids) {
         size_t const size = ids.size();
         size_t count = 0;
@@ -94,7 +101,8 @@ auto BazelNetworkReader::ReadGitTree(ArtifactDigest const& digest)
             }
             bool valid = std::all_of(
                 blobs.begin(), blobs.end(), [](ArtifactBlob const& blob) {
-                    return PathIsNonUpwards(*blob.ReadContent());
+                    auto const content = blob.ReadContent();
+                    return content != nullptr and PathIsNonUpwards(*content);
                 });
             if (not valid) {
                 return false;
@@ -104,9 +112,8 @@ auto BazelNetworkReader::ReadGitTree(ArtifactDigest const& digest)
         return true;
     };
 
-    std::string const content = *read_blob->ReadContent();
-    return GitRepo::ReadTreeData(content,
-                                 hash_function_.HashTreeData(content).Bytes(),
+    return GitRepo::ReadTreeData(*content,
+                                 hash_function_.HashTreeData(*content).Bytes(),
                                  check_symlinks,
                                  /*is_hex_id=*/false);
 }
@@ -122,7 +129,8 @@ auto BazelNetworkReader::DumpRawTree(Artifact::ObjectInfo const& info,
     }
 
     try {
-        return std::invoke(dumper, *read_blob->ReadContent());
+        auto const content = read_blob->ReadContent();
+        return content != nullptr and std::invoke(dumper, *content);
     } catch (...) {
         return false;
     }
@@ -236,11 +244,16 @@ auto BazelNetworkReader::BatchReadBlobs(
 
 auto BazelNetworkReader::Validate(ArtifactBlob const& blob) const noexcept
     -> bool {
+    auto const content = blob.ReadContent();
+    if (content == nullptr) {
+        return false;
+    }
+
     auto rehashed = blob.GetDigest().IsTree()
                         ? ArtifactDigestFactory::HashDataAs<ObjectType::Tree>(
-                              hash_function_, *blob.ReadContent())
+                              hash_function_, *content)
                         : ArtifactDigestFactory::HashDataAs<ObjectType::File>(
-                              hash_function_, *blob.ReadContent());
+                              hash_function_, *content);
     return rehashed == blob.GetDigest();
 }
 
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_response.cpp b/src/buildtool/execution_api/remote/bazel/bazel_response.cpp
index a0a621a4..f06a4275 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_response.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_response.cpp
@@ -67,7 +67,9 @@ auto BazelResponse::ReadStringBlob(bazel_re::Digest const& id) noexcept
     if (digest.has_value()) {
         auto reader = network_->CreateReader();
         if (auto blob = reader.ReadSingleBlob(*digest)) {
-            return *blob->ReadContent();
+            if (auto const content = blob->ReadContent()) {
+                return *content;
+            }
         }
     }
     Logger::Log(LogLevel::Warning,
@@ -239,8 +241,11 @@ auto BazelResponse::Populate() noexcept -> std::optional<std::string> {
     for (auto tree_blobs : reader.ReadIncrementally(&tree_digests)) {
         for (auto const& tree_blob : tree_blobs) {
             try {
-                auto tree = BazelMsgFactory::MessageFromString<bazel_re::Tree>(
-                    *tree_blob.ReadContent());
+                std::optional<bazel_re::Tree> tree;
+                if (auto const content = tree_blob.ReadContent()) {
+                    tree = BazelMsgFactory::MessageFromString<bazel_re::Tree>(
+                        *content);
+                }
                 if (not tree) {
                     return fmt::format(
                         "BazelResponse: failed to create Tree for {}",