rust/bt-common/src/debug_command.rs

// Copyright 2023 Google LLC
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Debug command traits and helpers for defining commands for integration
//! into a debug tool.
use log::LevelFilter;
use std::str::FromStr;
use std::sync::Arc;

#[must_use = "if unused the previous log level will immediately be restored"]
pub struct ScopedVerbosityGuard {
    prev_level: LevelFilter,
}

impl ScopedVerbosityGuard {
    pub fn new(scoped_level: LevelFilter) -> Self {
        let prev_level = log::max_level();
        log::set_max_level(scoped_level);
        Self { prev_level }
    }
}

impl Drop for ScopedVerbosityGuard {
    fn drop(&mut self) {
        log::set_max_level(self.prev_level);
    }
}

pub fn parse_level_filter(s: &str) -> Option<LevelFilter> {
    s.trim().parse().ok()
}

pub fn extract_verbosity_flags(args: Vec<String>) -> (Option<LevelFilter>, Vec<String>) {
    let mut level = None;
    let mut clean_args = Vec::new();
    let mut iter = args.into_iter().peekable();

    while let Some(arg) = iter.next() {
        match arg.as_str() {
            "-v" | "--verbose" => {
                level = Some(LevelFilter::Debug);
            }
            "-vv" | "--trace" => {
                level = Some(LevelFilter::Trace);
            }
            "-q" | "--quiet" => {
                level = Some(LevelFilter::Warn);
            }
            "--log-level" => {
                if let Some(next) = iter.peek() {
                    if let Some(lvl) = parse_level_filter(next) {
                        level = Some(lvl);
                        iter.next();
                        continue;
                    }
                }
                clean_args.push(arg);
            }
            _ if arg.starts_with("--log-level=") => {
                let val = &arg["--log-level=".len()..];
                if let Some(lvl) = parse_level_filter(val) {
                    level = Some(lvl);
                } else {
                    clean_args.push(arg);
                }
            }
            _ => {
                clean_args.push(arg);
            }
        }
    }

    (level, clean_args)
}

pub fn handle_verbose_command(args: &[String]) {
    match args.len() {
        0 => {
            println!("Current log level: {}", log::max_level());
        }
        1 => {
            if let Some(level) = parse_level_filter(&args[0]) {
                log::set_max_level(level);
                println!("Log level set to {}", level);
            } else {
                eprintln!(
                    "Invalid log level '{}'. Valid levels: off, error, warn, info, debug, trace",
                    args[0]
                );
            }
        }
        _ => {
            println!("Usage: verbose [off|error|warn|info|debug|trace]");
        }
    }
}

/// A CommandSet is a set of commands (usually an enum) that each represent an
/// action that can be performed.  i.e. 'list', 'volume' etc.  Each command can
/// take zero or more arguments and have zero or more flags.
/// Typically an Enum of commands would implement CommandSet trait.
pub trait CommandSet: FromStr + ::core::fmt::Display {
    /// Returns a vector of strings that are the commands supported by this.
    fn variants() -> Vec<String>;

    /// Returns a string listing the arguments that this command takes, in <>
    /// brackets
    fn arguments(&self) -> &'static str;

    /// Returns a string displaying the flags that this command supports, in []
    /// brackets
    fn flags(&self) -> &'static str;

    /// Returns a short description of this command
    fn desc(&self) -> &'static str;

    /// Help string for this variant (build from Display, arguments and flags by
    /// default)
    fn help_simple(&self) -> String {
        format!("{self} {} {} -- {}", self.flags(), self.arguments(), self.desc())
    }

    /// Possibly multi-line help string for all variants of this set.
    fn help_all() -> String {
        let mut help: String = Self::variants()
            .into_iter()
            .filter_map(|s| FromStr::from_str(&s).ok())
            .map(|s: Self| format!("{}\n", s))
            .collect();
        if !help.is_empty() {
            help.push('\n');
        }
        help.push_str("Global verbosity flags: [-v|--verbose] [-vv|--trace] [-q|--quiet] [--log-level=<level>]\n");
        help
    }
}

/// A top-level CLI command: either a framework meta-command (e.g. `verbose`)
/// or a profile-specific domain command.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CliCommand<T> {
    Verbose,
    Domain(T),
}

impl<T> From<T> for CliCommand<T> {
    fn from(cmd: T) -> Self {
        CliCommand::Domain(cmd)
    }
}

impl<T: CommandSet> CommandSet for CliCommand<T> {
    fn variants() -> Vec<String> {
        let mut variants = T::variants();
        variants.push("verbose".to_string());
        variants
    }

    fn arguments(&self) -> &'static str {
        match self {
            CliCommand::Domain(cmd) => cmd.arguments(),
            CliCommand::Verbose => "[level] ",
        }
    }

    fn flags(&self) -> &'static str {
        match self {
            CliCommand::Domain(cmd) => cmd.flags(),
            CliCommand::Verbose => "",
        }
    }

    fn desc(&self) -> &'static str {
        match self {
            CliCommand::Domain(cmd) => cmd.desc(),
            CliCommand::Verbose => "Get or set the logging verbosity level",
        }
    }
}

impl<T: ::core::fmt::Display> ::core::fmt::Display for CliCommand<T> {
    fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
        match self {
            CliCommand::Domain(cmd) => write!(f, "{cmd}"),
            CliCommand::Verbose => write!(f, "verbose"),
        }
    }
}

impl<T: FromStr> FromStr for CliCommand<T> {
    type Err = T::Err;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s == "verbose" {
            Ok(CliCommand::Verbose)
        } else {
            T::from_str(s).map(CliCommand::Domain)
        }
    }
}

/// Macro to help build CommandSets
#[macro_export]
macro_rules! gen_commandset {
    ($name:ident {
        $($variant:ident = ($val:expr, [$($flag:expr),*], [$($arg:expr),*], $help:expr)),*,
    }) => {
        /// Enum of all possible commands
        #[derive(PartialEq, Debug, Clone, Copy)]
        pub enum $name {
            $($variant,)*
        }

        impl CommandSet for $name {
            fn variants() -> Vec<String> {
                let mut variants = Vec::new();
                $(variants.push($val.to_string());)*
                variants
            }

            fn arguments(&self) -> &'static str {
                match self {
                    $(
                        $name::$variant => concat!($("<", $arg, "> ",)*),
                    )*
                }
            }

            fn flags(&self) -> &'static str {
                match self {
                    $(
                        $name::$variant => concat!($("[", $flag, "] ",)*),
                    )*
                }
            }

            fn desc(&self) -> &'static str {
                match self {
                    $(
                        $name::$variant => $help,
                    )*
                }
            }
        }

        impl ::core::fmt::Display for $name {
            fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
                match *self {
                    $($name::$variant => write!(f, $val),)*
                }
            }
        }

        impl ::std::str::FromStr for $name {
            type Err = ();

            fn from_str(s: &str) -> Result<$name, ()> {
                match s {
                    $($val => Ok($name::$variant),)*
                    _ => Err(()),
                }
            }
        }
    }
}

/// CommandRunner is used to perform a specific task based on the command set.
pub trait CommandRunner {
    type Set: CommandSet;
    type Error: ::std::error::Error;

    fn run_command(
        &self,
        cmd: Self::Set,
        args: Vec<String>,
    ) -> impl futures::Future<Output = Result<(), Self::Error>>;

    fn run(
        &self,
        cmd: impl Into<CliCommand<Self::Set>>,
        args: Vec<String>,
    ) -> impl futures::Future<Output = Result<(), Self::Error>> {
        let cmd = cmd.into();
        async move {
            match cmd {
                CliCommand::Verbose => {
                    handle_verbose_command(&args);
                    Ok(())
                }
                CliCommand::Domain(cmd) => {
                    let (scoped_level, clean_args) = extract_verbosity_flags(args);
                    let _guard = scoped_level.map(ScopedVerbosityGuard::new);
                    self.run_command(cmd, clean_args).await
                }
            }
        }
    }
}

impl<R: CommandRunner> CommandRunner for &R {
    type Set = R::Set;
    type Error = R::Error;

    fn run_command(
        &self,
        cmd: Self::Set,
        args: Vec<String>,
    ) -> impl futures::Future<Output = Result<(), Self::Error>> {
        (**self).run_command(cmd, args)
    }
}

impl<R: CommandRunner> CommandRunner for Arc<R> {
    type Set = R::Set;
    type Error = R::Error;

    fn run_command(
        &self,
        cmd: Self::Set,
        args: Vec<String>,
    ) -> impl futures::Future<Output = Result<(), Self::Error>> {
        (**self).run_command(cmd, args)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn gen_commandset_simple() {
        gen_commandset! {
            TestCmd {
                One = ("one", [], [], "First Command"),
                WithFlags = ("with-flags", ["-1","-2"], [], "Command with flags"),
                WithArgs = ("with-args", [], ["arg", "two"], "Command with args"),
                WithBoth = ("with-both", ["-w"], ["simple"], "Command with both flags and args"),
            }
        }

        let cmd: TestCmd = "one".parse().unwrap();
        assert_eq!(cmd, TestCmd::One);

        let cmd2: TestCmd = "with-flags".parse().unwrap();
        assert_eq!(cmd2.arguments(), "");
        assert_eq!(cmd2.flags(), "[-1] [-2] ");

        let verbose_cmd: CliCommand<TestCmd> = "verbose".parse().unwrap();
        assert_eq!(verbose_cmd, CliCommand::Verbose);
        assert_eq!(verbose_cmd.flags(), "");
        assert_eq!(verbose_cmd.to_string(), "verbose");

        let domain_cmd: CliCommand<TestCmd> = "one".parse().unwrap();
        assert_eq!(domain_cmd, CliCommand::Domain(TestCmd::One));

        assert!(!TestCmd::variants().contains(&"verbose".to_string()));
        assert!(CliCommand::<TestCmd>::variants().contains(&"verbose".to_string()));
        assert!(TestCmd::help_all().contains("Global verbosity flags:"));
        assert!(CliCommand::<TestCmd>::help_all().contains("Global verbosity flags:"));
    }

    #[test]
    fn test_extract_verbosity_flags() {
        // -v / --verbose
        let (lvl, args) = extract_verbosity_flags(vec!["-v".to_string()]);
        assert_eq!(lvl, Some(LevelFilter::Debug));
        assert!(args.is_empty());

        let (lvl, args) = extract_verbosity_flags(vec!["--verbose".to_string()]);
        assert_eq!(lvl, Some(LevelFilter::Debug));
        assert!(args.is_empty());

        // -vv / --trace
        let (lvl, args) = extract_verbosity_flags(vec!["-vv".to_string()]);
        assert_eq!(lvl, Some(LevelFilter::Trace));
        assert!(args.is_empty());

        let (lvl, args) = extract_verbosity_flags(vec!["--trace".to_string()]);
        assert_eq!(lvl, Some(LevelFilter::Trace));
        assert!(args.is_empty());

        // -q / --quiet
        let (lvl, args) = extract_verbosity_flags(vec!["-q".to_string()]);
        assert_eq!(lvl, Some(LevelFilter::Warn));
        assert!(args.is_empty());

        let (lvl, args) = extract_verbosity_flags(vec!["--quiet".to_string()]);
        assert_eq!(lvl, Some(LevelFilter::Warn));
        assert!(args.is_empty());

        // --log-level=<level> and --log-level <level> (case-insensitive)
        for (name, expected) in [
            ("off", LevelFilter::Off),
            ("error", LevelFilter::Error),
            ("warn", LevelFilter::Warn),
            ("info", LevelFilter::Info),
            ("debug", LevelFilter::Debug),
            ("trace", LevelFilter::Trace),
            ("DEBUG", LevelFilter::Debug),
            ("Info", LevelFilter::Info),
        ] {
            let (lvl, args) = extract_verbosity_flags(vec![format!("--log-level={name}")]);
            assert_eq!(lvl, Some(expected));
            assert!(args.is_empty());

            let (lvl, args) =
                extract_verbosity_flags(vec!["--log-level".to_string(), name.to_string()]);
            assert_eq!(lvl, Some(expected));
            assert!(args.is_empty());
        }

        // Mixed arguments
        let input = vec![
            "action".to_string(),
            "-v".to_string(),
            "--custom-flag".to_string(),
            "target".to_string(),
            "--log-level=warn".to_string(),
            "extra".to_string(),
        ];
        let (lvl, args) = extract_verbosity_flags(input);
        assert_eq!(lvl, Some(LevelFilter::Warn));
        assert_eq!(args, vec!["action", "--custom-flag", "target", "extra"]);

        // Invalid log level flag should not be extracted
        let (lvl, args) =
            extract_verbosity_flags(vec!["--log-level=invalid".to_string(), "arg".to_string()]);
        assert_eq!(lvl, None);
        assert_eq!(args, vec!["--log-level=invalid", "arg"]);
    }

    // Mutex to avoid parallel tests fighting over log levels
    static LOG_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());

    #[test]
    fn test_scoped_verbosity_guard() {
        let _lock = LOG_MUTEX.lock().unwrap();
        log::set_max_level(LevelFilter::Info);

        futures::executor::block_on(async {
            assert_eq!(log::max_level(), LevelFilter::Info);
            {
                let _guard = ScopedVerbosityGuard::new(LevelFilter::Trace);
                assert_eq!(log::max_level(), LevelFilter::Trace);
                // Yield across async boundary
                futures::future::ready(()).await;
                assert_eq!(log::max_level(), LevelFilter::Trace);
            }
            assert_eq!(log::max_level(), LevelFilter::Info);
        });
    }

    #[test]
    fn test_persistent_verbosity_and_handle_command() {
        let _lock = LOG_MUTEX.lock().unwrap();
        log::set_max_level(LevelFilter::Debug);

        // handle_verbose_command with 0 args
        handle_verbose_command(&[]);
        assert_eq!(log::max_level(), LevelFilter::Debug);

        // handle_verbose_command with 1 valid arg
        handle_verbose_command(&["trace".to_string()]);
        assert_eq!(log::max_level(), LevelFilter::Trace);

        handle_verbose_command(&["WARN".to_string()]);
        assert_eq!(log::max_level(), LevelFilter::Warn);

        // handle_verbose_command with invalid arg (should not change)
        handle_verbose_command(&["invalid_level".to_string()]);
        assert_eq!(log::max_level(), LevelFilter::Warn);

        // handle_verbose_command with > 1 args (should not change)
        handle_verbose_command(&["info".to_string(), "extra".to_string()]);
        assert_eq!(log::max_level(), LevelFilter::Warn);

        // Reset back to Info
        log::set_max_level(LevelFilter::Info);
    }

    #[test]
    fn test_command_runner_run() {
        gen_commandset! {
            RunnerCmd {
                Action = ("action", [], ["target"], "Perform action"),
            }
        }

        struct TestRunner {
            executed: std::sync::Mutex<Vec<(RunnerCmd, Vec<String>)>>,
        }

        impl CommandRunner for TestRunner {
            type Set = RunnerCmd;
            type Error = std::io::Error;

            fn run_command(
                &self,
                cmd: Self::Set,
                args: Vec<String>,
            ) -> impl futures::Future<Output = Result<(), Self::Error>> {
                self.executed.lock().unwrap().push((cmd, args));
                futures::future::ready(Ok(()))
            }
        }

        let runner = TestRunner { executed: std::sync::Mutex::new(Vec::new()) };
        let _lock = LOG_MUTEX.lock().unwrap();

        // Running Verbose command should update persistent verbosity and NOT
        // invoke run_command
        log::set_max_level(LevelFilter::Info);
        futures::executor::block_on(runner.run(CliCommand::Verbose, vec!["trace".to_string()]))
            .unwrap();
        assert_eq!(log::max_level(), LevelFilter::Trace);
        assert!(runner.executed.lock().unwrap().is_empty());

        // Running normal command with verbosity flag strips flag and runs
        // command
        let args = vec!["-v".to_string(), "my_target".to_string()];
        futures::executor::block_on(runner.run(RunnerCmd::Action, args)).unwrap();
        let executed = runner.executed.lock().unwrap().clone();
        assert_eq!(executed.len(), 1);
        assert_eq!(executed[0], (RunnerCmd::Action, vec!["my_target".to_string()]));

        // After run completes, the guard has restored the previous level
        assert_eq!(log::max_level(), LevelFilter::Trace);

        // Reset persistent verbosity back to Info
        log::set_max_level(LevelFilter::Info);
    }
}