docker-ctl 0.2.5

//! Control a container

use std::{
  ffi::OsString,
  process::{Command, ExitStatus, Stdio},
};

use shared_child::SharedChild;

use crate::{Error, Result};

/// Network mode, refer to the [docker documentation](https://docs.docker.com/engine/network/)
/// for an explanation on the different modes.
pub enum Network
{
  /// The default network driver.
  Bridge,
  /// Remove network isolation between the container and the Docker host.
  Host,
  /// Completely isolate a container from the host and other containers.
  None,
}

/// Setup IPC via shared memory between container and host.
pub enum Ipc
{
  /// No IPC
  None,
  /// IPC with Host
  Host
}

//   ____             __ _                       _
//  / ___|___  _ __  / _(_) __ _ _   _ _ __ __ _| |_ ___  _ __
// | |   / _ \| '_ \| |_| |/ _` | | | | '__/ _` | __/ _ \| '__|
// | |__| (_) | | | |  _| | (_| | |_| | | | (_| | || (_) | |
//  \____\___/|_| |_|_| |_|\__, |\__,_|_|  \__,_|\__\___/|_|
//                         |___/

fn generate_name(prefix: String) -> String
{
  use rand::Rng;
  const CHARSET: &[u8] = b"abcdefghijklmnopqrstuvwxyz0123456789";
  let mut rng = rand::rng();
  let rand_string: String = (0..10)
    .map(|_| {
      let idx = rng.random_range(0..CHARSET.len());
      CHARSET[idx] as char
    })
    .collect();
  format!("{}-{}", prefix, rand_string)
}

/// Configuration of the container
pub struct Configurator
{
  image: String,
  capture_stdio: bool,
  interactive: bool,
  tty: bool,
  daemon: bool,
  privileged: bool,
  x11_forwarding: bool,
  network: Network,
  mounts: Vec<(String, String)>,
  env: Vec<(String, String)>,
  username: Option<String>,
  work_directory: Option<String>,
  clean_up: bool,
  command: Vec<String>,
  name: String,
  port_mappings: Vec<(u16, u16)>,
  ipc: Ipc,
  ipc_size: Option<String>
}

impl Configurator
{
  /// Create the container
  pub fn create(self) -> Container
  {
    Container {
      config: self,
      process: None,
    }
  }
  /// Set the command for the docker
  pub fn set_command(mut self, command: impl IntoIterator<Item = impl Into<String>>) -> Self
  {
    self.command = command
      .into_iter()
      .map(|x| -> String { x.into() })
      .collect();
    self
  }
  /// Capture stdin/err/out
  pub fn set_capture_stdio(mut self, capture: bool) -> Self
  {
    self.capture_stdio = capture;
    self
  }
  /// Set the docker in interactive mode.
  pub fn set_interactive(mut self, interactive: bool) -> Self
  {
    self.interactive = interactive;
    self
  }
  /// Set the docker in TTY mode.
  pub fn set_tty(mut self, tty: bool) -> Self
  {
    self.tty = tty;
    self
  }
  /// Enable/disable privileged (see [docker documentation](https://docs.docker.com/engine/containers/run/#runtime-privilege-and-linux-capabilities)).
  #[deprecated(since = "0.2.4", note = "please use `set_privileged` instead")]
  pub fn set_privilged(mut self, privileged: bool) -> Self
  {
    self.privileged = privileged;
    self
  }
  /// Enable/disable privileged (see [docker documentation](https://docs.docker.com/engine/containers/run/#runtime-privilege-and-linux-capabilities)).
  pub fn set_privileged(mut self, privileged: bool) -> Self
  {
    self.privileged = privileged;
    self
  }
  /// Forward X11 to the container
  pub fn set_x11_forwarding(mut self, x11_forwarding: bool) -> Self
  {
    self.x11_forwarding = x11_forwarding;
    self
  }
  /// Set the network mode
  pub fn set_network(mut self, network: Network) -> Self
  {
    self.network = network;
    self
  }
  /// Set the IPC mode. Size is optional, for instance "1g"
  pub fn set_ipc(mut self, ipc: Ipc, size: Option<&str>) -> Self
  {
    self.ipc = ipc;
    self.ipc_size = size.map(|x| x.to_string());
    self
  }
  /// Mount a drive from the host to the container
  pub fn mount(mut self, host_dir: impl Into<String>, container_dir: impl Into<String>) -> Self
  {
    self.mounts.push((host_dir.into(), container_dir.into()));
    self
  }
  /// Set the username in the container, this is required for X11 forwarding
  pub fn set_username(mut self, username: impl Into<String>) -> Self
  {
    self.username = Some(username.into());
    self
  }
  /// Set an environment variable
  pub fn set_env_variable(mut self, key: impl Into<String>, variable: impl Into<String>) -> Self
  {
    self.env.push((key.into(), variable.into()));
    self
  }
  /// Copy a variable from host environment. Silently fails if the variable is not available in host.
  pub fn copy_env_variable_from_host(mut self, key: impl Into<String>) -> Self
  {
    let key = key.into();
    if let Ok(v) = std::env::var(&key)
    {
      self.env.push((key, v));
    }
    self
  }
  /// Set if the container should be cleaned up after execution (set --rm)
  pub fn set_clean_up(mut self, clean_up: bool) -> Self
  {
    self.clean_up = clean_up;
    self
  }
  /// set the name of the container
  pub fn set_name(mut self, name: impl Into<String>) -> Self
  {
    self.name = name.into();
    self
  }
  /// generate a name with the given prefix
  pub fn generate_name(mut self, prefix: impl Into<String>) -> Self
  {
    self.name = generate_name(prefix.into());
    self
  }
  /// Map a host port to a container port
  pub fn map_port(mut self, host_port: u16, container_port: u16) -> Self
  {
    self.port_mappings.push((host_port, container_port));
    self
  }
  /// Set the work directory in the container
  pub fn work_directory(mut self, working_directory: impl Into<String>) -> Self
  {
    self.work_directory = Some(working_directory.into());
    self
  }
  /// Convenience function for configuring the container with a value that
  /// comes from a result.
  pub fn with_ok<T,E>(self, value: Result<T, E>, f: impl FnOnce(Self, T) -> Self) -> Self
  {
    if let Ok(v) = value
    {
      f(self, v)
    } else {
      self
    }
  }
  /// Convenience function for configuring the container with a value that
  /// comes from an option.
  pub fn with_some<T>(self, value: Option<T>, f: impl FnOnce(Self, T) -> Self) -> Self
  {
    if let Some(v) = value
    {
      f(self, v)
    } else {
      self
    }
  }
  /// Convenience function for configuring the container with a condition.
  /// If `cond` is true, the closure `f` is executed.
  pub fn cond(self, cond: bool, f: impl FnOnce(Self) -> Self) -> Self
  {
    if cond
    {
      f(self)
    }
    else
    {
      self
    }
  }
  /// Convenience function for configuring the container with cond.
  /// If `cond` is true, the closure `f` is executed, otherwise `g` is executed.
  pub fn cond_or(
    self,
    cond: bool,
    f: impl FnOnce(Self) -> Self,
    g: impl FnOnce(Self) -> Self,
  ) -> Self
  {
    if cond
    {
      f(self)
    }
    else
    {
      g(self)
    }
  }
  /// If the option is set call the function `f`. Otherwise do nothing and return self.
  pub fn unwrap_option<T>(self, value: &Option<T>, f: impl FnOnce(Self, &T) -> Self) -> Self
  {
    match value
    {
      Some(value) => f(self, value),
      None => self,
    }
  }
  /// If the option is set call the function `f`. Otherwise do nothing and return self.
  pub fn unwrap_option_or<T>(
    self,
    value: &Option<T>,
    f: impl FnOnce(Self, &T) -> Self,
    g: impl FnOnce(Self) -> Self,
  ) -> Self
  {
    match value
    {
      Some(value) => f(self, value),
      None => g(self),
    }
  }

  /// Pull the image.
  pub fn pull(&self) -> Result<()>
  {
    let mut cmd = Command::new("docker");
    cmd.args(["pull", self.image.as_str()]);

    // Spawn
    let mut r = cmd.spawn()?;
    r.wait()?;
    Ok(())
  }
}

//   ____            _        _
//  / ___|___  _ __ | |_ __ _(_)_ __   ___ _ __
// | |   / _ \| '_ \| __/ _` | | '_ \ / _ \ '__|
// | |__| (_) | | | | || (_| | | | | |  __/ |
//  \____\___/|_| |_|\__\__,_|_|_| |_|\___|_|

/// Handle to a container
pub struct Container
{
  config: Configurator,
  process: Option<SharedChild>,
}

macro_rules! config_to_arg {
  ($check:expr, $args:ident, $arg:expr) => {
    if $check
    {
      $args.push($arg.into());
    }
  };
}

impl Container
{
  /// Call this function to configure a new container
  pub fn configure(image: impl Into<String>) -> Configurator
  {
    Configurator {
      image: image.into(),
      capture_stdio: false,
      interactive: false,
      tty: false,
      daemon: false,
      privileged: false,
      x11_forwarding: false,
      network: Network::Bridge,
      mounts: Default::default(),
      env: Default::default(),
      username: None,
      clean_up: true,
      name: generate_name("unknown".to_string()),
      command: Default::default(),
      port_mappings: Default::default(),
      work_directory: None,
      ipc: Ipc::None,
      ipc_size: None
    }
  }

  fn create_args(&mut self) -> Result<Vec<OsString>>
  {
    let mut args = Vec::<OsString>::new();
    args.push("run".into());
    args.push("--name".into());
    args.push(self.config.name.to_owned().into());

    config_to_arg!(self.config.daemon, args, "-d");
    config_to_arg!(self.config.tty, args, "-t");
    config_to_arg!(self.config.interactive, args, "-i");
    config_to_arg!(self.config.privileged, args, "--privileged");
    config_to_arg!(self.config.clean_up, args, "--rm");

    args.push("--network".into());
    args.push(
      match self.config.network
      {
        Network::Bridge => "bridge",
        Network::Host => "host",
        Network::None => "none",
      }
      .into(),
    );

    match self.config.ipc {
      Ipc::None=> {},
      Ipc::Host=> {
        args.push("--ipc".into());
        args.push("host".into());
        if let Some(ipc_size) = &self.config.ipc_size
        {
          args.push("--shm-size".into());
          args.push(ipc_size.into());
        }
      }
    }

    for (host_port, container_port) in self.config.port_mappings.iter()
    {
      args.push("-p".into());
      args.push(format!("{}:{}", host_port, container_port).into());
    }

    if self.config.x11_forwarding
    {
      args.push("-e".into());
      args.push(format!("DISPLAY={}", std::env::var("DISPLAY")?).into());
      args.push("-v".into());
      args.push(
        format!(
          "/home/{}/.Xauthority:/home/{}/.Xauthority",
          whoami::username(),
          self
            .config
            .username
            .as_ref()
            .ok_or(Error::MissingUsername)?
        )
        .into(),
      );
      args.push("-v".into());
      args.push("/tmp/.X11-unix:/tmp/.X11-unix".into());
      args.push("-h".into());
      args.push(hostname::get()?)
    }

    for (host_dir, container_dir) in self.config.mounts.iter()
    {
      args.push("-v".into());
      args.push(format!("{}:{}", host_dir, container_dir).into());
    }

    // Add environment
    for (env_name, env_value) in self.config.env.iter()
    {
      args.push("-e".into());
      args.push(format!("{}={}", env_name, env_value).into());
    }

    // Add the working directory
    if let Some(working_directory) = self.config.work_directory.as_ref()
    {
      args.push("-w".into());
      args.push(working_directory.into());
    }

    // Set the image
    args.push(self.config.image.to_owned().into());

    // Set the command arguments
    let mut command = self
      .config
      .command
      .iter()
      .map(|x| -> OsString { x.into() })
      .collect();
    args.append(&mut command);
    Ok(args)
  }

  /// Command line used to start the container
  pub fn command(&mut self) -> Result<Vec<OsString>>
  {
    let mut args = self.create_args()?;
    let mut r = Vec::<OsString>::new();
    r.push("docker".into());
    r.append(&mut args);
    Ok(r)
  }

  /// Call this function to start the container.
  pub fn start(&mut self) -> Result<()>
  {
    if self.is_running()?
    {
      return Err(Error::ContainerRunning);
    }

    // Prepare the process
    let mut cmd = Command::new("docker");
    cmd.args(self.create_args()?);

    if self.config.capture_stdio
    {
      cmd.stdin(Stdio::piped());
      cmd.stdout(Stdio::piped());
      cmd.stderr(Stdio::piped());
    }

    // Spawn
    self.process = Some(SharedChild::spawn(&mut cmd)?);
    Ok(())
  }
  /// Call this function to wait on the container
  pub fn wait(&self) -> Result<ExitStatus>
  {
    if let Some(process) = &self.process
    {
      Ok(process.wait()?)
    }
    else
    {
      Err(Error::ContainerNotRunning)
    }
  }
  /// Call this function to stop the container
  pub fn stop(&self) -> Result<()>
  {
    if self.process.is_some()
    {
      Command::new("docker")
        .args(["kill", self.config.name.to_owned().as_str()])
        .output()?;
      Ok(())
    }
    else
    {
      Err(Error::ContainerNotRunning)
    }
  }
  /// Check if running
  pub fn is_running(&self) -> Result<bool>
  {
    if let Some(p) = &self.process
      && p.try_wait()?.is_none()
      {
        return Ok(true);
      }
    Ok(false)
  }
  /// Get stdin
  pub fn take_stdin(&self) -> Result<std::process::ChildStdin>
  {
    self
        .process
        .as_ref()
        .ok_or(Error::ContainerNotRunning)?
        .take_stdin()
        .ok_or(Error::StdIONotPiped)
  }
  /// Get stdout
  pub fn take_stdout(&self) -> Result<std::process::ChildStdout>
  {
    self
        .process
        .as_ref()
        .ok_or(Error::ContainerNotRunning)?
        .take_stdout()
        .ok_or(Error::StdIONotPiped)
  }
  /// Get stderr
  pub fn take_stderr(&self) -> Result<std::process::ChildStderr>
  {
    self
        .process
        .as_ref()
        .ok_or(Error::ContainerNotRunning)?
        .take_stderr()
        .ok_or(Error::StdIONotPiped)
  }
}

impl Drop for Container
{
  fn drop(&mut self)
  {
    if self
      .is_running()
      .expect("to successfully query for running")
    {
      self.stop().expect("to successfully stop");
      self.wait().expect("to successfully wait");
    }
  }
}

#[cfg(test)]
mod tests
{
  use std::io::Read;

  use super::*;

  #[test]
  fn start_wait_stop_containers()
  {
    let mut container = Container::configure("alpine")
      .set_command(["sleep", "1"])
      .create();
    // Try to start and stop container
    container.start().expect("To start.");
    assert!(container.is_running().unwrap());
    container.stop().expect("To stop.");
    let wait_status = container.wait().unwrap();
    assert!(wait_status.success());

    assert!(!container.is_running().unwrap());

    // Run container and wait
    container.start().expect("To start.");
    assert!(container.is_running().unwrap());
    let wait_status = container.wait().unwrap();
    assert!(wait_status.success());
    assert_eq!(wait_status.code().unwrap(), 0);
    assert!(!container.is_running().unwrap());
  }
  #[test]
  fn interactive_containers()
  {
    use std::io::Write;

    let mut container = Container::configure("alpine")
      .set_interactive(true)
      .set_capture_stdio(true)
      .create();
    container.start().unwrap();
    let mut stdin = container.take_stdin().unwrap();
    stdin.write_all(b"echo Hello World!").unwrap();
    drop(stdin);
    let mut buf = vec![];
    container
      .take_stdout()
      .unwrap()
      .read_to_end(&mut buf)
      .unwrap();
    assert_eq!(buf, b"Hello World!\n");
  }
}