Skip to content

Commit

Permalink
fix: faster/more resiliant search for pcie/usb w/ u16 ids
Browse files Browse the repository at this point in the history
  • Loading branch information
zleyyij committed Feb 17, 2024
1 parent c04856b commit c07997d
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 62 deletions.
3 changes: 3 additions & 0 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
port = value;
}

let cache = UsbCache::new();
println!("{:#?}", cache.find("USB\\VID_1532&PID_008A&MI_01\\7&238AA5C2&1&0001"));

info!("Listening on port {}", port);
// run the app
let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{}", port))
Expand Down
79 changes: 29 additions & 50 deletions src/pcie/mod.rs
View file
Original file line number Diff line number Diff line change
@@ -1,16 +1,16 @@
use crate::NomError;
use nom::bytes::complete::{tag, take, take_until};
use nom::character::complete::char;
use nom::sequence::{delimited, preceded, terminated};
use nom::IResult;
use crate::NomError;

// the input file was obtained from https://pci-ids.ucw.cz/
const FILE_INPUT: &str = include_str!("./pci.ids.txt");

/// Vendors are at the root of the file
#[derive(PartialEq, Debug, Clone)]
pub struct Vendor {
pub id: String,
pub id: u16,
pub name: String,
pub devices: Vec<Device>,
}
Expand All @@ -19,7 +19,7 @@ pub struct Vendor {
/// and are marked with one tab before, the device ID, then two spaces and the device name
#[derive(PartialEq, Debug, Clone)]
pub struct Device {
pub id: String,
pub id: u16,
pub name: String,
pub subsystems: Vec<Subsystem>,
}
Expand All @@ -29,7 +29,7 @@ pub struct Device {
/// then two spaces, then the name of the subsystem
#[derive(PartialEq, Debug, Clone)]
pub struct Subsystem {
pub id: String,
pub id: u16,
pub name: String,
}

Expand Down Expand Up @@ -92,7 +92,7 @@ impl PcieCache {
/// Output is returned as a tuple of (`vendor`, `device`, `subsystem`)
fn parse_device_identifier<'a>(
input: &'a str,
) -> Result<(&'a str, &'a str, Option<&'a str>), NomError<'a>> {
) -> Result<(u16, u16, Option<u16>), NomError<'a>> {
// TODO: validate that ids are hex strings
let vid_combinator = delimited(tag("PCI\\VEN_"), take(4 as u8), char('&'))(input)?;
// https://learn.microsoft.com/en-us/windows-hardware/drivers/install/identifiers-for-pci-devices
Expand All @@ -103,7 +103,7 @@ fn parse_device_identifier<'a>(
ssid = Some(ssid_combinator.1);
}

Ok((vid_combinator.1, did_combinator.1, ssid))
Ok((u16::from_str_radix(vid_combinator.1, 16).unwrap(), u16::from_str_radix(did_combinator.1, 16).unwrap(), ssid.map(|s| u16::from_str_radix(s, 16).unwrap())))
}

/// Read the database from the file into memory
Expand Down Expand Up @@ -158,7 +158,7 @@ fn read_vendor(input: &str) -> IResult<&str, Vendor> {
Ok((
leftover,
Vendor {
id: String::from(vid),
id: u16::from_str_radix(vid, 16).unwrap(),
name: String::from(vname),
devices,
},
Expand Down Expand Up @@ -195,7 +195,7 @@ fn read_device(input: &str) -> IResult<&str, Device> {
Ok((
leftover,
Device {
id: String::from(did),
id: u16::from_str_radix(did, 16).unwrap(),
name: String::from(dname_combinator.1),
subsystems,
},
Expand All @@ -216,18 +216,15 @@ fn read_subsystem_line(input: &str) -> IResult<&str, Subsystem> {
Ok((
ss_name_combinator.0,
Subsystem {
id: String::from(ssid_combinator.1),
id: u16::from_str_radix(ssid_combinator.1, 16).unwrap(),
name: String::from(ss_name_combinator.1),
},
))
}

#[cfg(test)]
mod tests {
use crate::pcie::{
parse_device_identifier, read_device, read_subsystem_line, read_vendor, Device, Subsystem,
Vendor,
};
use crate::pcie::{parse_device_identifier, read_device, read_subsystem_line, read_vendor, Device, Subsystem, Vendor};

use super::{parse_pcie_db, read_header, PcieCache};

Expand All @@ -248,7 +245,7 @@ mod tests {
Ok((
"bat",
Subsystem {
id: String::from("0001"),
id: 0x0001,
name: String::from("foo bar")
}
))
Expand All @@ -263,7 +260,7 @@ mod tests {
Ok((
"\t0002",
Device {
id: String::from("0001"),
id: 0x0001,
name: String::from("foo bar"),
subsystems: vec![]
}
Expand All @@ -276,10 +273,10 @@ mod tests {
Ok((
"0002",
Device {
id: String::from("0001"),
id: 0x0001,
name: String::from("foo bar"),
subsystems: vec![Subsystem {
id: String::from("8008"),
id: 0x8008,
name: String::from("subsys")
}],
}
Expand All @@ -292,10 +289,10 @@ mod tests {
Ok((
"0002",
Device {
id: String::from("0001"),
id: 0x0001,
name: String::from("foo bar"),
subsystems: vec![Subsystem {
id: String::from("8008"),
id: 0x8008,
name: String::from("subsys")
}],
}
Expand All @@ -311,10 +308,10 @@ mod tests {
Ok((
"0002",
Vendor {
id: String::from("0001"),
id: 0x0001,
name: String::from("foo"),
devices: vec![Device {
id: String::from("000a"),
id: 0x000a,
name: String::from("bar"),
subsystems: vec![]
}]
Expand All @@ -332,35 +329,17 @@ mod tests {
#[test]
fn basic_parse_device_identifier() {
// https://learn.microsoft.com/en-us/windows-hardware/drivers/install/identifiers-for-pci-devices
assert_eq!(
parse_device_identifier("PCI\\VEN_1234&DEV_5678&SUBSYS_91230000&REV_00"),
Ok(("1234", "5678", Some("9123")))
);
assert_eq!(
parse_device_identifier("PCI\\VEN_1234&DEV_5678&SUBSYS_91230000"),
Ok(("1234", "5678", Some("9123")))
);
assert_eq!(
parse_device_identifier("PCI\\VEN_1234&DEV_5678&REV_00"),
Ok(("1234", "5678", None))
);
assert_eq!(
parse_device_identifier("PCI\\VEN_1234&DEV_5678"),
Ok(("1234", "5678", None))
);
assert_eq!(
parse_device_identifier("PCI\\VEN_1234&DEV_5678&CC_112200"),
Ok(("1234", "5678", None))
);
assert_eq!(
parse_device_identifier("PCI\\VEN_1234&DEV_5678&CC_1122"),
Ok(("1234", "5678", None))
);
assert_eq!(parse_device_identifier("PCI\\VEN_1234&DEV_5678&SUBSYS_91230000&REV_00"), Ok((0x1234, 0x5678, Some(0x9123))));
assert_eq!(parse_device_identifier("PCI\\VEN_1234&DEV_5678&SUBSYS_91230000"), Ok((0x1234, 0x5678, Some(0x9123))));
assert_eq!(parse_device_identifier("PCI\\VEN_1234&DEV_5678&REV_00"), Ok((0x1234, 0x5678, None)));
assert_eq!(parse_device_identifier("PCI\\VEN_1234&DEV_5678"), Ok((0x1234, 0x5678, None)));
assert_eq!(parse_device_identifier("PCI\\VEN_1234&DEV_5678&CC_112200"), Ok((0x1234, 0x5678, None)));
assert_eq!(parse_device_identifier("PCI\\VEN_1234&DEV_5678&CC_1122"), Ok((0x1234, 0x5678, None)));
}

// #[test]
// fn basic_find_device() {
// let cache = PcieCache::new();
// // println!("{:#?}", cache.find("PCI\\VEN_1022&DEV_1633&SUBSYS_14531022&REV_00\\3&2411E6FE&1&09").map(|t| t.1));
// }
#[test]
fn basic_find_device() {
let cache = PcieCache::new();
println!("{:#?}", cache.find("PCI\\VEN_1022&DEV_1633&SUBSYS_14531022&REV_00\\3&2411E6FE&1&09").map(|t| t.1));
}
}
24 changes: 12 additions & 12 deletions src/usb/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,14 @@ const INPUT_FILE: &[u8] = include_bytes!("usb.ids.txt");

#[derive(Clone, Debug, PartialEq)]
pub struct Vendor {
pub id: String,
pub id: u16,
pub name: String,
pub devices: Vec<Device>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Device {
pub id: String,
pub id: u16,
pub name: String,
}

Expand Down Expand Up @@ -70,13 +70,12 @@ impl UsbCache {
/// Input strings in the form of `USB\VID_1234&PID_5678\9479493` are assumed.
/// It returns a tuple, where the first value is the vendor id, and the second is the product id. This tuple contains substrings of the initial input string,
/// so handle lifetimes accordingly.
fn parse_device_identifier(device_string: &str) -> Result<(&str, &str), NomError> {
fn parse_device_identifier(device_string: &str) -> Result<(u16, u16), NomError> {
// https://learn.microsoft.com/en-us/windows-hardware/drivers/install/standard-usb-identifiers
// TODO: this does not fully support all formats of usb device identifiers
let vid_combinator = delimited(tag("USB\\VID_"), take(4 as u8), take(1 as u8))(device_string)?;
let pid_combinator = preceded(tag("PID_"), take(4 as u8))(vid_combinator.0)?;
// TODO: assert that the found values were actually valid hexadecimal strings
Ok((vid_combinator.1, pid_combinator.1))
Ok((u16::from_str_radix(vid_combinator.1, 16).unwrap(), u16::from_str_radix(pid_combinator.1, 16).unwrap()))
}

fn parse_usb_db() -> Vec<Vendor> {
Expand Down Expand Up @@ -137,7 +136,7 @@ fn read_vendor(input: &str) -> IResult<&str, Vendor> {
Ok((
leftover,
Vendor {
id: vid.to_string(),
id: u16::from_str_radix(vid, 16).unwrap(),
name: vname.to_string(),
devices,
},
Expand All @@ -153,7 +152,7 @@ fn read_device_line(input: &str) -> IResult<&str, Device> {
Ok((
combinator_output.0,
Device {
id: String::from(did_combinator_output.1),
id: u16::from_str_radix(did_combinator_output.1, 16).unwrap(),
name: String::from(dname),
},
))
Expand All @@ -170,7 +169,7 @@ mod tests {
let mock_device_string = "USB\\VID_1234&PID_5678\\9479493";
assert_eq!(
parse_device_identifier(mock_device_string),
Ok(("1234", "5678"))
Ok((0x1234, 0x5678))
);
}

Expand All @@ -187,10 +186,10 @@ mod tests {
fn basic_read_vendor() {
let mock_section = "1234 vendor_name\n\t5678 device_name\n9123";
let expected_output = Vendor {
id: String::from("1234"),
id: 0x1234,
name: String::from("vendor_name"),
devices: vec![Device {
id: String::from("5678"),
id: 0x5678,
name: String::from("device_name"),
}],
};
Expand All @@ -201,7 +200,7 @@ mod tests {
fn read_section_no_devices() {
let mock_section = "1234 vendor_name\n5678";
let expected_output = Vendor {
id: String::from("1234"),
id: 0x1234,
name: String::from("vendor_name"),
devices: vec![],
};
Expand All @@ -213,7 +212,7 @@ mod tests {
Ok((
"4567",
Device {
id: String::from("1234"),
id: 0x1234,
name: String::from("foo bar")
}
))
Expand All @@ -224,4 +223,5 @@ mod tests {
fn basic_parse_usbs() {
parse_usb_db();
}

}

0 comments on commit c07997d

Please sign in to comment.