Ambiguous digest problems with existing blob objects

crates/spfs/src/check.rs

@@ -230,7 +230,11 @@ where
         &self,
         partial: encoding::PartialDigest,
     ) -> Result<CheckItemResult> {
-        match self.repo.resolve_full_digest(&partial).await? {
+        match self
+            .repo
+            .resolve_full_digest(&partial, graph::PartialDigestType::Unknown)
+            .await?
+        {
             FoundDigest::Object(digest) => self
                 .check_object_digest(digest)
                 .await

crates/spfs/src/check_test.rs

@@ -210,6 +210,8 @@ impl CheckReporter for &DebugReporter {
 /// The check should complete successfully and report a missing payload.
 #[rstest]
 #[tokio::test]
+// This test just needs the config to not change while it is running.
+#[serial_test::serial(config)]
 async fn check_missing_annotation_payload(#[future] tmprepo: TempRepo) {
     init_logging();
     let tmprepo = tmprepo.await;

crates/spfs/src/graph/database.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // https://github.com/spkenv/spk
 
-use std::collections::VecDeque;
+use std::collections::{HashMap, VecDeque};
 use std::pin::Pin;
 use std::task::Poll;
 
@@ -160,7 +160,7 @@ pub enum DigestSearchCriteria {
 }
 
 /// The types of digests that can exist in a database.
-#[derive(PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum FoundDigest {
     Object(encoding::Digest),
     Payload(encoding::Digest),
@@ -195,6 +195,30 @@ impl FoundDigest {
     }
 }
 
+/// The types of items a partial digest can reference.
+#[derive(Clone, Copy)]
+pub enum PartialDigestType {
+    Object,
+    Payload,
+    Unknown,
+}
+
+impl PartialDigestType {
+    /// Return true if the `FoundDigest` is a different item type.
+    ///
+    /// Unknown always return false.
+    #[inline]
+    pub fn conflicts_with(&self, fd: &FoundDigest) -> bool {
+        match (self, fd) {
+            (PartialDigestType::Unknown, _)
+            | (PartialDigestType::Object, FoundDigest::Object(_))
+            | (PartialDigestType::Payload, FoundDigest::Payload(_)) => false,
+            (PartialDigestType::Object, FoundDigest::Payload(_))
+            | (PartialDigestType::Payload, FoundDigest::Object(_)) => true,
+        }
+    }
+}
+
 /// A read-only object database.
 #[async_trait::async_trait]
 pub trait DatabaseView: Sync + Send {
@@ -256,23 +280,65 @@ pub trait DatabaseView: Sync + Send {
 
     /// Resolve the complete item digest from a shortened one.
     ///
+    /// The type of item expected can be specified with `partial_digest_type`-
+    /// If `PartialDigestType::Unknown` is specified and both an object and
+    /// payload are found with the same digest, this will resolve to the payload
+    /// instead of fail with `AmbiguousReferenceError`, for interoperability
+    /// with repos containing legacy blob object files. Otherwise the type is
+    /// used to disambiguate in the unlikely case a non-blob object and payload
+    /// have the same digest.
+    ///
     /// By default this is an O(n) operation defined by the number of items.
     /// Other implementations may provide better results.
     ///
     /// # Errors
     /// - UnknownReferenceError: if the digest cannot be resolved
     /// - AmbiguousReferenceError: if the digest could point to multiple items
-    async fn resolve_full_digest(&self, partial: &encoding::PartialDigest) -> Result<FoundDigest> {
-        let options: Vec<_> = self
-            .find_digests(&crate::graph::DigestSearchCriteria::StartsWith(
-                partial.clone(),
-            ))
-            .try_collect()
-            .await?;
+    async fn resolve_full_digest(
+        &self,
+        partial: &encoding::PartialDigest,
+        partial_digest_type: PartialDigestType,
+    ) -> Result<FoundDigest> {
+        #[derive(Debug)]
+        struct UpgradeToPayload(FoundDigest);
+
+        impl UpgradeToPayload {
+            /// Replace the FoundDigest if fd is a Payload.
+            ///
+            /// If we observe both a legacy blob object and a payload with the
+            /// same digest we will only remember seeing the payload. This
+            /// method assumes it is only called with a FoundDigest that has the
+            /// same digest as the present one.
+            fn upgrade(&mut self, fd: FoundDigest) {
+                if matches!(fd, FoundDigest::Payload(_)) {
+                    self.0 = fd;
+                }
+            }
+        }
+
+        let mut options = HashMap::<_, UpgradeToPayload>::new();
+        let filter = crate::graph::DigestSearchCriteria::StartsWith(partial.clone());
+        let mut stream = self.find_digests(&filter);
+        while let Some(fd) = stream.try_next().await? {
+            if partial_digest_type.conflicts_with(&fd) {
+                continue;
+            }
+
+            // Hash on the raw digest to avoid double-counting legacy blobs and
+            // their payloads.
+            options
+                .entry(*fd.digest())
+                .and_modify(|utp| utp.upgrade(fd))
+                .or_insert_with(|| UpgradeToPayload(fd));
+        }
 
         match options.len() {
             0 => Err(Error::UnknownReference(partial.to_string())),
-            1 => Ok(options.into_iter().next().unwrap()),
+            1 => Ok(options
+                .into_iter()
+                .next()
+                .map(|(_, UpgradeToPayload(fd))| fd)
+                .unwrap()),
             _ => Err(Error::AmbiguousReference(partial.to_string())),
         }
     }
@@ -338,12 +404,34 @@ impl<T: Database> Database for &T {
 #[async_trait::async_trait]
 pub trait DatabaseExt: Send + Sync {
     /// Write an object to the database, for later retrieval.
+    ///
+    /// It is not permitted to write blob objects.
     async fn write_object<T: ObjectProto>(&self, obj: &FlatObject<T>) -> Result<()>;
+
+    /// Write an object to the database, for later retrieval.
+    ///
+    /// # Safety
+    ///
+    /// This function does not check the type of the object being written. It
+    /// is expected that the caller will not write a blob object except in
+    /// specific cases, such as in test code.
+    async unsafe fn write_object_unchecked<T: ObjectProto>(
+        &self,
+        obj: &FlatObject<T>,
+    ) -> Result<()>;
 }
 
 #[async_trait::async_trait]
 impl<T: DatabaseExt + Send + Sync> DatabaseExt for &T {
     async fn write_object<O: ObjectProto>(&self, obj: &FlatObject<O>) -> Result<()> {
         DatabaseExt::write_object(&**self, obj).await
     }
+
+    async unsafe fn write_object_unchecked<O: ObjectProto>(
+        &self,
+        obj: &FlatObject<O>,
+    ) -> Result<()> {
+        // Safety: transitive unsafe call
+        unsafe { DatabaseExt::write_object_unchecked(&**self, obj).await }
+    }
 }

crates/spfs/src/graph/mod.rs

@@ -34,6 +34,7 @@ pub use database::{
     DatabaseWalker,
     DigestSearchCriteria,
     FoundDigest,
+    PartialDigestType,
 };
 pub use entry::Entry;
 pub use kind::{HasKind, Kind, ObjectKind};

crates/spfs/src/proto/defs/database.proto

@@ -84,6 +84,16 @@ message WriteObjectResponse{
     }
 }
 
+message WriteObjectUncheckedRequest{
+    Object object = 1;
+}
+message WriteObjectUncheckedResponse{
+    oneof result {
+        Error error = 1;
+        Ok ok = 2;
+    }
+}
+
 message RemoveObjectRequest{
     Digest digest = 1;
 }
@@ -112,6 +122,7 @@ service DatabaseService {
     rpc IterObjects(IterObjectsRequest) returns (stream IterObjectsResponse);
     rpc WalkObjects(WalkObjectsRequest) returns (stream WalkObjectsResponse);
     rpc WriteObject(WriteObjectRequest) returns (WriteObjectResponse);
+    rpc WriteObjectUnchecked(WriteObjectUncheckedRequest) returns (WriteObjectUncheckedResponse);
     rpc RemoveObject(RemoveObjectRequest) returns (RemoveObjectResponse);
     rpc RemoveObjectIfOlderThan(RemoveObjectIfOlderThanRequest) returns (RemoveObjectIfOlderThanResponse);
 }

crates/spfs/src/proto/result.rs

@@ -120,6 +120,10 @@ rpc_result!(
     g::walk_objects_response::WalkObjectsItem
 );
 rpc_result!(g::WriteObjectResponse, g::write_object_response::Result);
+rpc_result!(
+    g::WriteObjectUncheckedResponse,
+    g::write_object_unchecked_response::Result
+);
 rpc_result!(g::RemoveObjectResponse, g::remove_object_response::Result);
 rpc_result!(
     g::RemoveObjectIfOlderThanResponse,

crates/spfs/src/resolve.rs

@@ -149,16 +149,20 @@ pub async fn compute_environment_manifest(
         .filter_map(|i| match i {
             tracking::EnvSpecItem::Digest(d) => Some(std::future::ready(Ok(*d)).boxed()),
             tracking::EnvSpecItem::PartialDigest(p) => Some(
-                repo.resolve_full_digest(p)
-                    .and_then(|found_digest| async move {
-                        match found_digest {
-                            graph::FoundDigest::Object(digest) => Ok(digest),
-                            graph::FoundDigest::Payload(_digest) => {
-                                Err("unexpected payload digest in environment spec".into())
-                            }
+                repo.resolve_full_digest(
+                    p,
+                    // by the error below, we only expect to find object digests
+                    graph::PartialDigestType::Object,
+                )
+                .and_then(|found_digest| async move {
+                    match found_digest {
+                        graph::FoundDigest::Object(digest) => Ok(digest),
+                        graph::FoundDigest::Payload(_digest) => {
+                            Err("unexpected payload digest in environment spec".into())
                         }
-                    })
-                    .boxed(),
+                    }
+                })
+                .boxed(),
             ),
             tracking::EnvSpecItem::TagSpec(t) => {
                 Some(repo.resolve_tag(t).map_ok(|t| t.target).boxed())

crates/spfs/src/server/database.rs

@@ -104,6 +104,19 @@ impl proto::database_service_server::DatabaseService for DatabaseService {
         Ok(Response::new(result))
     }
 
+    async fn write_object_unchecked(
+        &self,
+        request: Request<proto::WriteObjectUncheckedRequest>,
+    ) -> Result<Response<proto::WriteObjectUncheckedResponse>, Status> {
+        let request = request.into_inner();
+        let object = proto::handle_error!(request.object.try_into());
+        // Safety: in spirit this generated trait method would be unsafe, and
+        // this would be a transitive unsafe call.
+        unsafe { proto::handle_error!(self.repo.write_object_unchecked(&object).await) };
+        let result = proto::WriteObjectUncheckedResponse::ok(proto::Ok {});
+        Ok(Response::new(result))
+    }
+
     async fn remove_object(
         &self,
         request: Request<proto::RemoveObjectRequest>,

crates/spfs/src/storage/database_test.rs

@@ -6,6 +6,7 @@ use std::borrow::Cow;
 
 use rstest::rstest;
 
+use crate::encoding::PartialDigest;
 use crate::fixtures::*;
 use crate::graph;
 use crate::prelude::*;
@@ -40,3 +41,105 @@ async fn test_object_existence(
     let actual = tmprepo.has_object(digest).await;
     assert!(!actual, "object should not exist after being removed");
 }
+
+#[rstest]
+#[case::fs(tmprepo("fs"))]
+#[case::tar(tmprepo("tar"))]
+#[cfg_attr(feature = "server", case::rpc(tmprepo("rpc")))]
+#[tokio::test]
+async fn resolve_partial_digest_with_blob_object_file_present(
+    #[case]
+    #[future]
+    tmprepo: TempRepo,
+) {
+    let tmprepo = tmprepo.await;
+
+    let test_data = b"test data\n";
+
+    let payload = tmprepo
+        .commit_payload(Box::pin(test_data.as_slice()))
+        .await
+        .expect("failed to commit payload data");
+
+    let partial_digest = PartialDigest::from(&payload.as_bytes()[..8]);
+
+    // First test baseline behavior without blob object file present
+
+    // PartialDigestType::Unknown
+    {
+        let resolved = tmprepo
+            .resolve_full_digest(&partial_digest, graph::PartialDigestType::Unknown)
+            .await
+            .expect("failed to resolve partial digest");
+
+        assert_eq!(*resolved.digest(), payload);
+        assert!(matches!(resolved, graph::FoundDigest::Payload(_)));
+    }
+
+    // PartialDigestType::Payload
+    {
+        let resolved = tmprepo
+            .resolve_full_digest(&partial_digest, graph::PartialDigestType::Payload)
+            .await
+            .expect("failed to resolve partial digest");
+
+        assert_eq!(*resolved.digest(), payload);
+        assert!(matches!(resolved, graph::FoundDigest::Payload(_)));
+    }
+
+    // PartialDigestType::Object
+    {
+        let _ = tmprepo
+            .resolve_full_digest(&partial_digest, graph::PartialDigestType::Object)
+            .await
+            .expect_err("no object with this digest should exist");
+    }
+
+    // Write the blob object file
+    {
+        let blob = graph::Blob::new(payload, test_data.len() as u64);
+
+        // Safety: we are writing a blob object for the purposes of this test
+        unsafe {
+            tmprepo
+                .write_object_unchecked(&blob)
+                .await
+                .expect("failed to write blob object data");
+        }
+    }
+
+    // Then test behavior with blob object file present
+
+    // PartialDigestType::Unknown
+    {
+        let resolved = tmprepo
+            .resolve_full_digest(&partial_digest, graph::PartialDigestType::Unknown)
+            .await
+            .expect("failed to resolve partial digest");
+
+        assert_eq!(*resolved.digest(), payload);
+        assert!(matches!(resolved, graph::FoundDigest::Payload(_)));
+    }
+
+    // PartialDigestType::Payload
+    {
+        let resolved = tmprepo
+            .resolve_full_digest(&partial_digest, graph::PartialDigestType::Payload)
+            .await
+            .expect("failed to resolve partial digest");
+
+        assert_eq!(*resolved.digest(), payload);
+        assert!(matches!(resolved, graph::FoundDigest::Payload(_)));
+    }
+
+    // PartialDigestType::Object
+    {
+        let resolved = tmprepo
+            .resolve_full_digest(&partial_digest, graph::PartialDigestType::Object)
+            .await
+            .expect("failed to resolve partial digest");
+
+        assert_eq!(*resolved.digest(), payload);
+        assert!(matches!(resolved, graph::FoundDigest::Object(_)));
+    }
+}