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(())
 }