set env in exec mode

roboplc-cli/Cargo.lock

@@ -695,7 +695,7 @@ dependencies = [
 
 [[package]]
 name = "roboplc-cli"
-version = "0.4.5"
+version = "0.4.6"
 dependencies = [
  "ansi_term",
  "clap",

roboplc-cli/Cargo.toml

 [package]
 name = "roboplc-cli"
-version = "0.4.5"
+version = "0.4.6"
 edition = "2021"
 authors = ["Serhij S. <div@altertech.com>"]
 license = "Apache-2.0"

roboplc-cli/src/arguments.rs

-use std::path::PathBuf;
+use std::{collections::BTreeMap, path::PathBuf};
 
 use clap::{Parser, ValueEnum};
 
@@ -158,6 +158,12 @@ pub struct ExecCommand {
         help = "Force execute (ignore if other program is being executed)"
     )]
     pub force: bool,
+    #[clap(
+        short = 'e',
+        long,
+        help = "Environment variable to pass to the program, as NAME=VALUE"
+    )]
+    pub env: Vec<String>,
     #[arg(
         trailing_var_arg = true,
         allow_hyphen_values = true,
@@ -174,6 +180,7 @@ pub struct FlashExec {
     pub force: bool,
     pub run: bool,
     pub program_args: Vec<String>,
+    pub program_env: BTreeMap<String, String>,
 }
 
 impl From<FlashCommand> for FlashExec {
@@ -186,12 +193,26 @@ impl From<FlashCommand> for FlashExec {
             force: cmd.force,
             run: cmd.run,
             program_args: Vec::new(),
+            program_env: BTreeMap::new(),
         }
     }
 }
 
 impl From<ExecCommand> for FlashExec {
     fn from(cmd: ExecCommand) -> Self {
+        let program_env = cmd
+            .env
+            .iter()
+            .map(|s| {
+                let mut parts = s.splitn(2, '=');
+                let key = parts.next().unwrap().to_string();
+                let value = parts
+                    .next()
+                    .expect("No environment variable value set")
+                    .to_string();
+                (key, value)
+            })
+            .collect();
         Self {
             cargo: cmd.cargo,
             cargo_target: cmd.cargo_target,
@@ -200,6 +221,7 @@ impl From<ExecCommand> for FlashExec {
             force: cmd.force,
             run: false,
             program_args: cmd.args,
+            program_env,
         }
     }
 }

roboplc-cli/src/exec.rs

-use std::io::Write as _;
 use std::path::Path;
+use std::{collections::BTreeMap, io::Write as _};
 
 use colored::Colorize;
 #[cfg(not(target_os = "windows"))]
@@ -7,7 +7,6 @@ use tokio::io::AsyncReadExt as _;
 
 use futures_util::{SinkExt, StreamExt};
 use serde::{Deserialize, Serialize};
-use serde_json::json;
 use tokio_tungstenite::tungstenite::Message;
 
 pub fn exec(
@@ -16,11 +15,12 @@ pub fn exec(
     file: &Path,
     force: bool,
     args: Vec<String>,
+    env: BTreeMap<String, String>,
 ) -> Result<(), Box<dyn std::error::Error>> {
     tokio::runtime::Builder::new_current_thread()
         .enable_all()
         .build()?
-        .block_on(exec_remote(url, key, file, force, args))?;
+        .block_on(exec_remote(url, key, file, force, args, env))?;
     Ok(())
 }
 
@@ -31,6 +31,24 @@ enum Output {
     Terminated(i32),
 }
 
+#[derive(Serialize)]
+struct ExecPayload<'a> {
+    k: &'a str,
+    force: bool,
+    #[serde(skip_serializing_if = "Vec::is_empty")]
+    args: Vec<String>,
+    #[serde(skip_serializing_if = "BTreeMap::is_empty")]
+    env: BTreeMap<String, String>,
+    term: ExecTerm,
+}
+
+#[derive(Serialize)]
+struct ExecTerm {
+    width: usize,
+    height: usize,
+    name: String,
+}
+
 #[allow(clippy::too_many_lines)]
 async fn exec_remote(
     url: &str,
@@ -38,6 +56,7 @@ async fn exec_remote(
     file: &Path,
     force: bool,
     args: Vec<String>,
+    env: BTreeMap<String, String>,
 ) -> Result<(), Box<dyn std::error::Error>> {
     let (ws_uri, url_short) = if let Some(u) = url.strip_prefix("http://") {
         (format!("ws://{}/roboplc/api/ws.execute", u), u)
@@ -50,16 +69,17 @@ async fn exec_remote(
     println!();
     let (mut socket, _) = tokio_tungstenite::connect_async(&ws_uri).await?;
     let (width, height) = term_size::dimensions().ok_or("Failed to get terminal size")?;
-    let payload = json!({
-        "k": key,
-        "force": force,
-        "args": args,
-        "term": {
-            "width": width,
-            "height": height,
-            "name": std::env::var("TERM").unwrap_or("xterm-256color".to_string()),
+    let payload = ExecPayload {
+        k: key,
+        force,
+        args,
+        env,
+        term: ExecTerm {
+            width,
+            height,
+            name: std::env::var("TERM").unwrap_or("xterm-256color".to_string()),
         },
-    });
+    };
     socket
         .send(Message::Text(serde_json::to_string(&payload)?))
         .await?;

roboplc-cli/src/flashing.rs

 use std::{
+    collections::BTreeMap,
     env, fs,
     path::{Path, PathBuf},
 };
@@ -27,12 +28,13 @@ fn flash_file(
     run: bool,
     exec_only: bool,
     program_args: Vec<String>,
+    program_env: BTreeMap<String, String>,
 ) -> Result<(), Box<dyn std::error::Error>> {
     if !file.exists() {
         return Err(format!("File not found: {}", file.display()).into());
     }
     if exec_only {
-        return crate::exec::exec(url, key, file, force, program_args);
+        return crate::exec::exec(url, key, file, force, program_args, program_env);
     }
     if let Some(docker_img) = url.strip_prefix("docker://") {
         let tag = std::env::var("ROBOPLC_DOCKER_TAG").unwrap_or_else(|_| {
@@ -113,6 +115,7 @@ fn run_build_custom(
     file: &Path,
     exec_only: bool,
     program_args: Vec<String>,
+    program_env: BTreeMap<String, String>,
 ) -> Result<(), Box<dyn std::error::Error>> {
     println!("Remote: {}", url.yellow());
     println!("Build command line: {}", cmd.yellow());
@@ -128,7 +131,17 @@ fn run_build_custom(
     if !file.exists() {
         return Err(format!("File not found: {}", file.display()).into());
     }
-    flash_file(url, key, agent, file, force, run, exec_only, program_args)?;
+    flash_file(
+        url,
+        key,
+        agent,
+        file,
+        force,
+        run,
+        exec_only,
+        program_args,
+        program_env,
+    )?;
     Ok(())
 }
 
@@ -152,6 +165,7 @@ pub fn flash(
             opts.run,
             exec_only,
             opts.program_args,
+            opts.program_env,
         )?;
     } else if let Some(custom_cmd) = build_custom.command {
         run_build_custom(
@@ -166,6 +180,7 @@ pub fn flash(
                 .ok_or("Custom build command requires a file")?,
             exec_only,
             opts.program_args,
+            opts.program_env,
         )?;
     } else {
         let mut cargo_target: Option<String> = None;
@@ -240,6 +255,7 @@ pub fn flash(
             opts.run,
             exec_only,
             opts.program_args,
+            opts.program_env,
         )?;
     }
     report_ok()