clusterflux-public/crates/clusterflux-core/src/artifact.rs

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, TaskInstanceId, TenantId,
};

const MAX_ISSUED_DOWNLOAD_LINKS_PER_ARTIFACT: usize = 32;
const DOWNLOAD_LINK_TOMBSTONE_SECONDS: u64 = 15 * 60;
const MAX_ARTIFACT_METADATA_PER_PROCESS: usize = 256;

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ArtifactHandle {
    pub id: ArtifactId,
    pub digest: Digest,
    pub size_bytes: u64,
}

#[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, Copy, Debug)]
pub struct DownloadStreamRequest<'a> {
    pub context: &'a AuthContext,
    pub artifact: &'a ArtifactId,
    pub policy: &'a DownloadPolicy,
    pub presented_token_digest: &'a Digest,
    pub now_epoch_seconds: u64,
    pub limits: &'a ResourceLimits,
}

#[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 artifact_digest: Digest,
    pub artifact_size_bytes: u64,
    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: TaskInstanceId,
    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, PartialEq, Eq)]
pub struct ArtifactFlush {
    pub id: ArtifactId,
    pub tenant: TenantId,
    pub project: ProjectId,
    pub process: ProcessId,
    pub producer_task: TaskInstanceId,
    pub retaining_node: NodeId,
    pub digest: Digest,
    pub size: u64,
}

#[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, flush: ArtifactFlush) -> ArtifactMetadata {
        self.flush_metadata_bounded(flush, &BTreeSet::new())
            .expect("an unpinned artifact scope always has an evictable metadata entry")
    }

    pub fn flush_metadata_bounded(
        &mut self,
        flush: ArtifactFlush,
        pinned: &BTreeSet<ArtifactId>,
    ) -> Result<ArtifactMetadata, String> {
        let replacing_existing = self.artifacts.contains_key(&flush.id);
        while !replacing_existing
            && self
                .artifacts
                .values()
                .filter(|metadata| {
                    metadata.tenant == flush.tenant
                        && metadata.project == flush.project
                        && metadata.process == flush.process
                })
                .count()
                >= MAX_ARTIFACT_METADATA_PER_PROCESS
        {
            let candidate = self
                .artifacts
                .values()
                .filter(|metadata| {
                    metadata.tenant == flush.tenant
                        && metadata.project == flush.project
                        && metadata.process == flush.process
                        && !pinned.contains(&metadata.id)
                        && !self.issued_download_links.values().any(|issued| {
                            issued.link.artifact == metadata.id
                        })
                })
                .min_by_key(|metadata| metadata.flushed_epoch)
                .map(|metadata| metadata.id.clone())
                .ok_or_else(|| {
                    "artifact metadata retention limit reached and every retained object is pinned by active work, restart state, or a download"
                        .to_owned()
                })?;
            self.artifacts.remove(&candidate);
        }
        self.next_epoch += 1;
        let metadata = ArtifactMetadata {
            id: flush.id.clone(),
            tenant: flush.tenant,
            project: flush.project,
            process: flush.process,
            producer_task: flush.producer_task,
            producer_node: flush.retaining_node.clone(),
            digest: flush.digest,
            size: flush.size,
            flushed_epoch: self.next_epoch,
            retaining_nodes: BTreeSet::from([flush.retaining_node]),
            explicit_locations: Vec::new(),
            coordinator_has_large_bytes: false,
        };
        self.artifacts.insert(flush.id, metadata.clone());
        Ok(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 reconcile_node_retention(
        &mut self,
        node: &NodeId,
        retained_artifacts: &BTreeSet<ArtifactId>,
    ) {
        for (artifact, metadata) in &mut self.artifacts {
            if retained_artifacts.contains(artifact) {
                metadata.retaining_nodes.insert(node.clone());
            } else {
                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> {
        self.expire_download_links(now_epoch_seconds);
        if self
            .issued_download_links
            .values()
            .filter(|issued| issued.link.artifact == *artifact)
            .count()
            >= MAX_ISSUED_DOWNLOAD_LINKS_PER_ARTIFACT
        {
            return Err(DownloadError::Usage(
                "artifact download session limit reached; revoke a link or wait for expiry"
                    .to_owned(),
            ));
        }
        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(),
            artifact_digest: metadata.digest.clone(),
            artifact_size_bytes: metadata.size,
            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,
        request: DownloadStreamRequest<'_>,
        meter: &mut ResourceMeter,
    ) -> Result<ArtifactDownloadStream, DownloadError> {
        let DownloadStreamRequest {
            context,
            artifact,
            policy,
            presented_token_digest,
            now_epoch_seconds,
            limits,
        } = request;
        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)
    }

    pub fn expire_download_links(&mut self, now_epoch_seconds: u64) {
        self.issued_download_links.retain(|_, issued| {
            issued
                .link
                .expires_at_epoch_seconds
                .saturating_add(DOWNLOAD_LINK_TOMBSTONE_SECONDS)
                >= now_epoch_seconds
        });
    }

    pub fn retained_download_link_count(&self) -> usize {
        self.issued_download_links.len()
    }

    pub fn artifact_count(&self) -> usize {
        self.artifacts.len()
    }
}

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(ArtifactFlush {
            id: ArtifactId::from("artifact"),
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("project"),
            process: ProcessId::from("process"),
            producer_task: TaskInstanceId::from("task"),
            retaining_node: NodeId::from("node"),
            digest: Digest::sha256("bytes"),
            size: 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, TaskInstanceId::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 signed_node_retention_inventory_removes_garbage_collected_location() {
        let mut registry = registry_with_artifact();
        registry.reconcile_node_retention(&NodeId::from("node"), &BTreeSet::new());

        let metadata = registry.metadata(&ArtifactId::from("artifact")).unwrap();

        assert!(metadata.retaining_nodes.is_empty());
    }

    #[test]
    fn signed_node_retention_inventory_restores_readvertised_location() {
        let mut registry = registry_with_artifact();
        let node = NodeId::from("node");
        let artifact = ArtifactId::from("artifact");
        registry.garbage_collect_node(&node);

        registry.reconcile_node_retention(&node, &BTreeSet::from([artifact.clone()]));

        let metadata = registry.metadata(&artifact).unwrap();
        assert!(metadata.retaining_nodes.contains(&node));
    }

    #[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(
                    DownloadStreamRequest {
                        context: &other_actor,
                        artifact: &ArtifactId::from("artifact"),
                        policy: &policy,
                        presented_token_digest: &link.scoped_token_digest,
                        now_epoch_seconds: 11,
                        limits: &limits,
                    },
                    &mut meter,
                )
                .unwrap_err(),
            DownloadError::InvalidToken
        );
        assert_eq!(
            registry
                .open_download_stream(
                    DownloadStreamRequest {
                        context: &context,
                        artifact: &ArtifactId::from("artifact"),
                        policy: &DownloadPolicy { max_bytes: 99 },
                        presented_token_digest: &link.scoped_token_digest,
                        now_epoch_seconds: 11,
                        limits: &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(
                DownloadStreamRequest {
                    context: &context,
                    artifact: &ArtifactId::from("artifact"),
                    policy: &policy,
                    presented_token_digest: &expired.scoped_token_digest,
                    now_epoch_seconds: 16,
                    limits: &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(
                DownloadStreamRequest {
                    context: &context,
                    artifact: &ArtifactId::from("artifact"),
                    policy: &policy,
                    presented_token_digest: &active.scoped_token_digest,
                    now_epoch_seconds: 21,
                    limits: &limits,
                },
                &mut meter,
            )
            .unwrap_err();
        assert_eq!(error, DownloadError::Revoked);
    }

    #[test]
    fn download_session_history_is_count_and_time_bounded() {
        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 };
        for index in 0..MAX_ISSUED_DOWNLOAD_LINKS_PER_ARTIFACT {
            registry
                .create_download_link(
                    &context,
                    &ArtifactId::from("artifact"),
                    &policy,
                    &format!("bounded-{index}"),
                    10,
                    1,
                )
                .unwrap();
        }
        assert_eq!(
            registry.retained_download_link_count(),
            MAX_ISSUED_DOWNLOAD_LINKS_PER_ARTIFACT
        );
        assert!(matches!(
            registry.create_download_link(
                &context,
                &ArtifactId::from("artifact"),
                &policy,
                "over-limit",
                10,
                1,
            ),
            Err(DownloadError::Usage(_))
        ));

        registry.expire_download_links(12 + DOWNLOAD_LINK_TOMBSTONE_SECONDS);
        assert_eq!(registry.retained_download_link_count(), 0);
        registry
            .create_download_link(
                &context,
                &ArtifactId::from("artifact"),
                &policy,
                "after-expiry",
                12 + DOWNLOAD_LINK_TOMBSTONE_SECONDS,
                1,
            )
            .unwrap();
    }

    #[test]
    fn artifact_metadata_is_bounded_without_evicting_pins() {
        let mut registry = ArtifactRegistry::default();
        let mut pinned = BTreeSet::new();
        for index in 0..MAX_ARTIFACT_METADATA_PER_PROCESS {
            let id = ArtifactId::new(format!("artifact-{index}"));
            pinned.insert(id.clone());
            registry
                .flush_metadata_bounded(
                    ArtifactFlush {
                        id,
                        tenant: TenantId::from("tenant"),
                        project: ProjectId::from("project"),
                        process: ProcessId::from("process"),
                        producer_task: TaskInstanceId::new(format!("task-{index}")),
                        retaining_node: NodeId::from("node"),
                        digest: Digest::sha256(format!("content-{index}")),
                        size: 1,
                    },
                    &pinned,
                )
                .unwrap();
        }
        assert_eq!(registry.artifact_count(), MAX_ARTIFACT_METADATA_PER_PROCESS);
        let next = ArtifactFlush {
            id: ArtifactId::from("artifact-next"),
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("project"),
            process: ProcessId::from("process"),
            producer_task: TaskInstanceId::from("task-next"),
            retaining_node: NodeId::from("node"),
            digest: Digest::sha256("next"),
            size: 1,
        };
        assert!(registry
            .flush_metadata_bounded(next.clone(), &pinned)
            .unwrap_err()
            .contains("pinned"));

        pinned.remove(&ArtifactId::from("artifact-0"));
        registry.flush_metadata_bounded(next, &pinned).unwrap();
        assert_eq!(registry.artifact_count(), MAX_ARTIFACT_METADATA_PER_PROCESS);
        assert!(registry.metadata(&ArtifactId::from("artifact-0")).is_none());
        assert!(registry
            .metadata(&ArtifactId::from("artifact-next"))
            .is_some());
    }

    #[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(
                DownloadStreamRequest {
                    context: &context,
                    artifact: &ArtifactId::from("artifact"),
                    policy: &policy,
                    presented_token_digest: &link.scoped_token_digest,
                    now_epoch_seconds: 11,
                    limits: &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(
                DownloadStreamRequest {
                    context: &context,
                    artifact: &ArtifactId::from("artifact"),
                    policy: &policy,
                    presented_token_digest: &link.scoped_token_digest,
                    now_epoch_seconds: 11,
                    limits: &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(
                DownloadStreamRequest {
                    context: &context,
                    artifact: &ArtifactId::from("artifact"),
                    policy: &DownloadPolicy { max_bytes: 100 },
                    presented_token_digest: &Digest::sha256("guessed"),
                    now_epoch_seconds: 11,
                    limits: &limits,
                },
                &mut meter,
            )
            .unwrap_err();

        assert_eq!(error, DownloadError::InvalidToken);
    }
}