Skip to content

Commit

Permalink
Avoid potential panics due to indexing and blind unwraps
Browse files Browse the repository at this point in the history 
Restructure code and add docstrings
  • Loading branch information
@dormant-user
dormant-user committed Sep 28, 2024
1 parent 7a53b07 commit 1b8af06
Show file tree
Hide file tree
Showing 3 changed files with 121 additions and 39 deletions.
52 changes: 44 additions & 8 deletions src/legacy/disks.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,10 +121,16 @@ fn darwin_disks(lib_path: &str) -> Vec<HashMap<String, String>> {
}
};
let disks: Vec<&str> = output.lines().collect();
let disk_lines: Vec<&str> = disks.into_iter().filter(|&line| line.starts_with("/dev/disk")).collect();
let disk_lines: Vec<&str> = disks
.into_iter()
.filter(|&line| line.starts_with("/dev/disk"))
.collect();
let mut disk_info = Vec::new();
for line in disk_lines {
let device_id = line.split_whitespace().next().unwrap();
let device_id = line
.split_whitespace()
.next()
.unwrap_or_default();
if !is_physical_disk(lib_path, device_id) {
continue;
}
Expand All @@ -136,18 +142,41 @@ fn darwin_disks(lib_path: &str) -> Vec<HashMap<String, String>> {
return Vec::new();
}
};
let info_lines: Vec<&str> = disk_info_output.lines().collect();
let info_lines: Vec<&str> = disk_info_output
.lines()
.collect();
let mut disk_data = HashMap::new();
for info_line in info_lines {
if info_line.contains("Device / Media Name:") {
disk_data.insert("Name".to_string(), info_line.split(":").nth(1).unwrap().trim().to_string());
disk_data.insert(
"Name".to_string(),
info_line
.split(":")
.nth(1)
.unwrap_or_default()
.trim().to_string()
);
}
if info_line.contains("Disk Size:") {
let size_info = info_line.split(":").nth(1).unwrap().split("(").next().unwrap().trim().to_string();
disk_data.insert("Size".to_string(), size_info);
let size_info = info_line
.split(":")
.nth(1)
.unwrap_or_default()
.split("(")
.next()
.unwrap_or_default()
.trim()
.to_string();
disk_data.insert(
"Size".to_string(),
size_info
);
}
}
disk_data.insert("DeviceID".to_string(), device_id.to_string());
disk_data.insert(
"DeviceID".to_string(),
device_id.to_string()
);
disk_info.push(disk_data);
}
disk_info
Expand All @@ -168,7 +197,14 @@ fn reformat_windows(data: &mut HashMap<String, Value>) -> HashMap<String, String
let mut reformatted_data = HashMap::new();
reformatted_data.insert("Size".to_string(), squire::util::size_converter(size as u64));
reformatted_data.insert("Name".to_string(), model);
reformatted_data.insert("DeviceID".to_string(), data.get("DeviceID").unwrap().as_str().unwrap().to_string());
reformatted_data.insert(
"DeviceID".to_string(),
data.get("DeviceID")
.unwrap()
.as_str()
.unwrap()
.to_string()
);
reformatted_data
}

Expand Down
106 changes: 76 additions & 30 deletions src/legacy/gpu.rs
View file
Original file line number Diff line number Diff line change
@@ -1,34 +1,46 @@
use std::collections::HashMap;
use serde_json::Value;
use crate::squire;
use serde_json::Value;
use std::collections::HashMap;

pub fn get_gpu_info() -> Vec<HashMap<String, String>> {
let operating_system = std::env::consts::OS;
match operating_system {
"macos" => get_gpu_info_darwin("/usr/sbin/system_profiler"),
"linux" => get_gpu_info_linux("/usr/bin/lspci"),
"windows" => get_gpu_info_windows("C:\\Windows\\System32\\wbem\\wmic.exe"),
_ => {
log::error!("Unsupported operating system: {}", operating_system);
Vec::new()
}
}
}

/// Function to get GPU information for macOS machines.
///
/// # Arguments
///
/// * `lib_path` - The path to the library used to get GPU information.
///
/// # Returns
///
/// A `Vec` of `HashMap` containing GPU(s) information if successful, otherwise an empty `Vec`.
fn get_gpu_info_darwin(lib_path: &str) -> Vec<HashMap<String, String>> {
let result = squire::util::run_command(
let result: Result<String, String> = squire::util::run_command(
lib_path,
&["SPDisplaysDataType", "-json"],
true
true,
);
if result.is_err() {
return Vec::new();
let displays: Vec<Value>;
match result {
Ok(json_str) => {
match serde_json::from_str::<Value>(&json_str) {
Ok(json) => {
if let Some(d) = json.get("SPDisplaysDataType").and_then(Value::as_array) {
displays = d.to_vec();
} else {
log::error!("Key 'SPDisplaysDataType' not found or is not an array.");
return Vec::new();
}
}
Err(err) => {
log::error!("Failed to parse JSON: {}", err);
return Vec::new();
}
}
}
Err(err) => {
log::error!("Error retrieving result: {}", err);
return Vec::new();
}
}
let json: Value = serde_json::from_str(&result.unwrap()).unwrap();

let displays = json["SPDisplaysDataType"].as_array().unwrap();
let mut gpu_info = Vec::new();

let na = "N/A".to_string();
for display in displays {
if let Some(model) = display.get("sppci_model") {
Expand All @@ -37,7 +49,7 @@ fn get_gpu_info_darwin(lib_path: &str) -> Vec<HashMap<String, String>> {
"model".to_string(),
model.as_str()
.unwrap_or(na.as_str())
.to_string()
.to_string(),
);

// Handle cores
Expand Down Expand Up @@ -68,18 +80,26 @@ fn get_gpu_info_darwin(lib_path: &str) -> Vec<HashMap<String, String>> {
.unwrap_or(&na)
.to_string(),
);

gpu_info.push(info);
}
}
gpu_info
}

/// Function to get GPU information for Linux machines.
///
/// # Arguments
///
/// * `lib_path` - The path to the library used to get GPU information.
///
/// # Returns
///
/// A `Vec` of `HashMap` containing GPU(s) information if successful, otherwise an empty `Vec`.
fn get_gpu_info_linux(lib_path: &str) -> Vec<HashMap<String, String>> {
let result = squire::util::run_command(
lib_path,
&[],
true
true,
);
if result.is_err() {
return Vec::new();
Expand All @@ -96,19 +116,28 @@ fn get_gpu_info_linux(lib_path: &str) -> Vec<HashMap<String, String>> {
gpu_info
}

/// Function to get GPU information for Windows machines.
///
/// # Arguments
///
/// * `lib_path` - The path to the library used to get GPU information.
///
/// # Returns
///
/// A `Vec` of `HashMap` containing GPU(s) information if successful, otherwise an empty `Vec`.
fn get_gpu_info_windows(lib_path: &str) -> Vec<HashMap<String, String>> {
let result = squire::util::run_command(
lib_path,
&["path", "win32_videocontroller", "get", "Name,AdapterCompatibility", "/format:csv"],
true
true,
);
let output = match result {
Ok(output) => output.to_uppercase(),
Err(_) => {
return Vec::new();
}
};
let gpus_raw: Vec<&str> = output.lines().filter(|line| !line.trim().is_empty()).collect();
let gpus_raw: Vec<&str> = output.lines().filter(|line| !line.trim().is_empty()).collect();
if gpus_raw.is_empty() {
log::info!("No GPUs found!");
return vec![];
Expand All @@ -117,7 +146,6 @@ fn get_gpu_info_windows(lib_path: &str) -> Vec<HashMap<String, String>> {
let keys: Vec<String> = gpus_raw[0]
.trim()
.to_lowercase()
.replace("node", "node")
.replace("adaptercompatibility", "vendor")
.replace("name", "model")
.split(',')
Expand All @@ -129,11 +157,29 @@ fn get_gpu_info_windows(lib_path: &str) -> Vec<HashMap<String, String>> {
let mut gpu_info = Vec::new();
let key_len = keys.len();
for chunk in values.chunks(key_len) {
let mut map = std::collections::HashMap::new();
let mut map = HashMap::new();
for (key, value) in keys.iter().zip(chunk) {
map.insert(key.to_string(), value.to_string());
}
gpu_info.push(map);
}
gpu_info
}

/// OS-agnostic function to get GPU name.
///
/// # Returns
///
/// A `Vec` of `HashMap` containing GPU(s) information if successful, otherwise an empty `Vec`.
pub fn get_gpu_info() -> Vec<HashMap<String, String>> {
let operating_system = std::env::consts::OS;
match operating_system {
"macos" => get_gpu_info_darwin("/usr/sbin/system_profiler"),
"linux" => get_gpu_info_linux("/usr/bin/lspci"),
"windows" => get_gpu_info_windows("C:\\Windows\\System32\\wbem\\wmic.exe"),
_ => {
log::error!("Unsupported operating system: {}", operating_system);
Vec::new()
}
}
}
2 changes: 1 addition & 1 deletion src/resources/info.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ pub fn get_disks(disks: &Disks) -> Vec<HashMap<String, String>> {

fn get_gpu_info() -> Vec<String> {
let gpu_info = legacy::gpu::get_gpu_info();
log::info!("GPUs: {:?}", gpu_info);
log::debug!("GPUs: {:?}", gpu_info);
let mut gpus: Vec<String> = vec![];
for gpu in gpu_info {
if let Some(gpu_model) = gpu.get("model") {
Expand Down

0 comments on commit 1b8af06

Please sign in to comment.