diff --git a/changelog.d/+loading-env-from-inside-layer.added.md b/changelog.d/+loading-env-from-inside-layer.added.md new file mode 100644 index 00000000000..324dc2169f4 --- /dev/null +++ b/changelog.d/+loading-env-from-inside-layer.added.md @@ -0,0 +1 @@ +POC \ No newline at end of file diff --git a/mirrord/cli/src/execution.rs b/mirrord/cli/src/execution.rs index 7d3f1cef969..262568c1372 100644 --- a/mirrord/cli/src/execution.rs +++ b/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 diff --git a/mirrord/config/src/feature/env.rs b/mirrord/config/src/feature/env.rs index 30bff8bfca1..08391522347 100644 --- a/mirrord/config/src/feature/env.rs +++ b/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>, // `r#`: `override` is a Rust keyword. + + /// Remote environment will be applied during layer setup instead of CLI launch. + pub load_from_process: Option, } +pub const LOAD_ENV_FROM_PROCESS_FLAG: &str = "MIRRORD_LOAD_ENV_FROM_PROCESS"; + impl MirrordToggleableConfig for EnvFileConfig { fn disabled_config(context: &mut ConfigContext) -> Result { 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, }) } diff --git a/mirrord/intproxy/protocol/src/lib.rs b/mirrord/intproxy/protocol/src/lib.rs index d22f96e70af..ead2ac44761 100644 --- a/mirrord/intproxy/protocol/src/lib.rs +++ b/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), /// A response to layer's [`IncomingRequest`]. Incoming(IncomingResponse), + /// A response to layer's [`LayerToProxyMessage::GetEnv`]. + GetEnv(RemoteResult>), } /// 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>, + req_path = LayerToProxyMessage::GetEnv, + res_path = ProxyToLayerMessage::GetEnv, +); diff --git a/mirrord/intproxy/src/lib.rs b/mirrord/intproxy/src/lib.rs index 4aaf18d8ec9..eabd3ac3e33 100644 --- a/mirrord/intproxy/src/lib.rs +++ b/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)), } diff --git a/mirrord/intproxy/src/proxies/simple.rs b/mirrord/intproxy/src/proxies/simple.rs index caf797fc632..2f5d31a0d5e 100644 --- a/mirrord/intproxy/src/proxies/simple.rs +++ b/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>), } #[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 + } } } diff --git a/mirrord/layer/src/lib.rs b/mirrord/layer/src/lib.rs index 1bb72c93508..df1cd1acc94 100644 --- a/mirrord/layer/src/lib.rs +++ b/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 { + 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 diff --git a/mirrord/layer/src/setup.rs b/mirrord/layer/src/setup.rs index 98ea5efcb01..0636552f071 100644 --- a/mirrord/layer/src/setup.rs +++ b/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 }