Better block device handling (fix #14)

README.md

@@ -27,8 +27,7 @@ Dependencies will be handled for you if you install alma from AUR.
 sudo alma create /dev/disk/by-id/usb-Generic_USB_Flash_Disk-0:0
 ```
 
-This will wipe the entire disk and create a bootable installation of Arch Linux. As a precaution,
-this tool will refuse to work with drive paths which don't start with `/dev/disk/by-id/usb-`.
+This will wipe the entire disk and create a bootable installation of Arch Linux.
 
 After the installation is done you can either boot from it immediately or use `arch-chroot` to
 perform further customizations before your first boot.

src/block.rs

@@ -0,0 +1,67 @@
+use super::error::{Error, ErrorKind};
+use failure::ResultExt;
+use log::debug;
+use std::fs::read_to_string;
+use std::path::PathBuf;
+
+#[derive(Debug)]
+pub struct BlockDevice {
+    name: String,
+}
+
+impl BlockDevice {
+    pub fn from_path(path: PathBuf) -> Result<Self, Error> {
+        let real_path = path.canonicalize().context(ErrorKind::DeviceQuery)?;
+        let device_name = real_path
+            .file_name()
+            .and_then(|s| s.to_str())
+            .map(String::from)
+            .ok_or(Error::from(ErrorKind::InvalidDeviceName))?;
+
+        debug!(
+            "path: {:?}, real path: {:?}, device name: {:?}",
+            path, real_path, device_name
+        );
+
+        Ok(Self { name: device_name })
+    }
+
+    fn sys_path(&self) -> PathBuf {
+        let mut path = PathBuf::from("/sys/block");
+        path.push(self.name.clone());
+        path
+    }
+
+    pub fn removable(&self) -> Result<bool, Error> {
+        let mut path = self.sys_path();
+        path.push("removable");
+
+        debug!("Reading: {:?}", path);
+        let result = read_to_string(&path).context(ErrorKind::DeviceQuery)?;
+        debug!("{:?} -> {}", path, result);
+
+        Ok(result == "1\n")
+    }
+
+    pub fn device_path(&self) -> PathBuf {
+        let mut path = PathBuf::from("/dev");
+        path.push(self.name.clone());
+        path
+    }
+
+    pub fn partition_device_path(&self, index: u8) -> Result<PathBuf, Error> {
+        let name = if self.name.chars().rev().next().unwrap().is_digit(10) {
+            format!("{}p{}", self.name, index)
+        } else {
+            format!("{}{}", self.name, index)
+        };
+        let mut path = PathBuf::from("/dev");
+        path.push(name);
+
+        debug!("Partition {} for {} is in {:?}", index, self.name, path);
+        if !path.exists() {
+            return Err(ErrorKind::NoSuchPartition(index).into());
+        }
+        Ok(path)
+    }
+}

src/error.rs

@@ -8,10 +8,17 @@ pub struct Error {
 
 #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail)]
 pub enum ErrorKind {
-    #[fail(
-        display = "Provided Path does not point to a USB device. Make sure you specify a file matching /dev/disk/by-id/usb-*"
-    )]
-    NotUSB,
+    #[fail(display = "Error quering information about the block device")]
+    DeviceQuery,
+
+    #[fail(display = "Invalid device name")]
+    InvalidDeviceName,
+
+    #[fail(display = "The given block device is not removable")]
+    NotRemovableDevice,
+
+    #[fail(display = "Partition {} does not exist", _0)]
+    NoSuchPartition(u8),
 
     #[fail(display = "Could not find {}", _0)]
     NoTool(&'static str),

src/main.rs

@@ -1,6 +1,5 @@
-#[macro_use]
-extern crate log;
 extern crate failure;
+extern crate log;
 extern crate nix;
 extern crate simplelog;
 extern crate structopt;
@@ -8,6 +7,7 @@ extern crate tempfile;
 extern crate which;
 use nix::sys::signal;
 
+mod block;
 mod error;
 mod mountstack;
 mod process;
@@ -15,6 +15,7 @@ mod tool;
 
 use error::*;
 use failure::{Fail, ResultExt};
+use log::{debug, error, info, warn};
 use mountstack::{Filesystem, MountStack};
 use process::CommandExt;
 use simplelog::*;
@@ -70,7 +71,7 @@ struct ChrootCommand {
     disk: PathBuf,
 }
 
-fn create(command: &CreateCommand) -> Result<(), Error> {
+fn create(command: CreateCommand) -> Result<(), Error> {
     let sgdisk = Tool::find("sgdisk")?;
     let pacstrap = Tool::find("pacstrap")?;
     let arch_chroot = Tool::find("arch-chroot")?;
@@ -78,22 +79,20 @@ fn create(command: &CreateCommand) -> Result<(), Error> {
     let mkfat = Tool::find("mkfs.fat")?;
     let mkbtrfs = Tool::find("mkfs.btrfs")?;
 
-    if !(command.disk.starts_with("/dev/disk/by-id")
-        && (command
-            .disk
-            .file_name()
-            .and_then(|s| s.to_str())
-            .filter(|ref f| f.starts_with("usb-"))
-            .is_some()))
-    {
-        Err(ErrorKind::NotUSB)?;
+    let disk = block::BlockDevice::from_path(command.disk)?;
+
+    if !disk.removable()? {
+        Err(ErrorKind::NotRemovableDevice)?;
     }
 
     let mount_point = tempdir().context(ErrorKind::TmpDirError)?;
     let boot_point = mount_point.path().join("boot");
     let mut mount_stack = MountStack::new();
+    let disk_path = disk.device_path();
 
     info!("Partitioning the disk");
+    debug!("{:?}", disk_path);
+
     sgdisk
         .execute()
         .args(&[
@@ -104,27 +103,30 @@ fn create(command: &CreateCommand) -> Result<(), Error> {
             "--largest-new=3",
             "--typecode=1:EF02",
             "--typecode=2:EF00",
-        ]).arg(&command.disk)
+        ]).arg(&disk_path)
         .run(ErrorKind::Partitioning)?;
 
     thread::sleep(Duration::from_millis(1000));
 
     info!("Formatting filesystems");
+    let boot_partition = disk.partition_device_path(2)?;
     mkfat
         .execute()
         .arg("-F32")
-        .arg(format!("{}-part2", command.disk.display()))
+        .arg(&boot_partition)
         .run(ErrorKind::Formatting)?;
+
+    let root_partition = disk.partition_device_path(3)?;
     mkbtrfs
         .execute()
         .arg("-f")
-        .arg(format!("{}-part3", command.disk.display()))
+        .arg(&root_partition)
         .run(ErrorKind::Formatting)?;
 
     info!("Mounting filesystems to {}", mount_point.path().display());
     mount_stack
         .mount(
-            &PathBuf::from(format!("{}-part3", command.disk.display())),
+            &PathBuf::from(&root_partition),
             &mount_point.path(),
             Filesystem::Btrfs,
             Some("compress=lzo"),
@@ -133,12 +135,8 @@ fn create(command: &CreateCommand) -> Result<(), Error> {
     fs::create_dir(&boot_point).context(ErrorKind::CreateBoot)?;
 
     mount_stack
-        .mount(
-            &PathBuf::from(format!("{}-part2", command.disk.display())),
-            &boot_point,
-            Filesystem::Vfat,
-            None,
-        ).context(ErrorKind::Mounting)?;
+        .mount(&boot_partition, &boot_point, Filesystem::Vfat, None)
+        .context(ErrorKind::Mounting)?;
 
     info!("Bootstrapping system");
     pacstrap
@@ -191,7 +189,7 @@ fn create(command: &CreateCommand) -> Result<(), Error> {
         .execute()
         .arg(mount_point.path())
         .args(&["bash", "-c"])
-        .arg(format!("grub-install --target=i386-pc --boot-directory /boot {} && grub-install --target=x86_64-efi --efi-directory /boot --boot-directory /boot --removable &&  grub-mkconfig -o /boot/grub/grub.cfg", command.disk.display()))
+        .arg(format!("grub-install --target=i386-pc --boot-directory /boot {} && grub-install --target=x86_64-efi --efi-directory /boot --boot-directory /boot --removable &&  grub-mkconfig -o /boot/grub/grub.cfg", disk_path.display()))
         .run(ErrorKind::Bootloader)?;
 
     if command.interactive {
@@ -206,18 +204,12 @@ fn create(command: &CreateCommand) -> Result<(), Error> {
     Ok(())
 }
 
-fn chroot(command: &ChrootCommand) -> Result<(), Error> {
+fn chroot(command: ChrootCommand) -> Result<(), Error> {
     let arch_chroot = Tool::find("arch-chroot")?;
+    let disk = block::BlockDevice::from_path(command.disk)?;
 
-    if !(command.disk.starts_with("/dev/disk/by-id")
-        && (command
-            .disk
-            .file_name()
-            .and_then(|s| s.to_str())
-            .filter(|ref f| f.starts_with("usb-"))
-            .is_some()))
-    {
-        Err(ErrorKind::NotUSB)?;
+    if !disk.removable()? {
+        Err(ErrorKind::NotRemovableDevice)?;
     }
 
     let mount_point = tempdir().context(ErrorKind::TmpDirError)?;
@@ -225,21 +217,19 @@ fn chroot(command: &ChrootCommand) -> Result<(), Error> {
     let mut mount_stack = MountStack::new();
 
     info!("Mounting filesystems to {}", mount_point.path().display());
+    let root_partition = disk.partition_device_path(3)?;
     mount_stack
         .mount(
-            &PathBuf::from(format!("{}-part3", command.disk.display())),
+            &root_partition,
             &mount_point.path(),
             Filesystem::Btrfs,
             Some("compress=lzo"),
         ).context(ErrorKind::Mounting)?;
 
+    let boot_partition = disk.partition_device_path(2)?;
     mount_stack
-        .mount(
-            &PathBuf::from(format!("{}-part2", command.disk.display())),
-            &boot_point,
-            Filesystem::Vfat,
-            None,
-        ).context(ErrorKind::Mounting)?;
+        .mount(&boot_partition, &boot_point, Filesystem::Vfat, None)
+        .context(ErrorKind::Mounting)?;
 
     arch_chroot
         .execute()
@@ -273,8 +263,8 @@ fn main() {
     }
 
     let result = match app {
-        App::Create(command) => create(&command),
-        App::Chroot(command) => chroot(&command),
+        App::Create(command) => create(command),
+        App::Chroot(command) => chroot(command),
     };
 
     match result {

src/mountstack.rs

@@ -1,3 +1,4 @@
+use log::{debug, warn};
 use nix;
 use nix::mount::{mount, umount, MsFlags};
 use std::path::Path;

src/process.rs

@@ -1,5 +1,6 @@
 use super::error::*;
 use failure::{Fail, ResultExt};
+use log::error;
 use std::process::{Command, ExitStatus};
 use std::str;