Load env from layer (#2237)

* POC

* Changelog entry

* Changelog entry updated

* Docs + moved already-fetched flag to layer entirely

* Updated config schema

changelog.d/+loading-env-from-inside-layer.added.md

@@ -0,0 +1 @@
+Added configuration option `feature.env.load_from_process`, which allows for changing the way mirrord loads environment variables from the remote target.

mirrord-schema.json

@@ -463,6 +463,14 @@
             }
           ]
         },
+        "load_from_process": {
+          "title": "feature.env.load_from_process {#feature-env-load_from_process}",
+          "description": "Allows for changing the way mirrord loads remote environment variables. If set, the variables are fetched after the user application is started.\n\nThis setting is meant to resolve issues when using mirrord via the IntelliJ plugin on WSL and the remote environment contains a lot of variables.",
+          "type": [
+            "boolean",
+            "null"
+          ]
+        },
         "override": {
           "title": "feature.env.override {#feature-env-override}",
           "description": "Allows setting or overriding environment variables (locally) with a custom value.\n\nFor example, if the remote pod has an environment variable `REGION=1`, but this is an undesirable value, it's possible to use `override` to set `REGION=2` (locally) instead.",

mirrord/cli/src/execution.rs

@@ -146,9 +146,13 @@ impl MirrordExecution {
             .await
             .inspect_err(|_| analytics.set_error(AnalyticsError::AgentConnection))?;
 
-        let mut env_vars = Self::fetch_env_vars(config, &mut connection)
-            .await
-            .inspect_err(|_| analytics.set_error(AnalyticsError::EnvFetch))?;
+        let mut env_vars = if config.feature.env.load_from_process.unwrap_or(false) {
+            Default::default()
+        } else {
+            Self::fetch_env_vars(config, &mut connection)
+                .await
+                .inspect_err(|_| analytics.set_error(AnalyticsError::EnvFetch))?
+        };
 
         let lib_path: String = lib_path.to_string_lossy().into();
         // Set LD_PRELOAD/DYLD_INSERT_LIBRARIES

mirrord/config/src/feature/env.rs

@@ -69,6 +69,15 @@ pub struct EnvConfig {
     /// For example, if the remote pod has an environment variable `REGION=1`, but this is an
     /// undesirable value, it's possible to use `override` to set `REGION=2` (locally) instead.
     pub r#override: Option<HashMap<String, String>>, // `r#`: `override` is a Rust keyword.
+
+    /// ### feature.env.load_from_process {#feature-env-load_from_process}
+    ///
+    /// Allows for changing the way mirrord loads remote environment variables.
+    /// If set, the variables are fetched after the user application is started.
+    ///
+    /// This setting is meant to resolve issues when using mirrord via the IntelliJ plugin on WSL
+    /// and the remote environment contains a lot of variables.
+    pub load_from_process: Option<bool>,
 }
 
 impl MirrordToggleableConfig for EnvFileConfig {
@@ -81,6 +90,7 @@ impl MirrordToggleableConfig for EnvFileConfig {
                 .source_value(context)
                 .transpose()?
                 .or_else(|| Some(VecOrSingle::Single("*".to_owned()))),
+            load_from_process: None,
             r#override: None,
         })
     }

mirrord/intproxy/protocol/src/lib.rs

@@ -3,6 +3,7 @@
 //! internal proxy and the layer are shipped together in a single binary.
 
 use std::{
+    collections::HashMap,
     fmt,
     net::{IpAddr, SocketAddr},
 };
@@ -20,7 +21,7 @@ use mirrord_protocol::{
     },
     outgoing::SocketAddress,
     tcp::StealType,
-    FileRequest, FileResponse, Port, RemoteResult,
+    FileRequest, FileResponse, GetEnvVarsRequest, Port, RemoteResult,
 };
 
 #[cfg(feature = "codec")]
@@ -55,6 +56,8 @@ pub enum LayerToProxyMessage {
     OutgoingConnect(OutgoingConnectRequest),
     /// Requests related to incoming connections.
     Incoming(IncomingRequest),
+    /// Fetch environment variables from the target.
+    GetEnv(GetEnvVarsRequest),
 }
 
 /// Unique `layer <-> proxy` session identifier.
@@ -206,6 +209,8 @@ pub enum ProxyToLayerMessage {
     OutgoingConnect(RemoteResult<OutgoingConnectResponse>),
     /// A response to layer's [`IncomingRequest`].
     Incoming(IncomingResponse),
+    /// A response to layer's [`LayerToProxyMessage::GetEnv`].
+    GetEnv(RemoteResult<HashMap<String, String>>),
 }
 
 /// A response to layer's [`IncomingRequest`].
@@ -396,3 +401,10 @@ impl_request!(
     req_path = LayerToProxyMessage::Incoming => IncomingRequest::ConnMetadata,
     res_path = ProxyToLayerMessage::Incoming => IncomingResponse::ConnMetadata,
 );
+
+impl_request!(
+    req = GetEnvVarsRequest,
+    res = RemoteResult<HashMap<String, String>>,
+    req_path = LayerToProxyMessage::GetEnv,
+    res_path = ProxyToLayerMessage::GetEnv,
+);

mirrord/intproxy/src/lib.rs

@@ -326,6 +326,12 @@ impl IntProxy {
                 LogLevel::Error => tracing::error!("agent log: {}", log.message),
                 LogLevel::Warn => tracing::warn!("agent log: {}", log.message),
             },
+            DaemonMessage::GetEnvVarsResponse(res) => {
+                self.task_txs
+                    .simple
+                    .send(SimpleProxyMessage::GetEnvRes(res))
+                    .await
+            }
             other => {
                 return Err(IntProxyError::UnexpectedAgentMessage(other));
             }
@@ -372,6 +378,12 @@ impl IntProxy {
                     ))
                     .await
             }
+            LayerToProxyMessage::GetEnv(req) => {
+                self.task_txs
+                    .simple
+                    .send(SimpleProxyMessage::GetEnvReq(message_id, layer_id, req))
+                    .await
+            }
             other => return Err(IntProxyError::UnexpectedLayerMessage(other)),
         }
 

mirrord/intproxy/src/proxies/simple.rs

@@ -1,11 +1,13 @@
 //! The most basic proxying logic. Handles cases when the only job to do in the internal proxy is to
 //! pass requests and responses between the layer and the agent.
 
+use std::collections::HashMap;
+
 use mirrord_intproxy_protocol::{LayerId, MessageId, ProxyToLayerMessage};
 use mirrord_protocol::{
     dns::{GetAddrInfoRequest, GetAddrInfoResponse},
     file::{CloseDirRequest, CloseFileRequest, OpenDirResponse, OpenFileResponse},
-    ClientMessage, FileRequest, FileResponse,
+    ClientMessage, FileRequest, FileResponse, GetEnvVarsRequest, RemoteResult,
 };
 
 use crate::{
@@ -23,6 +25,8 @@ pub enum SimpleProxyMessage {
     AddrInfoRes(GetAddrInfoResponse),
     LayerForked(LayerForked),
     LayerClosed(LayerClosed),
+    GetEnvReq(MessageId, LayerId, GetEnvVarsRequest),
+    GetEnvRes(RemoteResult<HashMap<String, String>>),
 }
 
 #[derive(Clone, Copy, PartialEq, Eq, Hash)]
@@ -41,6 +45,8 @@ pub struct SimpleProxy {
     file_reqs: RequestQueue,
     /// For [`GetAddrInfoRequest`]s.
     addr_info_reqs: RequestQueue,
+    /// For [`GetEnvVarsRequest`]s.
+    get_env_reqs: RequestQueue,
 }
 
 impl BackgroundTask for SimpleProxy {
@@ -158,6 +164,22 @@ impl BackgroundTask for SimpleProxy {
                 SimpleProxyMessage::LayerForked(LayerForked { child, parent }) => {
                     self.remote_fds.clone_all(parent, child);
                 }
+                SimpleProxyMessage::GetEnvReq(message_id, layer_id, req) => {
+                    self.get_env_reqs.insert(message_id, layer_id);
+                    message_bus
+                        .send(ProxyMessage::ToAgent(ClientMessage::GetEnvVarsRequest(req)))
+                        .await;
+                }
+                SimpleProxyMessage::GetEnvRes(res) => {
+                    let (message_id, layer_id) = self.get_env_reqs.get()?;
+                    message_bus
+                        .send(ToLayer {
+                            message_id,
+                            message: ProxyToLayerMessage::GetEnv(res),
+                            layer_id,
+                        })
+                        .await
+                }
             }
         }
 

mirrord/layer/src/lib.rs

@@ -65,7 +65,15 @@
 extern crate alloc;
 extern crate core;
 
-use std::{cmp::Ordering, ffi::OsString, net::SocketAddr, panic, sync::OnceLock, time::Duration};
+use std::{
+    cmp::Ordering,
+    collections::{HashMap, HashSet},
+    ffi::OsString,
+    net::SocketAddr,
+    panic,
+    sync::OnceLock,
+    time::Duration,
+};
 
 use ctor::ctor;
 use error::{LayerError, Result};
@@ -81,12 +89,16 @@ use mirrord_config::{
 };
 use mirrord_intproxy_protocol::NewSessionRequest;
 use mirrord_layer_macro::{hook_fn, hook_guard_fn};
+use mirrord_protocol::{EnvVars, GetEnvVarsRequest};
 use proxy_connection::ProxyConnection;
 use setup::LayerSetup;
 use socket::SOCKETS;
 use tracing_subscriber::{fmt::format::FmtSpan, prelude::*};
 
-use crate::{debugger_ports::DebuggerPorts, detour::DetourGuard, load::LoadType};
+use crate::{
+    common::make_proxy_request_with_response, debugger_ports::DebuggerPorts, detour::DetourGuard,
+    load::LoadType,
+};
 
 mod common;
 mod debugger_ports;
@@ -279,11 +291,14 @@ fn init_tracing() {
 ///
 /// 1. [`tracing_subscriber`] or [`mirrord_console`];
 ///
-/// 2. Global [`STATE`];
+/// 2. Global [`SETUP`];
 ///
 /// 3. Global [`PROXY_CONNECTION`];
 ///
 /// 4. Replaces the [`libc`] calls with our hooks with [`enable_hooks`];
+///
+/// 5. Fetches remote environment from the agent (if enabled with
+/// [`EnvFileConfig::load_from_process`](mirrord_config::feature::env::EnvFileConfig::load_from_process)).
 fn layer_start(mut config: LayerConfig) {
     if config.target.path.is_none() {
         // Use localwithoverrides on targetless regardless of user config.
@@ -345,6 +360,64 @@ fn layer_start(mut config: LayerConfig) {
             .set(new_connection)
             .expect("setting PROXY_CONNECTION singleton")
     }
+
+    let fetch_env = setup().env_config().load_from_process.unwrap_or(false)
+        && !std::env::var(REMOTE_ENV_FETCHED)
+            .unwrap_or_default()
+            .parse::<bool>()
+            .unwrap_or(false);
+    if fetch_env {
+        let env = fetch_env_vars();
+        for (key, value) in env {
+            std::env::set_var(key, value);
+        }
+
+        std::env::set_var(REMOTE_ENV_FETCHED, "true");
+    }
+}
+
+/// Name of environment variable used to mark whether remote environment has already been fetched.
+const REMOTE_ENV_FETCHED: &str = "MIRRORD_REMOTE_ENV_FETCHED";
+
+/// Fetches remote environment from the agent.
+/// Uses [`SETUP`] and [`PROXY_CONNECTION`] globals.
+fn fetch_env_vars() -> HashMap<String, String> {
+    let (env_vars_exclude, env_vars_include) = match (
+        setup()
+            .env_config()
+            .exclude
+            .clone()
+            .map(|exclude| exclude.join(";")),
+        setup()
+            .env_config()
+            .include
+            .clone()
+            .map(|include| include.join(";")),
+    ) {
+        (Some(..), Some(..)) => {
+            panic!("invalid env config");
+        }
+        (Some(exclude), None) => (HashSet::from(EnvVars(exclude)), HashSet::new()),
+        (None, Some(include)) => (HashSet::new(), HashSet::from(EnvVars(include))),
+        (None, None) => (HashSet::new(), HashSet::from(EnvVars("*".to_owned()))),
+    };
+
+    if !env_vars_exclude.is_empty() || !env_vars_include.is_empty() {
+        let mut remote_env = make_proxy_request_with_response(GetEnvVarsRequest {
+            env_vars_filter: env_vars_exclude,
+            env_vars_select: env_vars_include,
+        })
+        .expect("failed to make request to proxy")
+        .expect("failed to fetch remote env");
+
+        if let Some(overrides) = setup().env_config().r#override.as_ref() {
+            remote_env.extend(overrides.iter().map(|(k, v)| (k.clone(), v.clone())));
+        }
+
+        remote_env
+    } else {
+        Default::default()
+    }
 }
 
 /// We need to hook execve syscall to allow mirrord-layer to be loaded with sip patch when loading

mirrord/layer/src/setup.rs

@@ -2,6 +2,7 @@ use std::{collections::HashSet, net::SocketAddr};
 
 use mirrord_config::{
     feature::{
+        env::EnvConfig,
         fs::FsConfig,
         network::{incoming::IncomingConfig, outgoing::OutgoingConfig},
     },
@@ -72,6 +73,10 @@ impl LayerSetup {
         }
     }
 
+    pub fn env_config(&self) -> &EnvConfig {
+        &self.config.feature.env
+    }
+
     pub fn fs_config(&self) -> &FsConfig {
         &self.config.feature.fs
     }