Public dry run dryrun-750f29790fe2

Source commit: 750f29790fe2046599fbe5b5d06fdb63d92fabff

Public tree identity: sha256:02326e08850bb287585789733ea5f3f153e1f7f3fd88afcc24e249107df91506

crates/disasmer-core/src/artifact.rs

@@ -0,0 +1,939 @@
+use std::collections::{BTreeMap, BTreeSet};
+
+use serde::{Deserialize, Serialize};
+use thiserror::Error;
+
+use crate::{
+    auth::same_tenant_project, Actor, ArtifactId, AuthContext, Digest, LimitError, LimitKind,
+    NodeId, ProcessId, ProjectId, ResourceLimits, ResourceMeter, Scope, TaskId, TenantId,
+};
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum StorageLocation {
+    RetainedNode(NodeId),
+    ExplicitStore(String),
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct RetentionPolicy {
+    pub best_effort_node_retention: bool,
+    pub max_download_bytes: u64,
+}
+
+impl Default for RetentionPolicy {
+    fn default() -> Self {
+        Self {
+            best_effort_node_retention: true,
+            max_download_bytes: 256 * 1024 * 1024,
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct DownloadPolicy {
+    pub max_bytes: u64,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct DownloadAction {
+    pub artifact: ArtifactId,
+    pub source: StorageLocation,
+    pub scoped_token_subject: String,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct DownloadLink {
+    pub artifact: ArtifactId,
+    pub source: StorageLocation,
+    pub url_path: String,
+    pub scoped_token_digest: Digest,
+    pub expires_at_epoch_seconds: u64,
+    pub tenant: TenantId,
+    pub project: ProjectId,
+    pub process: ProcessId,
+    pub actor: Actor,
+    pub max_bytes: u64,
+    pub policy_context_digest: Digest,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct IssuedDownloadLink {
+    pub link: DownloadLink,
+    pub revoked: bool,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ArtifactDownloadStream {
+    pub link: DownloadLink,
+    pub streamed_bytes: u64,
+}
+
+#[derive(Clone, Debug, Error, PartialEq, Eq)]
+pub enum DownloadError {
+    #[error("artifact does not exist")]
+    NotFound,
+    #[error("artifact is unavailable from current retention or explicit storage")]
+    Unavailable,
+    #[error("artifact download direct connectivity unavailable: {0}")]
+    DirectConnectivityUnavailable(String),
+    #[error("artifact download denied: {0}")]
+    Unauthorized(String),
+    #[error("artifact size {size} exceeds download limit {limit}")]
+    LimitExceeded { size: u64, limit: u64 },
+    #[error("download link token is invalid for this scoped artifact link")]
+    InvalidToken,
+    #[error("download link has expired")]
+    Expired,
+    #[error("download link has been revoked")]
+    Revoked,
+    #[error("download usage limit failed: {0}")]
+    Usage(String),
+}
+
+#[derive(Clone, Debug, Error, PartialEq, Eq)]
+#[error("artifact is unavailable because node-local unsynced bytes were lost")]
+pub struct ArtifactUnavailable;
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ArtifactMetadata {
+    pub id: ArtifactId,
+    pub tenant: TenantId,
+    pub project: ProjectId,
+    pub process: ProcessId,
+    pub producer_task: TaskId,
+    pub producer_node: NodeId,
+    pub digest: Digest,
+    pub size: u64,
+    pub flushed_epoch: u64,
+    pub retaining_nodes: BTreeSet<NodeId>,
+    pub explicit_locations: Vec<String>,
+    pub coordinator_has_large_bytes: bool,
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct ArtifactRegistry {
+    artifacts: BTreeMap<ArtifactId, ArtifactMetadata>,
+    issued_download_links: BTreeMap<Digest, IssuedDownloadLink>,
+    next_epoch: u64,
+}
+
+impl ArtifactRegistry {
+    pub fn flush_metadata(
+        &mut self,
+        id: ArtifactId,
+        tenant: TenantId,
+        project: ProjectId,
+        process: ProcessId,
+        producer_task: TaskId,
+        retaining_node: NodeId,
+        digest: Digest,
+        size: u64,
+    ) -> ArtifactMetadata {
+        self.next_epoch += 1;
+        let metadata = ArtifactMetadata {
+            id: id.clone(),
+            tenant,
+            project,
+            process,
+            producer_task,
+            producer_node: retaining_node.clone(),
+            digest,
+            size,
+            flushed_epoch: self.next_epoch,
+            retaining_nodes: BTreeSet::from([retaining_node]),
+            explicit_locations: Vec::new(),
+            coordinator_has_large_bytes: false,
+        };
+        self.artifacts.insert(id, metadata.clone());
+        metadata
+    }
+
+    pub fn sync_to_explicit_store(
+        &mut self,
+        artifact: &ArtifactId,
+        location: impl Into<String>,
+    ) -> Result<(), ArtifactUnavailable> {
+        let metadata = self
+            .artifacts
+            .get_mut(artifact)
+            .ok_or(ArtifactUnavailable)?;
+        metadata.explicit_locations.push(location.into());
+        Ok(())
+    }
+
+    pub fn garbage_collect_node(&mut self, node: &NodeId) {
+        for metadata in self.artifacts.values_mut() {
+            metadata.retaining_nodes.remove(node);
+        }
+    }
+
+    pub fn metadata(&self, artifact: &ArtifactId) -> Option<&ArtifactMetadata> {
+        self.artifacts.get(artifact)
+    }
+
+    pub fn download_action(
+        &self,
+        context: &AuthContext,
+        artifact: &ArtifactId,
+        policy: &DownloadPolicy,
+    ) -> Result<DownloadAction, DownloadError> {
+        let metadata = self
+            .artifacts
+            .get(artifact)
+            .ok_or(DownloadError::NotFound)?;
+        let scope = Scope {
+            tenant: metadata.tenant.clone(),
+            project: metadata.project.clone(),
+            process: Some(metadata.process.clone()),
+            task: Some(metadata.producer_task.clone()),
+            node: None,
+            artifact: Some(metadata.id.clone()),
+        };
+        let authz = same_tenant_project(context, &scope);
+        if !authz.allowed {
+            return Err(DownloadError::Unauthorized(authz.reason));
+        }
+        if metadata.size > policy.max_bytes {
+            return Err(DownloadError::LimitExceeded {
+                size: metadata.size,
+                limit: policy.max_bytes,
+            });
+        }
+        let source = metadata
+            .retaining_nodes
+            .iter()
+            .next()
+            .cloned()
+            .map(StorageLocation::RetainedNode)
+            .or_else(|| {
+                metadata
+                    .explicit_locations
+                    .first()
+                    .cloned()
+                    .map(StorageLocation::ExplicitStore)
+            })
+            .ok_or(DownloadError::Unavailable)?;
+
+        Ok(DownloadAction {
+            artifact: artifact.clone(),
+            source,
+            scoped_token_subject: format!(
+                "{}/{}/{}/{}",
+                metadata.tenant, metadata.project, metadata.process, metadata.id
+            ),
+        })
+    }
+
+    pub fn downloadable_size(
+        &self,
+        context: &AuthContext,
+        artifact: &ArtifactId,
+        policy: &DownloadPolicy,
+    ) -> Result<u64, DownloadError> {
+        self.download_action(context, artifact, policy)?;
+        let metadata = self
+            .artifacts
+            .get(artifact)
+            .ok_or(DownloadError::NotFound)?;
+        Ok(metadata.size)
+    }
+
+    pub fn create_download_link(
+        &mut self,
+        context: &AuthContext,
+        artifact: &ArtifactId,
+        policy: &DownloadPolicy,
+        token_nonce: &str,
+        now_epoch_seconds: u64,
+        ttl_seconds: u64,
+    ) -> Result<DownloadLink, DownloadError> {
+        let action = self.download_action(context, artifact, policy)?;
+        let metadata = self
+            .artifacts
+            .get(artifact)
+            .ok_or(DownloadError::NotFound)?;
+        let expires_at_epoch_seconds = now_epoch_seconds.saturating_add(ttl_seconds);
+        let policy_context_digest =
+            download_policy_context_digest(metadata, &action.source, policy);
+        let scoped_token_digest = Digest::from_parts([
+            b"artifact-download-token:v2".as_slice(),
+            action.scoped_token_subject.as_bytes(),
+            actor_subject(&context.actor).as_bytes(),
+            token_nonce.as_bytes(),
+            metadata.digest.as_str().as_bytes(),
+            metadata.size.to_string().as_bytes(),
+            policy_context_digest.as_str().as_bytes(),
+            expires_at_epoch_seconds.to_string().as_bytes(),
+        ]);
+        let link = DownloadLink {
+            artifact: artifact.clone(),
+            source: action.source,
+            url_path: format!(
+                "/artifacts/{}/{}/{}/{}",
+                metadata.tenant, metadata.project, metadata.process, metadata.id
+            ),
+            scoped_token_digest,
+            expires_at_epoch_seconds,
+            tenant: metadata.tenant.clone(),
+            project: metadata.project.clone(),
+            process: metadata.process.clone(),
+            actor: context.actor.clone(),
+            max_bytes: policy.max_bytes,
+            policy_context_digest,
+        };
+        self.issued_download_links.insert(
+            link.scoped_token_digest.clone(),
+            IssuedDownloadLink {
+                link: link.clone(),
+                revoked: false,
+            },
+        );
+        Ok(link)
+    }
+
+    pub fn revoke_download_link(
+        &mut self,
+        context: &AuthContext,
+        artifact: &ArtifactId,
+        presented_token_digest: &Digest,
+    ) -> Result<DownloadLink, DownloadError> {
+        let issued = self
+            .issued_download_links
+            .get(presented_token_digest)
+            .ok_or(DownloadError::InvalidToken)?;
+        if issued.link.artifact != *artifact || issued.link.actor != context.actor {
+            return Err(DownloadError::InvalidToken);
+        }
+        self.download_action(
+            context,
+            artifact,
+            &DownloadPolicy {
+                max_bytes: issued.link.max_bytes,
+            },
+        )?;
+
+        let issued = self
+            .issued_download_links
+            .get_mut(presented_token_digest)
+            .ok_or(DownloadError::InvalidToken)?;
+        issued.revoked = true;
+        Ok(issued.link.clone())
+    }
+
+    pub fn open_download_stream(
+        &self,
+        context: &AuthContext,
+        artifact: &ArtifactId,
+        policy: &DownloadPolicy,
+        presented_token_digest: &Digest,
+        now_epoch_seconds: u64,
+        limits: &ResourceLimits,
+        meter: &mut ResourceMeter,
+    ) -> Result<ArtifactDownloadStream, DownloadError> {
+        let issued = self
+            .issued_download_links
+            .get(presented_token_digest)
+            .ok_or(DownloadError::InvalidToken)?;
+        if issued.link.artifact != *artifact
+            || issued.link.max_bytes != policy.max_bytes
+            || issued.link.actor != context.actor
+        {
+            return Err(DownloadError::InvalidToken);
+        }
+        if issued.revoked {
+            return Err(DownloadError::Revoked);
+        }
+        if now_epoch_seconds > issued.link.expires_at_epoch_seconds {
+            return Err(DownloadError::Expired);
+        }
+        let action = self.download_action(context, artifact, policy)?;
+        if action.source != issued.link.source {
+            return Err(DownloadError::Unavailable);
+        }
+        let metadata = self
+            .artifacts
+            .get(artifact)
+            .ok_or(DownloadError::NotFound)?;
+        if download_policy_context_digest(metadata, &action.source, policy)
+            != issued.link.policy_context_digest
+        {
+            return Err(DownloadError::InvalidToken);
+        }
+        meter
+            .charge(limits, LimitKind::ArtifactDownloadBytes, 0)
+            .map_err(download_usage_error)?;
+        Ok(ArtifactDownloadStream {
+            link: issued.link.clone(),
+            streamed_bytes: 0,
+        })
+    }
+
+    pub fn stream_download_chunk(
+        &self,
+        stream: &mut ArtifactDownloadStream,
+        limits: &ResourceLimits,
+        meter: &mut ResourceMeter,
+        bytes: u64,
+    ) -> Result<(), DownloadError> {
+        let metadata = self
+            .artifacts
+            .get(&stream.link.artifact)
+            .ok_or(DownloadError::NotFound)?;
+        if !source_is_available(metadata, &stream.link.source) {
+            return Err(DownloadError::Unavailable);
+        }
+        stream.stream_chunk(limits, meter, bytes)
+    }
+}
+
+impl ArtifactDownloadStream {
+    pub fn stream_chunk(
+        &mut self,
+        limits: &ResourceLimits,
+        meter: &mut ResourceMeter,
+        bytes: u64,
+    ) -> Result<(), DownloadError> {
+        if self.streamed_bytes.saturating_add(bytes) > self.link.max_bytes {
+            return Err(DownloadError::LimitExceeded {
+                size: self.streamed_bytes.saturating_add(bytes),
+                limit: self.link.max_bytes,
+            });
+        }
+        meter
+            .charge(limits, LimitKind::ArtifactDownloadBytes, bytes)
+            .map_err(download_usage_error)?;
+        self.streamed_bytes += bytes;
+        Ok(())
+    }
+}
+
+fn download_usage_error(error: LimitError) -> DownloadError {
+    DownloadError::Usage(error.to_string())
+}
+
+fn download_policy_context_digest(
+    metadata: &ArtifactMetadata,
+    source: &StorageLocation,
+    policy: &DownloadPolicy,
+) -> Digest {
+    Digest::from_parts([
+        b"artifact-download-policy-context:v1".as_slice(),
+        metadata.tenant.as_str().as_bytes(),
+        metadata.project.as_str().as_bytes(),
+        metadata.process.as_str().as_bytes(),
+        metadata.id.as_str().as_bytes(),
+        metadata.digest.as_str().as_bytes(),
+        metadata.size.to_string().as_bytes(),
+        storage_location_key(source).as_bytes(),
+        policy.max_bytes.to_string().as_bytes(),
+    ])
+}
+
+fn actor_subject(actor: &Actor) -> String {
+    match actor {
+        Actor::User(id) => format!("user:{id}"),
+        Actor::Agent(id) => format!("agent:{id}"),
+        Actor::Node(id) => format!("node:{id}"),
+        Actor::Task(id) => format!("task:{id}"),
+    }
+}
+
+fn storage_location_key(source: &StorageLocation) -> String {
+    match source {
+        StorageLocation::RetainedNode(node) => format!("retained-node:{node}"),
+        StorageLocation::ExplicitStore(location) => format!("explicit-store:{location}"),
+    }
+}
+
+fn source_is_available(metadata: &ArtifactMetadata, source: &StorageLocation) -> bool {
+    match source {
+        StorageLocation::RetainedNode(node) => metadata.retaining_nodes.contains(node),
+        StorageLocation::ExplicitStore(location) => metadata.explicit_locations.contains(location),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeMap;
+
+    use crate::{Actor, LimitKind, ResourceLimits, ResourceMeter, UserId};
+
+    use super::*;
+
+    fn registry_with_artifact() -> ArtifactRegistry {
+        let mut registry = ArtifactRegistry::default();
+        registry.flush_metadata(
+            ArtifactId::from("artifact"),
+            TenantId::from("tenant"),
+            ProjectId::from("project"),
+            ProcessId::from("process"),
+            TaskId::from("task"),
+            NodeId::from("node"),
+            Digest::sha256("bytes"),
+            32,
+        );
+        registry
+    }
+
+    #[test]
+    fn flush_publishes_metadata_without_coordinator_bytes() {
+        let registry = registry_with_artifact();
+        let metadata = registry.metadata(&ArtifactId::from("artifact")).unwrap();
+
+        assert!(!metadata.coordinator_has_large_bytes);
+        assert_eq!(metadata.id, ArtifactId::from("artifact"));
+        assert_eq!(metadata.tenant, TenantId::from("tenant"));
+        assert_eq!(metadata.project, ProjectId::from("project"));
+        assert_eq!(metadata.process, ProcessId::from("process"));
+        assert_eq!(metadata.producer_task, TaskId::from("task"));
+        assert_eq!(metadata.producer_node, NodeId::from("node"));
+        assert_eq!(metadata.digest, Digest::sha256("bytes"));
+        assert_eq!(metadata.size, 32);
+        assert_eq!(metadata.flushed_epoch, 1);
+        assert!(metadata.retaining_nodes.contains(&NodeId::from("node")));
+        assert!(metadata.explicit_locations.is_empty());
+    }
+
+    #[test]
+    fn unsynced_node_loss_surfaces_as_unavailable() {
+        let mut registry = registry_with_artifact();
+        registry.garbage_collect_node(&NodeId::from("node"));
+        let context = AuthContext {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            actor: Actor::User(UserId::from("user")),
+        };
+
+        let error = registry
+            .download_action(
+                &context,
+                &ArtifactId::from("artifact"),
+                &DownloadPolicy { max_bytes: 100 },
+            )
+            .unwrap_err();
+
+        assert_eq!(error, DownloadError::Unavailable);
+    }
+
+    #[test]
+    fn explicit_user_storage_location_survives_node_retention_loss() {
+        let mut registry = registry_with_artifact();
+        registry
+            .sync_to_explicit_store(&ArtifactId::from("artifact"), "s3://bucket/app")
+            .unwrap();
+        registry.garbage_collect_node(&NodeId::from("node"));
+        let context = AuthContext {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            actor: Actor::User(UserId::from("user")),
+        };
+
+        let action = registry
+            .download_action(
+                &context,
+                &ArtifactId::from("artifact"),
+                &DownloadPolicy { max_bytes: 100 },
+            )
+            .unwrap();
+
+        assert_eq!(
+            action.source,
+            StorageLocation::ExplicitStore("s3://bucket/app".to_owned())
+        );
+        assert!(
+            !registry
+                .metadata(&ArtifactId::from("artifact"))
+                .unwrap()
+                .coordinator_has_large_bytes
+        );
+    }
+
+    #[test]
+    fn default_retention_policy_is_best_effort_node_retention() {
+        let policy = RetentionPolicy::default();
+
+        assert!(policy.best_effort_node_retention);
+        assert_eq!(policy.max_download_bytes, 256 * 1024 * 1024);
+    }
+
+    #[test]
+    fn cross_tenant_download_is_denied_even_with_known_artifact_id() {
+        let registry = registry_with_artifact();
+        let context = AuthContext {
+            tenant: TenantId::from("other"),
+            project: ProjectId::from("project"),
+            actor: Actor::User(UserId::from("user")),
+        };
+
+        let error = registry
+            .download_action(
+                &context,
+                &ArtifactId::from("artifact"),
+                &DownloadPolicy { max_bytes: 100 },
+            )
+            .unwrap_err();
+
+        assert!(matches!(error, DownloadError::Unauthorized(_)));
+    }
+
+    #[test]
+    fn cross_project_download_is_denied_even_with_known_artifact_id() {
+        let registry = registry_with_artifact();
+        let context = AuthContext {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("other-project"),
+            actor: Actor::User(UserId::from("user")),
+        };
+
+        let error = registry
+            .download_action(
+                &context,
+                &ArtifactId::from("artifact"),
+                &DownloadPolicy { max_bytes: 100 },
+            )
+            .unwrap_err();
+
+        assert!(matches!(error, DownloadError::Unauthorized(_)));
+    }
+
+    #[test]
+    fn download_link_is_not_created_when_artifact_is_unavailable_or_too_large() {
+        let mut registry = registry_with_artifact();
+        registry.garbage_collect_node(&NodeId::from("node"));
+        let context = AuthContext {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            actor: Actor::User(UserId::from("user")),
+        };
+
+        assert_eq!(
+            registry
+                .create_download_link(
+                    &context,
+                    &ArtifactId::from("artifact"),
+                    &DownloadPolicy { max_bytes: 100 },
+                    "nonce",
+                    10,
+                    60,
+                )
+                .unwrap_err(),
+            DownloadError::Unavailable
+        );
+
+        let mut registry = registry_with_artifact();
+        assert!(matches!(
+            registry.create_download_link(
+                &context,
+                &ArtifactId::from("artifact"),
+                &DownloadPolicy { max_bytes: 1 },
+                "nonce",
+                10,
+                60,
+            ),
+            Err(DownloadError::LimitExceeded { .. })
+        ));
+    }
+
+    #[test]
+    fn download_link_is_authenticated_scoped_and_not_guessable() {
+        let mut registry = registry_with_artifact();
+        let context = AuthContext {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            actor: Actor::User(UserId::from("user")),
+        };
+        let link = registry
+            .create_download_link(
+                &context,
+                &ArtifactId::from("artifact"),
+                &DownloadPolicy { max_bytes: 100 },
+                "nonce-a",
+                10,
+                60,
+            )
+            .unwrap();
+        let other = registry
+            .create_download_link(
+                &context,
+                &ArtifactId::from("artifact"),
+                &DownloadPolicy { max_bytes: 100 },
+                "nonce-b",
+                10,
+                60,
+            )
+            .unwrap();
+
+        assert_eq!(link.tenant, TenantId::from("tenant"));
+        assert_eq!(link.project, ProjectId::from("project"));
+        assert_eq!(link.process, ProcessId::from("process"));
+        assert_eq!(link.actor, Actor::User(UserId::from("user")));
+        assert_eq!(link.max_bytes, 100);
+        assert!(link.policy_context_digest.is_valid_sha256());
+        assert_eq!(link.expires_at_epoch_seconds, 70);
+        assert!(link
+            .url_path
+            .contains("/artifacts/tenant/project/process/artifact"));
+        assert_ne!(link.scoped_token_digest, other.scoped_token_digest);
+        assert!(matches!(link.source, StorageLocation::RetainedNode(_)));
+    }
+
+    #[test]
+    fn download_link_is_bound_to_actor_and_policy_context() {
+        let mut registry = registry_with_artifact();
+        let context = AuthContext {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            actor: Actor::User(UserId::from("user")),
+        };
+        let other_actor = AuthContext {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            actor: Actor::User(UserId::from("other-user")),
+        };
+        let policy = DownloadPolicy { max_bytes: 100 };
+        let link = registry
+            .create_download_link(
+                &context,
+                &ArtifactId::from("artifact"),
+                &policy,
+                "nonce",
+                10,
+                60,
+            )
+            .unwrap();
+        let limits = ResourceLimits {
+            limits: BTreeMap::from([(LimitKind::ArtifactDownloadBytes, 32)]),
+        };
+        let mut meter = ResourceMeter::default();
+
+        assert_eq!(
+            registry
+                .open_download_stream(
+                    &other_actor,
+                    &ArtifactId::from("artifact"),
+                    &policy,
+                    &link.scoped_token_digest,
+                    11,
+                    &limits,
+                    &mut meter,
+                )
+                .unwrap_err(),
+            DownloadError::InvalidToken
+        );
+        assert_eq!(
+            registry
+                .open_download_stream(
+                    &context,
+                    &ArtifactId::from("artifact"),
+                    &DownloadPolicy { max_bytes: 99 },
+                    &link.scoped_token_digest,
+                    11,
+                    &limits,
+                    &mut meter,
+                )
+                .unwrap_err(),
+            DownloadError::InvalidToken
+        );
+        assert_eq!(
+            registry
+                .revoke_download_link(
+                    &other_actor,
+                    &ArtifactId::from("artifact"),
+                    &link.scoped_token_digest,
+                )
+                .unwrap_err(),
+            DownloadError::InvalidToken
+        );
+    }
+
+    #[test]
+    fn download_link_expires_and_can_be_revoked() {
+        let mut registry = registry_with_artifact();
+        let context = AuthContext {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            actor: Actor::User(UserId::from("user")),
+        };
+        let policy = DownloadPolicy { max_bytes: 100 };
+        let limits = ResourceLimits {
+            limits: BTreeMap::from([(LimitKind::ArtifactDownloadBytes, 32)]),
+        };
+        let mut meter = ResourceMeter::default();
+        let expired = registry
+            .create_download_link(
+                &context,
+                &ArtifactId::from("artifact"),
+                &policy,
+                "expired",
+                10,
+                5,
+            )
+            .unwrap();
+
+        let error = registry
+            .open_download_stream(
+                &context,
+                &ArtifactId::from("artifact"),
+                &policy,
+                &expired.scoped_token_digest,
+                16,
+                &limits,
+                &mut meter,
+            )
+            .unwrap_err();
+        assert_eq!(error, DownloadError::Expired);
+
+        let active = registry
+            .create_download_link(
+                &context,
+                &ArtifactId::from("artifact"),
+                &policy,
+                "active",
+                20,
+                60,
+            )
+            .unwrap();
+        let revoked = registry
+            .revoke_download_link(
+                &context,
+                &ArtifactId::from("artifact"),
+                &active.scoped_token_digest,
+            )
+            .unwrap();
+        assert_eq!(revoked.scoped_token_digest, active.scoped_token_digest);
+
+        let error = registry
+            .open_download_stream(
+                &context,
+                &ArtifactId::from("artifact"),
+                &policy,
+                &active.scoped_token_digest,
+                21,
+                &limits,
+                &mut meter,
+            )
+            .unwrap_err();
+        assert_eq!(error, DownloadError::Revoked);
+    }
+
+    #[test]
+    fn download_stream_accounts_usage_before_and_during_streaming() {
+        let mut registry = registry_with_artifact();
+        let context = AuthContext {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            actor: Actor::User(UserId::from("user")),
+        };
+        let policy = DownloadPolicy { max_bytes: 100 };
+        let link = registry
+            .create_download_link(
+                &context,
+                &ArtifactId::from("artifact"),
+                &policy,
+                "nonce",
+                10,
+                60,
+            )
+            .unwrap();
+        let limits = ResourceLimits {
+            limits: BTreeMap::from([(LimitKind::ArtifactDownloadBytes, 32)]),
+        };
+        let mut meter = ResourceMeter::default();
+        let mut stream = registry
+            .open_download_stream(
+                &context,
+                &ArtifactId::from("artifact"),
+                &policy,
+                &link.scoped_token_digest,
+                11,
+                &limits,
+                &mut meter,
+            )
+            .unwrap();
+
+        stream.stream_chunk(&limits, &mut meter, 16).unwrap();
+        stream.stream_chunk(&limits, &mut meter, 16).unwrap();
+        assert!(matches!(
+            stream.stream_chunk(&limits, &mut meter, 1),
+            Err(DownloadError::Usage(_))
+        ));
+        assert_eq!(meter.used(&LimitKind::ArtifactDownloadBytes), 32);
+    }
+
+    #[test]
+    fn download_stream_fails_honestly_when_source_disappears_mid_stream() {
+        let mut registry = registry_with_artifact();
+        let context = AuthContext {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            actor: Actor::User(UserId::from("user")),
+        };
+        let policy = DownloadPolicy { max_bytes: 100 };
+        let link = registry
+            .create_download_link(
+                &context,
+                &ArtifactId::from("artifact"),
+                &policy,
+                "nonce",
+                10,
+                60,
+            )
+            .unwrap();
+        let limits = ResourceLimits {
+            limits: BTreeMap::from([(LimitKind::ArtifactDownloadBytes, 32)]),
+        };
+        let mut meter = ResourceMeter::default();
+        let mut stream = registry
+            .open_download_stream(
+                &context,
+                &ArtifactId::from("artifact"),
+                &policy,
+                &link.scoped_token_digest,
+                11,
+                &limits,
+                &mut meter,
+            )
+            .unwrap();
+
+        registry
+            .stream_download_chunk(&mut stream, &limits, &mut meter, 16)
+            .unwrap();
+        registry.garbage_collect_node(&NodeId::from("node"));
+
+        assert_eq!(
+            registry
+                .stream_download_chunk(&mut stream, &limits, &mut meter, 1)
+                .unwrap_err(),
+            DownloadError::Unavailable
+        );
+        assert_eq!(stream.streamed_bytes, 16);
+    }
+
+    #[test]
+    fn guessed_download_token_is_rejected() {
+        let registry = registry_with_artifact();
+        let context = AuthContext {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            actor: Actor::User(UserId::from("user")),
+        };
+        let limits = ResourceLimits {
+            limits: BTreeMap::from([(LimitKind::ArtifactDownloadBytes, 32)]),
+        };
+        let mut meter = ResourceMeter::default();
+
+        let error = registry
+            .open_download_stream(
+                &context,
+                &ArtifactId::from("artifact"),
+                &DownloadPolicy { max_bytes: 100 },
+                &Digest::sha256("guessed"),
+                11,
+                &limits,
+                &mut meter,
+            )
+            .unwrap_err();
+
+        assert_eq!(error, DownloadError::InvalidToken);
+    }
+}