1use std::borrow::Cow;
4use std::fmt;
5use std::str::FromStr;
6
7use oci_spec::distribution::Reference as OciDistributionReference;
8use schemars::{JsonSchema, Schema, SchemaGenerator, json_schema};
9use serde::{Deserialize, Deserializer, Serialize, Serializer};
10
11use super::hashes::{OCI_SHA256_DIGEST_ERROR_MESSAGE, OciSha256Digest};
12
13const LOCAL_IMAGE_REFERENCE_ERROR_MESSAGE: &str =
14 "local image reference must be a supported Agentics local repository with an explicit tag";
15const REGISTRY_IMAGE_REFERENCE_ERROR_MESSAGE: &str =
16 "registry image reference must include an explicit registry, repository, and tag";
17
18pub const SUPPORTED_LOCAL_AGENTICS_IMAGE_REPOSITORIES: &[&str] =
20 &["agentics-linux-arm64-cpu", "agentics-linux-arm64-cuda"];
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub struct LocalAgenticsImageReferenceError;
25
26impl fmt::Display for LocalAgenticsImageReferenceError {
27 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
29 f.write_str(LOCAL_IMAGE_REFERENCE_ERROR_MESSAGE)
30 }
31}
32
33impl std::error::Error for LocalAgenticsImageReferenceError {}
34
35#[derive(Debug, Clone, PartialEq, Eq, Hash)]
37pub struct LocalAgenticsImageReference {
38 original: String,
39 repository: String,
40 tag: String,
41}
42
43impl LocalAgenticsImageReference {
44 pub fn try_new(value: impl AsRef<str>) -> Result<Self, LocalAgenticsImageReferenceError> {
46 let value = value.as_ref();
47 if value.trim() != value || value.is_empty() || value.contains('/') || value.contains('@') {
48 return Err(LocalAgenticsImageReferenceError);
49 }
50 let Some((repository, tag)) = value.rsplit_once(':') else {
51 return Err(LocalAgenticsImageReferenceError);
52 };
53 if !SUPPORTED_LOCAL_AGENTICS_IMAGE_REPOSITORIES.contains(&repository)
54 || tag.is_empty()
55 || !tag
56 .chars()
57 .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '_' | '.' | '-'))
58 {
59 return Err(LocalAgenticsImageReferenceError);
60 }
61
62 Ok(Self {
63 original: value.to_string(),
64 repository: repository.to_string(),
65 tag: tag.to_string(),
66 })
67 }
68
69 pub fn as_str(&self) -> &str {
71 &self.original
72 }
73
74 pub fn repository(&self) -> &str {
76 &self.repository
77 }
78
79 pub fn tag(&self) -> &str {
81 &self.tag
82 }
83}
84
85impl fmt::Display for LocalAgenticsImageReference {
86 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88 f.write_str(self.as_str())
89 }
90}
91
92impl FromStr for LocalAgenticsImageReference {
93 type Err = LocalAgenticsImageReferenceError;
94
95 fn from_str(value: &str) -> Result<Self, Self::Err> {
97 Self::try_new(value)
98 }
99}
100
101impl Serialize for LocalAgenticsImageReference {
102 fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
104 where
105 S: Serializer,
106 {
107 serializer.serialize_str(self.as_str())
108 }
109}
110
111impl<'de> Deserialize<'de> for LocalAgenticsImageReference {
112 fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
114 where
115 D: Deserializer<'de>,
116 {
117 let value = String::deserialize(deserializer)?;
118 Self::try_new(&value).map_err(serde::de::Error::custom)
119 }
120}
121
122impl JsonSchema for LocalAgenticsImageReference {
123 fn inline_schema() -> bool {
125 true
126 }
127
128 fn schema_name() -> Cow<'static, str> {
130 "LocalAgenticsImageReference".into()
131 }
132
133 fn json_schema(_: &mut SchemaGenerator) -> Schema {
135 json_schema!({
136 "type": "string",
137 "pattern": "^(agentics-linux-arm64-cpu|agentics-linux-arm64-cuda):[A-Za-z0-9_.-]+$"
138 })
139 }
140}
141
142#[derive(Debug, Clone, PartialEq, Eq)]
144pub struct OciRegistryImageReferenceError(String);
145
146impl fmt::Display for OciRegistryImageReferenceError {
147 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
149 f.write_str(&self.0)
150 }
151}
152
153impl std::error::Error for OciRegistryImageReferenceError {}
154
155#[derive(Debug, Clone, PartialEq, Eq, Hash)]
157pub struct OciRegistryImageReference {
158 original: String,
159 parsed: OciDistributionReference,
160 digest: Option<OciSha256Digest>,
161 tag: String,
162}
163
164impl OciRegistryImageReference {
165 pub fn try_new(value: impl AsRef<str>) -> Result<Self, OciRegistryImageReferenceError> {
167 let value = value.as_ref();
168 if value.trim() != value
169 || value.is_empty()
170 || !has_explicit_registry(value)
171 || !has_explicit_tag(value)
172 {
173 return Err(OciRegistryImageReferenceError(
174 REGISTRY_IMAGE_REFERENCE_ERROR_MESSAGE.to_string(),
175 ));
176 }
177 let parsed = OciDistributionReference::from_str(value).map_err(|error| {
178 OciRegistryImageReferenceError(format!(
179 "{REGISTRY_IMAGE_REFERENCE_ERROR_MESSAGE}: {error}"
180 ))
181 })?;
182 let digest = parsed
183 .digest()
184 .map(|digest| {
185 OciSha256Digest::try_new(digest).map_err(|_| {
186 OciRegistryImageReferenceError(OCI_SHA256_DIGEST_ERROR_MESSAGE.to_string())
187 })
188 })
189 .transpose()?;
190 let Some(tag) = parsed.tag().map(ToOwned::to_owned) else {
191 return Err(OciRegistryImageReferenceError(
192 REGISTRY_IMAGE_REFERENCE_ERROR_MESSAGE.to_string(),
193 ));
194 };
195
196 Ok(Self {
197 original: value.to_string(),
198 parsed,
199 digest,
200 tag,
201 })
202 }
203
204 pub fn as_str(&self) -> &str {
206 &self.original
207 }
208
209 pub fn as_oci_reference(&self) -> &OciDistributionReference {
211 &self.parsed
212 }
213
214 pub fn policy_repository(&self) -> String {
216 format!("{}/{}", self.parsed.registry(), self.parsed.repository())
217 }
218
219 pub fn tag(&self) -> &str {
221 &self.tag
222 }
223
224 pub fn digest(&self) -> Option<&OciSha256Digest> {
226 self.digest.as_ref()
227 }
228}
229
230impl fmt::Display for OciRegistryImageReference {
231 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
233 f.write_str(self.as_str())
234 }
235}
236
237impl FromStr for OciRegistryImageReference {
238 type Err = OciRegistryImageReferenceError;
239
240 fn from_str(value: &str) -> Result<Self, Self::Err> {
242 Self::try_new(value)
243 }
244}
245
246impl Serialize for OciRegistryImageReference {
247 fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
249 where
250 S: Serializer,
251 {
252 serializer.serialize_str(self.as_str())
253 }
254}
255
256impl<'de> Deserialize<'de> for OciRegistryImageReference {
257 fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
259 where
260 D: Deserializer<'de>,
261 {
262 let value = String::deserialize(deserializer)?;
263 Self::try_new(&value).map_err(serde::de::Error::custom)
264 }
265}
266
267impl JsonSchema for OciRegistryImageReference {
268 fn inline_schema() -> bool {
270 true
271 }
272
273 fn schema_name() -> Cow<'static, str> {
275 "OciRegistryImageReference".into()
276 }
277
278 fn json_schema(_: &mut SchemaGenerator) -> Schema {
280 json_schema!({
281 "type": "string",
282 "pattern": "^[^/.:]+[.:][^/]*/[^\\s@:]+(/[^\\s@:]+)*:[^\\s@]+(@sha256:[0-9a-f]{64})?$"
283 })
284 }
285}
286
287#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
289#[serde(tag = "source", rename_all = "snake_case")]
290pub enum ChallengeImageReference {
291 Local {
292 reference: LocalAgenticsImageReference,
293 },
294 Registry {
295 reference: OciRegistryImageReference,
296 },
297}
298
299impl ChallengeImageReference {
300 pub fn docker_reference(&self) -> &str {
302 match self {
303 Self::Local { reference } => reference.as_str(),
304 Self::Registry { reference } => reference.as_str(),
305 }
306 }
307
308 pub fn policy_repository(&self) -> Cow<'_, str> {
310 match self {
311 Self::Local { reference } => Cow::Borrowed(reference.repository()),
312 Self::Registry { reference } => Cow::Owned(reference.policy_repository()),
313 }
314 }
315
316 pub fn tag(&self) -> &str {
318 match self {
319 Self::Local { reference } => reference.tag(),
320 Self::Registry { reference } => reference.tag(),
321 }
322 }
323
324 pub fn digest(&self) -> Option<&OciSha256Digest> {
326 match self {
327 Self::Local { .. } => None,
328 Self::Registry { reference } => reference.digest(),
329 }
330 }
331
332 pub fn is_local(&self) -> bool {
334 matches!(self, Self::Local { .. })
335 }
336}
337
338impl fmt::Display for ChallengeImageReference {
339 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
341 f.write_str(self.docker_reference())
342 }
343}
344
345fn has_explicit_registry(value: &str) -> bool {
347 let Some((registry, _)) = value.split_once('/') else {
348 return false;
349 };
350 registry == "localhost" || registry.contains('.') || registry.contains(':')
351}
352
353fn has_explicit_tag(value: &str) -> bool {
355 let image_without_digest = value
356 .split_once('@')
357 .map_or(value, |(reference, _digest)| reference);
358 let Some(tag_separator) = image_without_digest.rfind(':') else {
359 return false;
360 };
361 let slash = image_without_digest.rfind('/');
362 slash.is_none_or(|slash| tag_separator > slash)
363}
364
365#[cfg(test)]
366mod tests {
367 use super::{ChallengeImageReference, LocalAgenticsImageReference, OciRegistryImageReference};
368
369 #[test]
371 fn local_agentics_image_accepts_supported_tagged_images() {
372 let reference =
373 LocalAgenticsImageReference::try_new("agentics-linux-arm64-cpu:ubuntu26.04-local")
374 .expect("local Agentics image should parse");
375
376 assert_eq!(reference.repository(), "agentics-linux-arm64-cpu");
377 assert_eq!(reference.tag(), "ubuntu26.04-local");
378 assert_eq!(
379 reference.as_str(),
380 "agentics-linux-arm64-cpu:ubuntu26.04-local"
381 );
382 }
383
384 #[test]
386 fn local_agentics_image_rejects_non_local_or_unsupported_images() {
387 for invalid in [
388 "agentics-linux-arm64-cpu",
389 "ghcr.io/agentic-science/agentics-linux-arm64-cpu:ubuntu26.04-local",
390 "agentics-linux-arm64-cpu:ubuntu26.04-local@sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
391 "python:3.12-slim-bookworm",
392 " agentics-linux-arm64-cpu:ubuntu26.04-local",
393 ] {
394 assert!(
395 LocalAgenticsImageReference::try_new(invalid).is_err(),
396 "{invalid} should be rejected"
397 );
398 }
399 }
400
401 #[test]
403 fn registry_image_requires_explicit_registry_and_tag() {
404 for invalid in [
405 "agentics-linux-arm64-cpu:ubuntu26.04-local",
406 "ghcr.io/agentic-science/agentics-linux-arm64-cpu",
407 "busybox",
408 " ghcr.io/agentic-science/agentics-linux-arm64-cpu:ubuntu26.04-v0.1.0",
409 ] {
410 assert!(
411 OciRegistryImageReference::try_new(invalid).is_err(),
412 "{invalid} should be rejected"
413 );
414 }
415 }
416
417 #[test]
419 fn registry_image_accepts_digest_pinned_ghcr_references() {
420 let reference = OciRegistryImageReference::try_new(format!(
421 "ghcr.io/agentic-science/agentics-linux-arm64-cpu:ubuntu26.04-v0.1.0@sha256:{}",
422 "a".repeat(64)
423 ))
424 .expect("digest-pinned registry image should parse");
425
426 assert_eq!(reference.tag(), "ubuntu26.04-v0.1.0");
427 assert_eq!(
428 reference.policy_repository(),
429 "ghcr.io/agentic-science/agentics-linux-arm64-cpu"
430 );
431 assert!(reference.digest().is_some());
432 }
433
434 #[test]
436 fn registry_image_rejects_non_sha256_digests() {
437 let invalid = format!(
438 "ghcr.io/agentic-science/agentics-linux-arm64-cpu:ubuntu26.04-v0.1.0@sha512:{}",
439 "a".repeat(128)
440 );
441
442 assert!(OciRegistryImageReference::try_new(invalid).is_err());
443 }
444
445 #[test]
447 fn challenge_image_reference_serializes_with_source_tag() {
448 let reference = ChallengeImageReference::Local {
449 reference: LocalAgenticsImageReference::try_new(
450 "agentics-linux-arm64-cpu:ubuntu26.04-local",
451 )
452 .expect("local image should parse"),
453 };
454
455 let value = serde_json::to_value(reference).expect("image reference should serialize");
456
457 assert_eq!(
458 value,
459 serde_json::json!({
460 "source": "local",
461 "reference": "agentics-linux-arm64-cpu:ubuntu26.04-local"
462 })
463 );
464 }
465}