add: extension env-injector

Add extension to inject environment variables into snaps by it's command-chain
This extension takes snap options and transform then into environment variables

Signed-off-by: Lincoln Wallace <lincoln.wallace@canonical.com>

extensions/env-injector/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "env-exporter"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+http-body-util = "0.1.2"
+hyper = {git = "https://github.com/hyperium/hyper", branch = "master"}
+hyper-util = "0.1.6"
+hyperlocal = {git = "https://github.com/softprops/hyperlocal", branch = "main"}
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+tokio = { version = "1.35", features = ["full"] }
+
+
+[profile.release]
+panic = "abort"
+strip = true
+opt-level = "z"
+lto = true
+codegen-units = 1

extensions/env-injector/src/main.rs

@@ -0,0 +1,161 @@
+use http_body_util::{BodyExt, Full};
+use hyper::body::Bytes;
+use hyper::{Method, Request};
+use hyper_util::client::legacy::Client;
+use hyperlocal::{UnixClientExt, UnixConnector, Uri};
+use std::error::Error;
+use std::process::Command;
+use std::collections::HashMap;
+use std::path::Path;
+use std::io::BufRead;
+
+const SNAPD_SOCKET: &str = "/run/snapd-snap.socket";
+
+async fn snapdapi_req() -> Result<serde_json::Value, Box<dyn Error + Send + Sync>> {
+    let url: hyperlocal::Uri = Uri::new(SNAPD_SOCKET, "/v2/snapctl").into();
+
+    let client: Client<UnixConnector, Full<Bytes>> = Client::unix();
+
+    let snap_context = std::env::var("SNAP_CONTEXT")?;
+
+    let request_body = format!(
+        r#"{{"context-id":"{}","args":["get", "env", "envfile", "apps"]}}"#,
+        snap_context
+    );
+
+    let req: Request<Full<Bytes>> = Request::builder()
+        .method(Method::POST)
+        .uri(url)
+        .body(Full::from(request_body))?;
+
+    let mut res = client.request(req).await?;
+
+    let mut body: Vec<u8> = Vec::new();
+
+    while let Some(frame_result) = res.frame().await {
+        let frame = frame_result?;
+
+        if let Some(segment) = frame.data_ref() {
+            body.extend_from_slice(segment);
+        }
+    }
+
+    Ok(serde_json::from_slice(&body)?)
+}
+
+fn set_env_vars(app: &str, json: &serde_json::Value) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    let stdout_str = json["result"]["stdout"].as_str().ok_or("Invalid stdout")?;
+    let stdout_json: serde_json::Value = serde_json::from_str(stdout_str)?;
+
+    fn process_env(env: &serde_json::Value) -> HashMap<String, String> {
+        env.as_object()
+            .unwrap()
+            .iter()
+            .map(|(k, v)| {
+                let key = k.to_uppercase().replace("-", "_");
+                (key, v.to_string())
+            })
+            .collect()
+    }
+
+    if let Some(global_env) = stdout_json["env"].as_object() {
+        for (key, value) in process_env(&serde_json::Value::Object(global_env.clone())) {
+            std::env::set_var(key, value.trim_matches('"'));
+        }
+    }
+
+    if let Some(app_env) = stdout_json["apps"][app]["env"].as_object() {
+        for (key, value) in process_env(&serde_json::Value::Object(app_env.clone())) {
+            std::env::set_var(key, value.trim_matches('"'));
+        }
+    }
+
+    Ok(())
+}
+
+fn source_env_file(file_path: &str) -> std::io::Result<HashMap<String, String>> {
+    let path = Path::new(file_path);
+
+    if !path.exists() {
+        eprintln!("File does not exist: {}", file_path);
+        return Err(std::io::Error::new(std::io::ErrorKind::NotFound, "File does not exist"));
+    }
+
+    if !path.is_file() {
+        eprintln!("File is not readable: {}", file_path);
+        return Err(std::io::Error::new(std::io::ErrorKind::PermissionDenied, "File is not readable"));
+    }
+
+    let file = std::fs::File::open(file_path)?;
+    let reader = std::io::BufReader::new(file);
+
+    let mut env_vars = HashMap::new();
+    for line in reader.lines() {
+        let line = line?;
+        if !line.trim().is_empty() && !line.starts_with('#') {
+            if let Some((key, value)) = line.split_once('=') {
+                env_vars.insert(key.to_string(), value.to_string());
+            }
+        }
+    }
+
+    Ok(env_vars)
+}
+
+fn set_env_vars_from_file(app: &str, json: &serde_json::Value) ->  Result<(), Box<dyn std::error::Error + Send + Sync>>  {
+    // Extract the stdout JSON string and parse it
+    let stdout_str = json["result"]["stdout"].as_str().ok_or("Invalid stdout")?;
+    let stdout_json: serde_json::Value = serde_json::from_str(stdout_str)?;
+
+    // Source the global envfile first
+    if let Some(global_envfile) = stdout_json["envfile"].as_str() {
+        if let Ok(env_vars) = source_env_file(global_envfile) {
+            for (key, value) in env_vars {
+                std::env::set_var(key, value);
+            }
+        }
+    }
+
+    // Source the app-specific envfile
+    if let Some(app_envfile) = stdout_json["apps"][app]["envfile"].as_str() {
+        if let Ok(env_vars) = source_env_file(app_envfile) {
+            for (key, value) in env_vars {
+                std::env::set_var(key, value);
+            }
+        }
+    }
+
+    Ok(())
+}
+
+#[tokio::main]
+async fn run() -> Result<(), Box<dyn Error + Send + Sync>> {
+
+    let json = snapdapi_req().await?;
+
+    let app = std::env::var("env_alias")?;
+
+    set_env_vars_from_file(&app, &json)?;
+    set_env_vars(&app, &json)?;
+
+    Ok(())
+}
+
+fn main()-> Result<(), Box<dyn Error + Send + Sync>> {
+
+    let args: Vec<String> = std::env::args().collect();
+
+    if args.len() < 2 {
+        eprintln!("Usage: {} <app-path>", args[0]);
+        std::process::exit(1);
+    }
+
+    let command = args[1].clone();
+    let args = args[2..].to_vec();
+
+    run()?;
+
+    Command::new(command).args(args).status()?;
+
+    Ok(())
+}

schema/snapcraft.json

@@ -906,6 +906,7 @@
                             "uniqueItems": true,
                             "items": {
                                 "enum": [
+                                    "env-injector",
                                     "flutter-stable",
                                     "flutter-beta",
                                     "flutter-dev",

snapcraft/extensions/env_injector.py

@@ -0,0 +1,135 @@
+# -*- Mode:Python; indent-tabs-mode:nil; tab-width:4 -*-
+#
+# Copyright 2024 Canonical Ltd.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 3 as
+# published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+"""Extension to automatically set environment variables on snaps."""
+
+from typing import Any, Dict, Optional, Tuple
+
+from overrides import overrides
+
+from .extension import Extension, get_extensions_data_dir
+
+class EnvInjectorExtension(Extension):
+    """Extension to automatically set environment variables on snaps.
+
+    This extension allows you to transform snap options into environment
+    variables
+
+    It configures your application to run a command-chain that transforms the
+    snap options into environment variables automatically.
+
+    - To set global environment variables for all applications **inside** the snap:
+    
+    .. code-block:: shell
+        sudo snap set <snap-name> env.<key>=<value>
+
+    - To set environment variables for a specific application **inside** the snap:
+
+    .. code-block:: shell
+        sudo snap set <snap-name> apps.<app-name>.env.<key>=<value>
+
+    - To set environment file inside the snap:
+
+    .. code-block:: shell
+        sudo snap set <snap-name> env-file=<path-to-env-file>
+
+    """
+
+    @staticmethod
+    @overrides
+    def get_supported_bases() -> Tuple[str, ...]:
+        return ("core24",)
+
+
+    @staticmethod
+    @overrides
+    def get_supported_confinement() -> Tuple[str, ...]:
+        return ("strict", "devmode", "classic")
+
+
+    @staticmethod
+    @overrides
+    def is_experimental(base: Optional[str]) -> bool:
+        return True
+
+    @overrides
+    def get_root_snippet(self) -> Dict[str, Any]:
+        return {}
+
+
+    @overrides
+    def get_app_snippet(self, *, app_name: str) -> Dict[str, Any]:
+        """Return the app snippet to apply."""
+        return {
+            "command-chain": [ "bin/command-chain/env-exporter"],
+            "environment": {
+                "env_alias": f"{app_name}",
+            }
+        }
+
+    @overrides
+    def get_part_snippet(self, *, plugin_name: str) -> Dict[str, Any]:
+        return {
+             "build-environment": [
+                 { "VARS_INJECT_EXTENSION": "true" }
+             ]
+        }
+
+    @overrides
+    def get_parts_snippet(self) -> Dict[str, Any]:
+        toolchain = self.get_toolchain()
+        return {
+            f"env-injector/env-injector": {
+                "source": f"{get_extensions_data_dir()}/env-injector",
+                "plugin": "nil",
+                "build-snaps": [
+                    f"rustup",
+                ],
+                "build-packages": [
+                    f"musl-tools",
+                    f"upx-ucl", # for binary compression
+                ],
+                "override-build": f"""
+
+      rustup default stable
+      rustup target add {toolchain}
+
+      cargo build --target {toolchain} --release
+      mkdir -p $SNAPCRAFT_PART_INSTALL/bin/command-chain
+
+      cp target/{toolchain}/release/env-exporter $SNAPCRAFT_PART_INSTALL/bin/command-chain
+
+      # shrink the binary
+      upx --best --lzma target/{toolchain}/release/env-exporter
+      cp target/{toolchain}/release/env-exporter $SNAPCRAFT_PART_INSTALL/bin/command-chain/env-exporter-upx
+
+                """,
+
+            }
+        }
+
+    def get_toolchain(self):
+        # Dictionary mapping architecture names
+        toolchain = {
+            'amd64': 'x86_64-unknown-linux-musl',
+            'arm64': 'aarch64-unknown-linux-musl',
+            'armhf': 'arm-unknown-linux-musleabihf',
+            'riscv64': 'riscv64gc-unknown-linux-musl',
+            'ppc64el': 'powerpc64-unknown-linux-gnu',
+            's390x': 's390x-unknown-linux-musl',
+            'i386': 'i586-unknown-linux-musl',
+        }
+        return toolchain.get(self.arch)

snapcraft/extensions/registry.py

@@ -31,6 +31,7 @@
 from .ros2_jazzy_desktop import ROS2JazzyDesktopExtension
 from .ros2_jazzy_ros_base import ROS2JazzyRosBaseExtension
 from .ros2_jazzy_ros_core import ROS2JazzyRosCoreExtension
+from .env_injector import EnvInjectorExtension
 
 if TYPE_CHECKING:
     from .extension import Extension
@@ -49,6 +50,7 @@
     "ros2-jazzy-desktop": ROS2JazzyDesktopExtension,
     "kde-neon": KDENeon,
     "kde-neon-6": KDENeon6,
+    "env-injector": EnvInjectorExtension,
 }