1 // Copyright 2022 The ChromiumOS Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 /*
6 The struct must be named in non_camel and non_snake because we want to query the windows wmi
7 interface and conform to the windows naming convension.
8 */
9 #![allow(non_camel_case_types)]
10 #![allow(non_snake_case)]
11
12 use std::collections::HashMap;
13 use std::error::Error;
14 use std::rc::Rc;
15
16 use base::warn;
17 use serde::de::DeserializeOwned;
18 use serde::Deserialize;
19 use wmi::query::FilterValue;
20 use wmi::COMLibrary;
21 use wmi::WMIConnection;
22
23 const VIDEO_CONTROLLER_AVAILABILITY_ENABLED: i64 = 3;
24
25 #[derive(Deserialize, Debug)]
26 pub struct Win32_Processor {
27 pub Manufacturer: String,
28 pub Name: String,
29 pub NumberOfCores: u32,
30 pub NumberOfLogicalProcessors: u32,
31 pub ThreadCount: u32,
32 }
33
34 #[derive(Deserialize, Debug)]
35 pub struct Win32_VideoController {
36 pub Name: String,
37 // TODO(b/191406729): re-enable.
38 // pub AdapterRAM: u64,
39 pub DriverVersion: String,
40 pub Availability: u16,
41 pub Description: String,
42 }
43
44 #[derive(Deserialize, Debug)]
45 struct MSFT_Partition {
46 DriveLetter: String,
47 }
48
49 #[derive(Deserialize, Debug)]
50 struct MSFT_Disk {
51 __Path: String,
52 UniqueId: String,
53 }
54
55 #[derive(Deserialize, Debug)]
56 struct MSFT_DiskToPartition {}
57
58 #[derive(Deserialize, Debug)]
59 struct MSFT_PhysicalDisk {
60 FriendlyName: String,
61 MediaType: u16,
62 BusType: u16,
63 Size: u64,
64 UniqueId: String,
65 }
66
67 // Keep the formatting so that the debug output string matches the proto field
68 // values.
69 #[derive(Debug)]
70 pub enum MediaType {
71 UNKNOWN,
72 HDD,
73 SSD,
74 SCM,
75 }
76
77 impl From<u16> for MediaType {
from(value: u16) -> Self78 fn from(value: u16) -> Self {
79 match value {
80 3 => MediaType::HDD,
81 4 => MediaType::SSD,
82 5 => MediaType::SCM,
83 _ => MediaType::UNKNOWN,
84 }
85 }
86 }
87
88 // Keep the formatting so that the debug output string matches the proto field
89 // values.
90 #[derive(Debug)]
91 pub enum BusType {
92 UNKNOWN,
93 SCSI,
94 ATAPI,
95 ATA,
96 TYPE_1394,
97 SSA,
98 FIBRE_CHANNEL,
99 USB,
100 RAID,
101 ISCSI,
102 SAS,
103 SATA,
104 SD,
105 MMC,
106 FILE_BACKED_VIRTUAL,
107 STORAGE_SPACES,
108 NVME,
109 }
110
111 impl From<u16> for BusType {
from(value: u16) -> Self112 fn from(value: u16) -> Self {
113 match value {
114 1 => BusType::SCSI,
115 2 => BusType::ATAPI,
116 3 => BusType::ATA,
117 4 => BusType::TYPE_1394,
118 5 => BusType::SSA,
119 6 => BusType::FIBRE_CHANNEL,
120 7 => BusType::USB,
121 8 => BusType::RAID,
122 9 => BusType::ISCSI,
123 10 => BusType::SAS,
124 11 => BusType::SATA,
125 12 => BusType::SD,
126 13 => BusType::MMC,
127 15 => BusType::FILE_BACKED_VIRTUAL,
128 16 => BusType::STORAGE_SPACES,
129 17 => BusType::NVME,
130 _ => BusType::UNKNOWN,
131 }
132 }
133 }
134
135 // Friendly format for MSFT_PhysicalDisk.
136 // Also includes the cross-referenced partitions within the disk.
137 #[derive(Debug)]
138 pub struct PhysicalDisk {
139 pub Name: String,
140 pub MediaType: MediaType,
141 pub BusType: BusType,
142 pub Size: u64,
143 pub DriveLetters: Vec<String>,
144 }
145
146 #[derive(Deserialize, Debug)]
147 pub struct Win32_PhysicalMemory {
148 pub Capacity: u64,
149 pub ConfiguredClockSpeed: u32,
150 pub PartNumber: String,
151 }
152
153 #[derive(Clone, Deserialize, Debug)]
154 pub struct Win32_BaseBoard {
155 pub Manufacturer: String,
156 pub Product: String,
157 }
158
159 #[derive(Debug)]
160 pub struct WmiMetrics {
161 pub cpus: Vec<Win32_Processor>,
162 pub gpus: Vec<Win32_VideoController>,
163 pub disks: Vec<PhysicalDisk>,
164 pub mems: Vec<Win32_PhysicalMemory>,
165 pub motherboard: Option<Win32_BaseBoard>,
166 }
167
get_wmi_metrics() -> Result<WmiMetrics, Box<dyn Error>>168 pub fn get_wmi_metrics() -> Result<WmiMetrics, Box<dyn Error>> {
169 let com_con = Rc::new(COMLibrary::new()?);
170 let wmi_con = WMIConnection::new(Rc::clone(&com_con))?;
171
172 // Fetch WMI data, including all entries.
173 let cpus: Vec<Win32_Processor> = run_wmi_query(&wmi_con);
174 let disks = get_disks(Rc::clone(&com_con))?;
175 let mems: Vec<Win32_PhysicalMemory> = run_wmi_query(&wmi_con);
176 let motherboard: Option<Win32_BaseBoard> = run_wmi_query(&wmi_con).into_iter().next();
177 let gpus = get_gpus(&wmi_con);
178
179 let wmi_metrics = WmiMetrics {
180 cpus,
181 gpus,
182 disks,
183 mems,
184 motherboard,
185 };
186
187 Ok(wmi_metrics)
188 }
189
get_disks(com_con: Rc<COMLibrary>) -> Result<Vec<PhysicalDisk>, Box<dyn Error>>190 fn get_disks(com_con: Rc<COMLibrary>) -> Result<Vec<PhysicalDisk>, Box<dyn Error>> {
191 // For MSFT_PhysicalDisk, we need to connect with storage namespace.
192 let wmi_con = WMIConnection::with_namespace_path("Root\\Microsoft\\Windows\\Storage", com_con)?;
193 // First we get all instances of following classes:
194 // MSFT_Disk, MSFT_PhysicalDisk
195 // We use the WMI associator query to find mapping for each:
196 // MSFT_Disk -> MSFT_Partition (1:N)
197 // Then, we find the mapping from each:
198 // MSFT_Disk -> MSFT_PhysicalDisk (1:1)
199 // Finally, we construct each PhysicalDisk structure by combining the
200 // matched MSFT_PhysicalDisk and MSFT_Parition instances.
201 let msft_disks: Vec<MSFT_Disk> = run_wmi_query(&wmi_con);
202 let physical_disks: Vec<MSFT_PhysicalDisk> = run_wmi_query(&wmi_con);
203
204 let mut disks = Vec::with_capacity(physical_disks.len());
205 for msft_disk in msft_disks {
206 let partitions =
207 wmi_con.associators::<MSFT_Partition, MSFT_DiskToPartition>(&msft_disk.__Path)?;
208 let physical_disk = physical_disks
209 .iter()
210 .find(|d| d.UniqueId == msft_disk.UniqueId)
211 .ok_or("Could not find a matching MSFT_PhysicalDisk!")?;
212 disks.push(PhysicalDisk {
213 Name: physical_disk.FriendlyName.clone(),
214 MediaType: physical_disk.MediaType.into(),
215 BusType: physical_disk.BusType.into(),
216 Size: physical_disk.Size,
217 DriveLetters: partitions.into_iter().map(|p| p.DriveLetter).collect(),
218 });
219 }
220 Ok(disks)
221 }
222
get_gpus(wmi_con: &WMIConnection) -> Vec<Win32_VideoController>223 fn get_gpus(wmi_con: &WMIConnection) -> Vec<Win32_VideoController> {
224 // TODO(b/191406729): Fix the query once the AdapterRAM can be correctly
225 // queried.
226 let mut filters = HashMap::new();
227 filters.insert(
228 "Availability".to_string(),
229 FilterValue::Number(VIDEO_CONTROLLER_AVAILABILITY_ENABLED),
230 );
231 wmi_con
232 .filtered_query(&filters)
233 .map_err(|e| warn!("wmi query failed: {}", e))
234 .unwrap_or_default()
235 }
236
run_wmi_query<T>(wmi_con: &WMIConnection) -> Vec<T> where T: DeserializeOwned,237 fn run_wmi_query<T>(wmi_con: &WMIConnection) -> Vec<T>
238 where
239 T: DeserializeOwned,
240 {
241 wmi_con
242 .query()
243 .map_err(|e| warn!("wmi query failed: {}", e))
244 .unwrap_or_default()
245 }
246