Override remote env #25 (#134)

* Adding remote env var loading.

* Loading env vars from nginx.

* Loading env vars from remote proc works, but it doesn't load the correct string?

* Notes on how to proceed (nginx out).

* Working override env vars.

* Fix wrong pid value being used.

* Clean files.

* Use correct imagePullPolicy.

* Simplify env vars handling. Make it select all then filter.

* Fix launch.json env.

* Remove unused rust features.

* Put env vars override behind an enabled bool flag.

* Add configuration for setting up override env vars + filter.

* Refactor some layer initialization functions to take config, instead of select arguments. Changed name for env vars request/response messages.

* Update changelog. Renamed env vars config. Env vars overriding now filters out some keys by default, like PATH (user cant control these defaults).

* Fix comment.

* local config.

* Added config to change image pull policy from env.

* Fix clap deprecations.

* Remote env changes.

* Use clap 3, instead of specifying subversion.

* Update cargo.lock.

* Use std::env::set_var instead of libc::setenv (plus assert the var is properly set).

* Increase initial memory allocation for string buffer that receives env vars.

* Change mirrord to custom image for tests.

* Add remote env vars test.

* Make test app crash on incorrect env vars.

* Ignore unrelated tests to make CI go faster.

* Env vars now filter + select properly.

* Fix include string.

* Fix env vars test args.

* Added more tests for remote_env. Don't allow user to specify both include + exclude. Removed flag to enable/disable env vars, now controller by having values in the filters.

* Fix launch.json.

* Use proper test image.

* Fix remote_env tests using the wrong option.

* Fix comments.

* Support include * as an include all env vars.

* Log result of tests process.

* Failure test should panic.

* Less clutter in console for remote env tests.

* Debug why test that should panic (and panics), doesn't trigger the test panic (wow what a phrase).

* Return exit code -1 if we have an env var selected.

* Fix test for panicking remote_env. Re-enable other tests.

CHANGELOG.md

@@ -7,7 +7,12 @@ Check [Keep a Changelog](http://keepachangelog.com/) for recommendations on how
 
 ## [Unreleased]
 
+### Added
+
+- Added support for overriding a process' environment variables by setting `MIRRORD_OVERRIDE_ENV_VARS` to `true`. To filter out undesired variables, use the `MIRRORD_OVERRIDE_FILTER_ENV_VARS` configuration with arguments such as `FOO;BAR`.
+
 ### Changed
+
 - Removed `unwrap` from the `Future` that was waiting for Kube pod to spin up in `pod_api.rs`. (Fixes #110)
 - Speed up agent container image building by using a more specific base image.
 - CI: Remove building agent before building & running tests (duplicate)
@@ -23,17 +28,22 @@ Check [Keep a Changelog](http://keepachangelog.com/) for recommendations on how
 - Enable the blocking feature of the `reqwest` library 
 
 ## 2.2.1
+
 ### Changed
+
 - Compile universal binaries for MacOS. (Fixes #131)
 - E2E small improvements, removing sleeps. (Fixes #99)
 
 ## 2.2.0
+
 ### Added
+
 - File operations are now available behind the `MIRRORD_FILE_OPS` env variable, this means that mirrord now hooks into the following file functions: `open`, `fopen`, `fdopen`, `openat`, `read`, `fread`, `fileno`, `lseek`, and `write` to provide a mirrored file system.
-- Support for running x64 (Intel) binary on arm (Silicon) macOS using mirrord. This will download and use the x64 mirrord-layer binary when needed.   
+- Support for running x64 (Intel) binary on arm (Silicon) macOS using mirrord. This will download and use the x64 mirrord-layer binary when needed.
 - Add detours for fcntl/dup system calls, closes [#51](https://github.com/metalbear-co/mirrord/issues/51)
 
 ### Changed
+
 - Add graceful exit for library extraction logic in case of error.
 - Refactor the CI by splitting the building of mirrord-agent in a separate job and caching the agent image for E2E tests.
 - Update bug report template to apply to the latest version of mirrord.
@@ -47,13 +57,15 @@ Check [Keep a Changelog](http://keepachangelog.com/) for recommendations on how
 - Fix typos
 
 ## 2.1.0
+
 ### Added
 
 - Prompt user to update if their version is outdated in the VS Code extension or CLI.
 - Add support for docker runtime, closes [#95](https://github.com/metalbear-co/mirrord/issues/95).
 - Add a keep-alive to keep the agent-pod from exiting, closes [#63](https://github.com/metalbear-co/mirrord/issues/63)
 
 ## 2.0.4
+
 Complete refactor and re-write of everything.
 
 - The CLI/VSCode extension now use `mirrord-layer` which loads into debugged process using `LD_PRELOAD`/`DYLD_INSERT_LIBRARIES`.

mirrord-agent/src/main.rs

@@ -1,23 +1,26 @@
 #![feature(result_option_inspect)]
-#![feature(never_type)]
+#![feature(hash_drain_filter)]
 
 use std::{
     borrow::Borrow,
-    collections::HashSet,
+    collections::{HashMap, HashSet},
     hash::{Hash, Hasher},
     net::{Ipv4Addr, SocketAddrV4},
+    path::PathBuf,
 };
 
 use error::AgentError;
 use futures::SinkExt;
 use mirrord_protocol::{
     tcp::{DaemonTcp, LayerTcp},
-    ClientMessage, ConnectionID, DaemonCodec, DaemonMessage, FileRequest, FileResponse, Port,
+    ClientMessage, ConnectionID, DaemonCodec, DaemonMessage, FileError, FileRequest, FileResponse,
+    GetEnvVarsRequest, Port, ResponseError,
 };
 use tokio::{
+    io::AsyncReadExt,
     net::{TcpListener, TcpStream},
     select,
-    sync::mpsc::{self},
+    sync::mpsc,
 };
 use tokio_stream::StreamExt;
 use tracing::{debug, error, info, trace};
@@ -34,7 +37,7 @@ use cli::parse_args;
 use sniffer::{packet_worker, SnifferCommand, SnifferOutput};
 use util::{IndexAllocator, Subscriptions};
 
-use crate::file::file_worker;
+use crate::{file::file_worker, runtime::get_container_pid, sniffer::DEFAULT_RUNTIME};
 
 type PeerID = u32;
 
@@ -105,13 +108,95 @@ pub struct PeerMessage {
     peer_id: PeerID,
 }
 
+/// Helper function that loads the process' environment variables, and selects only those that were
+/// requested from `mirrord-layer` (ignores vars specified in `filter_env_vars`).
+///
+/// Returns an error if none of the requested environment variables were found.
+async fn select_env_vars(
+    environ_path: PathBuf,
+    filter_env_vars: HashSet<String>,
+    select_env_vars: HashSet<String>,
+) -> Result<HashMap<String, String>, ResponseError> {
+    debug!(
+        "select_env_vars -> environ_path {:#?} filter_env_vars {:#?} select_env_vars {:#?}",
+        environ_path, filter_env_vars, select_env_vars
+    );
+
+    let mut environ_file = tokio::fs::File::open(environ_path).await.map_err(|fail| {
+        ResponseError::FileOperation(FileError {
+            operation: "open".to_string(),
+            raw_os_error: fail.raw_os_error(),
+            kind: fail.kind().into(),
+        })
+    })?;
+
+    let mut raw_env_vars = String::with_capacity(8192);
+
+    // TODO: nginx doesn't play nice when we do this, it only returns a string that goes like
+    // "nginx -g daemon off;".
+    let read_amount = environ_file
+        .read_to_string(&mut raw_env_vars)
+        .await
+        .map_err(|fail| {
+            ResponseError::FileOperation(FileError {
+                operation: "read_to_string".to_string(),
+                raw_os_error: fail.raw_os_error(),
+                kind: fail.kind().into(),
+            })
+        })?;
+    debug!(
+        "select_env_vars -> read {:#?} bytes with pure ENV_VARS {:#?}",
+        read_amount, raw_env_vars
+    );
+
+    // TODO: These are env vars that should usually be ignored. Revisit this list if a user
+    // ever asks for a way to NOT filter out these.
+    let mut default_filter = HashSet::with_capacity(2);
+    default_filter.insert("PATH".to_string());
+    default_filter.insert("HOME".to_string());
+
+    let env_vars = raw_env_vars
+        // "DB=foo.db\0PORT=99\0HOST=\0PATH=/fake\0"
+        .split_terminator(char::from(0))
+        // ["DB=foo.db", "PORT=99", "HOST=", "PATH=/fake"]
+        .map(|key_and_value| key_and_value.split_terminator('=').collect::<Vec<_>>())
+        // [["DB", "foo.db"], ["PORT", "99"], ["HOST"], ["PATH", "/fake"]]
+        .filter_map(
+            |mut keys_and_values| match (keys_and_values.pop(), keys_and_values.pop()) {
+                (Some(value), Some(key)) => Some((key.to_string(), value.to_string())),
+                _ => None,
+            },
+        )
+        .filter(|(key, _)| !default_filter.contains(key))
+        // [("DB", "foo.db"), ("PORT", "99"), ("PATH", "/fake")]
+        .filter(|(key, _)| !filter_env_vars.contains(key))
+        // [("DB", "foo.db"), ("PORT", "99")]
+        .filter(|(key, _)| {
+            select_env_vars.is_empty()
+                || select_env_vars.contains("*")
+                || select_env_vars.contains(key)
+        })
+        // [("DB", "foo.db")]
+        .collect::<HashMap<_, _>>();
+
+    debug!("select_env_vars -> selected env vars found {:?}", env_vars);
+
+    if env_vars.is_empty() {
+        Err(ResponseError::NotFound)
+    } else {
+        Ok(env_vars)
+    }
+}
+
 async fn handle_peer_messages(
     // TODO: Possibly refactor `state` out to be more "independent", and live in its own worker
     // thread.
     state: &mut State,
     sniffer_command_tx: mpsc::Sender<SnifferCommand>,
     file_request_tx: mpsc::Sender<(PeerID, FileRequest)>,
     peer_message: PeerMessage,
+    container_id: &Option<String>,
+    container_runtime: &Option<String>,
 ) -> Result<(), AgentError> {
     match peer_message.client_message {
         ClientMessage::Tcp(LayerTcp::PortUnsubscribe(port)) => {
@@ -170,6 +255,37 @@ async fn handle_peer_messages(
                 .send((peer_message.peer_id, file_request))
                 .await?;
         }
+        ClientMessage::GetEnvVarsRequest(GetEnvVarsRequest {
+            env_vars_filter,
+            env_vars_select,
+        }) => {
+            debug!(
+                "ClientMessage::GetEnvVarsRequest peer id {:?} filter {:?} select {:?}",
+                peer_message.peer_id, env_vars_filter, env_vars_select
+            );
+
+            let container_runtime = container_runtime
+                .as_ref()
+                .map(String::as_str)
+                .unwrap_or(DEFAULT_RUNTIME);
+
+            let pid = match container_id {
+                Some(container_id) => get_container_pid(container_id, container_runtime).await,
+                None => Err(AgentError::NotFound(format!(
+                    "handle_peer_messages -> Container ID not specified {:#?} for runtime {:#?}!",
+                    container_id, container_runtime
+                ))),
+            }?;
+
+            let environ_path = PathBuf::from("/proc").join(pid.to_string()).join("environ");
+            let env_vars_result =
+                select_env_vars(environ_path, env_vars_filter, env_vars_select).await;
+
+            let peer = state.peers.get(&peer_message.peer_id).unwrap();
+            peer.channel
+                .send(DaemonMessage::GetEnvVarsResponse(env_vars_result))
+                .await?;
+        }
     }
 
     Ok(())
@@ -202,7 +318,7 @@ async fn peer_handler(
                 }
             },
             message = daemon_messages_rx.recv() => {
-                debug!("peer_handler -> daemon_messages_rx.recv received a message {:?}", message);
+                debug!("peer_handler -> daemon_messages_rx.recv received a message");
 
                 match message {
                     Some(message) => {
@@ -284,7 +400,13 @@ async fn start_agent() -> Result<(), AgentError> {
 
             },
             Some(peer_message) = peer_messages_rx.recv() => {
-                handle_peer_messages(&mut state, sniffer_command_tx.clone(), file_request_tx.clone(), peer_message).await?;
+                handle_peer_messages(&mut state,
+                    sniffer_command_tx.clone(),
+                    file_request_tx.clone(),
+                    peer_message,
+                    &args.container_id,
+                    &args.container_runtime,
+                ).await?;
 
             },
             Some((peer_id, file_response)) = file_response_rx.recv() => {

mirrord-cli/src/config.rs

@@ -0,0 +1,68 @@
+use clap::{Args, Parser, Subcommand};
+
+#[derive(Parser)]
+#[clap(author, version, about, long_about = None)]
+pub(super) struct Cli {
+    #[clap(subcommand)]
+    pub(super) commands: Commands,
+}
+
+#[derive(Subcommand)]
+pub(super) enum Commands {
+    Exec(ExecArgs),
+    Extract {
+        #[clap(value_parser)]
+        path: String,
+    },
+}
+
+#[derive(Args, Debug)]
+pub(super) struct ExecArgs {
+    /// Pod name to mirror.
+    #[clap(short, long, value_parser)]
+    pub pod_name: String,
+
+    /// Namespace of the pod to mirror. Defaults to "default".
+    #[clap(short = 'n', long, value_parser)]
+    pub pod_namespace: Option<String>,
+
+    /// Namespace to place agent in.
+    #[clap(short = 'a', long, value_parser)]
+    pub agent_namespace: Option<String>,
+
+    /// Agent log level
+    #[clap(short = 'l', long, value_parser)]
+    pub agent_log_level: Option<String>,
+
+    /// Agent image
+    #[clap(short = 'i', long, value_parser)]
+    pub agent_image: Option<String>,
+
+    /// Enable file hooking
+    #[clap(short = 'f', long, value_parser)]
+    pub enable_fs: bool,
+
+    /// The env vars to filter out
+    #[clap(short = 'x', long, value_parser)]
+    pub override_env_vars_exclude: Option<String>,
+
+    /// The env vars to select
+    #[clap(short = 's', long, value_parser)]
+    pub override_env_vars_include: Option<String>,
+
+    /// Binary to execute and mirror traffic into.
+    #[clap(value_parser)]
+    pub binary: String,
+
+    /// Agent TTL
+    #[clap(long, value_parser)]
+    pub agent_ttl: Option<u16>,
+
+    /// Accept/reject invalid certificates.
+    #[clap(short = 'c', long, value_parser)]
+    pub accept_invalid_certificates: bool,
+
+    /// Arguments to pass to the binary.
+    #[clap(value_parser)]
+    pub(super) binary_args: Vec<String>,
+}