aria_kit/backend/
model_path.rs

1//! Validates model paths, file types, ownership, and permissions before loading.
2
3use std::path::PathBuf;
4
5use crate::primitives::file_inspection::InspectedFile;
6
7use super::ModelError;
8
9pub fn validate_model_file(path: PathBuf) -> Result<PathBuf, ModelError> {
10    InspectedFile::inspect(path)
11        .and_then(InspectedFile::require_regular_file)
12        .and_then(InspectedFile::require_trusted)
13        .map_err(ModelError::Path)
14}
15
16pub fn validate_model_directory(
17    path: PathBuf,
18    required: &[&'static str],
19) -> Result<PathBuf, ModelError> {
20    let directory = InspectedFile::inspect(path)
21        .and_then(InspectedFile::require_directory)
22        .and_then(InspectedFile::require_trusted)
23        .map_err(ModelError::Path)?;
24
25    let entries = std::fs::read_dir(&directory).map_err(|source| ModelError::ReadDirectory {
26        path: directory.clone(),
27        source,
28    })?;
29    for entry in entries {
30        let entry = entry.map_err(|source| ModelError::ReadDirectory {
31            path: directory.clone(),
32            source,
33        })?;
34        let candidate = InspectedFile::inspect(entry.path()).map_err(ModelError::Path)?;
35        // Check regular files beside the weights; subdirectories are not traversed.
36        if candidate.is_regular_file() {
37            candidate.require_trusted().map_err(ModelError::Path)?;
38        }
39    }
40    if let Some(&file) = required.iter().find(|file| !directory.join(file).is_file()) {
41        return Err(ModelError::Missing { directory, file });
42    }
43    Ok(directory)
44}
45
46#[cfg(test)]
47mod tests {
48    use crate::primitives::file_inspection::InspectionError;
49
50    use std::fs;
51    use std::os::unix::fs::PermissionsExt;
52    use std::path::Path;
53
54    use tempfile::TempDir;
55
56    use super::*;
57
58    const SECURE_MODE: u32 = 0o644;
59    const GROUP_WRITABLE_MODE: u32 = 0o664;
60    const GROUP_WRITABLE_DIRECTORY_MODE: u32 = 0o775;
61
62    fn model_directory() -> TempDir {
63        let directory = TempDir::new().expect("test directory should be created");
64        write_model_file(directory.path(), "vocab.txt", SECURE_MODE);
65        directory
66    }
67
68    fn write_model_file(directory: &Path, name: &str, mode: u32) {
69        let path = directory.join(name);
70        fs::write(&path, []).expect("test model file should be written");
71        fs::set_permissions(&path, fs::Permissions::from_mode(mode))
72            .expect("test model file permissions should be set");
73    }
74
75    #[test]
76    fn a_directory_of_trusted_files_is_accepted() {
77        let directory = model_directory();
78
79        let validated = validate_model_directory(directory.path().to_path_buf(), &[])
80            .expect("a trusted model directory should be accepted");
81
82        assert_eq!(
83            validated,
84            fs::canonicalize(directory.path()).expect("the directory should canonicalize")
85        );
86    }
87
88    #[test]
89    fn a_nested_directory_is_left_alone() {
90        let directory = model_directory();
91        let nested = directory.path().join(".cache");
92        fs::create_dir(&nested).expect("nested directory should be created");
93        fs::set_permissions(
94            &nested,
95            fs::Permissions::from_mode(GROUP_WRITABLE_DIRECTORY_MODE),
96        )
97        .expect("nested directory permissions should be set");
98
99        assert!(validate_model_directory(directory.path().to_path_buf(), &[]).is_ok());
100    }
101
102    #[test]
103    fn a_directory_is_not_a_model_file() {
104        let directory = TempDir::new().expect("test directory should be created");
105
106        let error = validate_model_file(directory.path().to_path_buf())
107            .expect_err("a directory is not a model file");
108
109        assert!(matches!(
110            error,
111            ModelError::Path(InspectionError::NotRegularFile { .. })
112        ));
113    }
114
115    #[test]
116    fn a_missing_required_file_is_named() {
117        let directory = model_directory();
118
119        let error = validate_model_directory(
120            directory.path().to_path_buf(),
121            &["vocab.txt", "encoder-model.onnx"],
122        )
123        .expect_err("a directory without a required file must be rejected");
124
125        assert!(matches!(
126            error,
127            ModelError::Missing {
128                file: "encoder-model.onnx",
129                ..
130            }
131        ));
132    }
133
134    #[test]
135    fn a_group_writable_entry_is_rejected() {
136        let directory = model_directory();
137        write_model_file(directory.path(), "encoder-model.onnx", GROUP_WRITABLE_MODE);
138
139        let error = validate_model_directory(directory.path().to_path_buf(), &[])
140            .expect_err("a group-writable model file must be rejected");
141
142        assert!(matches!(
143            error,
144            ModelError::Path(InspectionError::WritableByOthers { actual_mode, .. })
145                if actual_mode == GROUP_WRITABLE_MODE
146        ));
147    }
148
149    #[test]
150    fn a_group_writable_model_file_is_rejected() {
151        let directory = TempDir::new().expect("test directory should be created");
152        write_model_file(directory.path(), "model.bin", GROUP_WRITABLE_MODE);
153
154        let error = validate_model_file(directory.path().join("model.bin"))
155            .expect_err("a group-writable model file must be rejected");
156
157        assert!(matches!(
158            error,
159            ModelError::Path(InspectionError::WritableByOthers { actual_mode, .. })
160                if actual_mode == GROUP_WRITABLE_MODE
161        ));
162    }
163
164    #[test]
165    fn a_group_writable_model_directory_is_rejected() {
166        let directory = model_directory();
167        fs::set_permissions(
168            directory.path(),
169            fs::Permissions::from_mode(GROUP_WRITABLE_DIRECTORY_MODE),
170        )
171        .expect("test directory permissions should be set");
172
173        let error = validate_model_directory(directory.path().to_path_buf(), &[])
174            .expect_err("a group-writable model directory must be rejected");
175
176        assert!(matches!(
177            error,
178            ModelError::Path(InspectionError::WritableByOthers { actual_mode, .. })
179                if actual_mode == GROUP_WRITABLE_DIRECTORY_MODE
180        ));
181    }
182
183    #[test]
184    fn a_regular_file_is_not_a_model_directory() {
185        let directory = TempDir::new().expect("test directory should be created");
186        write_model_file(directory.path(), "model.bin", SECURE_MODE);
187
188        let error = validate_model_directory(directory.path().join("model.bin"), &[])
189            .expect_err("a regular file must be rejected");
190
191        assert!(matches!(
192            error,
193            ModelError::Path(InspectionError::NotDirectory { .. })
194        ));
195    }
196}
aria_kit/backend/model_path.rs

aria_kit/backend/
model_path.rs
