diff --git a/cli/src/main.rs b/cli/src/main.rs index acb43cc9..482f3f26 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -18,7 +18,7 @@ use std::process::ExitCode; mod opts; use { - crate::opts::Opts, + crate::opts::*, clap::Parser, hyper::{client::Client, Body, Request}, std::{ @@ -35,11 +35,11 @@ use { /// Starts up a Viceroy server. /// /// Create a new server, bind it to an address, and serve responses until an error occurs. -pub async fn serve(opts: Opts) -> Result<(), Error> { +pub async fn serve(serve_args: ServeArgs) -> Result<(), Error> { // Load the wasm module into an execution context - let ctx = create_execution_context(&opts).await?; + let ctx = create_execution_context(&serve_args.shared(), true).await?; - let addr = opts.addr(); + let addr = serve_args.addr(); ViceroyService::new(ctx).serve(addr).await?; unreachable!() @@ -49,62 +49,62 @@ pub async fn serve(opts: Opts) -> Result<(), Error> { pub async fn main() -> ExitCode { // Parse the command-line options, exiting if there are any errors let opts = Opts::parse(); - install_tracing_subscriber(&opts); - if opts.run_mode() { - match run_wasm_main(opts).await { - Ok(_) => ExitCode::SUCCESS, - Err(_) => ExitCode::FAILURE, + let cmd = opts.command.unwrap_or(Commands::Serve(opts.serve)); + match cmd { + Commands::Run(run_args) => { + install_tracing_subscriber(0); + if let Ok(_) = run_wasm_main(run_args).await { + ExitCode::SUCCESS + } else { + ExitCode::FAILURE + } } - } else { - match { - tokio::select! { - _ = tokio::signal::ctrl_c() => { - Ok(()) - } - res = serve(opts) => { - if let Err(ref e) = res { - event!(Level::ERROR, "{}", e); + Commands::Serve(serve_args) => { + install_tracing_subscriber(serve_args.verbosity()); + match { + tokio::select! { + _ = tokio::signal::ctrl_c() => { + Ok(()) + } + res = serve(serve_args) => { + if let Err(ref e) = res { + event!(Level::ERROR, "{}", e); + } + res } - res } + } { + Ok(_) => ExitCode::SUCCESS, + Err(_) => ExitCode::FAILURE, } - } { - Ok(_) => ExitCode::SUCCESS, - Err(_) => ExitCode::FAILURE, } } } /// Execute a Wasm program in the Viceroy environment. -pub async fn run_wasm_main(opts: Opts) -> Result<(), anyhow::Error> { +pub async fn run_wasm_main(run_args: RunArgs) -> Result<(), anyhow::Error> { // Load the wasm module into an execution context - let ctx = create_execution_context(&opts).await?; - let program_name = match opts.input().file_stem() { + let ctx = create_execution_context(&run_args.shared(), false).await?; + let input = run_args.shared().input(); + let program_name = match input.file_stem() { Some(stem) => stem.to_string_lossy(), None => panic!("program cannot be a directory"), }; - ctx.run_main(&program_name, opts.wasm_args()).await + ctx.run_main(&program_name, run_args.wasm_args()).await } -fn install_tracing_subscriber(opts: &Opts) { +fn install_tracing_subscriber(verbosity: u8) { // Default to whatever a user provides, but if not set logging to work for // viceroy and viceroy-lib so that they can have output in the terminal if env::var("RUST_LOG").ok().is_none() { - match opts.verbosity() { - 0 => env::set_var("RUST_LOG", "viceroy=info,viceroy-lib=info"), - 1 => { - env::set_var("RUST_LOG", "viceroy=debug,viceroy-lib=debug"); - } - _ => { - env::set_var("RUST_LOG", "viceroy=trace,viceroy-lib=trace"); - } + match verbosity { + 0 => env::set_var("RUST_LOG", "viceroy=off,viceroy-lib=off"), + 1 => env::set_var("RUST_LOG", "viceroy=info,viceroy-lib=info"), + 2 => env::set_var("RUST_LOG", "viceroy=debug,viceroy-lib=debug"), + _ => env::set_var("RUST_LOG", "viceroy=trace,viceroy-lib=trace"), } } - // If the quiet flag is passed in, don't log anything (this should maybe - // just be a verbosity setting) - if opts.quiet() { - env::set_var("RUST_LOG", "viceroy=off,viceroy-lib=off"); - } + // Build a subscriber, using the default `RUST_LOG` environment variable for our filter. let builder = FmtSubscriber::builder() .with_writer(StdWriter::new()) @@ -185,12 +185,16 @@ impl<'a> MakeWriter<'a> for StdWriter { } } -async fn create_execution_context(opts: &Opts) -> Result { - let mut ctx = ExecuteCtx::new(opts.input(), opts.profiling_strategy(), opts.wasi_modules())? - .with_log_stderr(opts.log_stderr()) - .with_log_stdout(opts.log_stdout()); +async fn create_execution_context( + args: &SharedArgs, + check_backends: bool, +) -> Result { + let input = args.input(); + let mut ctx = ExecuteCtx::new(input, args.profiling_strategy(), args.wasi_modules())? + .with_log_stderr(args.log_stderr()) + .with_log_stdout(args.log_stdout()); - if let Some(config_path) = opts.config_path() { + if let Some(config_path) = args.config_path() { let config = FastlyConfig::from_file(config_path)?; let backends = config.backends(); let geolocation = config.geolocation(); @@ -214,7 +218,7 @@ async fn create_execution_context(opts: &Opts) -> Result, + + #[command(flatten)] + pub serve: ServeArgs, +} + +#[derive(Subcommand, Debug, Clone)] +pub enum Commands { + /// Run the wasm in a Viceroy server. This is the default if no subcommand + /// is given. + Serve(ServeArgs), + + /// Run the input wasm once and then exit. + Run(RunArgs), +} + +#[derive(Debug, Args, Clone)] +pub struct ServeArgs { /// The IP address that the service should be bound to. #[arg(long = "addr")] socket_addr: Option, + + /// Verbosity of logs for Viceroy. `-v` sets the log level to DEBUG and + /// `-vv` to TRACE. This option will not take effect if you set RUST_LOG + /// to a value before starting Viceroy + #[arg(short = 'v', action = clap::ArgAction::Count)] + verbosity: u8, + + #[command(flatten)] + shared: SharedArgs, +} + +#[derive(Args, Debug, Clone)] +pub struct RunArgs { + #[command(flatten)] + shared: SharedArgs, + + /// Args to pass along to the binary being executed. + #[arg(trailing_var_arg = true, allow_hyphen_values = true)] + wasm_args: Vec, +} + +#[derive(Args, Debug, Clone)] +pub struct SharedArgs { /// The path to the service's Wasm module. - #[arg(value_parser = check_module)] - input: PathBuf, + #[arg(value_parser = check_module, required=true)] + input: Option, /// The path to a TOML file containing `local_server` configuration. #[arg(short = 'C', long = "config")] config_path: Option, - /// [EXPERIMENTAL] Use Viceroy to run a module's _start function once, - /// rather than in a web server loop. This is experimental and the specific - /// interface for this is going to change in the near future. - #[arg(short = 'r', long = "run", default_value = "false", hide = true)] - run_mode: bool, /// Whether to treat stdout as a logging endpoint #[arg(long = "log-stdout", default_value = "false")] log_stdout: bool, /// Whether to treat stderr as a logging endpoint #[arg(long = "log-stderr", default_value = "false")] log_stderr: bool, - /// Verbosity of logs for Viceroy. `-v` sets the log level to DEBUG and - /// `-vv` to TRACE. This option will not take effect if you set RUST_LOG - /// to a value before starting Viceroy - #[arg(short = 'v', action = clap::ArgAction::Count)] - verbosity: u8, // Whether to enable wasmtime's builtin profiler. #[arg(long = "profiler", value_parser = check_wasmtime_profiler_mode)] profiler: Option, /// Set of experimental WASI modules to link against. #[arg(value_enum, long = "experimental_modules", required = false)] experimental_modules: Vec, - /// Don't log viceroy events to stdout or stderr - #[arg(short = 'q', long = "quiet", default_value = "false")] - quiet: bool, - /// [EXPERIMENTAL] Args to pass along to the binary being executed. This is - /// only used when run_mode=true - #[arg(trailing_var_arg = true, allow_hyphen_values = true, hide = true)] - wasm_args: Vec, } -impl Opts { +impl ServeArgs { /// The address that the service should be bound to. pub fn addr(&self) -> SocketAddr { self.socket_addr .unwrap_or_else(|| SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 7878)) } + /// Verbosity of logs for Viceroy. `-v` sets the log level to DEBUG and + /// `-vv` to TRACE. This option will not take effect if you set RUST_LOG + /// to a value before starting Viceroy + pub fn verbosity(&self) -> u8 { + self.verbosity + } + + pub fn shared(&self) -> &SharedArgs { + &self.shared + } +} + +impl RunArgs { + /// The arguments to pass to the underlying binary when run_mode=true + pub fn wasm_args(&self) -> &Vec { + &self.wasm_args + } + + pub fn shared(&self) -> &SharedArgs { + &self.shared + } +} + +impl SharedArgs { /// The path to the service's Wasm binary. - pub fn input(&self) -> &Path { - self.input.as_ref() + pub fn input(&self) -> PathBuf { + PathBuf::from(self.input.as_ref().unwrap()) } /// The path to a `local_server` configuration file. @@ -78,11 +129,6 @@ impl Opts { self.config_path.as_deref() } - /// Whether Viceroy should run the input once and then exit - pub fn run_mode(&self) -> bool { - self.run_mode - } - /// Whether to treat stdout as a logging endpoint pub fn log_stdout(&self) -> bool { self.log_stdout @@ -93,29 +139,11 @@ impl Opts { self.log_stderr } - /// Verbosity of logs for Viceroy. `-v` sets the log level to DEBUG and - /// `-vv` to TRACE. This option will not take effect if you set RUST_LOG - /// to a value before starting Viceroy - pub fn verbosity(&self) -> u8 { - self.verbosity - } - // Whether to enable wasmtime's builtin profiler. pub fn profiling_strategy(&self) -> ProfilingStrategy { self.profiler.unwrap_or(ProfilingStrategy::None) } - /// The arguments to pass to the underlying binary when run_mode=true - pub fn wasm_args(&self) -> &[String] { - self.wasm_args.as_ref() - } - - /// Prevents Viceroy from logging to stdout and stderr (note: any logs - /// emitted by the INPUT program will still go to stdout/stderr) - pub fn quiet(&self) -> bool { - self.quiet - } - // Set of experimental wasi modules to link against. pub fn wasi_modules(&self) -> HashSet { self.experimental_modules.iter().map(|x| x.into()).collect() @@ -164,11 +192,11 @@ impl From<&ExperimentalModule> for ExperimentalModuleArg { /// binary or text format. /// /// [opts]: struct.Opts.html -fn check_module(s: &str) -> Result { +fn check_module(s: &str) -> Result { let path = PathBuf::from(s); let contents = std::fs::read(&path)?; match wat::parse_bytes(&contents) { - Ok(_) => Ok(path), + Ok(_) => Ok(s.to_string()), _ => Err(Error::FileFormat), } } @@ -194,7 +222,7 @@ fn check_wasmtime_profiler_mode(s: &str) -> Result { #[cfg(test)] mod opts_tests { use { - super::Opts, + super::{Commands, Opts}, clap::{error::ErrorKind, Parser}, std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}, std::path::PathBuf, @@ -217,8 +245,11 @@ mod opts_tests { fn default_addr_works() -> TestResult { let empty_args = &["dummy-program-name", &test_file("minimal.wat")]; let opts = Opts::try_parse_from(empty_args)?; + let cmd = opts.command.unwrap_or(Commands::Serve(opts.serve)); let expected = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 7878); - assert_eq!(opts.addr(), expected); + if let Commands::Serve(serve_args) = cmd { + assert_eq!(serve_args.addr(), expected); + } Ok(()) } @@ -253,9 +284,12 @@ mod opts_tests { &test_file("minimal.wat"), ]; let opts = Opts::try_parse_from(args_with_ipv6_addr)?; + let cmd = opts.command.unwrap_or(Commands::Serve(opts.serve)); let addr_v6 = IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)); let expected = SocketAddr::new(addr_v6, 7878); - assert_eq!(opts.addr(), expected); + if let Commands::Serve(serve_args) = cmd { + assert_eq!(serve_args.addr(), expected); + } Ok(()) } @@ -358,16 +392,23 @@ mod opts_tests { /// Test that trailing arguments are collected successfully #[test] - fn trailing_args_are_collected() -> TestResult { + fn trailing_args_are_collected_in_run_mode() -> TestResult { let args = &[ "dummy-program-name", + "run", &test_file("minimal.wat"), "--", "--trailing-arg", "--trailing-arg-2", ]; let opts = Opts::try_parse_from(args)?; - assert_eq!(opts.wasm_args(), &["--trailing-arg", "--trailing-arg-2"]); + let cmd = opts.command.unwrap_or(Commands::Serve(opts.serve)); + if let Commands::Run(run_args) = cmd { + assert_eq!( + run_args.wasm_args(), + &["--trailing-arg", "--trailing-arg-2"] + ); + } Ok(()) } @@ -377,6 +418,7 @@ mod opts_tests { fn input_accepted_after_double_dash() -> TestResult { let args = &[ "dummy-program-name", + "run", "--", &test_file("minimal.wat"), "--trailing-arg", @@ -386,8 +428,17 @@ mod opts_tests { Ok(opts) => opts, res => panic!("unexpected result: {:?}", res), }; - assert_eq!(opts.input().to_str().unwrap(), &test_file("minimal.wat")); - assert_eq!(opts.wasm_args(), &["--trailing-arg", "--trailing-arg-2"]); + let cmd = opts.command.unwrap_or(Commands::Serve(opts.serve)); + if let Commands::Run(run_args) = cmd { + assert_eq!( + run_args.shared.input().to_str().unwrap(), + &test_file("minimal.wat") + ); + assert_eq!( + run_args.wasm_args(), + &["--trailing-arg", "--trailing-arg-2"] + ); + } Ok(()) } }