POC - load env from inside layer (#2234)

* POC

* Changelog entry

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

@@ -0,0 +1 @@
+POC

mirrord/cli/src/execution.rs

@@ -5,7 +5,7 @@ use std::{
 };
 
 use mirrord_analytics::{AnalyticsError, AnalyticsReporter};
-use mirrord_config::LayerConfig;
+use mirrord_config::{feature::env::LOAD_ENV_FROM_PROCESS_FLAG, LayerConfig};
 use mirrord_progress::Progress;
 use mirrord_protocol::{ClientMessage, DaemonMessage, EnvVars, GetEnvVarsRequest};
 #[cfg(target_os = "macos")]
@@ -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) {
+            HashMap::from([(LOAD_ENV_FROM_PROCESS_FLAG.to_string(), "true".to_string())])
+        } 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,8 +69,13 @@ 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.
+
+    /// Remote environment will be applied during layer setup instead of CLI launch.
+    pub load_from_process: Option<bool>,
 }
 
+pub const LOAD_ENV_FROM_PROCESS_FLAG: &str = "MIRRORD_LOAD_ENV_FROM_PROCESS";
+
 impl MirrordToggleableConfig for EnvFileConfig {
     fn disabled_config(context: &mut ConfigContext) -> Result<Self::Generated> {
         Ok(EnvConfig {
@@ -81,6 +86,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};
@@ -76,17 +84,21 @@ use load::ExecutableName;
 #[cfg(target_os = "macos")]
 use mirrord_config::feature::fs::FsConfig;
 use mirrord_config::{
-    feature::{fs::FsModeConfig, network::incoming::IncomingMode},
+    feature::{env::LOAD_ENV_FROM_PROCESS_FLAG, fs::FsModeConfig, network::incoming::IncomingMode},
     LayerConfig,
 };
 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;
@@ -345,6 +357,58 @@ fn layer_start(mut config: LayerConfig) {
             .set(new_connection)
             .expect("setting PROXY_CONNECTION singleton")
     }
+
+    let load_env = std::env::var(LOAD_ENV_FROM_PROCESS_FLAG)
+        .unwrap_or_default()
+        .parse()
+        .unwrap_or(false);
+    if load_env {
+        let env = fetch_env_vars();
+        for (key, value) in env {
+            std::env::set_var(key, value);
+        }
+
+        std::env::remove_var(LOAD_ENV_FROM_PROCESS_FLAG);
+    }
+}
+
+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
     }