use crate::error::VecboostError;
use hf_hub::{HFClientSync, HFRepositorySync, RepoTypeModel, split_id};
pub fn is_valid_hf_repo_id(repo_id: &str) -> bool {
if repo_id.is_empty() {
return false;
}
if repo_id.starts_with('/') || repo_id.ends_with('/') {
return false;
}
if repo_id.contains("..") || repo_id.contains("//") {
return false;
}
let segments: Vec<&str> = repo_id.split('/').collect();
if segments.len() > 2 {
return false;
}
segments.iter().all(|seg| {
!seg.is_empty()
&& *seg != "."
&& seg
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.')
})
}
pub(crate) fn build_hf_repo(
repo_id: &str,
) -> Result<HFRepositorySync<RepoTypeModel>, VecboostError> {
if !is_valid_hf_repo_id(repo_id) {
return Err(VecboostError::ModelLoadError(format!(
"Invalid HuggingFace repo ID '{}': must match 'organization/model-name' \
pattern with alphanumeric, dash, underscore, dot characters only",
repo_id
)));
}
let api = HFClientSync::new().map_err(|e| VecboostError::ModelLoadError(e.to_string()))?;
let (owner, name) = split_id(repo_id);
Ok(api.model(owner, name))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_is_valid_hf_repo_id_valid_two_segments() {
assert!(is_valid_hf_repo_id("BAAI/bge-m3"));
assert!(is_valid_hf_repo_id(
"sentence-transformers/all-MiniLM-L6-v2"
));
assert!(is_valid_hf_repo_id("org/model_name"));
assert!(is_valid_hf_repo_id("org/model.v2"));
}
#[test]
fn test_is_valid_hf_repo_id_valid_single_segment() {
assert!(is_valid_hf_repo_id("bert-base-uncased"));
assert!(is_valid_hf_repo_id("gpt2"));
assert!(is_valid_hf_repo_id("model_v1.2"));
}
#[test]
fn test_is_valid_hf_repo_id_rejects_empty() {
assert!(!is_valid_hf_repo_id(""));
}
#[test]
fn test_is_valid_hf_repo_id_rejects_path_traversal() {
assert!(!is_valid_hf_repo_id("../etc/passwd"));
assert!(!is_valid_hf_repo_id("org/../../etc/passwd"));
assert!(!is_valid_hf_repo_id("./model"));
assert!(!is_valid_hf_repo_id("org/.."));
}
#[test]
fn test_is_valid_hf_repo_id_rejects_leading_trailing_slash() {
assert!(!is_valid_hf_repo_id("/etc/passwd"));
assert!(!is_valid_hf_repo_id("org/model/"));
assert!(!is_valid_hf_repo_id("/"));
}
#[test]
fn test_is_valid_hf_repo_id_rejects_double_slash() {
assert!(!is_valid_hf_repo_id("org//model"));
assert!(!is_valid_hf_repo_id("//model"));
}
#[test]
fn test_is_valid_hf_repo_id_rejects_more_than_two_segments() {
assert!(!is_valid_hf_repo_id("org/sub/model"));
assert!(!is_valid_hf_repo_id("a/b/c/d"));
}
#[test]
fn test_is_valid_hf_repo_id_rejects_special_chars() {
assert!(!is_valid_hf_repo_id("org/model:name"));
assert!(!is_valid_hf_repo_id("org/model@v1"));
assert!(!is_valid_hf_repo_id("org/model name"));
assert!(!is_valid_hf_repo_id("org/model$evil"));
}
}