Add atomic durable job admission

crates/remora-admission/src/source.rs

@@ -0,0 +1,164 @@
+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<OpaqueId>,
+    pub allowed_credential_kinds: &'static [SourceCredentialKind],
+}
+
+pub(crate) fn validate_source(
+    source: SourceSubmissionDto,
+) -> Result<ValidatedSource, AdmissionRejection> {
+    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<Option<HuggingFaceSelectionDto>, 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<Vec<String>>) -> Result<Option<Vec<String>>, 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::<Result<BTreeSet<_>, _>>()?;
+    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, '.' | '_' | '-')
+}