Split Rust code into domain modules

src/model/transfer.rs

@@ -0,0 +1,507 @@
+use super::*;
+use sha2::{Digest, Sha256};
+use std::fs::{self, File, OpenOptions};
+use std::io::{Read, Write};
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+
+pub(crate) fn download_managed_artifact(
+    id: ManagedArtifactId,
+    models_path: &Path,
+    cancel: &AtomicBool,
+    verified_bytes: &AtomicU64,
+) -> Result<DownloadOutcome, String> {
+    download_artifact_with_cancel(
+        id.model(),
+        id.artifact(),
+        models_path,
+        cancel,
+        verified_bytes,
+    )
+}
+
+pub(crate) fn validate_managed_artifact(
+    id: ManagedArtifactId,
+    models_path: &Path,
+    cancel: &AtomicBool,
+    verified_bytes: &AtomicU64,
+) -> Result<DownloadOutcome, String> {
+    let model = id.model();
+    let artifact = id.artifact();
+    let destination = artifact.path(model, models_path);
+    let partial = artifact.partial_path(model, models_path);
+    let (path, promote) = if destination.exists() {
+        (destination.clone(), false)
+    } else if partial.exists() {
+        (partial.clone(), true)
+    } else {
+        return Err(format!("{} is not downloaded", artifact.label));
+    };
+
+    match verify(&path, artifact, model, cancel, verified_bytes) {
+        Ok(DownloadOutcome::Stopped) => return Ok(DownloadOutcome::Stopped),
+        Ok(DownloadOutcome::Complete) => {}
+        Err(error) => {
+            let marker = artifact.verification_path(model, models_path);
+            if let Err(remove_error) = fs::remove_file(marker)
+                && remove_error.kind() != std::io::ErrorKind::NotFound
+            {
+                return Err(format!(
+                    "{error}; could not remove checksum marker: {remove_error}"
+                ));
+            }
+            return Err(error);
+        }
+    }
+    if promote {
+        fs::rename(partial, destination).map_err(|error| error.to_string())?;
+    }
+    mark_verified(model, artifact, models_path)?;
+    Ok(DownloadOutcome::Complete)
+}
+
+pub(crate) fn delete_managed_artifact(
+    id: ManagedArtifactId,
+    models_path: &Path,
+) -> Result<(), String> {
+    let model = id.model();
+    let artifact = id.artifact();
+    for path in [
+        artifact.path(model, models_path),
+        artifact.partial_path(model, models_path),
+        artifact.verification_path(model, models_path),
+    ] {
+        match fs::remove_file(path) {
+            Ok(()) => {}
+            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
+            Err(error) => return Err(error.to_string()),
+        }
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+fn download_artifact(
+    model: ModelChoice,
+    artifact: &Artifact,
+    models_path: &Path,
+) -> Result<DownloadOutcome, String> {
+    download_artifact_with_cancel(
+        model,
+        artifact,
+        models_path,
+        &AtomicBool::new(false),
+        &AtomicU64::new(0),
+    )
+}
+
+fn download_artifact_with_cancel(
+    model: ModelChoice,
+    artifact: &Artifact,
+    models_path: &Path,
+    cancel: &AtomicBool,
+    verified_bytes: &AtomicU64,
+) -> Result<DownloadOutcome, String> {
+    if cancel.load(Ordering::Relaxed) {
+        return Ok(DownloadOutcome::Stopped);
+    }
+    let destination = artifact.path(model, models_path);
+    if artifact.is_installed(model, models_path) {
+        return Ok(DownloadOutcome::Complete);
+    }
+    if destination.exists() {
+        if verify(&destination, artifact, model, cancel, verified_bytes)?
+            == DownloadOutcome::Stopped
+        {
+            return Ok(DownloadOutcome::Stopped);
+        }
+        mark_verified(model, artifact, models_path)?;
+        return Ok(DownloadOutcome::Complete);
+    }
+
+    let directory = destination
+        .parent()
+        .ok_or_else(|| "model artifact path has no parent directory".to_owned())?;
+    fs::create_dir_all(directory).map_err(|error| error.to_string())?;
+    let partial = artifact.partial_path(model, models_path);
+    let partial_size = partial.metadata().map_or(0, |metadata| metadata.len());
+    if partial_size > artifact.size {
+        File::create(&partial).map_err(|error| error.to_string())?;
+    }
+    if partial.metadata().map_or(0, |metadata| metadata.len()) != artifact.size
+        && download_to_partial(artifact, &partial, cancel)? == DownloadOutcome::Stopped
+    {
+        return Ok(DownloadOutcome::Stopped);
+    }
+
+    if cancel.load(Ordering::Relaxed) {
+        return Ok(DownloadOutcome::Stopped);
+    }
+    match verify(&partial, artifact, model, cancel, verified_bytes) {
+        Ok(DownloadOutcome::Stopped) => return Ok(DownloadOutcome::Stopped),
+        Ok(DownloadOutcome::Complete) => {}
+        Err(error) => {
+            fs::remove_file(&partial).map_err(|remove_error| {
+                format!("{error}; could not remove partial file: {remove_error}")
+            })?;
+            return Err(error);
+        }
+    }
+    fs::rename(partial, destination).map_err(|error| error.to_string())?;
+    mark_verified(model, artifact, models_path)?;
+    Ok(DownloadOutcome::Complete)
+}
+
+fn mark_verified(
+    model: ModelChoice,
+    artifact: &Artifact,
+    models_path: &Path,
+) -> Result<(), String> {
+    fs::write(
+        artifact.verification_path(model, models_path),
+        artifact.sha256,
+    )
+    .map_err(|error| error.to_string())
+}
+
+fn verify(
+    path: &Path,
+    artifact: &Artifact,
+    model: ModelChoice,
+    cancel: &AtomicBool,
+    verified_bytes: &AtomicU64,
+) -> Result<DownloadOutcome, String> {
+    verified_bytes.store(0, Ordering::Relaxed);
+    let size = path.metadata().map_err(|error| error.to_string())?.len();
+    if size != artifact.size {
+        return Err(format!(
+            "{} has size {size}, expected {}",
+            path.display(),
+            artifact.size
+        ));
+    }
+
+    let mut file = File::open(path).map_err(|error| error.to_string())?;
+    let mut hasher = Sha256::new();
+    let mut buffer = vec![0; 1024 * 1024];
+    loop {
+        if cancel.load(Ordering::Relaxed) {
+            return Ok(DownloadOutcome::Stopped);
+        }
+        let count = file.read(&mut buffer).map_err(|error| error.to_string())?;
+        if count == 0 {
+            break;
+        }
+        hasher.update(&buffer[..count]);
+        verified_bytes.fetch_add(count as u64, Ordering::Relaxed);
+    }
+    let actual = hex(&hasher.finalize());
+    if actual != artifact.sha256 {
+        return Err(format!(
+            "Checksum verification failed for {}",
+            path.display()
+        ));
+    }
+    if let Some(support) = artifact.support {
+        crate::engine::validate_model_artifact(path, model, support)?;
+    }
+    Ok(DownloadOutcome::Complete)
+}
+
+fn hex(bytes: &[u8]) -> String {
+    const DIGITS: &[u8; 16] = b"0123456789abcdef";
+    let mut result = String::with_capacity(bytes.len() * 2);
+    for byte in bytes {
+        result.push(DIGITS[(byte >> 4) as usize] as char);
+        result.push(DIGITS[(byte & 0x0f) as usize] as char);
+    }
+    result
+}
+
+fn download_to_partial(
+    artifact: &Artifact,
+    partial: &Path,
+    cancel: &AtomicBool,
+) -> Result<DownloadOutcome, String> {
+    download_url_to_partial(&artifact.url(), partial, cancel)
+}
+
+fn download_url_to_partial(
+    url: &str,
+    partial: &Path,
+    cancel: &AtomicBool,
+) -> Result<DownloadOutcome, String> {
+    let offset = partial.metadata().map_or(0, |metadata| metadata.len());
+    let agent: ureq::Agent = ureq::Agent::config_builder()
+        .https_only(url.starts_with("https://"))
+        .build()
+        .into();
+    let mut request = agent.get(url);
+    if offset > 0 {
+        request = request.header("Range", format!("bytes={offset}-"));
+    }
+    let mut response = request
+        .call()
+        .map_err(|error| format!("Model download failed: {error}"))?;
+    let status = response.status().as_u16();
+    let append = offset > 0 && status == 206;
+    if offset > 0 && status != 200 && status != 206 {
+        return Err(format!(
+            "Model server returned HTTP {status} while resuming at byte {offset}"
+        ));
+    }
+    if append {
+        validate_content_range(&response, offset)?;
+    }
+
+    let mut output = OpenOptions::new()
+        .create(true)
+        .write(true)
+        .append(append)
+        .truncate(!append)
+        .open(partial)
+        .map_err(|error| error.to_string())?;
+    let mut body = response.body_mut().as_reader();
+    let mut buffer = vec![0; 1024 * 1024];
+    loop {
+        if cancel.load(Ordering::Relaxed) {
+            return Ok(DownloadOutcome::Stopped);
+        }
+        let count = body
+            .read(&mut buffer)
+            .map_err(|error| format!("Model download failed: {error}"))?;
+        if count == 0 {
+            break;
+        }
+        output
+            .write_all(&buffer[..count])
+            .map_err(|error| error.to_string())?;
+    }
+    Ok(DownloadOutcome::Complete)
+}
+
+fn validate_content_range(
+    response: &ureq::http::Response<ureq::Body>,
+    offset: u64,
+) -> Result<(), String> {
+    let expected = format!("bytes {offset}-");
+    let content_range = response
+        .headers()
+        .get("content-range")
+        .and_then(|value| value.to_str().ok())
+        .unwrap_or_default();
+    if content_range.starts_with(&expected) {
+        Ok(())
+    } else {
+        Err(format!(
+            "Model server returned an invalid Content-Range while resuming at byte {offset}"
+        ))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::net::TcpListener;
+    use std::thread;
+    use std::time::{SystemTime, UNIX_EPOCH};
+
+    #[test]
+    fn catalog_and_checksum_verification_are_explicit() {
+        assert_eq!(ModelChoice::from_id("glm-5.2"), Some(ModelChoice::Glm52));
+        assert!(ModelChoice::from_id("unknown").is_none());
+        assert_eq!(
+            ModelChoice::DeepSeekV4Flash.main_artifact().size,
+            86_720_111_488
+        );
+        assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448);
+        assert_eq!(ModelChoice::DeepSeekV4Flash.artifacts(true).count(), 2);
+        assert_eq!(ModelChoice::Glm52.artifacts(true).count(), 1);
+
+        let id = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_nanos();
+        let models_path = std::env::temp_dir().join(format!("ds4-server-models-{id}"));
+        let engine = engine_artifacts(ModelChoice::DeepSeekV4Flash, true, &models_path);
+        assert_eq!(
+            engine.model.file_name(),
+            Some(std::ffi::OsStr::new(FLASH.file_name))
+        );
+        assert_eq!(
+            engine.mtp.as_deref().and_then(Path::file_name),
+            Some(std::ffi::OsStr::new(FLASH_DSPARK.file_name))
+        );
+        let empty = Artifact {
+            label: "empty",
+            file_name: "empty",
+            repository: "",
+            size: 0,
+            sha256: "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+            support: None,
+        };
+        let partial = empty.partial_path(ModelChoice::DeepSeekV4Flash, &models_path);
+        fs::create_dir_all(partial.parent().unwrap()).unwrap();
+        fs::write(&partial, []).unwrap();
+        download_artifact(ModelChoice::DeepSeekV4Flash, &empty, &models_path).unwrap();
+        assert!(empty.is_installed(ModelChoice::DeepSeekV4Flash, &models_path));
+        assert!(!partial.exists());
+        assert_eq!(
+            fs::read_to_string(empty.verification_path(ModelChoice::DeepSeekV4Flash, &models_path))
+                .unwrap(),
+            empty.sha256
+        );
+        fs::remove_dir_all(models_path).unwrap();
+    }
+
+    #[test]
+    fn verification_reports_bytes_read() {
+        let path = std::env::temp_dir().join(format!(
+            "ds4-server-verify-{}",
+            SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap()
+                .as_nanos()
+        ));
+        fs::write(&path, b"abc").unwrap();
+        let artifact = Artifact {
+            label: "test model",
+            file_name: "unused",
+            repository: "unused",
+            size: 3,
+            sha256: "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad",
+            support: None,
+        };
+        let verified_bytes = AtomicU64::new(999);
+
+        assert_eq!(
+            verify(
+                &path,
+                &artifact,
+                ModelChoice::DeepSeekV4Flash,
+                &AtomicBool::new(false),
+                &verified_bytes,
+            )
+            .unwrap(),
+            DownloadOutcome::Complete
+        );
+        assert_eq!(verified_bytes.load(Ordering::Relaxed), 3);
+        fs::remove_file(path).unwrap();
+    }
+
+    #[test]
+    fn managed_artifact_inventory_and_delete_include_partial_files() {
+        let models_path = std::env::temp_dir().join(format!(
+            "ds4-server-manager-{}",
+            SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap()
+                .as_nanos()
+        ));
+        let id = ManagedArtifactId::DeepSeekV4Flash;
+        let partial = id.artifact().partial_path(id.model(), &models_path);
+        fs::create_dir_all(partial.parent().unwrap()).unwrap();
+        fs::write(&partial, b"part").unwrap();
+
+        let managed = managed_artifacts(&models_path)
+            .into_iter()
+            .find(|artifact| artifact.id == id)
+            .unwrap();
+        assert_eq!(managed.stored, 4);
+        assert_eq!(managed.state, ManagedArtifactState::Partial);
+
+        delete_managed_artifact(id, &models_path).unwrap();
+        assert!(!partial.exists());
+        fs::remove_dir_all(models_path).unwrap();
+    }
+
+    #[test]
+    fn restart_resumes_at_the_existing_partial_byte() {
+        let content = b"restart-resume works";
+        let offset = 8;
+        let directory = std::env::temp_dir().join(format!(
+            "ds4-server-resume-{}",
+            SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap()
+                .as_nanos()
+        ));
+        fs::create_dir_all(&directory).unwrap();
+        let partial = directory.join("model.gguf.part");
+        fs::write(&partial, &content[..offset]).unwrap();
+
+        let listener = match TcpListener::bind("127.0.0.1:0") {
+            Ok(listener) => listener,
+            Err(error) if error.kind() == std::io::ErrorKind::PermissionDenied => {
+                fs::remove_dir_all(directory).unwrap();
+                return;
+            }
+            Err(error) => panic!("could not start test server: {error}"),
+        };
+        let address = listener.local_addr().unwrap();
+        let server = thread::spawn(move || {
+            let (mut connection, _) = listener.accept().unwrap();
+            let mut request = [0; 2048];
+            let count = connection.read(&mut request).unwrap();
+            let request = String::from_utf8_lossy(&request[..count]).to_ascii_lowercase();
+            assert!(request.contains("range: bytes=8-"));
+            let remaining = &content[offset..];
+            write!(
+                connection,
+                "HTTP/1.1 206 Partial Content\r\nContent-Length: {}\r\nContent-Range: bytes {offset}-{}/{}\r\nConnection: close\r\n\r\n",
+                remaining.len(),
+                content.len() - 1,
+                content.len(),
+            )
+            .unwrap();
+            connection.write_all(remaining).unwrap();
+        });
+
+        let outcome = download_url_to_partial(
+            &format!("http://{address}/model.gguf"),
+            &partial,
+            &AtomicBool::new(false),
+        )
+        .unwrap();
+        server.join().unwrap();
+        assert_eq!(outcome, DownloadOutcome::Complete);
+        assert_eq!(fs::read(&partial).unwrap(), content);
+        fs::remove_dir_all(directory).unwrap();
+    }
+
+    #[test]
+    fn cancellation_keeps_the_partial_file_for_the_next_run() {
+        let directory = std::env::temp_dir().join(format!(
+            "ds4-server-cancel-{}",
+            SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap()
+                .as_nanos()
+        ));
+        let artifact = Artifact {
+            label: "test model",
+            file_name: "model.gguf",
+            repository: "unused",
+            size: 10,
+            sha256: "unused",
+            support: None,
+        };
+        let partial = artifact.partial_path(ModelChoice::DeepSeekV4Flash, &directory);
+        fs::create_dir_all(partial.parent().unwrap()).unwrap();
+        fs::write(&partial, b"part").unwrap();
+        let cancel = AtomicBool::new(true);
+
+        assert_eq!(
+            download_artifact_with_cancel(
+                ModelChoice::DeepSeekV4Flash,
+                &artifact,
+                &directory,
+                &cancel,
+                &AtomicU64::new(0),
+            )
+            .unwrap(),
+            DownloadOutcome::Stopped
+        );
+        assert_eq!(fs::read(&partial).unwrap(), b"part");
+        fs::remove_dir_all(directory).unwrap();
+    }
+}