Close compaction/publish TOCTOU with conditional-PUT on the generation head

crates/git-cache-domain/src/materializer.rs

@@ -13,9 +13,10 @@ use git_cache_disk::RepoLock;
 pub use git_cache_git::UploadPackProcess;
 use git_cache_objectstore::{
     generation_manifest_key, generation_manifest_prefix, pack_key, read_commit_manifest,
-    read_generation_manifest, read_json, read_repo_generation_head, write_commit_manifest,
-    write_json, write_ref_manifest, write_repo_generation_head, GenerationPublish,
-    PublishManifests,
+    read_generation_manifest, read_json, read_repo_generation_head,
+    read_repo_generation_head_versioned, write_commit_manifest, write_json, write_ref_manifest,
+    write_repo_generation_head_if_version_matches, GenerationPublish,
+    ObjectVersion, PublishManifests,
 };
 use serde::Serialize;
 use std::collections::{HashMap, HashSet};

crates/git-cache-domain/src/materializer/generations.rs

@@ -7,6 +7,7 @@ use tokio::task::JoinSet;
 
 const MAX_GENERATION_MANIFEST_SCAN_KEYS: usize = 10_000;
 const HYDRATE_PACK_DOWNLOAD_CONCURRENCY: usize = 4;
+const HEAD_CAS_MAX_ATTEMPTS: usize = 5;
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub struct CompactionReport {
@@ -586,14 +587,31 @@ impl Materializer {
                 .await?;
         }
 
-        let current_head = self.manifests().repo_head(repo).await?;
-        if current_head
-            .as_ref()
-            .map(|current| current.updated_at <= head.updated_at)
-            .unwrap_or(true)
-        {
-            self.manifests().write_repo_head(head).await?;
+        for _ in 0..HEAD_CAS_MAX_ATTEMPTS {
+            let current = self.manifests().repo_head_versioned(repo).await?;
+            let (current_head, version) = match &current {
+                Some((head, version)) => (Some(head), Some(version)),
+                None => (None, None),
+            };
+            if current_head
+                .map(|current| current.updated_at > head.updated_at)
+                .unwrap_or(false)
+            {
+                return Ok(());
+            }
+            if self
+                .manifests()
+                .write_repo_head_if_version_matches(head, version)
+                .await?
+            {
+                return Ok(());
+            }
         }
+        warn!(
+            %repo,
+            generation = %head.generation,
+            "generation head moved repeatedly during publish; leaving newer head in place"
+        );
         Ok(())
     }
 
@@ -623,7 +641,12 @@ impl Materializer {
         threshold: usize,
         dry_run: bool,
     ) -> CoreResult<Option<CompactionReport>> {
-        let Some(head) = self.manifests().repo_head(repo).await? else {
+        let outer_repo_lock = if dry_run {
+            None
+        } else {
+            Some(self.lock_repo(repo).await?)
+        };
+        let Some((head, head_version)) = self.manifests().repo_head_versioned(repo).await? else {
             return Ok(None);
         };
         let Some(head_manifest) = self.get_generation_manifest(repo, head.generation).await? else {
@@ -660,8 +683,8 @@ impl Materializer {
             }));
         }
 
+        let _repo_lock = outer_repo_lock;
         let repo_dir = self.ensure_repo_dir(repo).await?;
-        let _repo_lock = self.lock_repo(repo).await?;
         Box::pin(self.hydrate_generation(repo, &repo_dir, head.generation)).await?;
         self.state.git.repack_for_serving(&repo_dir).await?;
 
@@ -741,7 +764,18 @@ impl Materializer {
             };
             self.manifests().write_commit(&manifest).await?;
         }
-        self.manifests().write_repo_head(&new_head).await?;
+        if !self
+            .manifests()
+            .write_repo_head_if_version_matches(&new_head, Some(&head_version))
+            .await?
+        {
+            warn!(
+                %repo,
+                %new_generation,
+                "generation head changed during compaction; skipping cleanup of old packs"
+            );
+            return Ok(None);
+        }
 
         let retained_keys: HashSet<&str> = packs.iter().map(|pack| pack.key.as_str()).collect();
         let mut bytes_reclaimed = 0_u64;

crates/git-cache-domain/src/materializer/manifests.rs

@@ -63,8 +63,19 @@ impl<'a> ManifestStore<'a> {
         read_repo_generation_head(&*self.state.store, repo).await
     }
 
-    pub(super) async fn write_repo_head(&self, head: &RepoGenerationHead) -> CoreResult<()> {
-        write_repo_generation_head(&*self.state.store, head).await
+    pub(super) async fn repo_head_versioned(
+        &self,
+        repo: &RepoKey,
+    ) -> CoreResult<Option<(RepoGenerationHead, ObjectVersion)>> {
+        read_repo_generation_head_versioned(&*self.state.store, repo).await
+    }
+
+    pub(super) async fn write_repo_head_if_version_matches(
+        &self,
+        head: &RepoGenerationHead,
+        version: Option<&ObjectVersion>,
+    ) -> CoreResult<bool> {
+        write_repo_generation_head_if_version_matches(&*self.state.store, head, version).await
     }
 }
 

crates/git-cache-domain/src/materializer/tests/generation_tests.rs

@@ -1,5 +1,6 @@
 mod tests {
     use super::super::*;
+    use git_cache_objectstore::write_repo_generation_head;
 
     #[tokio::test]
     async fn publish_generation_links_delta_to_previous_generation() {

crates/git-cache-objectstore/src/lib.rs

@@ -14,18 +14,37 @@ pub use local::LocalObjectStore;
 pub use manifests::{
     acquire_lease, commit_manifest_key, generation_manifest_key, generation_manifest_prefix,
     lease_key, pack_key, read_commit_manifest, read_generation_manifest, read_json, read_lease,
-    read_ref_manifest, read_repo_generation_head, ref_manifest_key, repo_generation_head_key,
+    read_ref_manifest, read_repo_generation_head, read_repo_generation_head_versioned,
+    ref_manifest_key, repo_generation_head_key,
     write_commit_manifest, write_commit_manifest_if_absent,
     write_commit_manifest_if_absent_or_matches, write_generation_manifest,
     write_generation_manifest_if_absent, write_generation_manifest_if_absent_or_matches,
     write_json, write_json_if_absent, write_json_if_absent_or_matches, write_ref_manifest,
     write_ref_manifest_if_absent, write_ref_manifest_if_absent_or_matches,
-    write_repo_generation_head, GenerationPublish, LeaseManifest, PublishManifests,
+    write_repo_generation_head, write_repo_generation_head_if_version_matches, GenerationPublish,
+    LeaseManifest, PublishManifests,
 };
 
 #[cfg(feature = "s3")]
 pub use s3::S3ObjectStore;
 
+/// Opaque version token returned by `get_versioned` and consumed by
+/// `put_if_version_matches`. Backends choose the representation (S3 uses
+/// the object ETag, the local store uses a content digest); callers must
+/// treat it as opaque and only pass it back to the same store.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ObjectVersion(String);
+
+impl ObjectVersion {
+    pub(crate) fn new(token: impl Into<String>) -> Self {
+        Self(token.into())
+    }
+
+    pub(crate) fn token(&self) -> &str {
+        &self.0
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ObjectMeta {
     pub key: String,
@@ -38,6 +57,19 @@ pub trait ObjectStore: Send + Sync {
     async fn get(&self, key: &str) -> Result<Option<Bytes>>;
     async fn put(&self, key: &str, value: Bytes) -> Result<()>;
     async fn put_if_absent(&self, key: &str, value: Bytes) -> Result<bool>;
+
+    /// Fetch an object together with an opaque version token for a later
+    /// compare-and-swap via `put_if_version_matches`.
+    async fn get_versioned(&self, key: &str) -> Result<Option<(Bytes, ObjectVersion)>>;
+
+    /// Replace `key` only if the stored object still matches `version`.
+    /// Returns `Ok(false)` when the object changed or no longer exists.
+    async fn put_if_version_matches(
+        &self,
+        key: &str,
+        value: Bytes,
+        version: &ObjectVersion,
+    ) -> Result<bool>;
     async fn exists(&self, key: &str) -> Result<bool>;
     async fn delete(&self, key: &str) -> Result<()>;
     async fn list_prefix(&self, prefix: &str, max_keys: Option<usize>) -> Result<Vec<String>>;

crates/git-cache-objectstore/src/local.rs

@@ -1,9 +1,11 @@
-use crate::{validate_key, ObjectMeta, ObjectStore};
+use crate::{validate_key, ObjectMeta, ObjectStore, ObjectVersion};
 use async_trait::async_trait;
 use bytes::Bytes;
 use chrono::{DateTime, Utc};
 use git_cache_core::{GitCacheError, Result};
+use std::hash::{Hash, Hasher};
 use std::path::{Path, PathBuf};
+use std::sync::OnceLock;
 use std::time::{SystemTime, UNIX_EPOCH};
 use tokio::fs;
 use tokio::io::AsyncWriteExt;
@@ -80,6 +82,32 @@ impl ObjectStore for LocalObjectStore {
         }
     }
 
+    async fn get_versioned(&self, key: &str) -> Result<Option<(Bytes, ObjectVersion)>> {
+        let _guard = cas_lock().lock().await;
+        let Some(value) = self.get(key).await? else {
+            return Ok(None);
+        };
+        let version = content_version(&value);
+        Ok(Some((value, version)))
+    }
+
+    async fn put_if_version_matches(
+        &self,
+        key: &str,
+        value: Bytes,
+        version: &ObjectVersion,
+    ) -> Result<bool> {
+        let _guard = cas_lock().lock().await;
+        let Some(current) = self.get(key).await? else {
+            return Ok(false);
+        };
+        if content_version(&current) != *version {
+            return Ok(false);
+        }
+        self.put(key, value).await?;
+        Ok(true)
+    }
+
     async fn exists(&self, key: &str) -> Result<bool> {
         let path = self.object_path(key)?;
         match fs::metadata(path).await {
@@ -187,6 +215,20 @@ impl ObjectStore for LocalObjectStore {
     }
 }
 
+/// Serializes local read-compare-write sequences process-wide. The local
+/// store is single-node, so an in-process lock is sufficient to make
+/// `put_if_version_matches` atomic with respect to other CAS callers.
+fn cas_lock() -> &'static tokio::sync::Mutex<()> {
+    static LOCK: OnceLock<tokio::sync::Mutex<()>> = OnceLock::new();
+    LOCK.get_or_init(|| tokio::sync::Mutex::new(()))
+}
+
+fn content_version(value: &Bytes) -> ObjectVersion {
+    let mut hasher = std::hash::DefaultHasher::new();
+    value.as_ref().hash(&mut hasher);
+    ObjectVersion::new(format!("{}-{:016x}", value.len(), hasher.finish()))
+}
+
 fn allocate_temp_path(parent: &Path, final_path: &Path) -> Result<PathBuf> {
     let file_name = final_path
         .file_name()

crates/git-cache-objectstore/src/manifests.rs

@@ -1,4 +1,4 @@
-use crate::{validate_key, ObjectStore};
+use crate::{validate_key, ObjectStore, ObjectVersion};
 use bytes::Bytes;
 use chrono::{DateTime, Duration, Utc};
 use git_cache_core::{
@@ -204,6 +204,42 @@ where
     write_json(store, &repo_generation_head_key(&head.repo), head).await
 }
 
+pub async fn read_repo_generation_head_versioned<S>(
+    store: &S,
+    repo: &RepoKey,
+) -> Result<Option<(RepoGenerationHead, ObjectVersion)>>
+where
+    S: ObjectStore + ?Sized,
+{
+    let Some((value, version)) = store
+        .get_versioned(&repo_generation_head_key(repo))
+        .await?
+    else {
+        return Ok(None);
+    };
+    Ok(Some((serde_json::from_slice(&value)?, version)))
+}
+
+/// Compare-and-swap write of the generation head pointer. `version` is the
+/// token from `read_repo_generation_head_versioned`; `None` means the head
+/// is expected to be absent (first write). Returns `Ok(false)` when the
+/// stored head no longer matches the expectation.
+pub async fn write_repo_generation_head_if_version_matches<S>(
+    store: &S,
+    head: &RepoGenerationHead,
+    version: Option<&ObjectVersion>,
+) -> Result<bool>
+where
+    S: ObjectStore + ?Sized,
+{
+    let key = repo_generation_head_key(&head.repo);
+    let bytes = json_bytes(head)?;
+    match version {
+        Some(version) => store.put_if_version_matches(&key, bytes, version).await,
+        None => store.put_if_absent(&key, bytes).await,
+    }
+}
+
 pub async fn write_generation_manifest<S>(store: &S, manifest: &GenerationManifest) -> Result<()>
 where
     S: ObjectStore + ?Sized,

crates/git-cache-objectstore/src/s3.rs

@@ -1,4 +1,4 @@
-use crate::{validate_key, ObjectMeta, ObjectStore};
+use crate::{validate_key, ObjectMeta, ObjectStore, ObjectVersion};
 use async_trait::async_trait;
 use aws_sdk_s3::primitives::ByteStream;
 use aws_sdk_s3::types::{CompletedMultipartUpload, CompletedPart};
@@ -247,6 +247,60 @@ impl ObjectStore for S3ObjectStore {
         }
     }
 
+    async fn get_versioned(&self, key: &str) -> Result<Option<(Bytes, ObjectVersion)>> {
+        let s3_key = self.s3_key(key)?;
+        let output = match self
+            .client
+            .get_object()
+            .bucket(&self.bucket)
+            .key(&s3_key)
+            .send()
+            .await
+        {
+            Ok(output) => output,
+            Err(err) if is_not_found(&err) => return Ok(None),
+            Err(err) => return Err(s3_error("get_versioned", &s3_key, err)),
+        };
+
+        let e_tag = output
+            .e_tag()
+            .ok_or_else(|| {
+                GitCacheError::UpstreamUnavailable(format!(
+                    "s3 get `{s3_key}` returned no etag"
+                ))
+            })?
+            .to_string();
+        let body = output
+            .body
+            .collect()
+            .await
+            .map_err(|err| s3_error("read body", &s3_key, err))?;
+        Ok(Some((body.into_bytes(), ObjectVersion::new(e_tag))))
+    }
+
+    async fn put_if_version_matches(
+        &self,
+        key: &str,
+        value: Bytes,
+        version: &ObjectVersion,
+    ) -> Result<bool> {
+        let s3_key = self.s3_key(key)?;
+        match self
+            .client
+            .put_object()
+            .bucket(&self.bucket)
+            .key(&s3_key)
+            .if_match(version.token())
+            .body(ByteStream::new(value.into()))
+            .send()
+            .await
+        {
+            Ok(_) => Ok(true),
+            Err(err) if is_precondition_failed(&err) || is_not_found(&err) => Ok(false),
+            Err(err) => Err(s3_error("put_if_version_matches", &s3_key, err)),
+        }
+    }
+
     async fn exists(&self, key: &str) -> Result<bool> {
         let s3_key = self.s3_key(key)?;
         match self