1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

mod client;
mod server;

pub use client::*;
pub use server::*;

use tokio::process::{Child, Command};

/// Use [`Proxy::spawn`] to create a new proxy server and handler process.
#[derive(Default)]
pub struct Proxy {
  /// See [`Self::port`].
  pub port: Option<u16>,
  /// See [`Self::command`].
  pub command: Option<Command>,
}

impl Proxy {
  /// Create the handler command from the `argv[1..]`.
  /// For example if the command of the current process is `proxy node --help`
  /// then the handler command will be `node --help`.
  /// You can modify the handler command and pass it to [`Self::command`].
  /// # Examples
  /// ```
  /// use lambda_runtime_proxy::Proxy;
  /// use std::process::Stdio;
  ///
  /// #[tokio::main]
  /// async fn main {
  ///   // retrieve the default handler command
  ///   let mut command = Proxy::default_command();
  ///
  ///   // enhance the handler command
  ///   command
  ///     // override environment variables for the handler process
  ///     .env("KEY", "VALUE")
  ///     // pipe the stdout and stderr of the handler process
  ///     .stdout(Stdio::piped())
  ///     .stderr(Stdio::piped());
  ///
  ///   Proxy::default()
  ///     .command(command)
  ///     .spawn().await;
  /// }
  /// ```
  pub fn default_command() -> Command {
    let mut cmd = Command::new(std::env::args().nth(1).expect("Missing handler command"));
    cmd.args(std::env::args().skip(2));
    cmd
  }

  /// Set the port of the proxy server.
  /// If not set, the port will be read from the environment variable `AWS_LAMBDA_RUNTIME_PROXY_PORT`,
  /// or default to `3000`.
  pub fn port(mut self, port: u16) -> Self {
    self.port = Some(port);
    self
  }

  /// Set the command of the handler process.
  /// If not set, the command will be created using [`Self::default_command`].
  pub fn command(mut self, cmd: Command) -> Self {
    self.command = Some(cmd);
    self
  }

  /// Spawn the proxy server and the handler process.
  /// The handler process will be spawned with the environment variable `AWS_LAMBDA_RUNTIME_API`
  /// set to the address of the proxy server.
  pub async fn spawn(self) -> RunningProxy {
    let port = self
      .port
      .or_else(|| {
        std::env::var("AWS_LAMBDA_RUNTIME_PROXY_PORT")
          .ok()
          .and_then(|s| s.parse().ok())
      })
      .unwrap_or(3000);

    let mut command = self.command.unwrap_or_else(|| Self::default_command());
    command.env("AWS_LAMBDA_RUNTIME_API", format!("127.0.0.1:{}", port));

    let server = MockLambdaRuntimeApiServer::bind(port).await;

    // server is ready, spawn the real handler process
    let handler = command.spawn().expect("Failed to spawn handler process");

    RunningProxy { server, handler }
  }
}

/// Created by [`Proxy::spawn`].
pub struct RunningProxy {
  pub server: MockLambdaRuntimeApiServer,
  /// The lambda handler process.
  pub handler: Child,
}