leenfetch_core/modules/linux/info/
gpu.rs

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

leenfetch_core/modules/linux/info/
gpu.rs
