remora/crates/remora-admission/src/source.rs

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, '.' | '_' | '-')
}