spacedrive/core/src/volume.rs

use crate::{
	library::Library,
	prisma::volume::{self, *},
};

use serde::{Deserialize, Serialize};
use serde_with::{serde_as, DisplayFromStr};
use specta::Type;
use std::{fmt::Display, process::Command};
use sysinfo::{DiskExt, System, SystemExt};
use thiserror::Error;

#[derive(Serialize, Deserialize, Debug, Clone, Type)]
#[allow(clippy::upper_case_acronyms)]
pub enum DiskType {
	SSD,
	HDD,
	Removable,
}

impl Display for DiskType {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		f.write_str(match self {
			Self::SSD => "SSD",
			Self::HDD => "HDD",
			Self::Removable => "Removable",
		})
	}
}

#[serde_as]
#[derive(Serialize, Deserialize, Debug, Clone, Type)]
pub struct Volume {
	pub name: String,
	pub mount_point: String,
	#[specta(type = String)]
	#[serde_as(as = "DisplayFromStr")]
	pub total_capacity: u64,
	#[specta(type = String)]
	#[serde_as(as = "DisplayFromStr")]
	pub available_capacity: u64,
	pub is_removable: bool,
	pub disk_type: Option<DiskType>,
	pub file_system: Option<String>,
	pub is_root_filesystem: bool,
}

#[derive(Error, Debug)]
pub enum VolumeError {
	#[error("Database error: {0}")]
	DatabaseErr(#[from] prisma_client_rust::QueryError),
	#[error("FromUtf8Error: {0}")]
	FromUtf8Error(#[from] std::string::FromUtf8Error),
}

impl From<VolumeError> for rspc::Error {
	fn from(e: VolumeError) -> Self {
		rspc::Error::with_cause(rspc::ErrorCode::InternalServerError, e.to_string(), e)
	}
}

pub async fn save_volume(library: &Library) -> Result<(), VolumeError> {
	let volumes = get_volumes()?;

	// enter all volumes associate with this client add to db
	for volume in volumes {
		let params = vec![
			disk_type::set(volume.disk_type.map(|t| t.to_string())),
			filesystem::set(volume.file_system.clone()),
			total_bytes_capacity::set(volume.total_capacity.to_string()),
			total_bytes_available::set(volume.available_capacity.to_string()),
		];

		library
			.db
			.volume()
			.upsert(
				node_id_mount_point_name(
					library.node_local_id,
					volume.mount_point.to_string(),
					volume.name.to_string(),
				),
				volume::create(
					library.node_local_id,
					volume.name,
					volume.mount_point,
					params.clone(),
				),
				params,
			)
			.exec()
			.await?;
	}
	// cleanup: remove all unmodified volumes associate with this client

	Ok(())
}

// TODO: Error handling in this function
pub fn get_volumes() -> Result<Vec<Volume>, VolumeError> {
	System::new_all()
		.disks()
		.iter()
		.filter_map(|disk| {
			let mut total_capacity = disk.total_space();
			let mount_point = disk.mount_point().to_str().unwrap_or("/").to_string();
			let available_capacity = disk.available_space();
			let name = disk.name().to_str().unwrap_or("Volume").to_string();
			let is_removable = disk.is_removable();

			let file_system = String::from_utf8(disk.file_system().to_vec())
				.unwrap_or_else(|_| "Err".to_string());

			let disk_type = match disk.type_() {
				sysinfo::DiskType::SSD => DiskType::SSD,
				sysinfo::DiskType::HDD => DiskType::HDD,
				_ => DiskType::Removable,
			};

			if total_capacity < available_capacity && cfg!(target_os = "windows") {
				let mut caption = mount_point.clone();
				caption.pop();
				let wmic_process = Command::new("cmd")
					.args([
						"/C",
						&format!("wmic logical disk where Caption='{caption}' get Size"),
					])
					.output()
					.expect("failed to execute process");
				let wmic_process_output = String::from_utf8(wmic_process.stdout).ok()?;
				let parsed_size =
					wmic_process_output.split("\r\r\n").collect::<Vec<&str>>()[1].to_string();

				if let Ok(n) = parsed_size.trim().parse::<u64>() {
					total_capacity = n;
				}
			}

			(!mount_point.starts_with("/System")).then_some(Ok(Volume {
				name,
				is_root_filesystem: mount_point == "/",
				mount_point,
				total_capacity,
				available_capacity,
				is_removable,
				disk_type: Some(disk_type),
				file_system: Some(file_system),
			}))
		})
		.collect::<Result<Vec<_>, _>>()
}

// #[test]
// fn test_get_volumes() {
//   let volumes = get_volumes()?;
//   dbg!(&volumes);
//   assert!(volumes.len() > 0);
// }

// Adapted from: https://github.com/kimlimjustin/xplorer/blob/f4f3590d06783d64949766cc2975205a3b689a56/src-tauri/src/drives.rs