use std::collections::HashMap;
use std::ffi::OsString;
use std::fmt;
use std::fs::File;
use std::io::Read;
use std::path::{Path, PathBuf};
use crate::Tensor;
use crate::fs_guard::{MAX_MODEL_FILE_BYTES, open_leaf_no_follow};
use crate::interop::safetensors::{LoadError, SaveError, load_safetensors_f32_from_bytes};
#[cfg(unix)]
mod cache;
#[cfg(unix)]
mod dirfd;
#[cfg(unix)]
mod download;
#[cfg(all(test, unix))]
mod loopback_tests;
#[cfg(unix)]
mod sha256;
#[cfg(all(test, unix))]
mod test_http_server;
#[cfg(unix)]
mod transport;
#[cfg(unix)]
mod url;
const MODEL_FILE_NAME: &str = "model.safetensors";
#[non_exhaustive]
#[derive(Debug)]
pub enum ModelError {
CacheDirUnavailable,
InvalidComponent { kind: &'static str, value: String },
NotFound { name: String, version: String },
TooLarge {
name: String,
version: String,
len: u64,
max: u64,
},
EmptyModel { name: String, version: String },
Save(SaveError),
Load(LoadError),
InvalidUrl { reason: &'static str },
InvalidPin,
Http { status: u16 },
TooManyRedirects { max: u32 },
ContentLengthMismatch { expected: u64, received: u64 },
HashMismatch {
expected: Sha256Pin,
actual: Sha256Pin,
},
Network(std::io::Error),
Io(std::io::Error),
CacheMetadata { reason: &'static str },
StorageFull(std::io::Error),
PermissionDenied(std::io::Error),
}
impl fmt::Display for ModelError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
ModelError::CacheDirUnavailable => {
write!(
f,
"既定のモデルキャッシュディレクトリを解決できません（HOME／USERPROFILE 未設定）"
)
}
ModelError::InvalidComponent { kind, value } => {
write!(f, "不正な {kind}: {value:?}")
}
ModelError::NotFound { name, version } => {
write!(
f,
"モデルが見つかりません（name={name}・version={version}）"
)
}
ModelError::TooLarge {
name,
version,
len,
max,
} => write!(
f,
"モデルファイルが上限を超えています（name={name}・version={version}・{len} バイト > 上限 {max} バイト）"
),
ModelError::EmptyModel { name, version } => write!(
f,
"保存するテンソルが空です（name={name}・version={version}）"
),
ModelError::Save(e) => write!(f, "モデルの保存に失敗しました: {e}"),
ModelError::Load(e) => write!(f, "モデルのロードに失敗しました: {e}"),
ModelError::InvalidUrl { reason } => write!(f, "不正なダウンロード URL: {reason}"),
ModelError::InvalidPin => write!(
f,
"期待 SHA-256 の形式が不正です（64 桁の 16 進文字列が必要です）"
),
ModelError::Http { status } => {
write!(f, "HTTP 応答が成功ではありません（ステータス {status}）")
}
ModelError::TooManyRedirects { max } => {
write!(f, "リダイレクト回数が上限（{max}）を超えました")
}
ModelError::ContentLengthMismatch { expected, received } => write!(
f,
"受信サイズが Content-Length と一致しません（期待 {expected} バイト・受信 {received} バイト）"
),
ModelError::HashMismatch { expected, actual } => {
write!(
f,
"SHA-256 が一致しません（期待 {expected}・実際 {actual}）"
)
}
ModelError::Network(e) => write!(f, "ネットワーク取得に失敗しました: {e}"),
ModelError::Io(e) => write!(f, "モデルレジストリの I/O に失敗しました: {e}"),
ModelError::CacheMetadata { reason } => {
write!(f, "ダウンロードのキャッシュメタデータが不正です: {reason}")
}
ModelError::StorageFull(_) => {
write!(f, "ディスク容量が不足しているため書き込めません")
}
ModelError::PermissionDenied(_) => {
write!(f, "権限がないため書き込めません")
}
}
}
}
impl std::error::Error for ModelError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
ModelError::Save(e) => Some(e),
ModelError::Load(e) => Some(e),
ModelError::Io(e) => Some(e),
ModelError::Network(e) => Some(e),
ModelError::StorageFull(e) => Some(e),
ModelError::PermissionDenied(e) => Some(e),
_ => None,
}
}
}
#[derive(Clone, Copy, PartialEq, Eq)]
pub struct Sha256Pin([u8; 32]);
impl Sha256Pin {
pub fn from_hex(s: &str) -> Result<Sha256Pin, ModelError> {
let b = s.as_bytes();
if b.len() != 64 {
return Err(ModelError::InvalidPin);
}
let nibble = |c: u8| -> Result<u8, ModelError> {
match c {
b'0'..=b'9' => Ok(c - b'0'),
b'a'..=b'f' => Ok(c - b'a' + 10),
b'A'..=b'F' => Ok(c - b'A' + 10),
_ => Err(ModelError::InvalidPin),
}
};
let mut out = [0u8; 32];
for (i, pair) in b.as_chunks::<2>().0.iter().enumerate() {
out[i] = (nibble(pair[0])? << 4) | nibble(pair[1])?;
}
Ok(Sha256Pin(out))
}
}
impl fmt::Display for Sha256Pin {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
for b in self.0 {
write!(f, "{b:02x}")?;
}
Ok(())
}
}
impl fmt::Debug for Sha256Pin {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "Sha256Pin({self})")
}
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy)]
pub struct DownloadProgress {
pub received: u64,
pub total: Option<u64>,
}
#[non_exhaustive]
#[derive(Default)]
pub struct DownloadOptions<'a> {
expected_sha256: Option<Sha256Pin>,
progress: Option<&'a mut dyn FnMut(DownloadProgress)>,
}
impl<'a> DownloadOptions<'a> {
pub fn expected_sha256(mut self, pin: Sha256Pin) -> Self {
self.expected_sha256 = Some(pin);
self
}
pub fn progress(mut self, f: &'a mut dyn FnMut(DownloadProgress)) -> Self {
self.progress = Some(f);
self
}
}
impl fmt::Debug for DownloadOptions<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DownloadOptions")
.field("expected_sha256", &self.expected_sha256)
.field("progress", &self.progress.as_ref().map(|_| ".."))
.finish()
}
}
#[cfg(unix)]
const DOWNLOAD_MAX_REDIRECTS: u32 = 5;
fn validate_component(kind: &'static str, value: &str) -> Result<(), ModelError> {
let ok = !value.is_empty()
&& !value.starts_with('.')
&& value
.chars()
.all(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-'));
if ok {
Ok(())
} else {
Err(ModelError::InvalidComponent {
kind,
value: value.to_string(),
})
}
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
fn production_tmp_name() -> String {
let (pid, nanos, counter) = crate::fs_write::name_entropy();
format!(".{MODEL_FILE_NAME}.tmp-{pid}-{counter}-{nanos}")
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
fn proc_fd_path(dir: &File) -> PathBuf {
use std::os::unix::io::AsRawFd;
PathBuf::from(format!("/proc/self/fd/{}", dir.as_raw_fd()))
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
fn ensure_proc_fd_usable(proc_fd_dir: &Path) -> std::io::Result<()> {
std::fs::read_dir(proc_fd_dir).map(|_| ())
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
fn open_root_dir(root: &Path) -> std::io::Result<File> {
use std::os::unix::fs::OpenOptionsExt;
std::fs::OpenOptions::new()
.read(true)
.custom_flags(crate::fs_guard::open_flags::O_DIRECTORY)
.open(root)
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
fn open_dir_component(parent: &File, component: &str) -> std::io::Result<File> {
use crate::fs_guard::open_flags::{ELOOP, O_DIRECTORY, O_NOFOLLOW};
use std::io::{Error as IoError, ErrorKind};
use std::os::unix::fs::OpenOptionsExt;
const ENOTDIR: i32 = 20;
let path = proc_fd_path(parent).join(component);
match std::fs::create_dir(&path) {
Ok(()) => {}
Err(e) if e.kind() == ErrorKind::AlreadyExists => {}
Err(e) => return Err(e),
}
std::fs::OpenOptions::new()
.read(true)
.custom_flags(O_DIRECTORY | O_NOFOLLOW)
.open(&path)
.map_err(|e| match e.raw_os_error() {
Some(code) if code == ELOOP || code == ENOTDIR => IoError::new(
ErrorKind::InvalidInput,
"保存先の中間パスがシンボリックリンクまたはディレクトリではありません",
),
_ => e,
})
}
fn enforce_size_limit(name: &str, version: &str, len: u64, max: u64) -> Result<(), ModelError> {
if len > max {
Err(ModelError::TooLarge {
name: name.to_string(),
version: version.to_string(),
len,
max,
})
} else {
Ok(())
}
}
fn reserve_read_buffer(capacity_hint: usize) -> Result<Vec<u8>, ModelError> {
let mut bytes = Vec::new();
bytes
.try_reserve(capacity_hint)
.map_err(|_| ModelError::Io(std::io::Error::from(std::io::ErrorKind::OutOfMemory)))?;
Ok(bytes)
}
#[cfg(unix)]
fn verify_safetensors_before_commit(file: &File) -> Result<(), ModelError> {
use std::io::{Seek, SeekFrom};
let len = file.metadata().map_err(ModelError::Io)?.len();
let mut reader = file;
reader.seek(SeekFrom::Start(0)).map_err(ModelError::Io)?;
let mut bytes = reserve_read_buffer(usize::try_from(len).unwrap_or(usize::MAX))?;
reader
.take(MAX_MODEL_FILE_BYTES.saturating_add(1))
.read_to_end(&mut bytes)
.map_err(ModelError::Io)?;
crate::interop::layout::validate_from_bytes(&bytes).map_err(ModelError::Load)?;
Ok(())
}
#[cfg(unix)]
fn map_fs_error(e: std::io::Error) -> ModelError {
use std::io::ErrorKind;
match e.kind() {
ErrorKind::StorageFull | ErrorKind::QuotaExceeded => ModelError::StorageFull(e),
ErrorKind::PermissionDenied => ModelError::PermissionDenied(e),
_ => ModelError::Io(e),
}
}
fn default_cache_dir_from(home: Option<OsString>) -> Option<PathBuf> {
home.map(|h| PathBuf::from(h).join(".fandhe-ai").join("models"))
}
fn non_empty_os_string(value: Option<OsString>) -> Option<OsString> {
value.filter(|v| !v.is_empty())
}
fn select_home_var(
is_windows: bool,
home: Option<OsString>,
userprofile: Option<OsString>,
) -> Option<OsString> {
let home = non_empty_os_string(home);
let userprofile = non_empty_os_string(userprofile);
if is_windows {
userprofile.or(home)
} else {
home.or(userprofile)
}
}
pub struct ModelRegistry {
root: PathBuf,
}
impl ModelRegistry {
pub fn new() -> Result<Self, ModelError> {
let home = std::env::var_os("HOME");
let userprofile = std::env::var_os("USERPROFILE");
let selected = select_home_var(cfg!(windows), home, userprofile);
let root = default_cache_dir_from(selected).ok_or(ModelError::CacheDirUnavailable)?;
Ok(Self { root })
}
pub fn with_cache_dir(dir: impl Into<PathBuf>) -> Self {
Self { root: dir.into() }
}
pub fn cache_dir(&self) -> &Path {
&self.root
}
fn resolve_model_file_with_limit(
&self,
name: &str,
version: &str,
max: u64,
) -> Result<File, ModelError> {
validate_component("name", name)?;
validate_component("version", version)?;
let not_found = || ModelError::NotFound {
name: name.to_string(),
version: version.to_string(),
};
let canonical_root = match self.root.canonicalize() {
Ok(p) => p,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Err(not_found()),
Err(e) => return Err(ModelError::Io(e)),
};
let name_dir = self.root.join(name);
let version_dir = name_dir.join(version);
let leaf = version_dir.join(MODEL_FILE_NAME);
for component in [&name_dir, &version_dir] {
let meta = match std::fs::symlink_metadata(component) {
Ok(m) => m,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Err(not_found()),
Err(e) => return Err(ModelError::Io(e)),
};
let file_type = meta.file_type();
if file_type.is_symlink() || !file_type.is_dir() {
return Err(not_found());
}
}
let leaf_meta = match std::fs::symlink_metadata(&leaf) {
Ok(m) => m,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Err(not_found()),
Err(e) => return Err(ModelError::Io(e)),
};
if leaf_meta.file_type().is_symlink() || !leaf_meta.is_file() {
return Err(not_found());
}
let canonical_leaf = leaf.canonicalize().map_err(ModelError::Io)?;
if !canonical_leaf.starts_with(&canonical_root) {
return Err(not_found());
}
let file = match open_leaf_no_follow(&leaf) {
Ok(f) => f,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Err(not_found()),
#[cfg(any(
all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
),
target_os = "macos"
))]
Err(e) if e.raw_os_error() == Some(crate::fs_guard::open_flags::ELOOP) => {
return Err(not_found());
}
Err(e) => return Err(ModelError::Io(e)),
};
let open_meta = file.metadata().map_err(ModelError::Io)?;
if !open_meta.is_file() {
return Err(not_found());
}
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
if open_meta.dev() != leaf_meta.dev() || open_meta.ino() != leaf_meta.ino() {
return Err(not_found());
}
}
enforce_size_limit(name, version, open_meta.len(), max)?;
Ok(file)
}
fn resolve_model_file(&self, name: &str, version: &str) -> Result<File, ModelError> {
self.resolve_model_file_with_limit(name, version, MAX_MODEL_FILE_BYTES)
}
pub fn load(
&self,
name: &str,
version: &str,
) -> Result<HashMap<String, Tensor<f32>>, ModelError> {
self.load_with_limit(name, version, MAX_MODEL_FILE_BYTES)
}
fn load_with_limit(
&self,
name: &str,
version: &str,
max: u64,
) -> Result<HashMap<String, Tensor<f32>>, ModelError> {
let mut file = self.resolve_model_file_with_limit(name, version, max)?;
let limit = max.saturating_add(1);
let observed_len = file.metadata().map_err(ModelError::Io)?.len();
let capacity_hint = usize::try_from(observed_len.min(limit)).unwrap_or(usize::MAX);
let mut bytes = reserve_read_buffer(capacity_hint)?;
file.by_ref()
.take(limit)
.read_to_end(&mut bytes)
.map_err(ModelError::Io)?;
enforce_size_limit(name, version, bytes.len() as u64, max)?;
load_safetensors_f32_from_bytes(&bytes).map_err(ModelError::Load)
}
pub fn save(
&self,
name: &str,
version: &str,
tensors: &HashMap<String, Tensor<f32>>,
) -> Result<(), ModelError> {
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
{
self.save_with(
name,
version,
tensors,
MAX_MODEL_FILE_BYTES,
production_tmp_name,
|_: &Path| Ok(()),
)
}
#[cfg(not(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
)))]
{
let _ = (name, version, tensors);
Err(ModelError::Io(std::io::Error::from(
std::io::ErrorKind::Unsupported,
)))
}
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
fn save_with(
&self,
name: &str,
version: &str,
tensors: &HashMap<String, Tensor<f32>>,
max: u64,
mut next_tmp: impl FnMut() -> String,
before_commit: impl FnOnce(&Path) -> std::io::Result<()>,
) -> Result<(), ModelError> {
use crate::fs_write::{MAX_TMP_NAME_ATTEMPTS, create_new_with_retry, remove_own_tmp};
use crate::interop::safetensors::save_safetensors_f32_to_bytes;
use std::io::{Error as IoError, ErrorKind, Write};
use std::os::unix::fs::MetadataExt;
validate_component("name", name)?;
validate_component("version", version)?;
if tensors.is_empty() {
return Err(ModelError::EmptyModel {
name: name.to_string(),
version: version.to_string(),
});
}
let mut estimate: u64 = 0;
for t in tensors.values() {
match (t.numel() as u64)
.checked_mul(4)
.and_then(|b| estimate.checked_add(b))
{
Some(v) => estimate = v,
None => {
return Err(ModelError::TooLarge {
name: name.to_string(),
version: version.to_string(),
len: u64::MAX,
max,
});
}
}
}
enforce_size_limit(name, version, estimate, max)?;
let bytes = save_safetensors_f32_to_bytes(tensors, None).map_err(ModelError::Save)?;
enforce_size_limit(name, version, bytes.len() as u64, max)?;
{
let back = load_safetensors_f32_from_bytes(&bytes).map_err(ModelError::Load)?;
let consistent = back.len() == tensors.len()
&& tensors
.iter()
.all(|(k, t)| back.get(k).is_some_and(|b| b.shape() == t.shape()));
if !consistent {
return Err(ModelError::Io(IoError::new(
ErrorKind::InvalidData,
"エンコード結果の読み戻しがキー集合または shape と一致しません",
)));
}
}
ensure_proc_fd_usable(Path::new("/proc/self/fd")).map_err(ModelError::Io)?;
std::fs::create_dir_all(&self.root).map_err(ModelError::Io)?;
let root_dir = open_root_dir(&self.root).map_err(ModelError::Io)?;
let name_fd = open_dir_component(&root_dir, name).map_err(ModelError::Io)?;
let version_fd = open_dir_component(&name_fd, version).map_err(ModelError::Io)?;
let version_dir = proc_fd_path(&version_fd);
let (mut tmp_file, tmp_name) =
create_new_with_retry(&version_dir, &mut next_tmp, MAX_TMP_NAME_ATTEMPTS)
.map_err(ModelError::Io)?;
let tmp_path = version_dir.join(&tmp_name);
let leaf = version_dir.join(MODEL_FILE_NAME);
let committed: std::io::Result<()> = (|| {
tmp_file.write_all(&bytes)?;
let held = tmp_file.metadata()?;
let on_disk = std::fs::symlink_metadata(&tmp_path)?;
if !on_disk.file_type().is_file()
|| on_disk.dev() != held.dev()
|| on_disk.ino() != held.ino()
{
return Err(IoError::new(
ErrorKind::InvalidData,
"一時ファイルが書き込み中に差し替えられました",
));
}
before_commit(&tmp_path)?;
std::fs::rename(&tmp_path, &leaf)
})();
if let Err(e) = committed {
remove_own_tmp(&tmp_file, &tmp_path);
return Err(ModelError::Io(e));
}
let held = tmp_file.metadata().map_err(ModelError::Io)?;
let committed_file = match self.resolve_model_file_with_limit(name, version, max) {
Ok(f) => f,
Err(ModelError::NotFound { .. }) => {
return Err(ModelError::Io(IoError::new(
ErrorKind::InvalidData,
"保存直後のモデルファイルを解決できません",
)));
}
Err(e) => return Err(e),
};
let now = committed_file.metadata().map_err(ModelError::Io)?;
if now.dev() != held.dev() || now.ino() != held.ino() {
return Err(ModelError::Io(IoError::new(
ErrorKind::InvalidData,
"確定したファイルが書き込んだ実体と一致しません",
)));
}
Ok(())
}
pub fn download(&self, url: &str, name: &str, version: &str) -> Result<(), ModelError> {
self.download_with(url, name, version, DownloadOptions::default())
}
pub fn download_with(
&self,
url: &str,
name: &str,
version: &str,
options: DownloadOptions<'_>,
) -> Result<(), ModelError> {
validate_component("name", name)?;
validate_component("version", version)?;
#[cfg(unix)]
{
let url = url::validate_https_url(url)?;
let mut transport = transport::UreqTransport::new()?;
self.download_via(
&mut transport,
url,
name,
version,
options,
MAX_MODEL_FILE_BYTES,
)
}
#[cfg(not(unix))]
{
let _ = (url, options);
Err(ModelError::Io(std::io::Error::from(
std::io::ErrorKind::Unsupported,
)))
}
}
#[cfg(unix)]
fn download_via(
&self,
transport: &mut dyn transport::Transport,
url: url::ValidatedUrl,
name: &str,
version: &str,
options: DownloadOptions<'_>,
max_bytes: u64,
) -> Result<(), ModelError> {
let DownloadOptions {
expected_sha256,
progress,
} = options;
let root = self.root.clone();
let mut open_dir = || -> std::io::Result<std::os::fd::OwnedFd> {
let root_fd = dirfd::open_root(&root)?;
let name_fd = dirfd::open_or_create_dir(&root_fd, name)?;
dirfd::open_or_create_dir(&name_fd, version)
};
let mut lookup_dir = || -> std::io::Result<Option<std::os::fd::OwnedFd>> {
let Some(root_fd) = dirfd::open_root_existing(&root)? else {
return Ok(None);
};
let Some(name_fd) = dirfd::open_existing_dir(&root_fd, name)? else {
return Ok(None);
};
dirfd::open_existing_dir(&name_fd, version)
};
let cache_meta = lookup_dir()
.ok()
.flatten()
.and_then(|d| cache::read(&d))
.filter(|m| {
m.url == url.as_str()
&& m.etag.is_some()
&& expected_sha256.is_none_or(|pin| m.sha256 == pin)
});
let mut next_staging = || dirfd::production_staging_name(MODEL_FILE_NAME);
let mut next_tmp = || dirfd::production_tmp_name(MODEL_FILE_NAME);
let mut now = || {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |d| d.as_secs())
};
let job = download::Job {
url,
name,
version,
max_bytes,
max_redirects: DOWNLOAD_MAX_REDIRECTS,
deadline: std::time::Instant::now() + transport::OVERALL_DEADLINE,
expected: expected_sha256,
cache: cache_meta,
};
download::run(
transport,
job,
progress,
download::Hooks {
lookup_dir: &mut lookup_dir,
open_dir: &mut open_dir,
next_staging: &mut next_staging,
next_tmp: &mut next_tmp,
verify_before_commit: &mut verify_safetensors_before_commit,
now: &mut now,
},
)
}
pub fn available_models(&self) -> Vec<(String, Vec<String>)> {
let Ok(entries) = std::fs::read_dir(&self.root) else {
return Vec::new();
};
let mut result: Vec<(String, Vec<String>)> = Vec::new();
for entry in entries.flatten() {
let Ok(name) = entry.file_name().into_string() else {
continue;
};
if validate_component("name", &name).is_err() {
continue;
}
let Ok(file_type) = entry.file_type() else {
continue;
};
if !file_type.is_dir() {
continue;
}
let name_dir = entry.path();
let Ok(version_entries) = std::fs::read_dir(&name_dir) else {
continue;
};
let mut versions: Vec<String> = Vec::new();
for version_entry in version_entries.flatten() {
let Ok(version) = version_entry.file_name().into_string() else {
continue;
};
if validate_component("version", &version).is_err() {
continue;
}
if self.resolve_model_file(&name, &version).is_ok() {
versions.push(version);
}
}
if versions.is_empty() {
continue;
}
versions.sort();
result.push((name, versions));
}
result.sort_by(|a, b| a.0.cmp(&b.0));
result
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn reserve_read_buffer_maps_alloc_failure_to_out_of_memory() {
let err = reserve_read_buffer(usize::MAX).unwrap_err();
match err {
ModelError::Io(e) => {
assert_eq!(e.kind(), std::io::ErrorKind::OutOfMemory);
assert!(e.get_ref().is_none());
}
other => panic!("ModelError::Io を期待したが {other:?} だった"),
}
}
#[test]
fn reserve_read_buffer_succeeds_for_small_capacity() {
let bytes = reserve_read_buffer(16).expect("小さい確保要求は成功するべき");
assert!(bytes.capacity() >= 16);
assert_eq!(bytes.len(), 0);
}
#[test]
fn enforce_size_limit_accepts_within_bound() {
assert!(enforce_size_limit("mlp", "v1", 0, 100).is_ok());
assert!(enforce_size_limit("mlp", "v1", 100, 100).is_ok());
}
#[test]
fn enforce_size_limit_rejects_over_bound() {
let err = enforce_size_limit("mlp", "v1", 101, 100).unwrap_err();
match err {
ModelError::TooLarge {
name,
version,
len,
max,
} => {
assert_eq!(name, "mlp");
assert_eq!(version, "v1");
assert_eq!(len, 101);
assert_eq!(max, 100);
}
other => panic!("ModelError::TooLarge を期待したが {other:?} だった"),
}
}
#[test]
fn too_large_error_display_includes_sizes() {
let err = enforce_size_limit("mlp", "v1", 101, 100).unwrap_err();
let rendered = err.to_string();
assert!(rendered.contains("101"), "rendered={rendered:?}");
assert!(rendered.contains("100"), "rendered={rendered:?}");
}
#[test]
fn validate_component_accepts_ascii_allowlist() {
for value in ["a", "mlp", "v1", "1.2.3", "a_b-c.d", "A9"] {
assert!(
validate_component("name", value).is_ok(),
"許可されるべき値が拒否された: {value:?}"
);
}
}
#[test]
fn validate_component_rejects_traversal_and_unsafe_values() {
for value in [
"",
".",
"..",
"a/b",
"a\\b",
".hidden",
"a b",
"a\0b",
"/etc/passwd",
] {
assert!(
validate_component("name", value).is_err(),
"拒否されるべき値が許可された: {value:?}"
);
}
}
#[test]
fn default_cache_dir_from_home_appends_fixed_suffix() {
let home = Some(OsString::from("/home/x"));
let dir = default_cache_dir_from(home).unwrap();
assert_eq!(dir, PathBuf::from("/home/x/.fandhe-ai/models"));
}
#[test]
fn default_cache_dir_from_none_is_none() {
assert!(default_cache_dir_from(None).is_none());
}
#[test]
fn select_home_var_prefers_userprofile_on_windows() {
let home = Some(OsString::from("/home/x"));
let userprofile = Some(OsString::from(r"C:\Users\x"));
assert_eq!(
select_home_var(true, home.clone(), userprofile.clone()),
userprofile
);
assert_eq!(
select_home_var(true, home, None),
Some(OsString::from("/home/x"))
);
}
#[test]
fn select_home_var_prefers_home_on_non_windows() {
let home = Some(OsString::from("/home/x"));
let userprofile = Some(OsString::from(r"C:\Users\x"));
assert_eq!(
select_home_var(false, home.clone(), userprofile.clone()),
home
);
assert_eq!(
select_home_var(false, None, userprofile.clone()),
userprofile
);
}
#[test]
fn select_home_var_none_when_both_unset() {
assert_eq!(select_home_var(true, None, None), None);
assert_eq!(select_home_var(false, None, None), None);
}
#[test]
fn non_empty_os_string_filters_empty_value() {
assert_eq!(non_empty_os_string(Some(OsString::from(""))), None);
}
#[test]
fn non_empty_os_string_keeps_non_empty_value() {
let value = Some(OsString::from("/home/x"));
assert_eq!(non_empty_os_string(value.clone()), value);
}
#[test]
fn non_empty_os_string_none_stays_none() {
assert_eq!(non_empty_os_string(None), None);
}
#[test]
fn select_home_var_none_when_both_empty_strings() {
let empty = || Some(OsString::from(""));
assert_eq!(select_home_var(true, empty(), empty()), None);
assert_eq!(select_home_var(false, empty(), empty()), None);
assert!(default_cache_dir_from(select_home_var(true, empty(), empty())).is_none());
}
#[test]
fn select_home_var_falls_back_when_preferred_is_empty_string() {
let value = Some(OsString::from("/x"));
assert_eq!(
select_home_var(false, Some(OsString::from("")), value.clone()),
value
);
assert_eq!(
select_home_var(true, value.clone(), Some(OsString::from(""))),
value
);
}
struct TempDirGuard {
path: std::path::PathBuf,
}
impl TempDirGuard {
fn new(label: &str) -> Self {
use std::sync::atomic::{AtomicU64, Ordering};
static COUNTER: AtomicU64 = AtomicU64::new(0);
let pid = std::process::id();
for _ in 0..64 {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
let seq = COUNTER.fetch_add(1, Ordering::Relaxed);
let candidate = std::env::temp_dir().join(format!(
"fandhe-ai-model-registry-unit-{pid}-{nanos}-{seq}-{label}"
));
match std::fs::create_dir(&candidate) {
Ok(()) => return Self { path: candidate },
Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => continue,
Err(e) => panic!("一時ディレクトリを作成できない: {e}"),
}
}
panic!("一意な一時ディレクトリ名を 64 回試行しても確保できない");
}
fn path(&self) -> &Path {
&self.path
}
}
impl Drop for TempDirGuard {
fn drop(&mut self) {
let _ = std::fs::remove_dir_all(&self.path);
}
}
fn write_leaf(dir: &Path, contents: &[u8]) {
let version_dir = dir.join("mlp").join("v1");
std::fs::create_dir_all(&version_dir).unwrap();
std::fs::write(version_dir.join(MODEL_FILE_NAME), contents).unwrap();
}
#[cfg(any(
all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
),
target_os = "macos"
))]
#[test]
fn resolve_model_file_with_limit_rejects_over_bound() {
let dir = TempDirGuard::new("resolve-over-bound");
write_leaf(dir.path(), b"123456789");
let registry = ModelRegistry::with_cache_dir(dir.path());
let err = registry
.resolve_model_file_with_limit("mlp", "v1", 8)
.unwrap_err();
match err {
ModelError::TooLarge {
name,
version,
len,
max,
} => {
assert_eq!(name, "mlp");
assert_eq!(version, "v1");
assert_eq!(len, 9);
assert_eq!(max, 8);
}
other => panic!("ModelError::TooLarge を期待したが {other:?} だった"),
}
}
#[cfg(any(
all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
),
target_os = "macos"
))]
#[test]
fn resolve_model_file_with_limit_accepts_at_bound() {
let dir = TempDirGuard::new("resolve-at-bound");
write_leaf(dir.path(), b"12345678");
let registry = ModelRegistry::with_cache_dir(dir.path());
assert!(
registry
.resolve_model_file_with_limit("mlp", "v1", 8)
.is_ok()
);
}
#[cfg(any(
all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
),
target_os = "macos"
))]
#[test]
fn load_with_limit_rejects_over_bound() {
let dir = TempDirGuard::new("load-over-bound");
write_leaf(dir.path(), b"123456789");
let registry = ModelRegistry::with_cache_dir(dir.path());
let err = registry.load_with_limit("mlp", "v1", 8).unwrap_err();
assert!(
matches!(err, ModelError::TooLarge { len: 9, max: 8, .. }),
"ModelError::TooLarge {{ len: 9, max: 8, .. }} を期待したが {err:?} だった"
);
}
#[cfg(not(any(
all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
),
target_os = "macos"
)))]
#[test]
fn load_is_unsupported_on_non_supported_platform() {
let dir = TempDirGuard::new("unsupported");
write_leaf(dir.path(), b"12345678");
let registry = ModelRegistry::with_cache_dir(dir.path());
let err = registry
.resolve_model_file_with_limit("mlp", "v1", 8)
.unwrap_err();
assert!(
matches!(&err, ModelError::Io(e) if e.kind() == std::io::ErrorKind::Unsupported),
"Io(Unsupported) を期待したが {err:?} だった"
);
write_leaf(dir.path(), b"123456789");
let err = registry
.resolve_model_file_with_limit("mlp", "v1", 8)
.unwrap_err();
assert!(
matches!(&err, ModelError::Io(e) if e.kind() == std::io::ErrorKind::Unsupported),
"TooLarge ではなく Io(Unsupported) を期待したが {err:?} だった"
);
let err = registry.load_with_limit("mlp", "v1", 8).unwrap_err();
assert!(
matches!(&err, ModelError::Io(e) if e.kind() == std::io::ErrorKind::Unsupported),
"Io(Unsupported) を期待したが {err:?} だった"
);
}
fn sample_state() -> HashMap<String, Tensor<f32>> {
let mut m = HashMap::new();
m.insert(
"fc.weight".to_string(),
Tensor::new(
vec![0.1_f32, -0.2, 0.3, -0.4, f32::MIN_POSITIVE, 1e30],
&[2, 3],
)
.unwrap(),
);
m.insert(
"fc.bias".to_string(),
Tensor::new(vec![-0.0_f32, 7.5], &[2]).unwrap(),
);
m
}
#[cfg(any(
all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
),
target_os = "macos"
))]
fn bits(t: &Tensor<f32>) -> Vec<u32> {
t.contiguous()
.as_slice()
.expect("contiguous は連続")
.iter()
.map(|v| v.to_bits())
.collect()
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
fn tmp_leftovers(dir: &Path) -> Vec<String> {
std::fs::read_dir(dir)
.unwrap()
.flatten()
.map(|e| e.file_name().to_string_lossy().into_owned())
.filter(|n| n.starts_with(".model.safetensors.tmp-"))
.collect()
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
#[test]
fn save_then_load_roundtrips_bit_exact() {
let dir = TempDirGuard::new("save-roundtrip");
let registry = ModelRegistry::with_cache_dir(dir.path().join("root"));
let state = sample_state();
registry.save("mlp", "v1", &state).unwrap();
let back = registry.load("mlp", "v1").unwrap();
assert_eq!(back.len(), state.len());
for (k, t) in &state {
let b = back.get(k).expect("キーが往復するはず");
assert_eq!(b.shape(), t.shape());
assert_eq!(bits(b), bits(t), "{k}");
}
assert!(tmp_leftovers(&dir.path().join("root/mlp/v1")).is_empty());
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
#[test]
fn save_rejects_invalid_components_before_fs_access() {
let dir = TempDirGuard::new("save-invalid");
let root = dir.path().join("root");
let registry = ModelRegistry::with_cache_dir(&root);
let state = sample_state();
for bad in ["", ".", "..", ".hidden", "a/b", "a\\b", "a\0b", "a\nb"] {
assert!(
matches!(
registry.save(bad, "v1", &state),
Err(ModelError::InvalidComponent { kind: "name", .. })
),
"name={bad:?}"
);
assert!(
matches!(
registry.save("mlp", bad, &state),
Err(ModelError::InvalidComponent {
kind: "version",
..
})
),
"version={bad:?}"
);
}
assert!(!root.exists(), "検証失敗でルートを作ってはならない");
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
#[test]
fn save_rejects_empty_model_without_creating_root() {
let dir = TempDirGuard::new("save-empty");
let root = dir.path().join("root");
let registry = ModelRegistry::with_cache_dir(&root);
let err = registry.save("mlp", "v1", &HashMap::new()).unwrap_err();
assert!(matches!(err, ModelError::EmptyModel { .. }), "{err:?}");
assert!(!root.exists());
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
#[test]
fn save_with_limit_rejects_over_bound_without_creating_root() {
let dir = TempDirGuard::new("save-limit");
let root = dir.path().join("root");
let registry = ModelRegistry::with_cache_dir(&root);
let err = registry
.save_with("mlp", "v1", &sample_state(), 8, production_tmp_name, |_| {
Ok(())
})
.unwrap_err();
assert!(
matches!(err, ModelError::TooLarge { max: 8, .. }),
"{err:?}"
);
assert!(!root.exists());
let data_bytes = 8 * 4;
let err = registry
.save_with(
"mlp",
"v1",
&sample_state(),
data_bytes,
production_tmp_name,
|_| Ok(()),
)
.unwrap_err();
assert!(matches!(err, ModelError::TooLarge { .. }), "{err:?}");
assert!(!root.exists());
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
#[test]
fn save_failure_before_commit_keeps_existing_and_leaves_no_tmp() {
let dir = TempDirGuard::new("save-fail");
let registry = ModelRegistry::with_cache_dir(dir.path().join("root"));
let state = sample_state();
registry.save("mlp", "v1", &state).unwrap();
let mut other = HashMap::new();
other.insert("x".to_string(), Tensor::new(vec![1.0_f32], &[1]).unwrap());
let err = registry
.save_with(
"mlp",
"v1",
&other,
MAX_MODEL_FILE_BYTES,
production_tmp_name,
|_| Err(std::io::Error::other("注入した失敗")),
)
.unwrap_err();
assert!(matches!(err, ModelError::Io(_)), "{err:?}");
let back = registry.load("mlp", "v1").unwrap();
assert_eq!(back.len(), state.len());
for (k, t) in &state {
assert_eq!(bits(&back[k]), bits(t));
}
assert!(tmp_leftovers(&dir.path().join("root/mlp/v1")).is_empty());
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
#[test]
fn save_gives_up_after_tmp_name_collisions() {
let dir = TempDirGuard::new("save-collide");
let registry = ModelRegistry::with_cache_dir(dir.path().join("root"));
let state = sample_state();
registry.save("mlp", "v1", &state).unwrap();
let vdir = dir.path().join("root/mlp/v1");
std::fs::write(vdir.join(".collide"), b"keep").unwrap();
let err = registry
.save_with(
"mlp",
"v1",
&state,
MAX_MODEL_FILE_BYTES,
|| ".collide".to_string(),
|_| Ok(()),
)
.unwrap_err();
assert!(
matches!(&err, ModelError::Io(e) if e.kind() == std::io::ErrorKind::AlreadyExists),
"{err:?}"
);
assert_eq!(std::fs::read(vdir.join(".collide")).unwrap(), b"keep");
assert_eq!(registry.load("mlp", "v1").unwrap().len(), state.len());
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
#[test]
fn save_overwrites_existing_atomically() {
let dir = TempDirGuard::new("save-overwrite");
let registry = ModelRegistry::with_cache_dir(dir.path().join("root"));
registry.save("mlp", "v1", &sample_state()).unwrap();
let mut next = HashMap::new();
next.insert(
"w".to_string(),
Tensor::new(vec![3.0_f32, 4.0], &[2]).unwrap(),
);
registry.save("mlp", "v1", &next).unwrap();
let back = registry.load("mlp", "v1").unwrap();
assert_eq!(back.len(), 1);
assert_eq!(bits(&back["w"]), bits(&next["w"]));
assert_eq!(
registry.available_models(),
vec![("mlp".to_string(), vec!["v1".to_string()])]
);
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
#[test]
fn proc_fd_precheck_rejects_missing_proc_and_accepts_real_proc() {
let dir = TempDirGuard::new("save-noproc");
assert!(ensure_proc_fd_usable(&dir.path().join("no-such-proc")).is_err());
assert!(ensure_proc_fd_usable(Path::new("/proc/self/fd")).is_ok());
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
#[test]
fn save_rejects_symlinked_name_dir_without_writing_outside() {
let dir = TempDirGuard::new("save-symlink");
let root = dir.path().join("root");
let outside = dir.path().join("outside");
std::fs::create_dir_all(&root).unwrap();
std::fs::create_dir_all(&outside).unwrap();
std::os::unix::fs::symlink(&outside, root.join("mlp")).unwrap();
let registry = ModelRegistry::with_cache_dir(&root);
let err = registry.save("mlp", "v1", &sample_state()).unwrap_err();
assert!(
matches!(&err, ModelError::Io(e) if e.kind() == std::io::ErrorKind::InvalidInput),
"{err:?}"
);
assert_eq!(std::fs::read_dir(&outside).unwrap().count(), 0);
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
#[test]
fn save_does_not_write_outside_when_version_dir_swapped_to_symlink() {
let dir = TempDirGuard::new("save-swap");
let root = dir.path().join("root");
let outside = dir.path().join("outside");
std::fs::create_dir_all(&outside).unwrap();
let registry = ModelRegistry::with_cache_dir(&root);
let vdir = root.join("mlp").join("v1");
let moved = root.join("mlp").join("v1-moved");
let outside_c = outside.clone();
let (vdir_c, moved_c) = (vdir.clone(), moved.clone());
let result = registry.save_with(
"mlp",
"v1",
&sample_state(),
MAX_MODEL_FILE_BYTES,
production_tmp_name,
move |_| {
std::fs::rename(&vdir_c, &moved_c)?;
std::os::unix::fs::symlink(&outside_c, &vdir_c)
},
);
assert!(result.is_err(), "差し替えを検出して Err のはず");
assert_eq!(
std::fs::read_dir(&outside).unwrap().count(),
0,
"ルート外へ書き込んではならない"
);
}
#[cfg(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
))]
#[test]
fn save_does_not_write_outside_when_name_dir_swapped_to_symlink() {
let dir = TempDirGuard::new("save-swap-name");
let root = dir.path().join("root");
let outside = dir.path().join("outside");
std::fs::create_dir_all(outside.join("v1")).unwrap();
let registry = ModelRegistry::with_cache_dir(&root);
let ndir = root.join("mlp");
let moved = root.join("mlp-moved");
let (outside_c, ndir_c, moved_c) = (outside.clone(), ndir.clone(), moved.clone());
let result = registry.save_with(
"mlp",
"v1",
&sample_state(),
MAX_MODEL_FILE_BYTES,
production_tmp_name,
move |_| {
std::fs::rename(&ndir_c, &moved_c)?;
std::os::unix::fs::symlink(&outside_c, &ndir_c)
},
);
assert!(result.is_err(), "差し替えを検出して Err のはず");
assert_eq!(std::fs::read_dir(outside.join("v1")).unwrap().count(), 0);
}
#[cfg(not(all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
)))]
#[test]
fn save_is_unsupported_on_non_supported_platform() {
let registry = ModelRegistry::with_cache_dir("unused-root-for-unsupported-save");
let err = registry.save("mlp", "v1", &sample_state()).unwrap_err();
assert!(
matches!(&err, ModelError::Io(e) if e.kind() == std::io::ErrorKind::Unsupported),
"{err:?}"
);
}
#[test]
fn new_error_variants_display_without_path() {
let e = ModelError::EmptyModel {
name: "mlp".into(),
version: "v1".into(),
};
let text = format!("{e} {e:?}");
assert!(text.contains("mlp") && text.contains("v1"));
assert!(!text.contains('/'));
use std::error::Error as _;
assert!(e.source().is_none());
}
#[test]
fn sha256_pin_from_hex_accepts_64_hex_either_case() {
let lower = "0123456789abcdef".repeat(4);
let upper = lower.to_uppercase();
let a = Sha256Pin::from_hex(&lower).unwrap();
assert_eq!(a, Sha256Pin::from_hex(&upper).unwrap());
assert_eq!(a.to_string(), lower);
assert!(format!("{a:?}").contains(&lower));
}
#[test]
fn sha256_pin_from_hex_rejects_malformed() {
let ok = "a".repeat(64);
for bad in [
"a".repeat(63),
"a".repeat(65),
String::new(),
format!("0x{}", "a".repeat(62)),
format!(" {}", "a".repeat(63)),
format!("{}g", "a".repeat(63)),
format!("{}\n", "a".repeat(63)),
] {
assert!(
matches!(Sha256Pin::from_hex(&bad), Err(ModelError::InvalidPin)),
"{bad:?}"
);
}
assert!(Sha256Pin::from_hex(&ok).is_ok());
}
#[test]
fn download_error_displays_never_contain_url_material() {
let secret = "secret-token-xyz";
let errs = [
ModelError::InvalidUrl { reason: "r" },
ModelError::InvalidPin,
ModelError::Http { status: 404 },
ModelError::TooManyRedirects { max: 5 },
ModelError::ContentLengthMismatch {
expected: 1,
received: 2,
},
ModelError::HashMismatch {
expected: Sha256Pin([1; 32]),
actual: Sha256Pin([2; 32]),
},
ModelError::Network(std::io::Error::other("net")),
ModelError::CacheMetadata { reason: "r" },
ModelError::StorageFull(std::io::Error::other("full")),
ModelError::PermissionDenied(std::io::Error::other("perm")),
];
for e in errs {
let text = format!("{e} {e:?}");
assert!(
!text.contains(secret) && !text.contains("https://"),
"{text}"
);
}
}
#[cfg(unix)]
fn tmp_root(tag: &str) -> PathBuf {
let (pid, nanos, counter) = crate::fs_write::name_entropy();
std::env::temp_dir().join(format!("fandhe-dlreg-{tag}-{pid}-{counter}-{nanos}"))
}
#[cfg(unix)]
#[test]
fn download_rejects_bad_input_before_touching_fs() {
let root = tmp_root("reject");
let registry = ModelRegistry::with_cache_dir(&root);
for url in [
"http://example.com/m",
"https://user:pw@example.com/m",
"file:///etc/passwd",
"https://example.com/\n",
"ftp://example.com/m",
"https://user@example.com/m",
] {
let err = registry.download(url, "m", "v1").unwrap_err();
assert!(matches!(err, ModelError::InvalidUrl { .. }), "{err:?}");
let text = err.to_string();
assert!(!text.contains("pw@") && !text.contains("etc/passwd"));
}
let err = registry
.download("https://example.com/m", "../x", "v1")
.unwrap_err();
assert!(matches!(err, ModelError::InvalidComponent { .. }));
assert!(!root.exists(), "拒否時にキャッシュルートを作らない");
}
#[cfg(unix)]
#[test]
fn bad_url_and_malformed_pin_never_connect() {
use super::test_http_server::{MockReply, MockServer};
let root = tmp_root("noconnect");
let registry = ModelRegistry::with_cache_dir(&root);
let srv = MockServer::start(vec![MockReply::Stall]);
let p = srv.port();
for url in [
format!("http://127.0.0.1:{p}/m"),
"file:///etc/passwd".to_string(),
format!("ftp://127.0.0.1:{p}/m"),
format!("https://u:secretpw@127.0.0.1:{p}/m"),
] {
let err = registry.download(&url, "m", "v1").unwrap_err();
assert!(matches!(err, ModelError::InvalidUrl { .. }), "{err:?}");
let text = err.to_string();
assert!(!text.contains("secretpw") && !text.contains("127.0.0.1"));
}
for bad in ["", "zz", &"g".repeat(64), &"0".repeat(63)] {
assert!(matches!(
Sha256Pin::from_hex(bad),
Err(ModelError::InvalidPin)
));
}
assert_eq!(srv.connection_count(), 0, "接続前に拒否される");
assert!(!root.exists(), "拒否時にキャッシュルートを作らない");
}
#[cfg(unix)]
struct Scripted {
steps: std::collections::VecDeque<(u16, Option<&'static str>, Vec<u8>)>,
conditional: Vec<Option<String>>,
}
#[cfg(unix)]
impl Scripted {
fn new(steps: Vec<(u16, Option<&'static str>, Vec<u8>)>) -> Self {
Self {
steps: steps.into(),
conditional: Vec::new(),
}
}
fn one(body: Vec<u8>) -> Self {
Self::new(vec![(200, None, body)])
}
}
#[cfg(unix)]
impl transport::Transport for Scripted {
fn get(
&mut self,
_u: &url::ValidatedUrl,
_d: std::time::Instant,
if_none_match: Option<&str>,
) -> Result<transport::RawResponse, ModelError> {
self.conditional.push(if_none_match.map(String::from));
let (status, etag, body) = self.steps.pop_front().expect("scripted step");
Ok(transport::RawResponse {
status,
location: None,
content_length: Some(body.len() as u64),
content_encoding: None,
etag: etag.map(String::from),
body: Box::new(std::io::Cursor::new(body)),
})
}
}
#[cfg(unix)]
fn dir_names(dir: &Path) -> Vec<String> {
let mut v: Vec<String> = std::fs::read_dir(dir)
.unwrap()
.flatten()
.map(|e| e.file_name().to_string_lossy().into_owned())
.collect();
v.sort();
v
}
#[cfg(unix)]
fn dl(
registry: &ModelRegistry,
t: &mut Scripted,
opts: DownloadOptions<'_>,
) -> Result<(), ModelError> {
let u = url::validate_https_url("https://example.com/m").unwrap();
registry.download_via(t, u, "m", "v1", opts, MAX_MODEL_FILE_BYTES)
}
#[cfg(any(
all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
),
target_os = "macos"
))]
#[test]
fn download_via_publishes_and_load_reads_back() {
let root = tmp_root("via");
let registry = ModelRegistry::with_cache_dir(&root);
let state = sample_state();
let bytes =
crate::interop::safetensors::save_safetensors_f32_to_bytes(&state, None).unwrap();
dl(
®istry,
&mut Scripted::one(bytes.clone()),
DownloadOptions::default(),
)
.unwrap();
assert_eq!(
std::fs::read(root.join("m/v1").join(MODEL_FILE_NAME)).unwrap(),
bytes
);
assert_eq!(
dir_names(&root.join("m/v1")),
vec!["download.json".to_string(), MODEL_FILE_NAME.to_string()]
);
let loaded = registry.load("m", "v1").unwrap();
assert_eq!(loaded.len(), state.len());
for (k, v) in &state {
assert_eq!(bits(&loaded[k]), bits(v), "{k}");
}
assert_eq!(registry.available_models().len(), 1);
let _ = std::fs::remove_dir_all(&root);
}
#[cfg(unix)]
#[test]
fn download_via_rejects_invalid_safetensors_and_keeps_existing_model() {
let root = tmp_root("via-invalid");
let registry = ModelRegistry::with_cache_dir(&root);
let good =
crate::interop::safetensors::save_safetensors_f32_to_bytes(&sample_state(), None)
.unwrap();
let run = |body: Vec<u8>| {
dl(
®istry,
&mut Scripted::one(body),
DownloadOptions::default(),
)
};
run(good.clone()).unwrap();
let mut bad_offsets = good.clone();
bad_offsets.truncate(bad_offsets.len() - 4);
for body in [
b"<html>404 not found</html>".to_vec(),
vec![0xff; 64],
bad_offsets,
] {
let e = run(body).unwrap_err();
assert!(matches!(e, ModelError::Load(_)), "got {e:?}");
assert_eq!(
std::fs::read(root.join("m/v1").join(MODEL_FILE_NAME)).unwrap(),
good,
"既存の正常なモデルが残る"
);
assert_eq!(
dir_names(&root.join("m/v1")),
vec!["download.json".to_string(), MODEL_FILE_NAME.to_string()],
"一時ファイルなし"
);
}
let _ = std::fs::remove_dir_all(&root);
}
#[cfg(unix)]
#[test]
fn download_via_accepts_f16_safetensors_but_load_stays_f32_only() {
let root = tmp_root("via-f16");
let registry = ModelRegistry::with_cache_dir(&root);
let header = br#"{"w":{"dtype":"F16","shape":[2],"data_offsets":[0,4]}}"#;
let mut bytes = (header.len() as u64).to_le_bytes().to_vec();
bytes.extend_from_slice(header);
bytes.extend_from_slice(&[0u8; 4]);
dl(
®istry,
&mut Scripted::one(bytes),
DownloadOptions::default(),
)
.unwrap();
assert!(root.join("m/v1").join(MODEL_FILE_NAME).is_file());
#[cfg(any(
all(
target_os = "linux",
any(target_arch = "x86_64", target_arch = "aarch64")
),
target_os = "macos"
))]
assert!(matches!(
registry.load("m", "v1"),
Err(ModelError::Load(LoadError::UnsupportedDtype { .. }))
));
let _ = std::fs::remove_dir_all(&root);
}
#[cfg(unix)]
#[test]
fn download_via_second_fetch_uses_conditional_get_and_reuses_on_304() {
let root = tmp_root("via-304");
let registry = ModelRegistry::with_cache_dir(&root);
let bytes =
crate::interop::safetensors::save_safetensors_f32_to_bytes(&sample_state(), None)
.unwrap();
let mut t = Scripted::new(vec![
(200, Some("\"v1\""), bytes.clone()),
(304, None, Vec::new()),
]);
dl(®istry, &mut t, DownloadOptions::default()).unwrap();
dl(®istry, &mut t, DownloadOptions::default()).unwrap();
assert_eq!(t.conditional, vec![None, Some("\"v1\"".to_string())]);
assert_eq!(
std::fs::read(root.join("m/v1").join(MODEL_FILE_NAME)).unwrap(),
bytes
);
let other = Sha256Pin::from_hex(&"1".repeat(64)).unwrap();
let mut t2 = Scripted::new(vec![(200, Some("\"v1\""), bytes.clone())]);
let e = dl(
®istry,
&mut t2,
DownloadOptions::default().expected_sha256(other),
)
.unwrap_err();
assert!(matches!(e, ModelError::HashMismatch { .. }));
assert_eq!(t2.conditional, vec![None]);
let _ = std::fs::remove_dir_all(&root);
}
#[cfg(unix)]
#[test]
fn download_permission_denied_is_typed() {
use std::os::unix::fs::{MetadataExt, PermissionsExt};
let root = tmp_root("perm");
let registry = ModelRegistry::with_cache_dir(&root);
let dir = root.join("m/v1");
std::fs::create_dir_all(&dir).unwrap();
if std::fs::metadata(&dir).unwrap().uid() == 0 {
let _ = std::fs::remove_dir_all(&root);
return;
}
std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o500)).unwrap();
let bytes =
crate::interop::safetensors::save_safetensors_f32_to_bytes(&sample_state(), None)
.unwrap();
let r = dl(
®istry,
&mut Scripted::one(bytes),
DownloadOptions::default(),
);
std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700)).unwrap();
let e = r.unwrap_err();
assert!(matches!(e, ModelError::PermissionDenied(_)), "got {e:?}");
assert!(!e.to_string().contains(root.to_str().unwrap()));
let _ = std::fs::remove_dir_all(&root);
}
#[cfg(unix)]
#[test]
fn map_fs_error_assigns_typed_variants() {
use std::io::{Error, ErrorKind};
assert!(matches!(
map_fs_error(Error::from(ErrorKind::StorageFull)),
ModelError::StorageFull(_)
));
assert!(matches!(
map_fs_error(Error::from(ErrorKind::QuotaExceeded)),
ModelError::StorageFull(_)
));
assert!(matches!(
map_fs_error(Error::from(ErrorKind::PermissionDenied)),
ModelError::PermissionDenied(_)
));
assert!(matches!(
map_fs_error(Error::from(ErrorKind::NotFound)),
ModelError::Io(_)
));
}
#[test]
fn new_error_variants_display_without_path_and_chain_source() {
use std::error::Error as _;
let full = ModelError::StorageFull(std::io::Error::other("/secret/path"));
let perm = ModelError::PermissionDenied(std::io::Error::other("/secret/path"));
let meta = ModelError::CacheMetadata { reason: "r" };
for e in [&full, &perm, &meta] {
assert!(!e.to_string().contains("/secret"));
}
assert!(full.source().is_some() && perm.source().is_some());
assert!(meta.source().is_none());
}
}