Display available extensions in the help text

With recent changes we are able to execute user-provided extensions
through the program. However, discoverability is arguably lacking,
because nitrocli provides no insight into what extensions are available
to begin with.
This patch changes this state of affairs by listing available extensions
in the help text.

4 files changed, 114 insertions, 7 deletions

diff --git a/src/commands.rs b/src/commands.rs
index 549ebec..7612b76 100644
--- a/src/commands.rs
+++ b/src/commands.rs
@@ -8,6 +8,7 @@ use std::convert::TryFrom as _;
 use std::env;
 use std::ffi;
 use std::fmt;
+use std::fs;
 use std::io;
 use std::mem;
 use std::ops;
@@ -33,6 +34,8 @@ use crate::output;
 use crate::pinentry;
 use crate::Context;
 
+const NITROCLI_EXT_PREFIX: &str = "nitrocli-";
+
 /// Set `libnitrokey`'s log level based on the execution context's verbosity.
 fn set_log_level(ctx: &mut Context<'_>) {
   let log_lvl = match ctx.config.verbosity {
@@ -1087,15 +1090,52 @@ pub fn pws_status(ctx: &mut Context<'_>, all: bool) -> anyhow::Result<()> {
   })
 }
 
+/// Find and list all available extensions.
+///
+/// The logic used in this function should use the same criteria as
+/// `resolve_extension`.
+pub(crate) fn discover_extensions(path_var: &ffi::OsStr) -> anyhow::Result<Vec<String>> {
+  let dirs = env::split_paths(path_var);
+  let mut commands = Vec::new();
+
+  for dir in dirs {
+    match fs::read_dir(&dir) {
+      Ok(entries) => {
+        for entry in entries {
+          let entry = entry?;
+          let path = entry.path();
+          if path.is_file() {
+            let name = entry.file_name();
+            let file = name.to_string_lossy();
+            if file.starts_with(NITROCLI_EXT_PREFIX) {
+              let mut file = file.into_owned();
+              file.replace_range(..NITROCLI_EXT_PREFIX.len(), "");
+              commands.push(file);
+            }
+          }
+        }
+      }
+      Err(ref err) if err.kind() == io::ErrorKind::NotFound => (),
+      x => x
+        .map(|_| ())
+        .with_context(|| format!("Failed to iterate entries of directory {}", dir.display()))?,
+    }
+  }
+  Ok(commands)
+}
+
 /// Resolve an extension provided by name to an actual path.
 ///
 /// Extensions are (executable) files that have the "nitrocli-" prefix
 /// and are discoverable via the `PATH` environment variable.
+///
+/// The logic used in this function should use the same criteria as
+/// `discover_extensions`.
 pub(crate) fn resolve_extension(
   path_var: &ffi::OsStr,
   ext_name: &ffi::OsStr,
 ) -> anyhow::Result<path::PathBuf> {
-  let mut bin_name = ffi::OsString::from("nitrocli-");
+  let mut bin_name = ffi::OsString::from(NITROCLI_EXT_PREFIX);
   bin_name.push(ext_name);
 
   for dir in env::split_paths(path_var) {
diff --git a/src/main.rs b/src/main.rs
index e7a7d2f..c8b73dc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -69,6 +69,10 @@ use std::io;
 use std::process;
 use std::str;
 
+use structopt::clap::ErrorKind;
+use structopt::clap::SubCommand;
+use structopt::StructOpt;
+
 const NITROCLI_BINARY: &str = "NITROCLI_BINARY";
 const NITROCLI_MODEL: &str = "NITROCLI_MODEL";
 const NITROCLI_USB_PATH: &str = "NITROCLI_USB_PATH";
@@ -105,15 +109,44 @@ impl fmt::Display for DirectExitError {
 impl error::Error for DirectExitError {}
 
 /// Parse the command-line arguments and execute the selected command.
-fn handle_arguments(ctx: &mut Context<'_>, args: Vec<String>) -> anyhow::Result<()> {
-  use structopt::StructOpt;
-
-  match args::Args::from_iter_safe(args.iter()) {
+fn handle_arguments(ctx: &mut Context<'_>, argv: Vec<String>) -> anyhow::Result<()> {
+  match args::Args::from_iter_safe(argv.iter()) {
     Ok(args) => {
       ctx.config.update(&args);
       args.cmd.execute(ctx)
     }
-    Err(err) => {
+    Err(mut err) => {
+      if err.kind == ErrorKind::HelpDisplayed {
+        // For the convenience of the user we'd like to list the
+        // available extensions in the help text. At the same time, we
+        // don't want to unconditionally iterate through PATH (which may
+        // contain directories with loads of files that need scanning)
+        // for every command invoked. So we do that listing only if a
+        // help text is actually displayed.
+        let path = ctx.path.clone().unwrap_or_else(ffi::OsString::new);
+        if let Ok(extensions) = commands::discover_extensions(&path) {
+          let mut clap = args::Args::clap();
+          for name in extensions {
+            // Because of clap's brain dead API, we see no other way
+            // but to leak the string we created here. That's okay,
+            // though, because we exit in a moment anyway.
+            let about = Box::leak(format!("Run the {} extension", name).into_boxed_str());
+            clap = clap.subcommand(
+              SubCommand::with_name(&name)
+                // Use some magic number here that causes all
+                // extensions to be listed after all other
+                // subcommands.
+                .display_order(1000)
+                .about(about as &'static str),
+            );
+          }
+          // At this point we are *pretty* sure that repeated invocation
+          // will result in another error. So should be fine to unwrap
+          // here.
+          err = clap.get_matches_from_safe(argv.iter()).unwrap_err();
+        }
+      }
+
       if err.use_stderr() {
         Err(err.into())
       } else {
diff --git a/src/tests/extensions.rs b/src/tests/extensions.rs
index a295949..d546e8a 100644
--- a/src/tests/extensions.rs
+++ b/src/tests/extensions.rs
@@ -9,6 +9,36 @@ use std::fs;
 use super::*;
 
 #[test]
+fn no_extensions_to_discover() -> anyhow::Result<()> {
+  let exts = crate::commands::discover_extensions(&ffi::OsString::new())?;
+  assert!(exts.is_empty(), "{:?}", exts);
+  Ok(())
+}
+
+#[test]
+fn extension_discovery() -> anyhow::Result<()> {
+  let dir1 = tempfile::tempdir()?;
+  let dir2 = tempfile::tempdir()?;
+
+  {
+    let ext1_path = dir1.path().join("nitrocli-ext1");
+    let ext2_path = dir1.path().join("nitrocli-ext2");
+    let ext3_path = dir2.path().join("nitrocli-super-1337-extensions111one");
+    let _ext1 = fs::File::create(&ext1_path)?;
+    let _ext2 = fs::File::create(&ext2_path)?;
+    let _ext3 = fs::File::create(&ext3_path)?;
+
+    let path = env::join_paths(&[dir1.path(), dir2.path()])?;
+    let mut exts = crate::commands::discover_extensions(&path)?;
+    // We can't assume a fixed ordering of extensions, because that is
+    // platform/file system dependent. So sort here to fix it.
+    exts.sort();
+    assert_eq!(exts, vec!["ext1", "ext2", "super-1337-extensions111one"]);
+  }
+  Ok(())
+}
+
+#[test]
 fn resolve_extensions() -> anyhow::Result<()> {
   let dir1 = tempfile::tempdir()?;
   let dir2 = tempfile::tempdir()?;
diff --git a/src/tests/run.rs b/src/tests/run.rs
index e4bbb28..1dae166 100644
--- a/src/tests/run.rs
+++ b/src/tests/run.rs
@@ -300,7 +300,11 @@ print("success")
   }
 
   let path = ext_dir.path().as_os_str().to_os_string();
-  let out = Nitrocli::new().path(path).handle(&["ext"])?;
+  // Make sure that the extension appears in the help text.
+  let out = Nitrocli::new().path(&path).handle(&["--help"])?;
+  assert!(out.contains("ext            Run the ext extension\n"), out);
+  // And, of course, that we can invoke it.
+  let out = Nitrocli::new().path(&path).handle(&["ext"])?;
   assert_eq!(out, "success\n");
   Ok(())
 }