From e23fac0d89ef20ec9acf5be7ccabaf42496c2396 Mon Sep 17 00:00:00 2001 From: Lincoln Wallace Date: Wed, 17 Jul 2024 11:22:19 -0300 Subject: [PATCH] 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 --- extensions/env-injector/Cargo.toml | 21 ++++ extensions/env-injector/src/main.rs | 161 +++++++++++++++++++++++++++ schema/snapcraft.json | 1 + snapcraft/extensions/env_injector.py | 135 ++++++++++++++++++++++ snapcraft/extensions/registry.py | 2 + 5 files changed, 320 insertions(+) create mode 100644 extensions/env-injector/Cargo.toml create mode 100644 extensions/env-injector/src/main.rs create mode 100644 snapcraft/extensions/env_injector.py diff --git a/extensions/env-injector/Cargo.toml b/extensions/env-injector/Cargo.toml new file mode 100644 index 00000000000..94a73715ca1 --- /dev/null +++ b/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 diff --git a/extensions/env-injector/src/main.rs b/extensions/env-injector/src/main.rs new file mode 100644 index 00000000000..61480805810 --- /dev/null +++ b/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> { + let url: hyperlocal::Uri = Uri::new(SNAPD_SOCKET, "/v2/snapctl").into(); + + let client: Client> = 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> = Request::builder() + .method(Method::POST) + .uri(url) + .body(Full::from(request_body))?; + + let mut res = client.request(req).await?; + + let mut body: Vec = 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> { + 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 { + 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> { + 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> { + // 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> { + + 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> { + + let args: Vec = std::env::args().collect(); + + if args.len() < 2 { + eprintln!("Usage: {} ", 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(()) +} diff --git a/schema/snapcraft.json b/schema/snapcraft.json index 9da873f1849..f0890b80f3b 100644 --- a/schema/snapcraft.json +++ b/schema/snapcraft.json @@ -906,6 +906,7 @@ "uniqueItems": true, "items": { "enum": [ + "env-injector", "flutter-stable", "flutter-beta", "flutter-dev", diff --git a/snapcraft/extensions/env_injector.py b/snapcraft/extensions/env_injector.py new file mode 100644 index 00000000000..d5fad95f380 --- /dev/null +++ b/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 . + +"""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 env.= + + - To set environment variables for a specific application **inside** the snap: + + .. code-block:: shell + sudo snap set apps..env.= + + - To set environment file inside the snap: + + .. code-block:: shell + sudo snap set 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) \ No newline at end of file diff --git a/snapcraft/extensions/registry.py b/snapcraft/extensions/registry.py index 42c60af1d6d..64906dad276 100644 --- a/snapcraft/extensions/registry.py +++ b/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, }