use std::collections::BTreeSet; use remora_contracts::{ HttpSourceProjectionDto, HuggingFaceSelectionDto, HuggingFaceSourceProjectionDto, ObjectStoreSourceProjectionDto, SourceProjectionDto, SourceSubmissionDto, }; use remora_credentials::SourceCredentialKind; use remora_domain::{BucketName, MemberPath, OpaqueId, TargetStoreId}; use serde_json::Value; use url::Url; use crate::AdmissionRejection; pub(crate) struct ValidatedSource { pub projection: SourceProjectionDto, pub private_input: Value, pub credential_id: Option, pub allowed_credential_kinds: &'static [SourceCredentialKind], } pub(crate) fn validate_source( source: SourceSubmissionDto, ) -> Result { let private_input = serde_json::to_value(&source).map_err(|_| AdmissionRejection::InvalidSource)?; match source { SourceSubmissionDto::Http(source) => { if source.url.len() > 8192 { return Err(AdmissionRejection::InvalidSource); } let parsed = Url::parse(&source.url).map_err(|_| AdmissionRejection::InvalidSource)?; if parsed.scheme() != "https" || parsed.host_str().is_none() || parsed.port_or_known_default() != Some(443) || !parsed.username().is_empty() || parsed.password().is_some() || parsed.fragment().is_some() { return Err(AdmissionRejection::InvalidSource); } let result_path = MemberPath::parse(&source.result_path) .map_err(|_| AdmissionRejection::InvalidSource)?; let mut display = parsed; display.set_query(None); display.set_fragment(None); let credential_id = source.credential_id; Ok(ValidatedSource { projection: SourceProjectionDto::Http(HttpSourceProjectionDto { display_url: display.to_string(), result_path: result_path.canonical(), credential_id: credential_id.clone(), }), private_input, credential_id, allowed_credential_kinds: &[ SourceCredentialKind::BasicAuth, SourceCredentialKind::BearerToken, SourceCredentialKind::HttpHeaders, ], }) } SourceSubmissionDto::ObjectStore(source) => { TargetStoreId::parse(&source.store).map_err(|_| AdmissionRejection::InvalidSource)?; BucketName::parse(&source.bucket).map_err(|_| AdmissionRejection::InvalidSource)?; if source.key.is_empty() || source.key.len() > 4096 || source.key.chars().any(char::is_control) || source .version_id .as_ref() .is_some_and(|version| version.is_empty() || version.len() > 1024) { return Err(AdmissionRejection::InvalidSource); } let result_path = MemberPath::parse(&source.result_path) .map_err(|_| AdmissionRejection::InvalidSource)?; let credential_id = source.credential_id; Ok(ValidatedSource { projection: SourceProjectionDto::ObjectStore(ObjectStoreSourceProjectionDto { store: source.store, bucket: source.bucket, key: source.key, version_id: source.version_id, result_path: result_path.canonical(), credential_id: credential_id.clone(), }), private_input, credential_id, allowed_credential_kinds: &[SourceCredentialKind::ObjectStore], }) } SourceSubmissionDto::HuggingFace(source) => { if !valid_repository_id(&source.repository_id) || source.revision.is_empty() || source.revision.len() > 256 || source.revision.chars().any(char::is_control) { return Err(AdmissionRejection::InvalidSource); } let selection = source .selection .map(canonical_selection) .transpose()? .flatten(); let credential_id = source.credential_id; Ok(ValidatedSource { projection: SourceProjectionDto::HuggingFace(HuggingFaceSourceProjectionDto { repository_type: source.repository_type, repository_id: source.repository_id, revision: source.revision, selection, credential_id: credential_id.clone(), }), private_input, credential_id, allowed_credential_kinds: &[SourceCredentialKind::HuggingFace], }) } } } fn canonical_selection( selection: HuggingFaceSelectionDto, ) -> Result, AdmissionRejection> { let paths = canonical_paths(selection.paths)?; let prefixes = canonical_paths(selection.prefixes)?; if paths.is_none() && prefixes.is_none() { return Err(AdmissionRejection::InvalidSource); } Ok(Some(HuggingFaceSelectionDto { paths, prefixes })) } fn canonical_paths(paths: Option>) -> Result>, AdmissionRejection> { let Some(paths) = paths else { return Ok(None); }; if paths.is_empty() || paths.len() > 1024 { return Err(AdmissionRejection::InvalidSource); } let canonical = paths .into_iter() .map(|path| { MemberPath::parse(&path) .map(|path| path.canonical()) .map_err(|_| AdmissionRejection::InvalidSource) }) .collect::, _>>()?; Ok(Some(canonical.into_iter().collect())) } fn valid_repository_id(value: &str) -> bool { let Some((owner, name)) = value.split_once('/') else { return false; }; !owner.is_empty() && !name.is_empty() && !name.contains('/') && owner.chars().all(repository_character) && name.chars().all(repository_character) } fn repository_character(character: char) -> bool { character.is_ascii_alphanumeric() || matches!(character, '.' | '_' | '-') }