yggdrasil/userspace/sysutils/src/lspci.rs

use std::{fmt, fs, io, path::PathBuf, process::ExitCode, str::FromStr};

use clap::Parser;
use pci_ids::FromId;

#[derive(Debug, Parser)]
struct Args {
    #[clap(short, help = "Display numeric IDs only")]
    numeric: bool,
    #[clap(short, help = "Print all the information about the devices")]
    verbose: bool,
}

struct PciAddress {
    bus: u8,
    device: u8,
    function: u8,
}

struct DeviceId {
    vendor: u16,
    device: u16,
}

struct Device {
    path: PathBuf,
    address: PciAddress,
}

enum Resource {
    Io(u16),
    Memory32(u32),
    Memory64(u64),
}

struct PciRegion {
    index: usize,
    resource: Resource
}

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("{0}")]
    Io(#[from] io::Error),
    #[error("Invalid PCI ID")]
    InvalidId,
}

impl FromStr for PciAddress {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut elements = s.split(':');
        let bus = elements
            .next()
            .and_then(|p| u8::from_str_radix(p, 16).ok())
            .ok_or(Error::InvalidId)?;
        let device = elements
            .next()
            .and_then(|p| u8::from_str_radix(p, 16).ok())
            .ok_or(Error::InvalidId)?;
        let function = elements
            .next()
            .and_then(|p| u8::from_str_radix(p, 16).ok())
            .ok_or(Error::InvalidId)?;
        if elements.next().is_some() {
            return Err(Error::InvalidId);
        }
        Ok(Self {
            bus,
            device,
            function,
        })
    }
}

impl FromStr for DeviceId {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let (vendor, device) = s.split_once(':').ok_or(Error::InvalidId)?;
        let vendor = u16::from_str_radix(vendor, 16).map_err(|_| Error::InvalidId)?;
        let device = u16::from_str_radix(device, 16).map_err(|_| Error::InvalidId)?;
        Ok(Self { vendor, device })
    }
}

impl FromStr for Resource {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let (ty, addr) = s.split_once(':').ok_or(Error::InvalidId)?;
        let addr = addr.strip_prefix("0x").ok_or(Error::InvalidId)?;
        match ty {
            "pio" => {
                let addr = u16::from_str_radix(addr, 16).map_err(|_| Error::InvalidId)?;
                Ok(Self::Io(addr))
            },
            "m32" => {
                let addr = u32::from_str_radix(addr, 16).map_err(|_| Error::InvalidId)?;
                Ok(Self::Memory32(addr))
            },
            "m64" => {
                let addr = u64::from_str_radix(addr, 16).map_err(|_| Error::InvalidId)?;
                Ok(Self::Memory64(addr))
            },
            _ => Err(Error::InvalidId)
        }
    }
}

impl FromStr for PciRegion {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let (bar, resource) = s.split_once(':').ok_or(Error::InvalidId)?;
        let index = usize::from_str(bar).map_err(|_| Error::InvalidId)?;
        let resource = Resource::from_str(resource)?;
        Ok(Self {
            index, resource
        })
    }
}

impl fmt::Display for PciAddress {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{:02x}:{:02x}.{:x}",
            self.bus, self.device, self.function
        )
    }
}

impl fmt::Display for DeviceId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:04x}:{:04x}", self.vendor, self.device)
    }
}

impl fmt::Display for Resource {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Self::Io(addr) => write!(f, "I/O at {addr:#x}"),
            Self::Memory32(addr) => write!(f, "Memory (32-bit) at {addr:#x}"),
            Self::Memory64(addr) => write!(f, "Memory (64-bit) at {addr:#x}")
        }
    }
}

fn list_devices() -> Result<Vec<Device>, Error> {
    let mut entries = Vec::new();
    for entry in fs::read_dir("/sys/bus/pci")? {
        let Ok(entry) = entry else { continue };
        let file_name = entry.file_name();
        let Some(name) = file_name.to_str() else {
            continue;
        };

        if let Ok(address) = PciAddress::from_str(name) {
            entries.push(Device {
                path: entry.path(),
                address,
            });
        }
    }

    Ok(entries)
}

fn read_device_id(device: &Device) -> Result<DeviceId, Error> {
    let str = fs::read_to_string(device.path.join("id"))?;
    DeviceId::from_str(str.trim())
}

fn read_device_regions(device: &Device) -> Result<Vec<PciRegion>, Error> {
    let str = fs::read_to_string(device.path.join("resources"))?;
    let mut regions = Vec::new();
    for line in str.split('\n') {
        let line = line.trim();
        let Ok(region) = PciRegion::from_str(line) else { continue };
        regions.push(region);
    }
    Ok(regions)
}

fn vendor_id(args: &Args, id: u16) -> String {
    if args.numeric {
        format!("{id:04x}")
    } else {
        pci_ids::Vendor::from_id(id).map(|v| v.name().into()).unwrap_or_else(|| format!("{id:04x}"))
    }
}

fn device_id(args: &Args, vid: u16, pid: u16) -> String {
    if args.numeric {
        format!("{pid:04x}")
    } else {
        pci_ids::Device::from_vid_pid(vid, pid).map(|v| v.name().into()).unwrap_or_else(|| format!("{pid:04x}"))
    }
}

fn print_device(args: &Args, device: &Device) -> Result<(), Error> {
    let id = read_device_id(device)?;

    print!("{}: ", device.address);

    let vid = vendor_id(args, id.vendor);
    let pid = device_id(args, id.vendor, id.device);
    let device_name = if args.numeric {
        format!("{vid}:{pid}")
    }else {
        format!("[{vid}] {pid}")
    };

    println!("{device_name}");

    if args.verbose {
        for region in read_device_regions(device)? {
            println!("  BAR {}: {}", region.index, region.resource);
        }

    }

    Ok(())
}

fn run(args: &Args) -> Result<(), Error> {
    let devices = list_devices()?;
    for device in devices {
        if let Err(error) = print_device(args, &device) {
            eprintln!("{}: {error}", device.address);
        }
    }
    Ok(())
}

fn main() -> ExitCode {
    let args = Args::parse();

    match run(&args) {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("{error}");
            ExitCode::FAILURE
        }
    }
}