Use PCI bus and device numbers in block test util

Bug: 36558
Change-Id: I77548e604f294f1bed223f781fdf855f0b5ac05e

src/bin/virtio_block_test_util.rs

diff --git a/src/bin/virtio_block_test_util.rs b/src/bin/virtio_block_test_util.rs
index d1175d4..fdc2319 100644
--- a/src/bin/virtio_block_test_util.rs
+++ b/src/bin/virtio_block_test_util.rs
@@ -4,13 +4,13 @@
 #![deny(warnings)]
 
 use libc;
-use std::fs::{self, File, OpenOptions};
+use std::fs::{File, OpenOptions};
 use std::io::{Error, ErrorKind, Read, Seek, SeekFrom, Write};
 use std::os::unix::io::AsRawFd;
-use std::path::Path;
+use std::path::PathBuf;
 use structopt::StructOpt;
 
-const DEV_DIR: &str = "/dev";
+const PCI_DIR: &str = "/dev/disk/by-path";
 
 // Ioctl requests are comprised of 32 bits:
 // [31:30]  Access mode
@@ -40,7 +40,8 @@
 #[derive(StructOpt, Debug)]
 struct Config {
     block_size: u32,
-    block_count: u64,
+    pci_bus: u8,
+    pci_device: u8,
     #[structopt(subcommand)]
     cmd: Command,
 }
@@ -48,42 +49,45 @@
 #[derive(StructOpt, Debug)]
 enum Command {
     #[structopt(name = "check")]
-    Check,
+    Check { block_count: u64 },
     #[structopt(name = "read")]
     Read { offset: u64, expected: u8 },
     #[structopt(name = "write")]
     Write { offset: u64, value: u8 },
 }
 
-fn find_block_device(block_size: u32, block_count: u64, write: bool) -> std::io::Result<File> {
-    let dir = Path::new(DEV_DIR);
-    if !dir.is_dir() {
+fn open_block_device(pci_bus: u8, pci_device: u8, write: bool) -> std::io::Result<File> {
+    let mut path = PathBuf::from(PCI_DIR);
+    path.push(format!(
+        "virtio-pci-0000:{:02}:{:02}.0",
+        pci_bus, pci_device
+    ));
+    OpenOptions::new().read(true).write(write).open(path)
+}
+
+fn check(block_dev: &File, block_size: u32, block_count: u64) -> std::io::Result<()> {
+    let actual_block_size = block_dev_sector_size(block_dev);
+    let actual_block_count = block_dev_size(block_dev);
+
+    if block_size != actual_block_size {
         return Err(Error::new(
             ErrorKind::Other,
-            format!("{} is not a directory", DEV_DIR),
+            format!(
+                "Incorrect size: {} (expected {}).",
+                actual_block_size, block_size
+            ),
         ));
     }
-    for entry in fs::read_dir(dir)?.collect::<std::io::Result<Vec<_>>>()? {
-        if !entry
-            .file_name()
-            .to_str()
-            .map_or(false, |f| f.starts_with("vd"))
-        {
-            continue;
-        }
-        let file = OpenOptions::new()
-            .read(true)
-            .write(write)
-            .open(entry.path())?;
-        if block_size != block_dev_sector_size(&file) {
-            continue;
-        }
-        if block_count != block_dev_size(&file) {
-            continue;
-        }
-        return Ok(file);
+    if block_count != actual_block_count {
+        return Err(Error::new(
+            ErrorKind::Other,
+            format!(
+                "Incorrect count: {} (expected {}).",
+                actual_block_count, block_count
+            ),
+        ));
     }
-    Err(Error::new(ErrorKind::NotFound, "Block device not found"))
+    Ok(())
 }
 
 fn read_block(
@@ -121,9 +125,9 @@
     } else {
         false
     };
-    let mut block_dev = find_block_device(config.block_size, config.block_count, write)?;
+    let mut block_dev = open_block_device(config.pci_bus, config.pci_device, write)?;
     let result = match config.cmd {
-        Command::Check => Ok(()),
+        Command::Check { block_count } => check(&block_dev, config.block_size, block_count),
         Command::Read { offset, expected } => {
             read_block(&mut block_dev, config.block_size, offset, expected)
         }