Skip to main content

leenfetch_core/modules/linux/info/
gpu.rs

1#![allow(
2    clippy::collapsible_if,
3    clippy::explicit_auto_deref,
4    clippy::unused_enumerate_index,
5    clippy::useless_format
6)]
7
8use once_cell::sync::Lazy;
9use std::collections::HashMap;
10use std::fs;
11use std::path::Path;
12use std::process::Command;
13
14pub fn get_gpus() -> Vec<String> {
15    let mut entries = collect_from_sysfs();
16    if entries.is_empty() {
17        entries = collect_from_lspci();
18    }
19
20    if entries.is_empty() {
21        return vec!["Unknown GPU".to_string()];
22    }
23
24    entries
25        .into_iter()
26        .enumerate()
27        .map(|(_, info)| format!("{info}"))
28        .collect()
29}
30
31fn collect_from_sysfs() -> Vec<String> {
32    collect_from_sysfs_root(Path::new("/sys/class/drm"))
33}
34
35fn collect_from_sysfs_root(root: &Path) -> Vec<String> {
36    let Ok(read_dir) = fs::read_dir(root) else {
37        return Vec::new();
38    };
39    let mut out = Vec::new();
40
41    for entry in read_dir.flatten() {
42        let name = entry.file_name();
43        let name = name.to_string_lossy();
44        if !name.starts_with("card") || name.contains('-') {
45            continue;
46        }
47
48        let device_dir = entry.path().join("device");
49        if !device_dir.is_dir() {
50            continue;
51        }
52
53        if let Some(line) = describe_device(&device_dir) {
54            out.push(line);
55        }
56    }
57
58    out
59}
60
61fn describe_device(device_dir: &Path) -> Option<String> {
62    let vendor_hex = read_trimmed(device_dir.join("vendor")).and_then(|s| normalize_hex(&s));
63    let device_hex = read_trimmed(device_dir.join("device")).and_then(|s| normalize_hex(&s));
64    let driver_path = device_dir.join("driver");
65    let driver = fs::read_link(&driver_path)
66        .ok()
67        .and_then(|path| path.file_name().map(|s| s.to_string_lossy().into_owned()))
68        .or_else(|| {
69            read_trimmed(&driver_path).and_then(|raw| {
70                Path::new(&raw)
71                    .file_name()
72                    .map(|s| s.to_string_lossy().into_owned())
73            })
74        });
75
76    let vendor_id = vendor_hex
77        .as_deref()
78        .and_then(|hex| u16::from_str_radix(hex, 16).ok());
79    let device_id = device_hex
80        .as_deref()
81        .and_then(|hex| u16::from_str_radix(hex, 16).ok());
82
83    let db = pci_database();
84    let vendor_name = vendor_id.and_then(|id| {
85        db.as_ref()
86            .and_then(|db| db.vendors.get(&id).cloned())
87            .map(|s| s.trim().to_string())
88    });
89    let device_name: Option<String> =
90        if let (Some(vendor_id), Some(device_id)) = (vendor_id, device_id) {
91            db.as_ref()
92                .and_then(|db| db.devices.get(&(vendor_id, device_id)).cloned())
93                .and_then(|name| {
94                    name.split_once('[')
95                        .and_then(|(_, rest)| rest.split_once(']'))
96                        .map(|(inside, _)| inside.trim().to_owned())
97                })
98        } else {
99            None
100        };
101
102    let mut label = match (vendor_name, device_name) {
103        (Some(vendor), Some(model)) => {
104            format!("{} {}", vendor.replace(" Corporation", ""), model)
105        }
106        (Some(vendor), _) => vendor,
107        (_, Some(model)) => model,
108        _ => {
109            if let Some(driver) = driver.as_deref() {
110                driver.to_string()
111            } else {
112                format!(
113                    "GPU [{}:{}]",
114                    vendor_hex.as_deref().unwrap_or("????"),
115                    device_hex.as_deref().unwrap_or("????")
116                )
117            }
118        }
119    };
120
121    if let Some(role) = classify_gpu(vendor_id, driver.as_deref()) {
122        label.push_str(&format!(" [{}]", role));
123    }
124
125    Some(label)
126}
127
128fn collect_from_lspci() -> Vec<String> {
129    let output = Command::new("lspci")
130        .arg("-mm")
131        .output()
132        .ok()
133        .and_then(|o| String::from_utf8(o.stdout).ok())
134        .unwrap_or_default();
135
136    let mut gpus = Vec::new();
137
138    for line in output.lines() {
139        if !(line.contains("\"VGA") || line.contains("\"3D") || line.contains("\"Display")) {
140            continue;
141        }
142
143        let parts: Vec<&str> = line.split('"').collect();
144        if parts.len() < 6 {
145            continue;
146        }
147
148        let vendor = parts[3].trim();
149        let model = parts[5].trim();
150        if vendor.is_empty() && model.is_empty() {
151            continue;
152        }
153
154        let mut label = format!("{vendor} {model}").trim().to_string();
155        let role = classify_gpu_from_name(&label);
156        if let Some(role) = role {
157            label.push_str(&format!(" [{}]", role));
158        }
159
160        gpus.push(label);
161    }
162
163    gpus
164}
165
166fn read_trimmed<P: AsRef<Path>>(path: P) -> Option<String> {
167    let contents = fs::read_to_string(path).ok()?;
168    let trimmed = contents.trim();
169    if trimmed.is_empty() {
170        None
171    } else {
172        Some(trimmed.to_string())
173    }
174}
175
176fn normalize_hex(value: &str) -> Option<String> {
177    let trimmed = value
178        .trim()
179        .trim_start_matches("0x")
180        .trim_start_matches("0X");
181    if trimmed.is_empty() {
182        return None;
183    }
184    u16::from_str_radix(trimmed, 16)
185        .ok()
186        .map(|v| format!("{:04X}", v))
187}
188
189fn classify_gpu(vendor: Option<u16>, driver: Option<&str>) -> Option<&'static str> {
190    match vendor {
191        Some(0x8086) => Some("Integrated"),
192        Some(0x1002) | Some(0x1022) | Some(0x10DE) => Some("Discrete"),
193        Some(0x1234) | Some(0x1AF4) | Some(0x1B36) => Some("Virtual"),
194        Some(0x1A03) => Some("BMC"),
195        _ => match driver {
196            Some("i915") | Some("xe") => Some("Integrated"),
197            Some("amdgpu") | Some("radeon") | Some("nvidia") => Some("Discrete"),
198            Some("vc4") | Some("v3d") => Some("Integrated"),
199            Some("virtio-pci") | Some("bochs-drm") => Some("Virtual"),
200            _ => None,
201        },
202    }
203}
204
205fn classify_gpu_from_name(name: &str) -> Option<&'static str> {
206    let lower = name.to_ascii_lowercase();
207    if lower.contains("intel") {
208        Some("Integrated")
209    } else if lower.contains("nvidia") || lower.contains("geforce") || lower.contains("radeon") {
210        Some("Discrete")
211    } else if lower.contains("virtio") || lower.contains("qxl") || lower.contains("vmware") {
212        Some("Virtual")
213    } else {
214        None
215    }
216}
217
218struct PciDatabase {
219    vendors: HashMap<u16, String>,
220    devices: HashMap<(u16, u16), String>,
221}
222
223static PCI_DB: Lazy<Option<PciDatabase>> = Lazy::new(load_pci_database);
224
225fn pci_database() -> &'static Option<PciDatabase> {
226    &*PCI_DB
227}
228
229fn load_pci_database() -> Option<PciDatabase> {
230    if let Ok(custom) = std::env::var("LEENFETCH_PCI_IDS") {
231        if let Ok(contents) = fs::read_to_string(&custom) {
232            return Some(parse_pci_ids(&contents));
233        }
234    }
235
236    for candidate in ["/usr/share/hwdata/pci.ids", "/usr/share/misc/pci.ids"] {
237        if let Ok(contents) = fs::read_to_string(candidate) {
238            return Some(parse_pci_ids(&contents));
239        }
240    }
241
242    None
243}
244
245fn parse_pci_ids(contents: &str) -> PciDatabase {
246    let mut vendors = HashMap::new();
247    let mut devices = HashMap::new();
248    let mut current_vendor = None;
249
250    for line in contents.lines() {
251        let line = line.trim_end();
252        if line.is_empty() || line.starts_with('#') {
253            continue;
254        }
255
256        if let Some(rest) = line.strip_prefix('\t') {
257            if rest.starts_with('\t') {
258                continue;
259            }
260            if let (Some(vendor_id), Some(device_id)) = (current_vendor, parse_id(rest)) {
261                devices.insert((vendor_id, device_id), rest[4..].trim().to_string());
262            }
263        } else if let Some(vendor_id) = parse_id(line) {
264            vendors.insert(vendor_id, line[4..].trim().to_string());
265            current_vendor = Some(vendor_id);
266        }
267    }
268
269    PciDatabase { vendors, devices }
270}
271
272fn parse_id(line: &str) -> Option<u16> {
273    if line.len() < 4 {
274        return None;
275    }
276    u16::from_str_radix(&line[0..4], 16).ok()
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use crate::test_utils::EnvLock;
283    use std::fs;
284    use std::time::{SystemTime, UNIX_EPOCH};
285
286    #[test]
287    fn test_normalize_hex() {
288        assert_eq!(normalize_hex("0x10de"), Some("10DE".into()));
289        assert_eq!(normalize_hex("10DE"), Some("10DE".into()));
290        assert_eq!(normalize_hex(""), None);
291    }
292
293    #[test]
294    fn test_classify_gpu_by_vendor_and_driver() {
295        assert_eq!(classify_gpu(Some(0x8086), None), Some("Integrated"));
296        assert_eq!(classify_gpu(Some(0x10DE), None), Some("Discrete"));
297        assert_eq!(classify_gpu(Some(0x1234), None), Some("Virtual"));
298        assert_eq!(classify_gpu(None, Some("virtio-pci")), Some("Virtual"));
299        assert_eq!(classify_gpu(None, Some("unknown")), None);
300    }
301
302    #[test]
303    fn test_classify_from_name() {
304        assert_eq!(
305            classify_gpu_from_name("Intel UHD Graphics"),
306            Some("Integrated")
307        );
308        assert_eq!(
309            classify_gpu_from_name("NVIDIA GeForce RTX 3050 Mobile"),
310            Some("Discrete")
311        );
312        assert_eq!(classify_gpu_from_name("QXL GPU"), Some("Virtual"));
313        assert_eq!(classify_gpu_from_name("Mystery Adapter"), None);
314    }
315
316    #[test]
317    fn describe_device_falls_back_to_hex_ids() {
318        let unique = SystemTime::now()
319            .duration_since(UNIX_EPOCH)
320            .unwrap()
321            .as_nanos();
322        let temp = std::env::temp_dir().join(format!("leenfetch_gpu_unknown_{unique}"));
323        let device_dir = temp.join("card1/device");
324        fs::create_dir_all(&device_dir).unwrap();
325        fs::write(device_dir.join("vendor"), "0xFFFF\n").unwrap();
326        fs::write(device_dir.join("device"), "0xEEEE\n").unwrap();
327        fs::write(device_dir.join("driver"), "virtio-pci\n").unwrap();
328
329        let line = super::describe_device(&device_dir).expect("device string");
330        assert!(
331            line.contains("FFFF") || line.contains("Illegal Vendor ID"),
332            "unexpected output: {line}"
333        );
334        assert!(
335            line.contains("[Virtual]"),
336            "expected virtual classification tag: {line}"
337        );
338
339        fs::remove_dir_all(&temp).unwrap();
340    }
341
342    #[test]
343    fn test_collect_from_sysfs_formatting() {
344        let unique = SystemTime::now()
345            .duration_since(UNIX_EPOCH)
346            .unwrap()
347            .as_nanos();
348        let temp = std::env::temp_dir().join(format!("leenfetch_gpu_test_{unique}"));
349        fs::create_dir_all(temp.join("card0/device")).unwrap();
350        fs::write(temp.join("card0/device/vendor"), "0x8086\n").unwrap();
351        fs::write(temp.join("card0/device/device"), "0x9A60\n").unwrap();
352        fs::write(temp.join("card0/device/driver"), "i915\n").unwrap();
353
354        let database = "\
3558086  Intel Corporation
356\t9A60  Alder Lake-P GT1 [UHD Graphics]
357";
358        let db_path = temp.join("pci.ids");
359        fs::write(&db_path, database).unwrap();
360        let env_lock = EnvLock::acquire(&["LEENFETCH_PCI_IDS"]);
361        env_lock.set_var("LEENFETCH_PCI_IDS", db_path.to_str().unwrap());
362
363        let result = super::collect_from_sysfs_root(temp.as_path());
364        assert_eq!(result, vec!["Intel UHD Graphics [Integrated]"]);
365
366        env_lock.remove_var("LEENFETCH_PCI_IDS");
367        drop(env_lock);
368        fs::remove_dir_all(temp).unwrap();
369    }
370}