GithubFollow/One-KVM
1
0
Fork 0
mirror of https://github.com/mofeng-git/One-KVM.git synced 2026-06-14 11:42:02 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactor: 删除部分多余的代码和注释

Browse Source
This commit is contained in:
mofeng-git
2026-05-01 17:31:04 +08:00
parent 74035f8e12
commit d8e7de74a6
165 changed files with 2960 additions and 9917 deletions
Download Patch File Download Diff File Expand all files Collapse all files

228
src/msd/controller.rs
View File

@@ -1,11 +1,3 @@
//! MSD Controller
//!
//! Manages the mass storage device lifecycle including:
//! - Image mounting and unmounting
//! - Virtual drive management
//! - State tracking
//! - Image downloads from URL
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
@@ -14,41 +6,25 @@ use tokio_util::sync::CancellationToken;
use tracing::{debug, info, warn};
use super::image::ImageManager;
use super::monitor::{MsdHealthMonitor, MsdHealthStatus};
use super::monitor::MsdHealthMonitor;
use super::types::{DownloadProgress, DownloadStatus, DriveInfo, ImageInfo, MsdMode, MsdState};
use crate::error::{AppError, Result};
use crate::otg::{MsdFunction, MsdLunConfig, OtgService};
/// MSD Controller
pub struct MsdController {
/// OTG Service reference
otg_service: Arc<OtgService>,
/// MSD function manager (provided by OtgService)
msd_function: RwLock<Option<MsdFunction>>,
/// Current state
state: RwLock<MsdState>,
/// Images storage path
images_path: PathBuf,
/// Ventoy directory path
ventoy_dir: PathBuf,
/// Virtual drive path
drive_path: PathBuf,
/// Event bus for broadcasting state changes (optional)
events: tokio::sync::RwLock<Option<Arc<crate::events::EventBus>>>,
/// Active downloads (download_id -> CancellationToken)
downloads: Arc<RwLock<HashMap<String, CancellationToken>>>,
/// Operation mutex lock (prevents concurrent operations)
operation_lock: Arc<RwLock<()>>,
/// Health monitor for error tracking and recovery
monitor: Arc<MsdHealthMonitor>,
}
impl MsdController {
/// Create new MSD controller
///
/// # Parameters
/// * `otg_service` - OTG service for gadget management
/// * `msd_dir` - Base directory for MSD storage
pub fn new(otg_service: Arc<OtgService>, msd_dir: impl Into<PathBuf>) -> Self {
let msd_dir = msd_dir.into();
let images_path = msd_dir.join("images");
@@ -68,11 +44,9 @@ impl MsdController {
}
}
/// Initialize the MSD controller
pub async fn init(&self) -> Result<()> {
info!("Initializing MSD controller");
// 1. Ensure images directory exists
if let Err(e) = std::fs::create_dir_all(&self.images_path) {
warn!("Failed to create images directory: {}", e);
}
@@ -80,20 +54,16 @@ impl MsdController {
warn!("Failed to create ventoy directory: {}", e);
}
// 2. Get active MSD function from OtgService
info!("Fetching MSD function from OtgService");
let msd_func = self.otg_service.msd_function().await.ok_or_else(|| {
AppError::Internal("MSD function is not active in OtgService".to_string())
})?;
// 3. Store function handle
*self.msd_function.write().await = Some(msd_func);
// 4. Update state
let mut state = self.state.write().await;
state.available = true;
// 5. Check for existing virtual drive
if self.drive_path.exists() {
if let Ok(metadata) = std::fs::metadata(&self.drive_path) {
state.drive_info = Some(DriveInfo {
@@ -114,17 +84,14 @@ impl MsdController {
Ok(())
}
/// Get current MSD state
pub async fn state(&self) -> MsdState {
self.state.read().await.clone()
}
/// Set event bus for broadcasting state changes
pub async fn set_event_bus(&self, events: std::sync::Arc<crate::events::EventBus>) {
*self.events.write().await = Some(events);
}
/// Publish an event to the event bus
async fn publish_event(&self, event: crate::events::SystemEvent) {
if let Some(ref bus) = *self.events.read().await {
bus.publish(event);
@@ -137,43 +104,21 @@ impl MsdController {
}
}
/// Check if MSD is available
pub async fn is_available(&self) -> bool {
self.state.read().await.available
}
/// Connect an image file
///
/// # Parameters
/// * `image` - Image info to mount
/// * `cdrom` - Mount as CD-ROM (read-only, removable)
/// * `read_only` - Mount as read-only
pub async fn connect_image(
&self,
image: &ImageInfo,
cdrom: bool,
read_only: bool,
) -> Result<()> {
// Acquire operation lock to prevent concurrent operations
let _op_guard = self.operation_lock.write().await;
let mut state = self.state.write().await;
if !state.available {
let err = AppError::Internal("MSD not available".to_string());
self.monitor
.report_error("MSD not available", "not_available")
.await;
return Err(err);
}
self.assert_can_connect(&state).await?;
if state.connected {
return Err(AppError::Internal(
"Already connected. Disconnect first.".to_string(),
));
}
// Verify image exists
if !image.path.exists() {
let error_msg = format!("Image file not found: {}", image.path.display());
self.monitor
@@ -182,29 +127,12 @@ impl MsdController {
return Err(AppError::Internal(error_msg));
}
// Configure LUN
let config = if cdrom {
MsdLunConfig::cdrom(image.path.clone())
} else {
MsdLunConfig::disk(image.path.clone(), read_only)
};
let gadget_path = self.active_gadget_path().await?;
if let Some(ref msd) = *self.msd_function.read().await {
if let Err(e) = msd.configure_lun_async(&gadget_path, 0, &config).await {
let error_msg = format!("Failed to configure LUN: {}", e);
self.monitor
.report_error(&error_msg, "configfs_error")
.await;
return Err(e);
}
} else {
let err = AppError::Internal("MSD function not initialized".to_string());
self.monitor
.report_error("MSD function not initialized", "not_initialized")
.await;
return Err(err);
}
self.configure_lun_now(&config).await?;
state.connected = true;
state.mode = MsdMode::Image;
@@ -215,42 +143,19 @@ impl MsdController {
image.name, cdrom, read_only
);
// Release the lock before publishing events
drop(state);
drop(_op_guard);
// Report recovery if we were in an error state
if self.monitor.is_error().await {
self.monitor.report_recovered().await;
}
self.mark_device_info_dirty().await;
self.finish_connect_success().await;
Ok(())
}
/// Connect the virtual drive
pub async fn connect_drive(&self) -> Result<()> {
// Acquire operation lock to prevent concurrent operations
let _op_guard = self.operation_lock.write().await;
let mut state = self.state.write().await;
if !state.available {
let err = AppError::Internal("MSD not available".to_string());
self.monitor
.report_error("MSD not available", "not_available")
.await;
return Err(err);
}
self.assert_can_connect(&state).await?;
if state.connected {
return Err(AppError::Internal(
"Already connected. Disconnect first.".to_string(),
));
}
// Check drive exists
if !self.drive_path.exists() {
let err =
AppError::Internal("Virtual drive not initialized. Call init first.".to_string());
@@ -260,25 +165,8 @@ impl MsdController {
return Err(err);
}
// Configure LUN as read-write disk
let config = MsdLunConfig::disk(self.drive_path.clone(), false);
let gadget_path = self.active_gadget_path().await?;
if let Some(ref msd) = *self.msd_function.read().await {
if let Err(e) = msd.configure_lun_async(&gadget_path, 0, &config).await {
let error_msg = format!("Failed to configure LUN: {}", e);
self.monitor
.report_error(&error_msg, "configfs_error")
.await;
return Err(e);
}
} else {
let err = AppError::Internal("MSD function not initialized".to_string());
self.monitor
.report_error("MSD function not initialized", "not_initialized")
.await;
return Err(err);
}
self.configure_lun_now(&config).await?;
state.connected = true;
state.mode = MsdMode::Drive;
@@ -286,23 +174,57 @@ impl MsdController {
info!("Connected virtual drive: {}", self.drive_path.display());
// Release the lock before publishing event
drop(state);
drop(_op_guard);
// Report recovery if we were in an error state
if self.monitor.is_error().await {
self.monitor.report_recovered().await;
}
self.mark_device_info_dirty().await;
self.finish_connect_success().await;
Ok(())
}
/// Disconnect current storage
async fn assert_can_connect(&self, state: &MsdState) -> Result<()> {
if !state.available {
self.monitor
.report_error("MSD not available", "not_available")
.await;
return Err(AppError::Internal("MSD not available".to_string()));
}
if state.connected {
return Err(AppError::Internal(
"Already connected. Disconnect first.".to_string(),
));
}
Ok(())
}
async fn configure_lun_now(&self, config: &MsdLunConfig) -> Result<()> {
let gadget_path = self.active_gadget_path().await?;
let msd_hold = self.msd_function.read().await;
let Some(ref msd) = *msd_hold else {
self.monitor
.report_error("MSD function not initialized", "not_initialized")
.await;
return Err(AppError::Internal(
"MSD function not initialized".to_string(),
));
};
if let Err(e) = msd.configure_lun_async(&gadget_path, 0, config).await {
let error_msg = format!("Failed to configure LUN: {}", e);
self.monitor
.report_error(&error_msg, "configfs_error")
.await;
return Err(e);
}
Ok(())
}
async fn finish_connect_success(&self) {
if self.monitor.is_error().await {
self.monitor.report_recovered().await;
}
self.mark_device_info_dirty().await;
}
pub async fn disconnect(&self) -> Result<()> {
// Acquire operation lock to prevent concurrent operations
let _op_guard = self.operation_lock.write().await;
let mut state = self.state.write().await;
@@ -323,7 +245,6 @@ impl MsdController {
info!("Disconnected storage");
// Release the lock before publishing events
drop(state);
drop(_op_guard);
@@ -332,41 +253,31 @@ impl MsdController {
Ok(())
}
/// Get images storage path
pub fn images_path(&self) -> &PathBuf {
&self.images_path
}
/// Get ventoy directory path
pub fn ventoy_dir(&self) -> &PathBuf {
&self.ventoy_dir
}
/// Get virtual drive path
pub fn drive_path(&self) -> &PathBuf {
&self.drive_path
}
/// Check if currently connected
pub async fn is_connected(&self) -> bool {
self.state.read().await.connected
}
/// Get current mode
pub async fn mode(&self) -> MsdMode {
self.state.read().await.mode.clone()
}
/// Update drive info
pub async fn update_drive_info(&self, info: DriveInfo) {
let mut state = self.state.write().await;
state.drive_info = Some(info);
}
/// Start downloading an image from URL
///
/// Returns the download_id that can be used to track or cancel the download.
/// Progress is reported via MsdDownloadProgress events.
pub async fn download_image(
&self,
url: String,
@@ -375,18 +286,15 @@ impl MsdController {
let download_id = uuid::Uuid::new_v4().to_string();
let cancel_token = CancellationToken::new();
// Register download
{
let mut downloads = self.downloads.write().await;
downloads.insert(download_id.clone(), cancel_token.clone());
}
// Extract filename for initial response
let display_filename = filename
.clone()
.unwrap_or_else(|| url.rsplit('/').next().unwrap_or("download").to_string());
// Create initial progress
let initial_progress = DownloadProgress {
download_id: download_id.clone(),
url: url.clone(),
@@ -398,7 +306,6 @@ impl MsdController {
error: None,
};
// Publish started event
self.publish_event(crate::events::SystemEvent::MsdDownloadProgress {
download_id: download_id.clone(),
url: url.clone(),
@@ -410,18 +317,15 @@ impl MsdController {
})
.await;
// Clone what we need for the spawned task
let images_path = self.images_path.clone();
let events = self.events.read().await.clone();
let downloads = self.downloads.clone();
let download_id_clone = download_id.clone();
let url_clone = url.clone();
// Spawn download task
tokio::spawn(async move {
let manager = ImageManager::new(images_path);
// Create progress callback
let events_for_callback = events.clone();
let download_id_for_callback = download_id_clone.clone();
let url_for_callback = url_clone.clone();
@@ -443,18 +347,15 @@ impl MsdController {
}
};
// Run download
let result = manager
.download_from_url(&url_clone, filename, progress_callback)
.await;
// Remove from active downloads
{
let mut downloads_guard = downloads.write().await;
downloads_guard.remove(&download_id_clone);
}
// Publish completion event
match result {
Ok(image_info) => {
if let Some(ref bus) = events {
@@ -489,7 +390,6 @@ impl MsdController {
Ok(initial_progress)
}
/// Cancel an active download
pub async fn cancel_download(&self, download_id: &str) -> Result<()> {
let mut downloads = self.downloads.write().await;
@@ -505,12 +405,6 @@ impl MsdController {
}
}
/// Get list of active download IDs
pub async fn active_downloads(&self) -> Vec<String> {
let downloads = self.downloads.read().await;
downloads.keys().cloned().collect()
}
async fn active_gadget_path(&self) -> Result<PathBuf> {
self.otg_service
.gadget_path()
@@ -518,16 +412,13 @@ impl MsdController {
.ok_or_else(|| AppError::Internal("OTG gadget path is not available".to_string()))
}
/// Shutdown the controller
pub async fn shutdown(&self) -> Result<()> {
info!("Shutting down MSD controller");
// 1. Disconnect if connected
if let Err(e) = self.disconnect().await {
warn!("Error disconnecting during shutdown: {}", e);
}
// 2. Clear local state
*self.msd_function.write().await = None;
let mut state = self.state.write().await;
@@ -537,27 +428,9 @@ impl MsdController {
Ok(())
}
/// Get the health monitor reference
pub fn monitor(&self) -> &Arc<MsdHealthMonitor> {
&self.monitor
}
/// Get current health status
pub async fn health_status(&self) -> MsdHealthStatus {
self.monitor.status().await
}
/// Check if the MSD is healthy
pub async fn is_healthy(&self) -> bool {
self.monitor.is_healthy().await
}
}
impl Drop for MsdController {
fn drop(&mut self) {
// Cleanup is handled by OtgGadgetManager when the gadget is torn down
// Individual controllers don't need to cleanup the ConfigFS
}
}
#[cfg(test)]
@@ -573,7 +446,6 @@ mod tests {
let controller = MsdController::new(otg_service, &msd_dir);
// Check that MSD is not initialized (msd_function is None)
let state = controller.state().await;
assert!(!state.available);
assert!(controller.images_path.ends_with("images"));

183
src/msd/image.rs
View File

@@ -1,53 +1,37 @@
//! Image file manager
//!
//! Handles ISO/IMG image file operations:
//! - List available images
//! - Upload new images
//! - Delete images
//! - Metadata management
//! - Download from URL
use futures::StreamExt;
use time::OffsetDateTime;
use std::fs::{self, File};
use std::io::{self, Read, Write};
use std::fs;
#[cfg(test)]
use std::io::Write;
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
use time::OffsetDateTime;
use tokio::io::AsyncWriteExt;
use tracing::info;
use super::types::ImageInfo;
use crate::error::{AppError, Result};
/// Maximum image size (32 GB)
const MAX_IMAGE_SIZE: u64 = 32 * 1024 * 1024 * 1024;
/// Progress report throttle interval (milliseconds)
const PROGRESS_THROTTLE_MS: u64 = 200;
/// Progress report throttle bytes threshold (512 KB)
const PROGRESS_THROTTLE_BYTES: u64 = 512 * 1024;
/// Image Manager
pub struct ImageManager {
/// Images storage directory
images_path: PathBuf,
}
impl ImageManager {
/// Create a new image manager
pub fn new(images_path: PathBuf) -> Self {
Self { images_path }
}
/// Ensure images directory exists
pub fn ensure_dir(&self) -> Result<()> {
fs::create_dir_all(&self.images_path)
.map_err(|e| AppError::Internal(format!("Failed to create images directory: {}", e)))?;
Ok(())
}
/// List all available images
pub fn list(&self) -> Result<Vec<ImageInfo>> {
self.ensure_dir()?;
@@ -68,19 +52,16 @@ impl ImageManager {
}
}
// Sort by creation time (newest first)
images.sort_by(|a, b| b.created_at.cmp(&a.created_at));
Ok(images)
}
/// Get image info from path
fn get_image_info(&self, path: &Path) -> Option<ImageInfo> {
let metadata = fs::metadata(path).ok()?;
let name = path.file_name()?.to_string_lossy().to_string();
// Use filename hash as ID (stable across restarts)
let id = format!("{:x}", md5_hash(&name));
let id = stable_image_id_from_filename(&name);
let created_at = metadata
.created()
@@ -101,7 +82,6 @@ impl ImageManager {
})
}
/// Get image by ID
pub fn get(&self, id: &str) -> Result<ImageInfo> {
for image in self.list()? {
if image.id == id {
@@ -111,24 +91,21 @@ impl ImageManager {
Err(AppError::NotFound(format!("Image not found: {}", id)))
}
/// Get image by name
pub fn get_by_name(&self, name: &str) -> Result<ImageInfo> {
let path = self.images_path.join(name);
self.get_image_info(&path)
.ok_or_else(|| AppError::NotFound(format!("Image not found: {}", name)))
}
/// Create a new image from bytes
pub fn create(&self, name: &str, data: &[u8]) -> Result<ImageInfo> {
#[cfg(test)]
fn create(&self, name: &str, data: &[u8]) -> Result<ImageInfo> {
self.ensure_dir()?;
// Validate name
let name = sanitize_filename(name);
if name.is_empty() {
return Err(AppError::Internal("Invalid filename".to_string()));
}
// Check size
if data.len() as u64 > MAX_IMAGE_SIZE {
return Err(AppError::Internal(format!(
"Image too large. Maximum size: {} GB",
@@ -136,7 +113,6 @@ impl ImageManager {
)));
}
// Write file
let path = self.images_path.join(&name);
if path.exists() {
return Err(AppError::Internal(format!(
@@ -145,11 +121,10 @@ impl ImageManager {
)));
}
let mut file = File::create(&path)
let mut file = fs::File::create(&path)
.map_err(|e| AppError::Internal(format!("Failed to create image file: {}", e)))?;
file.write_all(data).map_err(|e| {
// Try to clean up on error
let _ = fs::remove_file(&path);
AppError::Internal(format!("Failed to write image data: {}", e))
})?;
@@ -159,55 +134,6 @@ impl ImageManager {
self.get_by_name(&name)
}
/// Create a new image from a file stream (for chunked uploads)
pub fn create_from_stream<R: Read>(
&self,
name: &str,
reader: &mut R,
expected_size: Option<u64>,
) -> Result<ImageInfo> {
self.ensure_dir()?;
let name = sanitize_filename(name);
if name.is_empty() {
return Err(AppError::Internal("Invalid filename".to_string()));
}
if let Some(size) = expected_size {
if size > MAX_IMAGE_SIZE {
return Err(AppError::Internal(format!(
"Image too large. Maximum size: {} GB",
MAX_IMAGE_SIZE / 1024 / 1024 / 1024
)));
}
}
let path = self.images_path.join(&name);
if path.exists() {
return Err(AppError::Internal(format!(
"Image already exists: {}",
name
)));
}
// Create file and copy data
let mut file = File::create(&path)
.map_err(|e| AppError::Internal(format!("Failed to create image file: {}", e)))?;
let bytes_written = io::copy(reader, &mut file).map_err(|e| {
let _ = fs::remove_file(&path);
AppError::Internal(format!("Failed to write image data: {}", e))
})?;
info!("Created image: {} ({} bytes)", name, bytes_written);
self.get_by_name(&name)
}
/// Create a new image from an async multipart field (streaming, memory-efficient)
///
/// This method streams data directly to disk without buffering the entire file in memory,
/// making it suitable for large files (multi-GB ISOs).
pub async fn create_from_multipart_field(
&self,
name: &str,
@@ -220,12 +146,10 @@ impl ImageManager {
return Err(AppError::Internal("Invalid filename".to_string()));
}
// Use a temporary file during upload
let temp_name = format!(".upload_{}", uuid::Uuid::new_v4());
let temp_path = self.images_path.join(&temp_name);
let final_path = self.images_path.join(&name);
// Check if final file already exists
if final_path.exists() {
return Err(AppError::Internal(format!(
"Image already exists: {}",
@@ -233,23 +157,19 @@ impl ImageManager {
)));
}
// Create temp file
let mut file = tokio::fs::File::create(&temp_path)
.await
.map_err(|e| AppError::Internal(format!("Failed to create temp file: {}", e)))?;
let mut bytes_written: u64 = 0;
// Stream chunks directly to disk
while let Some(chunk) = field
.chunk()
.await
.map_err(|e| AppError::Internal(format!("Failed to read upload chunk: {}", e)))?
{
// Check size limit
bytes_written += chunk.len() as u64;
if bytes_written > MAX_IMAGE_SIZE {
// Cleanup and return error
drop(file);
let _ = tokio::fs::remove_file(&temp_path).await;
return Err(AppError::Internal(format!(
@@ -258,19 +178,16 @@ impl ImageManager {
)));
}
// Write chunk to file
file.write_all(&chunk)
.await
.map_err(|e| AppError::Internal(format!("Failed to write chunk: {}", e)))?;
}
// Flush and close file
file.flush()
.await
.map_err(|e| AppError::Internal(format!("Failed to flush file: {}", e)))?;
drop(file);
// Move temp file to final location
tokio::fs::rename(&temp_path, &final_path)
.await
.map_err(|e| {
@@ -286,7 +203,6 @@ impl ImageManager {
self.get_by_name(&name)
}
/// Delete an image by ID
pub fn delete(&self, id: &str) -> Result<()> {
let image = self.get(id)?;
@@ -297,45 +213,6 @@ impl ImageManager {
Ok(())
}
/// Delete an image by name
pub fn delete_by_name(&self, name: &str) -> Result<()> {
let path = self.images_path.join(name);
if !path.exists() {
return Err(AppError::NotFound(format!("Image not found: {}", name)));
}
fs::remove_file(&path)
.map_err(|e| AppError::Internal(format!("Failed to delete image: {}", e)))?;
info!("Deleted image: {}", name);
Ok(())
}
/// Get total storage used
pub fn used_space(&self) -> u64 {
self.list()
.map(|images| images.iter().map(|i| i.size).sum())
.unwrap_or(0)
}
/// Check if storage has space for new image
pub fn has_space(&self, size: u64) -> bool {
// For now, just check against max size
// In the future, could check disk space
size <= MAX_IMAGE_SIZE
}
/// Download image from URL with progress callback
///
/// # Arguments
/// * `url` - The URL to download from
/// * `filename` - Optional custom filename (extracted from URL or Content-Disposition if not provided)
/// * `progress_callback` - Callback function called with (bytes_downloaded, total_bytes)
///
/// # Returns
/// * `Ok(ImageInfo)` - The downloaded image info
/// * `Err(AppError)` - If download fails
pub async fn download_from_url<F>(
&self,
url: &str,
@@ -347,20 +224,17 @@ impl ImageManager {
{
self.ensure_dir()?;
// Validate URL
let parsed_url = reqwest::Url::parse(url)
.map_err(|e| AppError::BadRequest(format!("Invalid URL: {}", e)))?;
info!("Starting download from: {}", url);
// Create HTTP client with timeout
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(3600)) // 1 hour timeout for large files
.timeout(std::time::Duration::from_secs(3600))
.connect_timeout(std::time::Duration::from_secs(30))
.build()
.map_err(|e| AppError::Internal(format!("Failed to create HTTP client: {}", e)))?;
// Send HEAD request first to get content info
let head_response = client
.head(url)
.send()
@@ -380,7 +254,6 @@ impl ImageManager {
.and_then(|v| v.to_str().ok())
.and_then(|s| s.parse::<u64>().ok());
// Check file size
if let Some(size) = total_size {
if size > MAX_IMAGE_SIZE {
return Err(AppError::BadRequest(format!(
@@ -391,11 +264,9 @@ impl ImageManager {
}
}
// Determine filename
let final_filename = if let Some(name) = filename {
sanitize_filename(&name)
} else {
// Try Content-Disposition header first
let from_header = head_response
.headers()
.get(reqwest::header::CONTENT_DISPOSITION)
@@ -405,7 +276,6 @@ impl ImageManager {
if let Some(name) = from_header {
sanitize_filename(&name)
} else {
// Fall back to URL path
let path = parsed_url.path();
let name = path.rsplit('/').next().unwrap_or("download");
let name = urlencoding::decode(name).unwrap_or_else(|_| name.into());
@@ -419,7 +289,6 @@ impl ImageManager {
));
}
// Check if file already exists
let final_path = self.images_path.join(&final_filename);
if final_path.exists() {
return Err(AppError::BadRequest(format!(
@@ -428,11 +297,9 @@ impl ImageManager {
)));
}
// Create temporary file for download
let temp_filename = format!(".download_{}", uuid::Uuid::new_v4());
let temp_path = self.images_path.join(&temp_filename);
// Start actual download
let response = client
.get(url)
.send()
@@ -446,7 +313,6 @@ impl ImageManager {
)));
}
// Get actual content length from response (may differ from HEAD)
let content_length = response
.headers()
.get(reqwest::header::CONTENT_LENGTH)
@@ -454,19 +320,16 @@ impl ImageManager {
.and_then(|s| s.parse::<u64>().ok())
.or(total_size);
// Create temp file
let mut file = tokio::fs::File::create(&temp_path)
.await
.map_err(|e| AppError::Internal(format!("Failed to create temp file: {}", e)))?;
// Stream download with progress (throttled)
let mut stream = response.bytes_stream();
let mut downloaded: u64 = 0;
let mut last_report_time = Instant::now();
let mut last_reported_bytes: u64 = 0;
let throttle_interval = Duration::from_millis(PROGRESS_THROTTLE_MS);
// Report initial progress
progress_callback(0, content_length);
while let Some(chunk_result) = stream.next().await {
@@ -474,14 +337,12 @@ impl ImageManager {
chunk_result.map_err(|e| AppError::Internal(format!("Download error: {}", e)))?;
file.write_all(&chunk).await.map_err(|e| {
// Cleanup on error
let _ = std::fs::remove_file(&temp_path);
AppError::Internal(format!("Failed to write data: {}", e))
})?;
downloaded += chunk.len() as u64;
// Throttled progress reporting: report if enough time or bytes have passed
let now = Instant::now();
let time_elapsed = now.duration_since(last_report_time) >= throttle_interval;
let bytes_elapsed = downloaded - last_reported_bytes >= PROGRESS_THROTTLE_BYTES;
@@ -493,18 +354,15 @@ impl ImageManager {
}
}
// Always report final progress
if downloaded != last_reported_bytes {
progress_callback(downloaded, content_length);
}
// Ensure all data is flushed
file.flush()
.await
.map_err(|e| AppError::Internal(format!("Failed to flush file: {}", e)))?;
drop(file);
// Verify downloaded size
let metadata = tokio::fs::metadata(&temp_path)
.await
.map_err(|e| AppError::Internal(format!("Failed to read file metadata: {}", e)))?;
@@ -520,7 +378,6 @@ impl ImageManager {
}
}
// Move temp file to final location
tokio::fs::rename(&temp_path, &final_path)
.await
.map_err(|e| {
@@ -534,35 +391,29 @@ impl ImageManager {
metadata.len()
);
// Return image info
self.get_by_name(&final_filename)
}
/// Get images storage path
pub fn images_path(&self) -> &PathBuf {
&self.images_path
}
}
/// Simple hash function for generating stable IDs
fn md5_hash(s: &str) -> u64 {
fn stable_image_id_from_filename(name: &str) -> String {
let mut hash: u64 = 0;
for (i, byte) in s.bytes().enumerate() {
for (i, byte) in name.bytes().enumerate() {
hash = hash.wrapping_add((byte as u64).wrapping_mul((i as u64).wrapping_add(1)));
hash = hash.wrapping_mul(31);
}
hash
format!("{:x}", hash)
}
/// Sanitize filename to prevent path traversal
fn sanitize_filename(name: &str) -> String {
let name = name.trim();
let name = name.replace(['/', '\\', '\0', ':', '*', '?', '"', '<', '>', '|'], "_");
// Remove leading dots (hidden files)
let name = name.trim_start_matches('.');
// Limit length
if name.len() > 255 {
name[..255].to_string()
} else {
@@ -570,17 +421,10 @@ fn sanitize_filename(name: &str) -> String {
}
}
/// Extract filename from Content-Disposition header
fn extract_filename_from_content_disposition(header: &str) -> Option<String> {
// Handle both:
// Content-Disposition: attachment; filename="example.iso"
// Content-Disposition: attachment; filename*=UTF-8''example.iso
// Try filename* first (RFC 5987)
if let Some(pos) = header.find("filename*=") {
let start = pos + 10;
let value = &header[start..];
// Format: charset'language'value
if let Some(quote_start) = value.find("''") {
let encoded = value[quote_start + 2..].split(';').next()?;
let decoded = urlencoding::decode(encoded.trim()).ok()?;
@@ -591,7 +435,6 @@ fn extract_filename_from_content_disposition(header: &str) -> Option<String> {
}
}
// Try filename next
if let Some(pos) = header.find("filename=") {
let start = pos + 9;
let value = &header[start..];
@@ -613,7 +456,7 @@ mod tests {
#[test]
fn test_sanitize_filename() {
assert_eq!(sanitize_filename("test.iso"), "test.iso");
assert_eq!(sanitize_filename("../test.iso"), "_test.iso"); // .. becomes empty after trim_start_matches('.')
assert_eq!(sanitize_filename("../test.iso"), "_test.iso");
assert_eq!(sanitize_filename("test/file.iso"), "test_file.iso");
assert_eq!(sanitize_filename(".hidden.iso"), "hidden.iso");
}

19
src/msd/mod.rs
View File

@@ -1,19 +1,3 @@
//! MSD (Mass Storage Device) module
//!
//! Provides virtual USB storage functionality with two modes:
//! - Image mounting: Mount ISO/IMG files for system installation
//! - Ventoy drive: Bootable exFAT drive for multiple ISO files
//!
//! Architecture:
//! ```text
//! Web API --> MSD Controller --> ConfigFS Mass Storage --> Target PC
//! |
//! ┌──────┴──────┐
//! │ │
//! Image Manager Ventoy Drive
//! (ISO/IMG) (Bootable exFAT)
//! ```
pub mod controller;
pub mod image;
pub mod monitor;
@@ -22,12 +6,11 @@ pub mod ventoy_drive;
pub use controller::MsdController;
pub use image::ImageManager;
pub use monitor::{MsdHealthMonitor, MsdHealthStatus, MsdMonitorConfig};
pub use monitor::MsdHealthMonitor;
pub use types::{
DownloadProgress, DownloadStatus, DriveFile, DriveInfo, DriveInitRequest, ImageDownloadRequest,
ImageInfo, MsdConnectRequest, MsdMode, MsdState,
};
pub use ventoy_drive::VentoyDrive;
// Re-export from otg module for backward compatibility
pub use crate::otg::{MsdFunction, MsdLunConfig};

90
src/msd/monitor.rs
View File

@@ -1,99 +1,46 @@
//! MSD (Mass Storage Device) health monitoring
//!
//! This module provides health monitoring for MSD operations, including:
//! - ConfigFS operation error tracking
//! - Image mount/unmount error tracking
//! - Error state tracking
//! - Log throttling to prevent log flooding
use std::sync::atomic::{AtomicU32, Ordering};
use tokio::sync::RwLock;
use tracing::{info, warn};
use crate::utils::LogThrottler;
/// MSD health status
const LOG_THROTTLE_SECS: u64 = 5;
#[derive(Debug, Clone, PartialEq, Default)]
pub enum MsdHealthStatus {
/// Device is healthy and operational
pub(crate) enum MsdHealthStatus {
#[default]
Healthy,
/// Device has an error
Error {
/// Human-readable error reason
reason: String,
/// Error code for programmatic handling
error_code: String,
},
}
/// MSD health monitor configuration
#[derive(Debug, Clone)]
pub struct MsdMonitorConfig {
/// Log throttle interval in seconds
pub log_throttle_secs: u64,
}
impl Default for MsdMonitorConfig {
fn default() -> Self {
Self {
log_throttle_secs: 5,
}
}
}
/// MSD health monitor
///
/// Monitors MSD operation health and manages error state.
pub struct MsdHealthMonitor {
/// Current health status
status: RwLock<MsdHealthStatus>,
/// Log throttler to prevent log flooding
throttler: LogThrottler,
/// Error count (for tracking)
error_count: AtomicU32,
/// Last error code (for change detection)
last_error_code: RwLock<Option<String>>,
}
impl MsdHealthMonitor {
/// Create a new MSD health monitor with the specified configuration
pub fn new(config: MsdMonitorConfig) -> Self {
let throttle_secs = config.log_throttle_secs;
pub fn with_defaults() -> Self {
Self {
status: RwLock::new(MsdHealthStatus::Healthy),
throttler: LogThrottler::with_secs(throttle_secs),
throttler: LogThrottler::with_secs(LOG_THROTTLE_SECS),
error_count: AtomicU32::new(0),
last_error_code: RwLock::new(None),
}
}
/// Create a new MSD health monitor with default configuration
pub fn with_defaults() -> Self {
Self::new(MsdMonitorConfig::default())
}
/// Report an error from MSD operations
///
/// This method is called when an MSD operation fails. It:
/// 1. Updates the health status
/// 2. Logs the error (with throttling)
/// 3. Updates in-memory error state
///
/// # Arguments
///
/// * `reason` - Human-readable error description
/// * `error_code` - Error code for programmatic handling
pub async fn report_error(&self, reason: &str, error_code: &str) {
let count = self.error_count.fetch_add(1, Ordering::Relaxed) + 1;
// Check if error code changed
let error_changed = {
let last = self.last_error_code.read().await;
last.as_ref().map(|s| s.as_str()) != Some(error_code)
};
// Log with throttling (always log if error type changed)
let throttle_key = format!("msd_{}", error_code);
if error_changed || self.throttler.should_log(&throttle_key) {
warn!(
@@ -102,29 +49,21 @@ impl MsdHealthMonitor {
);
}
// Update last error code
*self.last_error_code.write().await = Some(error_code.to_string());
// Update status
*self.status.write().await = MsdHealthStatus::Error {
reason: reason.to_string(),
error_code: error_code.to_string(),
};
}
/// Report that the MSD has recovered from error
///
/// This method is called when an MSD operation succeeds after errors.
/// It resets the error state.
pub async fn report_recovered(&self) {
let prev_status = self.status.read().await.clone();
// Only report recovery if we were in an error state
if prev_status != MsdHealthStatus::Healthy {
let error_count = self.error_count.load(Ordering::Relaxed);
info!("MSD recovered after {} errors", error_count);
// Reset state
self.error_count.store(0, Ordering::Relaxed);
self.throttler.clear_all();
*self.last_error_code.write().await = None;
@@ -132,29 +71,25 @@ impl MsdHealthMonitor {
}
}
/// Get the current health status
pub async fn status(&self) -> MsdHealthStatus {
#[cfg(test)]
pub(crate) async fn status(&self) -> MsdHealthStatus {
self.status.read().await.clone()
}
/// Get the current error count
pub fn error_count(&self) -> u32 {
#[cfg(test)]
pub(crate) fn error_count(&self) -> u32 {
self.error_count.load(Ordering::Relaxed)
}
/// Check if the monitor is in an error state
pub async fn is_error(&self) -> bool {
matches!(*self.status.read().await, MsdHealthStatus::Error { .. })
}
/// Check if the monitor is healthy
pub async fn is_healthy(&self) -> bool {
#[cfg(test)]
pub(crate) async fn is_healthy(&self) -> bool {
matches!(*self.status.read().await, MsdHealthStatus::Healthy)
}
/// Reset the monitor to healthy state without publishing events
///
/// This is useful during initialization.
pub async fn reset(&self) {
self.error_count.store(0, Ordering::Relaxed);
*self.last_error_code.write().await = None;
@@ -162,7 +97,6 @@ impl MsdHealthMonitor {
self.throttler.clear_all();
}
/// Get the current error message if in error state
pub async fn error_message(&self) -> Option<String> {
match &*self.status.read().await {
MsdHealthStatus::Error { reason, .. } => Some(reason.clone()),
@@ -212,13 +146,11 @@ mod tests {
async fn test_report_recovered() {
let monitor = MsdHealthMonitor::with_defaults();
// First report an error
monitor
.report_error("Image not found", "image_not_found")
.await;
assert!(monitor.is_error().await);
// Then report recovery
monitor.report_recovered().await;
assert!(monitor.is_healthy().await);
assert_eq!(monitor.error_count(), 0);

66
src/msd/types.rs
View File

@@ -1,42 +1,28 @@
//! MSD data types and structures
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use time::OffsetDateTime;
/// MSD operating mode
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
#[derive(Default)]
pub enum MsdMode {
/// No storage connected
#[default]
None,
/// Image file mounted (ISO/IMG)
Image,
/// Virtual drive (FAT32) connected
Drive,
}
/// Image file metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageInfo {
/// Unique image ID
pub id: String,
/// Display name
pub name: String,
/// File path on disk
#[serde(skip_serializing)]
pub path: PathBuf,
/// File size in bytes
pub size: u64,
/// Creation timestamp
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
}
impl ImageInfo {
/// Create new image info
pub fn new(id: String, name: String, path: PathBuf, size: u64) -> Self {
Self {
id,
@@ -47,7 +33,6 @@ impl ImageInfo {
}
}
/// Format size for display
pub fn size_display(&self) -> String {
const KB: u64 = 1024;
const MB: u64 = KB * 1024;
@@ -65,18 +50,12 @@ impl ImageInfo {
}
}
/// MSD state information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MsdState {
/// Whether MSD feature is available
pub available: bool,
/// Current mode
pub mode: MsdMode,
/// Whether storage is connected to target
pub connected: bool,
/// Currently mounted image (if mode is Image)
pub current_image: Option<ImageInfo>,
/// Virtual drive info (if mode is Drive)
pub drive_info: Option<DriveInfo>,
}
@@ -92,24 +71,17 @@ impl Default for MsdState {
}
}
/// Virtual drive information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DriveInfo {
/// Drive size in bytes
pub size: u64,
/// Used space in bytes
pub used: u64,
/// Free space in bytes
pub free: u64,
/// Whether drive is initialized
pub initialized: bool,
/// Drive file path
#[serde(skip_serializing)]
pub path: PathBuf,
}
impl DriveInfo {
/// Create new drive info
pub fn new(path: PathBuf, size: u64) -> Self {
Self {
size,
@@ -121,92 +93,60 @@ impl DriveInfo {
}
}
/// File entry in virtual drive
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DriveFile {
/// File name
pub name: String,
/// Relative path from drive root
pub path: String,
/// File size in bytes (0 for directories)
pub size: u64,
/// Whether this is a directory
pub is_dir: bool,
/// Last modified timestamp
#[serde(with = "time::serde::rfc3339::option")]
pub modified: Option<OffsetDateTime>,
}
/// MSD connect request
#[derive(Debug, Clone, Deserialize)]
pub struct MsdConnectRequest {
/// Connection mode: "image" or "drive"
pub mode: MsdMode,
/// Image ID to mount (required for image mode)
pub image_id: Option<String>,
/// Mount as CD-ROM (optional, defaults based on image type)
#[serde(default)]
pub cdrom: Option<bool>,
/// Mount as read-only
#[serde(default)]
pub read_only: Option<bool>,
}
/// Virtual drive init request
#[derive(Debug, Clone, Deserialize)]
pub struct DriveInitRequest {
/// Drive size in megabytes (defaults to 16GB)
#[serde(default = "default_drive_size")]
pub size_mb: u32,
/// Optional custom path for Ventoy installation
pub ventoy_path: Option<String>,
}
fn default_drive_size() -> u32 {
16 * 1024 // 16GB
16 * 1024
}
/// Image download request
#[derive(Debug, Clone, Deserialize)]
pub struct ImageDownloadRequest {
/// URL to download from
pub url: String,
/// Optional custom filename
pub filename: Option<String>,
}
/// Download status
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DownloadStatus {
/// Download has started
Started,
/// Download is in progress
InProgress,
/// Download completed successfully
Completed,
/// Download failed
Failed,
}
/// Download progress information
#[derive(Debug, Clone, Serialize)]
pub struct DownloadProgress {
/// Unique download ID
pub download_id: String,
/// Source URL
pub url: String,
/// Target filename
pub filename: String,
/// Bytes downloaded so far
pub bytes_downloaded: u64,
/// Total file size (None if unknown)
pub total_bytes: Option<u64>,
/// Progress percentage (0.0 - 100.0, None if total unknown)
pub progress_pct: Option<f32>,
/// Download status
pub status: DownloadStatus,
/// Error message if failed
pub error: Option<String>,
}
@@ -220,7 +160,7 @@ mod tests {
"test".into(),
"test.iso".into(),
PathBuf::from("/tmp/test.iso"),
1024 * 1024 * 1024 * 2, // 2 GB
1024 * 1024 * 1024 * 2,
);
assert!(info.size_display().contains("GB"));
}

123
src/msd/ventoy_drive.rs
View File

@@ -1,8 +1,3 @@
//! Ventoy Virtual Drive
//!
//! Replaces FAT32 VirtualDrive with a Ventoy bootable image.
//! Provides a bootable USB with exFAT data partition for ISO files.
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::RwLock;
@@ -13,33 +8,20 @@ use ventoy_img::{FileInfo as VentoyFileInfo, VentoyError, VentoyImage};
use super::types::{DriveFile, DriveInfo};
use crate::error::{AppError, Result};
/// Chunk size for streaming reads (64 KB)
const STREAM_CHUNK_SIZE: usize = 64 * 1024;
/// Minimum drive size (1 GB) - Ventoy requires space for boot partition
const MIN_DRIVE_SIZE_MB: u32 = 1024;
/// Maximum drive size (128 GB)
const MAX_DRIVE_SIZE_MB: u32 = 128 * 1024;
/// Default drive label
const DEFAULT_LABEL: &str = "ONE-KVM";
/// Ventoy Drive Manager
///
/// Thread-safe wrapper around VentoyImage providing async file operations.
/// Uses spawn_blocking for all ventoy-img-rs operations since they are synchronous.
/// Uses RwLock to allow concurrent read operations while serializing writes.
pub struct VentoyDrive {
/// Drive image path
path: PathBuf,
/// RwLock for concurrent reads, exclusive writes
/// (ventoy-img-rs operations are synchronous and not thread-safe)
lock: Arc<RwLock<()>>,
}
impl VentoyDrive {
/// Create new Ventoy drive manager
pub fn new(path: PathBuf) -> Self {
Self {
path,
@@ -47,40 +29,32 @@ impl VentoyDrive {
}
}
/// Check if drive image exists
pub fn exists(&self) -> bool {
self.path.exists()
}
/// Get drive path
pub fn path(&self) -> &PathBuf {
&self.path
}
/// Initialize a new Ventoy drive image
///
/// Creates a bootable Ventoy image with the specified size.
/// The image includes boot partitions and an exFAT data partition.
pub async fn init(&self, size_mb: u32) -> Result<DriveInfo> {
let size_mb = size_mb.clamp(MIN_DRIVE_SIZE_MB, MAX_DRIVE_SIZE_MB);
let size_str = format!("{}M", size_mb);
let path = self.path.clone();
let _lock = self.lock.write().await; // Write lock for initialization
let _lock = self.lock.write().await;
info!("Creating {} MB Ventoy drive at {}", size_mb, path.display());
// Run Ventoy creation in blocking task
let info = tokio::task::spawn_blocking(move || {
VentoyImage::create(&path, &size_str, DEFAULT_LABEL).map_err(ventoy_to_app_error)?;
// Get file metadata for DriveInfo
let metadata = std::fs::metadata(&path)
.map_err(|e| AppError::Internal(format!("Failed to read drive metadata: {}", e)))?;
Ok::<DriveInfo, AppError>(DriveInfo {
size: metadata.len(),
used: 0,
free: metadata.len(), // Approximate - exFAT overhead not calculated
free: metadata.len(),
initialized: true,
path,
})
@@ -92,20 +66,18 @@ impl VentoyDrive {
Ok(info)
}
/// Get drive information
pub async fn info(&self) -> Result<DriveInfo> {
if !self.exists() {
return Err(AppError::Internal("Drive not initialized".to_string()));
}
let path = self.path.clone();
let _lock = self.lock.read().await; // Read lock for info query
let _lock = self.lock.read().await;
tokio::task::spawn_blocking(move || {
let metadata = std::fs::metadata(&path)
.map_err(|e| AppError::Internal(format!("Failed to read drive metadata: {}", e)))?;
// Open image to get file list and calculate used space
let image = VentoyImage::open(&path).map_err(ventoy_to_app_error)?;
let files = image.list_files_recursive().map_err(ventoy_to_app_error)?;
@@ -116,7 +88,6 @@ impl VentoyDrive {
.map(|f| f.size)
.sum();
// Note: This is approximate since we don't have exact exFAT overhead
let size = metadata.len();
let free = size.saturating_sub(used);
@@ -132,7 +103,6 @@ impl VentoyDrive {
.map_err(|e| AppError::Internal(format!("Task join error: {}", e)))?
}
/// List files at a given path (or root if empty/"/")
pub async fn list_files(&self, dir_path: &str) -> Result<Vec<DriveFile>> {
if !self.exists() {
return Err(AppError::Internal("Drive not initialized".to_string()));
@@ -140,7 +110,7 @@ impl VentoyDrive {
let path = self.path.clone();
let dir_path = dir_path.to_string();
let _lock = self.lock.read().await; // Read lock for listing
let _lock = self.lock.read().await;
tokio::task::spawn_blocking(move || {
let image = VentoyImage::open(&path).map_err(ventoy_to_app_error)?;
@@ -161,9 +131,6 @@ impl VentoyDrive {
.map_err(|e| AppError::Internal(format!("Task join error: {}", e)))?
}
/// Write a file to the drive from multipart upload (streaming)
///
/// Streams the file directly into the Ventoy image's exFAT partition.
pub async fn write_file_from_multipart_field(
&self,
file_path: &str,
@@ -173,12 +140,10 @@ impl VentoyDrive {
return Err(AppError::Internal("Drive not initialized".to_string()));
}
// First, stream to a temporary file (to get the size)
let temp_dir = self.path.parent().unwrap_or(Path::new("/tmp"));
let temp_name = format!(".upload_ventoy_{}", uuid::Uuid::new_v4());
let temp_path = temp_dir.join(&temp_name);
// Stream upload to temp file
let mut temp_file = tokio::fs::File::create(&temp_path)
.await
.map_err(|e| AppError::Internal(format!("Failed to create temp file: {}", e)))?;
@@ -201,23 +166,16 @@ impl VentoyDrive {
.map_err(|e| AppError::Internal(format!("Failed to flush temp file: {}", e)))?;
drop(temp_file);
// Now copy from temp file to Ventoy image
let path = self.path.clone();
let file_path = file_path.to_string();
let temp_path_clone = temp_path.clone();
let _lock = self.lock.write().await; // Write lock for file write
let _lock = self.lock.write().await;
let result = tokio::task::spawn_blocking(move || {
let mut image = VentoyImage::open(&path).map_err(ventoy_to_app_error)?;
// Use add_file_to_path which handles streaming internally
image
.add_file_to_path(
&temp_path_clone,
&file_path,
true, // create_parents
true, // overwrite
)
.add_file_to_path(&temp_path_clone, &file_path, true, true)
.map_err(ventoy_to_app_error)?;
Ok::<(), AppError>(())
@@ -225,14 +183,13 @@ impl VentoyDrive {
.await
.map_err(|e| AppError::Internal(format!("Task join error: {}", e)))?;
// Cleanup temp file
let _ = tokio::fs::remove_file(&temp_path).await;
result?;
Ok(bytes_written)
}
/// Read a file from the drive (for download)
#[cfg(test)]
pub async fn read_file(&self, file_path: &str) -> Result<Vec<u8>> {
if !self.exists() {
return Err(AppError::Internal("Drive not initialized".to_string()));
@@ -240,7 +197,7 @@ impl VentoyDrive {
let path = self.path.clone();
let file_path = file_path.to_string();
let _lock = self.lock.read().await; // Read lock for file read
let _lock = self.lock.read().await;
tokio::task::spawn_blocking(move || {
let image = VentoyImage::open(&path).map_err(ventoy_to_app_error)?;
@@ -251,10 +208,6 @@ impl VentoyDrive {
.map_err(|e| AppError::Internal(format!("Task join error: {}", e)))?
}
/// Get file information without reading content
///
/// Returns file size, name, and other metadata.
/// Returns None if the file doesn't exist.
pub async fn get_file_info(&self, file_path: &str) -> Result<Option<DriveFile>> {
if !self.exists() {
return Err(AppError::Internal("Drive not initialized".to_string()));
@@ -262,7 +215,7 @@ impl VentoyDrive {
let path = self.path.clone();
let file_path_owned = file_path.to_string();
let _lock = self.lock.read().await; // Read lock for file info
let _lock = self.lock.read().await;
let info = tokio::task::spawn_blocking(move || {
let image = VentoyImage::open(&path).map_err(ventoy_to_app_error)?;
@@ -282,10 +235,6 @@ impl VentoyDrive {
}))
}
/// Read a file from the drive as a stream (for large file downloads)
///
/// Returns an async channel receiver that yields chunks of file data.
/// This avoids loading the entire file into memory.
pub async fn read_file_stream(
&self,
file_path: &str,
@@ -297,7 +246,6 @@ impl VentoyDrive {
return Err(AppError::Internal("Drive not initialized".to_string()));
}
// First, get the file size
let file_info = self
.get_file_info(file_path)
.await?
@@ -315,15 +263,12 @@ impl VentoyDrive {
let file_path_owned = file_path.to_string();
let lock = self.lock.clone();
// Create a channel for streaming data
let (tx, rx) =
tokio::sync::mpsc::channel::<std::result::Result<bytes::Bytes, std::io::Error>>(8);
// Spawn blocking task to read and send chunks
tokio::task::spawn_blocking(move || {
// Hold read lock for the entire read operation
let rt = tokio::runtime::Handle::current();
let _lock = rt.block_on(lock.read()); // Read lock for streaming
let _lock = rt.block_on(lock.read());
let image = match VentoyImage::open(&path) {
Ok(img) => img,
@@ -333,10 +278,8 @@ impl VentoyDrive {
}
};
// Create a channel writer that sends chunks
let mut chunk_writer = ChannelWriter::new(tx.clone(), rt.clone());
// Stream the file through the writer
if let Err(e) = image.read_file_to_writer(&file_path_owned, &mut chunk_writer) {
let _ = rt.block_on(tx.send(Err(std::io::Error::other(e.to_string()))));
}
@@ -345,7 +288,6 @@ impl VentoyDrive {
Ok((file_size, rx))
}
/// Create a directory
pub async fn mkdir(&self, dir_path: &str) -> Result<()> {
if !self.exists() {
return Err(AppError::Internal("Drive not initialized".to_string()));
@@ -353,7 +295,7 @@ impl VentoyDrive {
let path = self.path.clone();
let dir_path = dir_path.to_string();
let _lock = self.lock.write().await; // Write lock for mkdir
let _lock = self.lock.write().await;
tokio::task::spawn_blocking(move || {
let mut image = VentoyImage::open(&path).map_err(ventoy_to_app_error)?;
@@ -366,7 +308,6 @@ impl VentoyDrive {
.map_err(|e| AppError::Internal(format!("Task join error: {}", e)))?
}
/// Delete a file or directory
pub async fn delete(&self, path_to_delete: &str) -> Result<()> {
if !self.exists() {
return Err(AppError::Internal("Drive not initialized".to_string()));
@@ -374,12 +315,11 @@ impl VentoyDrive {
let path = self.path.clone();
let path_to_delete = path_to_delete.to_string();
let _lock = self.lock.write().await; // Write lock for delete
let _lock = self.lock.write().await;
tokio::task::spawn_blocking(move || {
let mut image = VentoyImage::open(&path).map_err(ventoy_to_app_error)?;
// Use recursive delete to handle directories
image
.remove_recursive(&path_to_delete)
.map_err(ventoy_to_app_error)
@@ -389,7 +329,6 @@ impl VentoyDrive {
}
}
/// Convert VentoyError to AppError
fn ventoy_to_app_error(err: VentoyError) -> AppError {
match err {
VentoyError::Io(e) => AppError::Io(e),
@@ -405,7 +344,6 @@ fn ventoy_to_app_error(err: VentoyError) -> AppError {
}
}
/// Convert VentoyFileInfo to DriveFile
fn ventoy_file_to_drive_file(info: VentoyFileInfo, parent_path: &str) -> DriveFile {
let full_path = if parent_path.is_empty() || parent_path == "/" {
format!("/{}", info.name)
@@ -418,13 +356,10 @@ fn ventoy_file_to_drive_file(info: VentoyFileInfo, parent_path: &str) -> DriveFi
path: full_path,
size: info.size,
is_dir: info.is_directory,
modified: None, // Ventoy FileInfo doesn't include timestamps
modified: None,
}
}
/// A writer that sends chunks to an async channel
///
/// This bridges the sync Write trait with async channels for streaming.
struct ChannelWriter {
tx: tokio::sync::mpsc::Sender<std::result::Result<bytes::Bytes, std::io::Error>>,
rt: tokio::runtime::Handle,
@@ -484,7 +419,6 @@ impl std::io::Write for ChannelWriter {
impl Drop for ChannelWriter {
fn drop(&mut self) {
// Flush any remaining data when the writer is dropped
let _ = self.flush_buffer();
}
}
@@ -496,16 +430,13 @@ mod tests {
use std::sync::OnceLock;
use tempfile::TempDir;
/// Path to ventoy resources directory
static RESOURCE_DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../ventoy-img-rs/resources");
/// Initialize ventoy resources once
fn init_ventoy_resources() -> bool {
static INIT: OnceLock<bool> = OnceLock::new();
*INIT.get_or_init(|| {
let resource_path = std::path::Path::new(RESOURCE_DIR);
// Decompress xz files if needed
let core_xz = resource_path.join("core.img.xz");
let core_img = resource_path.join("core.img");
if core_xz.exists() && !core_img.exists() {
@@ -524,7 +455,6 @@ mod tests {
}
}
// Initialize resources
if let Err(e) = ventoy_img::resources::init_resources(resource_path) {
eprintln!("Failed to init ventoy resources: {}", e);
return false;
@@ -534,7 +464,6 @@ mod tests {
})
}
/// Decompress xz file using system command
fn decompress_xz(src: &std::path::Path, dst: &std::path::Path) -> std::io::Result<()> {
let output = Command::new("xz")
.args(["-d", "-k", "-c", src.to_str().unwrap()])
@@ -551,7 +480,6 @@ mod tests {
Ok(())
}
/// Ensure resources are initialized, skip test if failed
fn ensure_resources() -> bool {
if !init_ventoy_resources() {
eprintln!("Skipping test: ventoy resources not available");
@@ -602,15 +530,12 @@ mod tests {
let drive_path = temp_dir.path().join("test_ventoy.img");
let drive = VentoyDrive::new(drive_path.clone());
// Initialize drive
drive.init(MIN_DRIVE_SIZE_MB).await.unwrap();
// Write a test file
let test_content = b"Hello, Ventoy!";
let test_file_path = temp_dir.path().join("test.txt");
std::fs::write(&test_file_path, test_content).unwrap();
// Add file to drive using ventoy-img directly
let path = drive.path().clone();
tokio::task::spawn_blocking(move || {
let mut image = VentoyImage::open(&path).unwrap();
@@ -619,7 +544,6 @@ mod tests {
.await
.unwrap();
// Read file from drive
let read_data = drive.read_file("/test.txt").await.unwrap();
assert_eq!(read_data, test_content);
}
@@ -633,18 +557,14 @@ mod tests {
let drive_path = temp_dir.path().join("test_ventoy.img");
let drive = VentoyDrive::new(drive_path.clone());
// Initialize drive
drive.init(MIN_DRIVE_SIZE_MB).await.unwrap();
// Create a directory
drive.mkdir("/mydir").await.unwrap();
// Write a test file
let test_content = b"Test file content for info check";
let test_file_path = temp_dir.path().join("info_test.txt");
std::fs::write(&test_file_path, test_content).unwrap();
// Add file to drive
let path = drive.path().clone();
tokio::task::spawn_blocking(move || {
let mut image = VentoyImage::open(&path).unwrap();
@@ -653,7 +573,6 @@ mod tests {
.await
.unwrap();
// Test get_file_info for file
let file_info = drive.get_file_info("/info_test.txt").await.unwrap();
assert!(file_info.is_some());
let file_info = file_info.unwrap();
@@ -661,14 +580,12 @@ mod tests {
assert_eq!(file_info.size, test_content.len() as u64);
assert!(!file_info.is_dir);
// Test get_file_info for directory
let dir_info = drive.get_file_info("/mydir").await.unwrap();
assert!(dir_info.is_some());
let dir_info = dir_info.unwrap();
assert_eq!(dir_info.name, "mydir");
assert!(dir_info.is_dir);
// Test get_file_info for non-existent file
let not_found = drive.get_file_info("/nonexistent.txt").await.unwrap();
assert!(not_found.is_none());
}
@@ -682,16 +599,13 @@ mod tests {
let drive_path = temp_dir.path().join("test_ventoy.img");
let drive = VentoyDrive::new(drive_path.clone());
// Initialize drive
drive.init(MIN_DRIVE_SIZE_MB).await.unwrap();
// Create test data that spans multiple chunks (>64KB)
let test_size = 200 * 1024; // 200 KB
let test_size = 200 * 1024;
let test_content: Vec<u8> = (0..test_size).map(|i| (i % 256) as u8).collect();
let test_file_path = temp_dir.path().join("large_file.bin");
std::fs::write(&test_file_path, &test_content).unwrap();
// Add file to drive
let path = drive.path().clone();
let file_path_clone = test_file_path.clone();
tokio::task::spawn_blocking(move || {
@@ -701,18 +615,15 @@ mod tests {
.await
.unwrap();
// Stream read the file
let (file_size, mut rx) = drive.read_file_stream("/large_file.bin").await.unwrap();
assert_eq!(file_size, test_size as u64);
// Collect all chunks
let mut received_data = Vec::new();
while let Some(chunk_result) = rx.recv().await {
let chunk = chunk_result.expect("Chunk should not be an error");
received_data.extend_from_slice(&chunk);
}
// Verify data matches
assert_eq!(received_data.len(), test_content.len());
assert_eq!(received_data, test_content);
}
@@ -726,15 +637,12 @@ mod tests {
let drive_path = temp_dir.path().join("test_ventoy.img");
let drive = VentoyDrive::new(drive_path.clone());
// Initialize drive
drive.init(MIN_DRIVE_SIZE_MB).await.unwrap();
// Create a small test file
let test_content = b"Small file for streaming test";
let test_file_path = temp_dir.path().join("small.txt");
std::fs::write(&test_file_path, test_content).unwrap();
// Add file to drive
let path = drive.path().clone();
tokio::task::spawn_blocking(move || {
let mut image = VentoyImage::open(&path).unwrap();
@@ -743,18 +651,15 @@ mod tests {
.await
.unwrap();
// Stream read the file
let (file_size, mut rx) = drive.read_file_stream("/small.txt").await.unwrap();
assert_eq!(file_size, test_content.len() as u64);
// Collect all chunks
let mut received_data = Vec::new();
while let Some(chunk_result) = rx.recv().await {
let chunk = chunk_result.expect("Chunk should not be an error");
received_data.extend_from_slice(&chunk);
}
// Verify data matches
assert_eq!(received_data.as_slice(), test_content);
}
}