libshpool 0.11.0

// Copyright 2023 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{
    collections::hash_map::DefaultHasher,
    env, fs,
    hash::{Hash, Hasher},
    io,
    path::PathBuf,
};

use anyhow::{anyhow, Context};
use clap::{Parser, Subcommand};
pub use hooks::Hooks;
use parking_lot::{Mutex, MutexGuard};
use tracing::error;
use tracing_subscriber::{fmt::format::FmtSpan, prelude::*};

mod attach;
mod common;
mod config;
mod config_watcher;
mod consts;
mod daemon;
mod daemonize;
mod detach;
mod duration;
mod events;
mod exe;
mod hooks;
mod kill;
mod list;
mod protocol;
mod session_restore;
mod set_log_level;
mod template;
mod test_hooks;
mod tty;
mod user;
mod var;

/// The command line arguments that shpool expects.
/// These can be directly parsed with clap or manually
/// constructed in order to present some other user
/// interface.
///
/// NOTE: You must check `version()` and handle it yourself
/// if it is set. Clap won't do a good job with its
/// automatic version support for a library.
#[derive(Parser, Debug, Default)]
#[clap(author, about)]
pub struct Args {
    #[clap(
        short,
        long,
        action,
        long_help = "The file to write logs to

In most modes logs are discarded by default, but if shpool is
running in daemon mode, the logs will go to stderr by default."
    )]
    pub log_file: Option<String>,

    #[clap(
        short,
        long,
        action = clap::ArgAction::Count,
        help = "Show more in logs, may be provided multiple times",
    )]
    pub verbose: u8,

    #[clap(
        short,
        long,
        action,
        long_help = "The path for the unix socket to listen on

This defaults to $XDG_RUNTIME_DIR/shpool/shpool.socket or ~/.local/run/shpool/shpool.socket
if XDG_RUNTIME_DIR is unset.

This flag gets overridden by systemd socket activation when
the daemon is launched by systemd."
    )]
    pub socket: Option<String>,

    #[clap(short, long, action, help = "a toml file containing configuration")]
    pub config_file: Option<String>,

    #[clap(short, long, action, help = "automatically launch a daemon if one is not running")]
    pub daemonize: bool,

    #[clap(short = 'D', long, action, help = "do not automatically launch a daemon")]
    pub no_daemonize: bool,

    #[clap(subcommand)]
    pub command: Commands,

    // A hidden field rather than using the #[non_exhaustive] attribute
    // allows users to build this struct using the default value plus
    // modifications, while the #[non_exhaustive] attribute would not.
    // See https://rust-lang.github.io/rfcs/2008-non-exhaustive.html#functional-record-updates
    // (the attribute behaves as if there is implicitly a field like this
    // that is private).
    #[doc(hidden)]
    pub __non_exhaustive: (),
}

/// The subcommds that shpool supports.
#[derive(Subcommand, Debug, Default)]
#[non_exhaustive]
pub enum Commands {
    #[clap(about = "Print version")]
    #[default]
    Version,

    #[clap(about = "Starts running a daemon that holds a pool of shells")]
    Daemon,

    #[clap(about = "Creates or attaches to an existing shell session")]
    #[non_exhaustive]
    Attach {
        #[clap(short, long, help = "If a tty is already attached to the session, detach it first")]
        force: bool,
        #[clap(
            short = 'b',
            long,
            help = "Create/attach the session and immediately detach (use with --force to detach any existing client first)"
        )]
        background: bool,
        #[clap(
            long,
            long_help = "Automatically kill the session after the given time

This option only applies when first creating a session, it is ignored on
reattach.

The duration can be specified either in a colon seperated format
of the form dd:hh:mm:ss where any prefix may be left off (i.e. '01:00:30:00'
for 1 day and 30 minutes or '10:45:00' for 10 hours and 45 minutes), or
using a number with a trailing letter to indicate time unit
(i.e. '3d', '19h', or '5s')."
        )]
        ttl: Option<String>,
        #[clap(
            short,
            long,
            long_help = "A command to run instead of the user's default shell

The command is broken up into a binary to invoke and a list of arguments to
pass to the binary using the shell-words crate."
        )]
        cmd: Option<String>,
        #[clap(
            short,
            long,
            long_help = "The directory to start the shell in.

$HOME by default. Use '.' for pwd."
        )]
        dir: Option<String>,
        #[clap(
            short,
            long,
            long_help = "A command to run in the shell immediately on startup

This differs from the --cmd flag in that it is run
directly in the shell rather than replacing the shell. Think of
it as running `source cmd` rather than `exec cmd`. The main
usecase is to be able to automatically enter some useful context
such as a particular directory with a python virtual environment
already set up for example.

--start-cmd supports templates in the same way that session names do,
so you can parameterize your command with shpool variables. For example,
you could do --start-cmd 'cd /workspace/path/prefix/{workspace}' to have
your session start in your workspace root depending on a 'workspace' var.
See 'shpool help var' for more details."
        )]
        start_cmd: Option<String>,
        #[clap(long_help = "The name of the shell session to create or attach to

Session names support templates, and if a session name changes based on a
variable change, 'shpool attach' will automatically hang up and re-dial the
new session name. See 'shpool help var' for more details.")]
        name: String,
    },

    #[clap(about = "Make the given session detach from shpool

This does not close the shell. If no session name is provided
$SHPOOL_SESSION_NAME will be used if it is present in the
environment.")]
    #[non_exhaustive]
    Detach {
        #[clap(help = "sessions to detach")]
        sessions: Vec<String>,
    },

    #[clap(about = "Subscribe to the daemon's push-event stream

Connects to the events socket and writes each event (one JSON object
per line) to stdout, flushing after every line so the stream is
pipeline-friendly (e.g. `shpool events | jq`). See EVENTS.md for
details.")]
    Events,

    #[clap(about = "Kill the given sessions

This detaches the session if it is attached and kills the underlying
shell with a SIGHUP followed by a SIGKILL if the shell fails to exit
quickly enough. If no session name is provided $SHPOOL_SESSION_NAME
will be used if it is present in the environment.")]
    #[non_exhaustive]
    Kill {
        #[clap(help = "sessions to kill")]
        sessions: Vec<String>,
    },

    #[clap(about = "lists all the running shell sessions")]
    #[non_exhaustive]
    List {
        #[clap(short, long, help = "Output as JSON, includes extra fields")]
        json: bool,
    },

    #[clap(about = "Dynamically change daemon log level

This command changes the log level of the shpool daemon without
restarting. It may to useful if the daemon gets into a state that
needs debugging, but would be clobbered by a restart.")]
    #[non_exhaustive]
    SetLogLevel {
        #[clap(help = "new log level")]
        level: shpool_protocol::LogLevel,
    },

    #[clap(about = "Manipulate template variables

shpool session names can include {variables} which are resolved via
an environment stored globally in the shpool daemon. This command
manipulates that environment.

The main usecase for templated session names is the ability to switch
multiple shpool sessions to new targets at the same time. For example,
you might have a `shpool attach -f '{workspace}-edit'` session and
a `shpool attach -f '{workspace}-term'` session. To switch both
sessions from the fun-feature workspace to the key-bugfix workspace,
you could just do `shpool var set workspace key-bugfix`.
")]
    #[non_exhaustive]
    Var {
        #[clap(subcommand)]
        command: VarCommands,
    },
}

/// The subcommds of the var command.
#[derive(Subcommand, Debug)]
#[non_exhaustive]
pub enum VarCommands {
    #[clap(about = "List the variables

This command dumps out the whole variable list with
both vars and values in a JSON object using vars as keys.")]
    List {
        #[clap(short, long, help = "Output as JSON")]
        json: bool,
    },
    #[clap(about = "Get a variable

This returns the raw value of the given variable.")]
    #[non_exhaustive]
    Get { var: String },
    #[clap(about = "Set a variable

This updates the value of the given variable.")]
    #[non_exhaustive]
    Set { var: String, val: String },
    #[clap(about = "Unset a variable

This removes the given variable from the environment.")]
    #[non_exhaustive]
    Unset { var: String },
}

impl Args {
    /// Version indicates if the wrapping binary must display the
    /// version then exit.
    pub fn version(&self) -> bool {
        matches!(self.command, Commands::Version)
    }
}

// Copied from the tracing-subscriber crate. This is public in
// a future version of the crate, but for now we don't have
// access to it. If tracing-subscriber is 0.3.19 or better,
// it is worth checking to see if we can rip this out.
#[derive(Debug)]
pub struct MutexGuardWriter<'a, W>(MutexGuard<'a, W>);
impl<W> io::Write for MutexGuardWriter<'_, W>
where
    W: io::Write,
{
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }

    #[inline]
    fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
        self.0.write_vectored(bufs)
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.0.write_all(buf)
    }

    #[inline]
    fn write_fmt(&mut self, fmt: std::fmt::Arguments<'_>) -> io::Result<()> {
        self.0.write_fmt(fmt)
    }
}

struct LogWriterBuilder {
    log_file: Option<Mutex<fs::File>>,
    is_daemon: bool,
}

impl<'writer> tracing_subscriber::fmt::MakeWriter<'writer> for LogWriterBuilder {
    type Writer = Box<dyn io::Write + 'writer>;

    fn make_writer(&'writer self) -> Self::Writer {
        if let Some(log_file) = &self.log_file {
            Box::new(MutexGuardWriter(log_file.lock()))
        } else if self.is_daemon {
            Box::new(io::stderr())
        } else {
            Box::new(io::empty())
        }
    }
}

/// Run the shpool tool with the given arguments. If hooks is provided,
/// inject the callbacks into the daemon.
///
/// # Safety
///
/// Callers MUST NOT call this routine after forking a thread.
/// internally shpool will double-fork to daemonize, and doing so with
/// multiple threads running is unsafe. Really this only applies when
/// consts::AUTODAEMONIZE_VAR is set to true in the environment, but
/// we still need to mark this entrypoint unsafe to be sound.
pub unsafe fn run(
    args: Args,
    hooks: Option<Box<dyn hooks::Hooks + Send + Sync>>,
) -> anyhow::Result<()> {
    match (&args.command, env::var(consts::SENTINEL_FLAG_VAR).as_deref()) {
        (Commands::Daemon, Ok("prompt")) => {
            println!("{}", consts::PROMPT_SENTINEL);
            std::process::exit(0);
        }
        (Commands::Daemon, Ok("startup")) => {
            println!("{}", consts::STARTUP_SENTINEL);
            std::process::exit(0);
        }
        _ => {}
    }

    let log_level_filter = if args.verbose == 0 {
        tracing_subscriber::filter::LevelFilter::INFO
    } else if args.verbose == 1 {
        tracing_subscriber::filter::LevelFilter::DEBUG
    } else {
        tracing_subscriber::filter::LevelFilter::TRACE
    };
    let (log_level_layer, log_level_handle) =
        tracing_subscriber::reload::Layer::new(log_level_filter);

    let log_writer_builder = LogWriterBuilder {
        log_file: if let Some(lf) = &args.log_file {
            Some(Mutex::new(fs::File::create(lf).context("unable to create log file")?))
        } else {
            None
        },
        is_daemon: matches!(args.command, Commands::Daemon),
    };
    tracing_subscriber::registry::Registry::default()
        .with(log_level_layer)
        .with(
            tracing_subscriber::fmt::layer()
                .with_ansi(false)
                .with_thread_ids(true)
                .with_target(false)
                .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
                .with_writer(log_writer_builder),
        )
        .init();

    let mut runtime_dir = match env::var("XDG_RUNTIME_DIR") {
        Ok(runtime_dir) => PathBuf::from(runtime_dir),
        Err(_) => PathBuf::from(env::var("HOME").context("no XDG_RUNTIME_DIR or HOME")?)
            .join(".local")
            .join("run"),
    }
    .join("shpool");
    fs::create_dir_all(&runtime_dir).context("ensuring runtime dir exists")?;

    let socket = match &args.socket {
        Some(s) => {
            // The user can reasonably expect that if they provide seperate
            // sockets for differnt shpool instances to run on, they won't
            // stomp on one another. To respect this expectation we need to
            // namespace the rest of the runtime data if they provide a socket
            // name. A short hash is probably good enough.
            let mut hasher = DefaultHasher::new();
            s.hash(&mut hasher);
            let hash = hasher.finish();
            runtime_dir = runtime_dir.join(format!("{hash:x}"));

            PathBuf::from(s)
        }
        None => runtime_dir.join("shpool.socket"),
    };

    let config_manager = config::Manager::new(args.config_file.as_deref())?;

    if !config_manager.get().nodaemonize.unwrap_or(false) || args.daemonize {
        let arg0 = env::args().next().ok_or(anyhow!("arg0 missing"))?;
        if !args.no_daemonize && !matches!(args.command, Commands::Daemon) {
            daemonize::maybe_launch_daemon(&config_manager, &args, arg0, &socket)?;
        }
    }

    #[cfg(feature = "test_hooks")]
    if let Ok(test_hook_sock) = std::env::var("SHPOOL_TEST_HOOK_SOCKET_PATH") {
        log::info!("spawning test hook sock at {}", test_hook_sock);
        test_hooks::TEST_HOOK_SERVER.set_socket_path(test_hook_sock.clone());
        std::thread::spawn(|| {
            test_hooks::TEST_HOOK_SERVER.start();
        });
        log::info!("waiting for test hook connection");
        test_hooks::TEST_HOOK_SERVER.wait_for_connect()?;
    }

    let res: anyhow::Result<()> = match args.command {
        Commands::Version => return Err(anyhow!("wrapper binary must handle version")),
        Commands::Daemon => daemon::run(
            config_manager,
            runtime_dir,
            hooks.unwrap_or(Box::new(NoopHooks {})),
            log_level_handle,
            socket,
        ),
        Commands::Attach { force, background, ttl, cmd, dir, start_cmd, name } => {
            attach::run(socket, config_manager, name, force, background, ttl, cmd, dir, start_cmd)
        }
        Commands::Detach { sessions } => detach::run(sessions, socket),
        Commands::Kill { sessions } => kill::run(sessions, socket),
        Commands::List { json } => list::run(socket, json),
        Commands::SetLogLevel { level } => set_log_level::run(level, socket),
        Commands::Var { command } => var::run(socket, command),
        Commands::Events => events::run(&crate::daemon::events::socket_path(&socket)),
    };

    if let Err(err) = res {
        error!("{:?}", err);
        std::process::exit(1);
    }

    Ok(())
}

struct NoopHooks {}
impl hooks::Hooks for NoopHooks {}