Files
2026-02-22 14:06:01 -03:00

447 lines
19 KiB
Rust

use std::path::Path;
use std::process::Command;
use std::fs::File;
use std::io::Read;
use std::io::Write;
use std::time::Duration;
// Constantes para download e validação do OllamaSetup.exe
const OLLAMA_SETUP_URL: &str = "https://ollama.com/download/OllamaSetup.exe";
const MAX_RETRIES: u32 = 3;
const DOWNLOAD_TIMEOUT_SECS: u64 = 300;
// Hash SHA256 esperado do OllamaSetup.exe
// Pode ser sobrescrito via variável de ambiente OLLAMA_SETUP_EXE_SHA256
// IMPORTANTE: Atualizar este valor com o hash SHA256 oficial do instalador
const DEFAULT_EXPECTED_OLLAMA_SETUP_SHA256: &str = "AC9C74B41A8303E5566C5AD3ED6EF0E0123D58AEBC87D3109973F02A38D3FF61";
/// Calcula o hash SHA256 de um arquivo
/// Usa chunks de 8KB para leitura eficiente (padrão do projeto)
fn calculate_sha256(file_path: &Path) -> Result<String, String> {
use sha2::{Sha256, Digest};
let mut file = File::open(file_path)
.map_err(|e| format!("Failed to open file for hash calculation: {}", e))?;
let mut hasher = Sha256::new();
let mut buffer = vec![0u8; 8192]; // Buffer de 8KB para leitura em chunks
loop {
let bytes_read = file.read(&mut buffer)
.map_err(|e| format!("Failed to read file for hash: {}", e))?;
if bytes_read == 0 {
break;
}
hasher.update(&buffer[..bytes_read]);
}
let hash = hasher.finalize();
Ok(format!("{:x}", hash))
}
/// Faz download de um arquivo com retry logic e backoff exponencial
fn download_with_retry(url: &str, dest_path: &Path) -> Result<(), String> {
use reqwest::blocking::Client;
// Limpar arquivo parcial se existir antes de iniciar
if dest_path.exists() {
std::fs::remove_file(dest_path)
.map_err(|e| format!("Failed to remove existing file before download: {}", e))?;
}
let client = Client::builder()
.timeout(Duration::from_secs(DOWNLOAD_TIMEOUT_SECS))
.build()
.map_err(|e| format!("Failed to create HTTP client: {}", e))?;
let mut last_error = None;
for attempt in 1..=MAX_RETRIES {
println!("cargo:warning=Download attempt {}/{}: {}", attempt, MAX_RETRIES, url);
// Limpar arquivo parcial antes de cada tentativa (exceto na primeira, já limpo acima)
if attempt > 1 && dest_path.exists() {
std::fs::remove_file(dest_path)
.map_err(|e| format!("Failed to remove partial file before retry: {}", e))?;
}
// Garantir que o diretório pai existe
if let Some(parent) = dest_path.parent() {
std::fs::create_dir_all(parent)
.map_err(|e| format!("Failed to create parent directory: {}", e))?;
}
match client.get(url).send() {
Ok(response) => {
if !response.status().is_success() {
let err_msg = format!("HTTP error: {}", response.status());
println!("cargo:warning={}", err_msg);
last_error = Some(err_msg);
if attempt < MAX_RETRIES {
let backoff_secs = 2_u64.pow(attempt - 1); // 1s, 2s, 4s
println!("cargo:warning=Waiting {} seconds before retry...", backoff_secs);
std::thread::sleep(Duration::from_secs(backoff_secs));
}
continue;
}
// Obter bytes do response
match response.bytes() {
Ok(bytes) => {
// Criar arquivo e escrever conteúdo
match File::create(dest_path) {
Ok(mut dest_file) => {
match dest_file.write_all(&bytes) {
Ok(_) => {
drop(dest_file); // Fechar arquivo antes de validar
println!("cargo:warning=Download completed successfully");
return Ok(());
}
Err(e) => {
let err_msg = format!("Failed to write file: {}", e);
println!("cargo:warning={}", err_msg);
last_error = Some(err_msg);
// Limpar arquivo parcial
let _ = std::fs::remove_file(dest_path);
if attempt < MAX_RETRIES {
let backoff_secs = 2_u64.pow(attempt - 1);
println!("cargo:warning=Waiting {} seconds before retry...", backoff_secs);
std::thread::sleep(Duration::from_secs(backoff_secs));
}
continue;
}
}
}
Err(e) => {
let err_msg = format!("Failed to create destination file: {}", e);
println!("cargo:warning={}", err_msg);
last_error = Some(err_msg);
if attempt < MAX_RETRIES {
let backoff_secs = 2_u64.pow(attempt - 1);
println!("cargo:warning=Waiting {} seconds before retry...", backoff_secs);
std::thread::sleep(Duration::from_secs(backoff_secs));
}
continue;
}
}
}
Err(e) => {
let err_msg = format!("Failed to read response bytes: {}", e);
println!("cargo:warning={}", err_msg);
last_error = Some(err_msg);
// Limpar arquivo parcial
let _ = std::fs::remove_file(dest_path);
if attempt < MAX_RETRIES {
let backoff_secs = 2_u64.pow(attempt - 1);
println!("cargo:warning=Waiting {} seconds before retry...", backoff_secs);
std::thread::sleep(Duration::from_secs(backoff_secs));
}
continue;
}
}
}
Err(e) => {
let err_msg = format!("Download failed: {}", e);
println!("cargo:warning={}", err_msg);
last_error = Some(err_msg);
if attempt < MAX_RETRIES {
let backoff_secs = 2_u64.pow(attempt - 1);
println!("cargo:warning=Waiting {} seconds before retry...", backoff_secs);
std::thread::sleep(Duration::from_secs(backoff_secs));
}
}
}
}
Err(format!(
"Failed to download {} after {} attempts. Last error: {}",
url,
MAX_RETRIES,
last_error.unwrap_or_else(|| "Unknown error".to_string())
))
}
/// Valida o hash SHA256 de um arquivo contra o hash esperado
fn validate_sha256(file_path: &Path, expected_hash: &str) -> Result<(), String> {
println!("cargo:warning=Validating SHA256 hash...");
let calculated_hash = calculate_sha256(file_path)?;
let calculated_lower = calculated_hash.to_lowercase();
let expected_lower = expected_hash.to_lowercase();
if calculated_lower != expected_lower {
// Remover arquivo corrompido
let _ = std::fs::remove_file(file_path);
return Err(format!(
"SHA256 hash mismatch! Expected: {}, Got: {}. Corrupted file removed.",
expected_hash,
calculated_hash
));
}
println!("cargo:warning=SHA256 hash validated successfully");
Ok(())
}
/// Obtém o hash SHA256 esperado (variável de ambiente ou constante padrão)
fn get_expected_sha256() -> String {
std::env::var("OLLAMA_SETUP_EXE_SHA256")
.unwrap_or_else(|_| DEFAULT_EXPECTED_OLLAMA_SETUP_SHA256.to_string())
}
/// Baixa e valida o instalador do Ollama
fn download_and_validate_ollama_setup(ollama_path: &Path) -> Result<(), String> {
let expected_hash = get_expected_sha256();
// Se o hash esperado está vazio, falhar imediatamente
if expected_hash.is_empty() {
panic!(
"OLLAMA_SETUP_EXE_SHA256 não configurado! \
Configure a variável de ambiente OLLAMA_SETUP_EXE_SHA256 ou \
atualize DEFAULT_EXPECTED_OLLAMA_SETUP_SHA256 no build.rs"
);
}
// Verificar se arquivo já existe
if ollama_path.exists() {
println!("cargo:warning=OllamaSetup.exe already exists, validating hash...");
// Validar hash do arquivo existente
match validate_sha256(ollama_path, &expected_hash) {
Ok(_) => {
println!("cargo:warning=Existing OllamaSetup.exe is valid, skipping download");
return Ok(());
}
Err(e) => {
println!("cargo:warning=Existing file failed validation: {}", e);
println!("cargo:warning=Re-downloading...");
// Continuar para baixar novamente
}
}
}
// Fazer download com retry
println!("cargo:warning=Downloading Ollama installer...");
download_with_retry(OLLAMA_SETUP_URL, ollama_path)?;
// Validar hash após download
validate_sha256(ollama_path, &expected_hash)?;
println!("cargo:warning=Ollama installer downloaded and validated successfully");
Ok(())
}
fn main() {
// Criar estrutura de diretórios para bundle
let bundle_dir = Path::new("bundle");
let ollama_dir = bundle_dir.join("ollama");
let embeddings_dir = bundle_dir.join("models/embeddings");
std::fs::create_dir_all(&ollama_dir).ok();
std::fs::create_dir_all(&embeddings_dir).ok();
#[cfg(target_os = "windows")]
{
println!("cargo:warning=Preparing bundle resources for Windows...");
// 1. Baixar e validar instalador do Ollama com SHA256
let ollama_path = ollama_dir.join("OllamaSetup.exe");
match download_and_validate_ollama_setup(&ollama_path) {
Ok(_) => {}
Err(e) => {
panic!("cargo:error=Failed to download and validate OllamaSetup.exe: {}", e);
}
}
// 2. Baixar modelo de embeddings ONNX (apenas se não existir)
let model_url = "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/onnx/model.onnx";
let model_path = embeddings_dir.join("all-MiniLM-L6-v2.onnx");
if !model_path.exists() {
println!("cargo:warning=Downloading ONNX embedding model (this may take a while, ~90MB)...");
let status = Command::new("powershell")
.args(&[
"-Command",
&format!(
"if (!(Test-Path '{}')) {{ Invoke-WebRequest -Uri '{}' -OutFile '{}' -UseBasicParsing }}",
model_path.display(),
model_url,
model_path.display()
)
])
.status();
if let Ok(exit_status) = status {
if exit_status.success() {
println!("cargo:warning=ONNX model downloaded successfully");
} else {
println!("cargo:warning=Failed to download ONNX model (non-zero exit code)");
}
} else {
println!("cargo:warning=Failed to download ONNX model (command failed)");
}
} else {
println!("cargo:warning=ONNX model already exists, skipping download");
}
// 3. Baixar tokenizer
let tokenizer_url = "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/tokenizer.json";
let tokenizer_path = embeddings_dir.join("tokenizer.json");
if !tokenizer_path.exists() {
println!("cargo:warning=Downloading tokenizer...");
let status = Command::new("powershell")
.args(&[
"-Command",
&format!(
"if (!(Test-Path '{}')) {{ Invoke-WebRequest -Uri '{}' -OutFile '{}' -UseBasicParsing }}",
tokenizer_path.display(),
tokenizer_url,
tokenizer_path.display()
)
])
.status();
if let Ok(exit_status) = status {
if exit_status.success() {
println!("cargo:warning=Tokenizer downloaded successfully");
} else {
println!("cargo:warning=Failed to download tokenizer (non-zero exit code)");
}
} else {
println!("cargo:warning=Failed to download tokenizer (command failed)");
}
} else {
println!("cargo:warning=Tokenizer already exists, skipping download");
}
// 4. Baixar e extrair ONNX Runtime DLL
let ort_url = "https://github.com/microsoft/onnxruntime/releases/download/v1.20.1/onnxruntime-win-x64-1.20.1.zip";
let ort_zip_path = embeddings_dir.join("onnxruntime.zip");
let ort_dll_path = embeddings_dir.join("onnxruntime.dll");
if !ort_dll_path.exists() {
if !ort_zip_path.exists() {
println!("cargo:warning=Downloading ONNX Runtime library (~50MB)...");
let status = Command::new("powershell")
.args(&[
"-Command",
&format!(
"if (!(Test-Path '{}')) {{ Invoke-WebRequest -Uri '{}' -OutFile '{}' -UseBasicParsing }}",
ort_zip_path.display(),
ort_url,
ort_zip_path.display()
)
])
.status();
if let Ok(exit_status) = status {
if !exit_status.success() {
println!("cargo:warning=Failed to download ONNX Runtime (non-zero exit code)");
} else {
// Aguardar um pouco para garantir que o download foi concluído
std::thread::sleep(std::time::Duration::from_secs(1));
}
} else {
println!("cargo:warning=Failed to download ONNX Runtime (command failed)");
}
}
// Extrair DLL do zip se o zip existe e está completo
if ort_zip_path.exists() && !ort_dll_path.exists() {
// Verificar se o arquivo zip não está vazio e tem tamanho mínimo razoável (~40MB)
let min_zip_size = 40 * 1024 * 1024; // 40MB
let zip_valid = if let Ok(metadata) = std::fs::metadata(&ort_zip_path) {
let size = metadata.len();
if size < 1000 {
println!("cargo:warning=ONNX Runtime zip file appears to be incomplete (too small: {} bytes), will be downloaded at runtime", size);
false
} else if size < min_zip_size {
println!("cargo:warning=ONNX Runtime zip file appears to be incomplete ({} bytes, expected at least {} bytes), will be re-downloaded", size, min_zip_size);
// Tentar remover ZIP corrompido para forçar re-download
let _ = std::fs::remove_file(&ort_zip_path);
false
} else {
true
}
} else {
println!("cargo:warning=Failed to get metadata for ONNX Runtime zip file");
false
};
if zip_valid {
println!("cargo:warning=Extracting ONNX Runtime DLL from zip...");
// Usar zip crate para extrair (mais confiável que PowerShell)
match extract_onnx_dll_from_zip(&ort_zip_path, &ort_dll_path) {
Ok(_) => {
println!("cargo:warning=ONNX Runtime DLL extracted successfully");
// Remover zip após extração bem-sucedida
let _ = std::fs::remove_file(&ort_zip_path);
}
Err(e) => {
println!("cargo:warning=Failed to extract ONNX Runtime DLL: {}. ZIP may be corrupted, will be downloaded at runtime", e);
// Remover ZIP corrompido para forçar re-download em runtime
let _ = std::fs::remove_file(&ort_zip_path);
// Não falhar o build se a extração falhar - será baixado em runtime se necessário
}
}
}
}
} else {
println!("cargo:warning=ONNX Runtime DLL already exists, skipping download");
}
}
tauri_build::build()
}
/// Extrai onnxruntime.dll do arquivo zip
fn extract_onnx_dll_from_zip(zip_path: &Path, dll_path: &Path) -> Result<(), String> {
use zip::ZipArchive;
let file = File::open(zip_path)
.map_err(|e| format!("Failed to open zip file: {}", e))?;
let mut archive = ZipArchive::new(file)
.map_err(|e| format!("Failed to read zip archive: {}", e))?;
// Procurar pela DLL no zip
for i in 0..archive.len() {
let mut file = archive.by_index(i)
.map_err(|e| format!("Failed to read zip entry {}: {}", i, e))?;
let name = file.name().to_string();
// Procurar por onnxruntime.dll (pode estar em subdiretórios)
if name.ends_with("onnxruntime.dll") {
// Criar diretório de destino se não existir
if let Some(parent) = dll_path.parent() {
std::fs::create_dir_all(parent)
.map_err(|e| format!("Failed to create destination directory: {}", e))?;
}
// Extrair arquivo
let mut outfile = File::create(dll_path)
.map_err(|e| format!("Failed to create output file: {}", e))?;
std::io::copy(&mut file, &mut outfile)
.map_err(|e| format!("Failed to copy file: {}", e))?;
return Ok(());
}
}
Err("onnxruntime.dll not found in zip archive".to_string())
}