systemprompt_loader/bundle/source/oci/
mod.rs

1//! OCI distribution transport for services bundles.
2//!
3//! Only the parts of the distribution spec a bundle needs are implemented:
4//! a manifest GET, the Bearer challenge dance, and a single blob pull whose
5//! `mediaType` is
6//! [`BUNDLE_MEDIA_TYPE`](systemprompt_models::services::bundle::BUNDLE_MEDIA_TYPE).
7//! A manifest carrying zero or several
8//! such layers is refused rather than guessed at, because picking one would
9//! make which bytes an instance runs depend on registry ordering.
10//!
11//! The registry scheme is `https` unless the host is loopback or listed in
12//! the trusted-host escape hatch, and every constructed URL goes through the
13//! shared SSRF guard.
14//!
15//! Copyright (c) systemprompt.io — Business Source License 1.1.
16//! See <https://systemprompt.io> for licensing details.
17
18pub mod auth;
19pub mod pull;
20pub mod push;
21
22use std::path::Path;
23use std::str::FromStr;
24
25use systemprompt_models::net::{trusted_http_hosts_from_env, validate_outbound_url_with_trust};
26use systemprompt_models::profile::OciReference;
27
28use super::{BundleFetcher, FetchedBundle, MAX_BUNDLE_BYTES, RemoteRef};
29use crate::bundle::error::{BundleError, BundleResult};
30
31pub use push::push_bundle;
32
33#[derive(Debug)]
34pub struct RegistryClient {
35    pub client: reqwest::Client,
36    pub name: String,
37    pub reference: OciReference,
38    pub secret: Option<String>,
39}
40
41impl RegistryClient {
42    pub fn new(
43        name: &str,
44        reference: &str,
45        secret: Option<String>,
46        client: reqwest::Client,
47    ) -> BundleResult<Self> {
48        let reference = OciReference::from_str(reference)
49            .map_err(|e| BundleError::policy(format!("source {name}: {e}")))?;
50        Ok(Self {
51            client,
52            name: name.to_owned(),
53            reference,
54            secret,
55        })
56    }
57
58    pub fn url(&self, path: &str) -> BundleResult<url::Url> {
59        let host = &self.reference.registry;
60        let trusted = trusted_http_hosts_from_env();
61        let bare = host.split(':').next().unwrap_or(host);
62        let plain = bare == "localhost"
63            || bare == "127.0.0.1"
64            || trusted.iter().any(|t| t.eq_ignore_ascii_case(bare));
65        let scheme = if plain { "http" } else { "https" };
66        let raw = format!("{scheme}://{host}/v2/{}{path}", self.reference.repository);
67        validate_outbound_url_with_trust(&raw, &trusted)
68            .map_err(|e| BundleError::fetch(&self.name, e))
69    }
70
71    #[must_use]
72    pub fn manifest_ref(&self) -> String {
73        self.reference.digest.clone().unwrap_or_else(|| {
74            self.reference
75                .tag
76                .clone()
77                .unwrap_or_else(|| "latest".to_owned())
78        })
79    }
80
81    pub async fn send(
82        &self,
83        build: impl Fn(&reqwest::Client) -> reqwest::RequestBuilder + Send,
84    ) -> BundleResult<reqwest::Response> {
85        let first = auth::apply_credential(build(&self.client), self.secret.as_deref())
86            .send()
87            .await
88            .map_err(|e| BundleError::fetch(&self.name, e))?;
89
90        if first.status() != reqwest::StatusCode::UNAUTHORIZED {
91            return Ok(first);
92        }
93
94        let header = first
95            .headers()
96            .get(reqwest::header::WWW_AUTHENTICATE)
97            .and_then(|v| v.to_str().ok())
98            .unwrap_or_default()
99            .to_owned();
100        let challenge = auth::parse_challenge(&header).ok_or_else(|| BundleError::Auth {
101            source_name: self.name.clone(),
102        })?;
103        let token =
104            auth::fetch_token(&self.client, &challenge, self.secret.as_deref(), &self.name).await?;
105
106        let retried = build(&self.client)
107            .bearer_auth(token)
108            .send()
109            .await
110            .map_err(|e| BundleError::fetch(&self.name, e))?;
111        if retried.status() == reqwest::StatusCode::UNAUTHORIZED
112            || retried.status() == reqwest::StatusCode::FORBIDDEN
113        {
114            return Err(BundleError::Auth {
115                source_name: self.name.clone(),
116            });
117        }
118        Ok(retried)
119    }
120}
121
122#[derive(Debug)]
123pub struct OciFetcher {
124    registry: RegistryClient,
125    max_bytes: u64,
126}
127
128impl OciFetcher {
129    pub fn new(
130        name: &str,
131        reference: &str,
132        secret: Option<String>,
133        client: reqwest::Client,
134    ) -> BundleResult<Self> {
135        Ok(Self {
136            registry: RegistryClient::new(name, reference, secret, client)?,
137            max_bytes: MAX_BUNDLE_BYTES,
138        })
139    }
140
141    #[must_use]
142    pub const fn with_max_bytes(mut self, max_bytes: u64) -> Self {
143        self.max_bytes = max_bytes;
144        self
145    }
146}
147
148impl BundleFetcher for OciFetcher {
149    async fn head(&self) -> BundleResult<RemoteRef> {
150        let (digest, _manifest) = pull::get_manifest(&self.registry).await?;
151        Ok(RemoteRef { digest })
152    }
153
154    async fn fetch(&self, into: &Path) -> BundleResult<FetchedBundle> {
155        let (_digest, manifest) = pull::get_manifest(&self.registry).await?;
156        pull::pull_bundle_layer(&self.registry, &manifest, into, self.max_bytes).await
157    }
158}
systemprompt_loader/bundle/source/oci/mod.rs

systemprompt_loader/bundle/source/oci/
mod.rs
