From 1386c1dc4c8bc96baff6820690a9cd19d57a931b Mon Sep 17 00:00:00 2001
From: C0ffeeCode <ritters_werth@outlook.com>
Date: Mon, 7 Oct 2024 16:22:45 +0200
Subject: [PATCH] Feature: Allow setting environment variables

---
 src/channel.rs       |  6 +++++-
 src/channel/shell.rs | 34 ++++++++++++++++++----------------
 src/connection.rs    |  4 ++++
 3 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/src/channel.rs b/src/channel.rs
index fda158a..7fb79d4 100644
--- a/src/channel.rs
+++ b/src/channel.rs
@@ -26,12 +26,14 @@ pub struct Channel {
     peer_window_size: u32,
     max_packet_size: u32,
     read_thread: Option<JoinHandle<()>>,
+    env: Vec<(String, String)>,
 }
 
 #[derive(Debug)]
 pub enum ChannelRequest {
     Pty(PtyConfig),
     Shell,
+    Env(String, String),
 }
 
 impl Channel {
@@ -52,6 +54,7 @@ impl Channel {
             peer_window_size,
             max_packet_size,
             read_thread: None,
+            env: Vec::new(),
         }
     }
 
@@ -68,11 +71,12 @@ impl Channel {
     }
 
     pub fn handle_request(&mut self, request: ChannelRequest) {
+        debug!("Channel Request: {:?}", request);
         match request {
             ChannelRequest::Pty(ref pty) => self.setup_tty(pty),
             ChannelRequest::Shell => self.setup_shell(),
+            ChannelRequest::Env(key, value) => self.env.push((key, value)),
         }
-        debug!("Channel Request: {:?}", request);
     }
 
     /// Writes `data` **to** this channel;
diff --git a/src/channel/shell.rs b/src/channel/shell.rs
index c7c2a1d..bb461a2 100644
--- a/src/channel/shell.rs
+++ b/src/channel/shell.rs
@@ -21,7 +21,7 @@ impl Channel {
         match self.pty.as_ref() {
             Some((_, tty_path)) => with_tty(tty_path),
             None => {
-                let pipes = without_tty();
+                let pipes = self.without_tty();
                 #[cfg(unix)]
                 use crate::sys::non_blockify_reader;
                 non_blockify_reader(pipes.stdout.get_ref());
@@ -30,23 +30,25 @@ impl Channel {
             },
         }
     }
-}
 
-fn without_tty() -> PipeContainer {
-    let proc = unsafe {
-        process::Command::new("/bin/sh")    
-            .stdin(Stdio::piped())
-            .stdout(Stdio::piped())
-            .stderr(Stdio::piped())
-            .pre_exec(sys::before_exec)
-            .spawn()
-            .unwrap()
-    };
+    fn without_tty(&mut self) -> PipeContainer {
+        let proc = unsafe {
+            process::Command::new("/bin/sh")
+                // .env_clear()
+                .envs(self.env.drain(..))
+                .stdin(Stdio::piped())
+                .stdout(Stdio::piped())
+                .stderr(Stdio::piped())
+                .pre_exec(sys::before_exec)
+                .spawn()
+                .unwrap()
+        };
 
-    PipeContainer {
-        stdin: proc.stdin.unwrap(),
-        stdout: BufReader::with_capacity(1, proc.stdout.unwrap()),
-        stderr: BufReader::with_capacity(1, proc.stderr.unwrap()),
+        PipeContainer {
+            stdin: proc.stdin.unwrap(),
+            stdout: BufReader::with_capacity(1, proc.stdout.unwrap()),
+            stderr: BufReader::with_capacity(1, proc.stderr.unwrap()),
+        }
     }
 }
 
diff --git a/src/connection.rs b/src/connection.rs
index ef3700e..8a44d26 100644
--- a/src/connection.rs
+++ b/src/connection.rs
@@ -356,6 +356,10 @@ impl Connection {
                 modes: reader.read_string()?,
             })),
             "shell" => Some(ChannelRequest::Shell),
+            "env" => Some(ChannelRequest::Env(
+                String::from_utf8(reader.read_string()?).unwrap(), // TODO: error handling
+                String::from_utf8(reader.read_string()?).unwrap(),
+            )),
             _ => None,
         };