From febd51c3184bd991505161d498c935e9df497008 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Smolarek?=
 <34063647+Razz4780@users.noreply.github.com>
Date: Tue, 12 Dec 2023 19:17:29 +0100
Subject: [PATCH] Fix intproxy timeouts with lengthy build tools (#2111)

* Making proxy connection from skipped processes

* Making connection from sip_only binaries

* Changelog entry
---
 changelog.d/2101.fixed.md |  1 +
 mirrord/layer/src/lib.rs  | 35 ++++++++++++++++++++++++++++++-----
 mirrord/layer/src/load.rs | 14 ++++++++------
 3 files changed, 39 insertions(+), 11 deletions(-)
 create mode 100644 changelog.d/2101.fixed.md

diff --git a/changelog.d/2101.fixed.md b/changelog.d/2101.fixed.md
new file mode 100644
index 00000000000..8777528e74a
--- /dev/null
+++ b/changelog.d/2101.fixed.md
@@ -0,0 +1 @@
+Fixed an issue with internal proxy timing out when the user application spawns lengthy build processes.
\ No newline at end of file
diff --git a/mirrord/layer/src/lib.rs b/mirrord/layer/src/lib.rs
index d3c17759ac6..b76e903e168 100644
--- a/mirrord/layer/src/lib.rs
+++ b/mirrord/layer/src/lib.rs
@@ -65,7 +65,7 @@
 extern crate alloc;
 extern crate core;
 
-use std::{cmp::Ordering, ffi::OsString, panic, sync::OnceLock, time::Duration};
+use std::{cmp::Ordering, ffi::OsString, net::SocketAddr, panic, sync::OnceLock, time::Duration};
 
 use ctor::ctor;
 use error::{LayerError, Result};
@@ -184,16 +184,39 @@ fn layer_pre_initialization() -> Result<(), LayerError> {
         }
     }
 
-    match given_process.load_type(config) {
-        LoadType::Full(config) => layer_start(*config),
+    match given_process.load_type(&config) {
+        LoadType::Full => layer_start(config),
         #[cfg(target_os = "macos")]
-        LoadType::SIPOnly(config) => sip_only_layer_start(*config, patch_binaries),
-        LoadType::Skip => {}
+        LoadType::SIPOnly => sip_only_layer_start(config, patch_binaries),
+        LoadType::Skip => load_only_layer_start(&config),
     }
 
     Ok(())
 }
 
+/// Initialize a new session with the internal proxy.
+/// Sets [`PROXY_CONNECTION`].
+fn load_only_layer_start(config: &LayerConfig) {
+    let address: SocketAddr = config
+        .connect_tcp
+        .as_ref()
+        .expect("missing internal proxy address")
+        .parse()
+        .expect("failed to parse internal proxy address");
+
+    let new_connection =
+        ProxyConnection::new(address, NewSessionRequest::New, Duration::from_secs(5))
+            .expect("failed to initialize proxy connection");
+
+    unsafe {
+        // SAFETY
+        // Called only from library constructor.
+        PROXY_CONNECTION
+            .set(new_connection)
+            .expect("setting PROXY_CONNECTION singleton")
+    }
+}
+
 /// The one true start of mirrord-layer.
 ///
 /// Calls [`layer_pre_initialization`], which runs mirrord-layer.
@@ -325,6 +348,8 @@ fn layer_start(mut config: LayerConfig) {
 /// mirrord-layer on a process where specified to skip with MIRRORD_SKIP_PROCESSES
 #[cfg(target_os = "macos")]
 fn sip_only_layer_start(mut config: LayerConfig, patch_binaries: Vec<String>) {
+    load_only_layer_start(&config);
+
     let mut hook_manager = HookManager::default();
 
     unsafe { exec_utils::enable_execve_hook(&mut hook_manager, patch_binaries) };
diff --git a/mirrord/layer/src/load.rs b/mirrord/layer/src/load.rs
index ed64480df7f..bb15fa1e1a6 100644
--- a/mirrord/layer/src/load.rs
+++ b/mirrord/layer/src/load.rs
@@ -97,7 +97,7 @@ impl ExecutableName {
     }
 
     /// Determine the [`LoadType`] for this process.
-    pub fn load_type(&self, config: LayerConfig) -> LoadType {
+    pub fn load_type(&self, config: &LayerConfig) -> LoadType {
         let skip_processes = config
             .skip_processes
             .as_ref()
@@ -106,12 +106,12 @@ impl ExecutableName {
 
         if self.should_load(skip_processes, config.skip_build_tools) {
             trace!("Loading into process: {self}.");
-            LoadType::Full(Box::new(config))
+            LoadType::Full
         } else {
             #[cfg(target_os = "macos")]
             if sip::is_sip_only(self) {
                 trace!("Loading into process: {self}, but only hooking exec/spawn.");
-                return LoadType::SIPOnly(Box::new(config));
+                return LoadType::SIPOnly;
             }
 
             trace!("Not loading into process: {self}.");
@@ -146,12 +146,14 @@ mod sip {
 /// Load Type of mirrord-layer
 pub enum LoadType {
     /// Mirrord is loaded fully and layer should connect to agent
-    Full(Box<LayerConfig>),
+    Full,
+
     /// Only load sip patch required hooks
     #[cfg(target_os = "macos")]
-    SIPOnly(Box<LayerConfig>),
+    SIPOnly,
 
-    /// Skip on current process
+    /// Skip on current process, make only a dummy connection to the internal proxy (to prevent
+    /// timeouts)
     Skip,
 }