use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct RuntimeOverrides {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub wasm: Option<WasmOverrides>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub docker: Option<DockerOverrides>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub native: Option<NativeOverrides>,
}
impl RuntimeOverrides {
pub fn new() -> Self {
Self::default()
}
pub fn with_wasm(mut self, wasm: WasmOverrides) -> Self {
self.wasm = Some(wasm);
self
}
pub fn with_docker(mut self, docker: DockerOverrides) -> Self {
self.docker = Some(docker);
self
}
pub fn with_native(mut self, native: NativeOverrides) -> Self {
self.native = Some(native);
self
}
pub fn is_empty(&self) -> bool {
self.wasm.is_none() && self.docker.is_none() && self.native.is_none()
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct WasmOverrides {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub stack_size: Option<usize>,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub wasi_capabilities: HashMap<String, bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub fuel_limit: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub epoch_interruption: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_memory_pages: Option<u32>,
#[serde(default)]
pub debug_info: bool,
}
impl WasmOverrides {
pub fn new() -> Self {
Self::default()
}
pub fn with_stack_size(mut self, size: usize) -> Self {
self.stack_size = Some(size);
self
}
pub fn with_wasi_capability(mut self, capability: impl Into<String>, enabled: bool) -> Self {
self.wasi_capabilities.insert(capability.into(), enabled);
self
}
pub fn enable_capability(self, capability: impl Into<String>) -> Self {
self.with_wasi_capability(capability, true)
}
pub fn disable_capability(self, capability: impl Into<String>) -> Self {
self.with_wasi_capability(capability, false)
}
pub fn with_fuel_limit(mut self, limit: u64) -> Self {
self.fuel_limit = Some(limit);
self
}
pub fn with_epoch_interruption(mut self) -> Self {
self.epoch_interruption = Some(true);
self
}
pub fn with_max_memory_pages(mut self, pages: u32) -> Self {
self.max_memory_pages = Some(pages);
self
}
pub fn with_debug_info(mut self) -> Self {
self.debug_info = true;
self
}
pub fn is_capability_enabled(&self, capability: &str) -> Option<bool> {
self.wasi_capabilities.get(capability).copied()
}
pub fn stack_size_or_default(&self) -> usize {
self.stack_size.unwrap_or(1024 * 1024) }
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct DockerOverrides {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub image: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub extra_args: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub entrypoint: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub command: Option<Vec<String>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub user: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub gpus: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub platform: Option<String>,
#[serde(default)]
pub privileged: bool,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub security_opt: Vec<String>,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub sysctls: HashMap<String, String>,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub labels: HashMap<String, String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub restart: Option<String>,
#[serde(default = "default_true")]
pub rm: bool,
#[serde(default)]
pub init: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub hostname: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ipc: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub pid: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub cap_add: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub cap_drop: Vec<String>,
}
fn default_true() -> bool {
true
}
impl Default for DockerOverrides {
fn default() -> Self {
Self {
image: None,
extra_args: Vec::new(),
entrypoint: None,
command: None,
user: None,
gpus: None,
platform: None,
privileged: false,
security_opt: Vec::new(),
sysctls: HashMap::new(),
labels: HashMap::new(),
restart: None,
rm: true, init: false,
hostname: None,
ipc: None,
pid: None,
cap_add: Vec::new(),
cap_drop: Vec::new(),
}
}
}
impl DockerOverrides {
pub fn new() -> Self {
Self::default()
}
pub fn with_image(mut self, image: impl Into<String>) -> Self {
self.image = Some(image.into());
self
}
pub fn with_extra_arg(mut self, arg: impl Into<String>) -> Self {
self.extra_args.push(arg.into());
self
}
pub fn with_entrypoint(mut self, entrypoint: impl Into<String>) -> Self {
self.entrypoint = Some(entrypoint.into());
self
}
pub fn with_command(mut self, command: Vec<String>) -> Self {
self.command = Some(command);
self
}
pub fn with_user(mut self, user: impl Into<String>) -> Self {
self.user = Some(user.into());
self
}
pub fn with_gpus(mut self, gpus: impl Into<String>) -> Self {
self.gpus = Some(gpus.into());
self
}
pub fn with_all_gpus(self) -> Self {
self.with_gpus("all")
}
pub fn with_platform(mut self, platform: impl Into<String>) -> Self {
self.platform = Some(platform.into());
self
}
pub fn privileged(mut self) -> Self {
self.privileged = true;
self
}
pub fn with_security_opt(mut self, opt: impl Into<String>) -> Self {
self.security_opt.push(opt.into());
self
}
pub fn with_no_new_privileges(self) -> Self {
self.with_security_opt("no-new-privileges")
}
pub fn with_sysctl(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.sysctls.insert(key.into(), value.into());
self
}
pub fn with_label(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.labels.insert(key.into(), value.into());
self
}
pub fn with_restart(mut self, policy: impl Into<String>) -> Self {
self.restart = Some(policy.into());
self
}
pub fn keep_container(mut self) -> Self {
self.rm = false;
self
}
pub fn with_init(mut self) -> Self {
self.init = true;
self
}
pub fn with_hostname(mut self, hostname: impl Into<String>) -> Self {
self.hostname = Some(hostname.into());
self
}
pub fn add_capability(mut self, cap: impl Into<String>) -> Self {
self.cap_add.push(cap.into());
self
}
pub fn drop_capability(mut self, cap: impl Into<String>) -> Self {
self.cap_drop.push(cap.into());
self
}
pub fn drop_all_capabilities(self) -> Self {
self.drop_capability("ALL")
}
pub fn to_docker_args(&self) -> Vec<String> {
let mut args = Vec::new();
if self.rm {
args.push("--rm".to_string());
}
if self.init {
args.push("--init".to_string());
}
if self.privileged {
args.push("--privileged".to_string());
}
if let Some(ref user) = self.user {
args.push("--user".to_string());
args.push(user.clone());
}
if let Some(ref gpus) = self.gpus {
args.push("--gpus".to_string());
args.push(gpus.clone());
}
if let Some(ref platform) = self.platform {
args.push("--platform".to_string());
args.push(platform.clone());
}
if let Some(ref entrypoint) = self.entrypoint {
args.push("--entrypoint".to_string());
args.push(entrypoint.clone());
}
if let Some(ref hostname) = self.hostname {
args.push("--hostname".to_string());
args.push(hostname.clone());
}
if let Some(ref ipc) = self.ipc {
args.push("--ipc".to_string());
args.push(ipc.clone());
}
if let Some(ref pid) = self.pid {
args.push("--pid".to_string());
args.push(pid.clone());
}
if let Some(ref restart) = self.restart {
args.push("--restart".to_string());
args.push(restart.clone());
}
for opt in &self.security_opt {
args.push("--security-opt".to_string());
args.push(opt.clone());
}
for (key, value) in &self.sysctls {
args.push("--sysctl".to_string());
args.push(format!("{}={}", key, value));
}
for (key, value) in &self.labels {
args.push("--label".to_string());
args.push(format!("{}={}", key, value));
}
for cap in &self.cap_add {
args.push("--cap-add".to_string());
args.push(cap.clone());
}
for cap in &self.cap_drop {
args.push("--cap-drop".to_string());
args.push(cap.clone());
}
args.extend(self.extra_args.clone());
args
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct NativeOverrides {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub working_dir: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub shell: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub path_additions: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub run_as: Option<String>,
#[serde(default)]
pub clear_env: bool,
#[serde(default = "default_true")]
pub inherit_env: bool,
}
impl Default for NativeOverrides {
fn default() -> Self {
Self {
working_dir: None,
shell: None,
path_additions: Vec::new(),
run_as: None,
clear_env: false,
inherit_env: true, }
}
}
impl NativeOverrides {
pub fn new() -> Self {
Self::default()
}
pub fn with_working_dir(mut self, dir: impl Into<String>) -> Self {
self.working_dir = Some(dir.into());
self
}
pub fn with_shell(mut self, shell: impl Into<String>) -> Self {
self.shell = Some(shell.into());
self
}
pub fn with_path_addition(mut self, path: impl Into<String>) -> Self {
self.path_additions.push(path.into());
self
}
pub fn with_run_as(mut self, user: impl Into<String>) -> Self {
self.run_as = Some(user.into());
self
}
pub fn with_clear_env(mut self) -> Self {
self.clear_env = true;
self.inherit_env = false;
self
}
pub fn without_inherit_env(mut self) -> Self {
self.inherit_env = false;
self
}
pub fn shell_or_default(&self) -> &str {
self.shell.as_deref().unwrap_or_else(|| {
if cfg!(windows) {
"cmd.exe"
} else {
"/bin/sh"
}
})
}
pub fn build_path(&self, existing_path: Option<&str>) -> String {
let separator = if cfg!(windows) { ";" } else { ":" };
let additions = self.path_additions.join(separator);
match (additions.is_empty(), existing_path) {
(true, Some(p)) => p.to_string(),
(true, None) => String::new(),
(false, Some(p)) if self.inherit_env => format!("{}{}{}",additions, separator, p),
(false, _) => additions,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_runtime_overrides_builder() {
let overrides = RuntimeOverrides::new()
.with_wasm(WasmOverrides::new().with_fuel_limit(1000))
.with_docker(DockerOverrides::new().with_image("python:3.11"));
assert!(overrides.wasm.is_some());
assert!(overrides.docker.is_some());
assert!(!overrides.is_empty());
}
#[test]
fn test_wasm_overrides() {
let wasm = WasmOverrides::new()
.with_stack_size(2 * 1024 * 1024)
.with_fuel_limit(100_000)
.enable_capability("filesystem")
.disable_capability("network")
.with_debug_info();
assert_eq!(wasm.stack_size, Some(2 * 1024 * 1024));
assert_eq!(wasm.fuel_limit, Some(100_000));
assert_eq!(wasm.is_capability_enabled("filesystem"), Some(true));
assert_eq!(wasm.is_capability_enabled("network"), Some(false));
assert!(wasm.debug_info);
}
#[test]
fn test_docker_overrides() {
let docker = DockerOverrides::new()
.with_image("python:3.11-slim")
.with_user("1000:1000")
.with_no_new_privileges()
.drop_all_capabilities()
.add_capability("NET_BIND_SERVICE")
.with_label("app", "skill-engine");
assert_eq!(docker.image, Some("python:3.11-slim".to_string()));
assert_eq!(docker.user, Some("1000:1000".to_string()));
assert!(docker.security_opt.contains(&"no-new-privileges".to_string()));
assert!(docker.cap_drop.contains(&"ALL".to_string()));
assert!(docker.cap_add.contains(&"NET_BIND_SERVICE".to_string()));
}
#[test]
fn test_docker_args() {
let docker = DockerOverrides::new()
.with_user("1000:1000")
.with_all_gpus()
.with_init()
.with_no_new_privileges();
let args = docker.to_docker_args();
assert!(args.contains(&"--rm".to_string()));
assert!(args.contains(&"--init".to_string()));
assert!(args.contains(&"--user".to_string()));
assert!(args.contains(&"--gpus".to_string()));
assert!(args.contains(&"--security-opt".to_string()));
}
#[test]
fn test_native_overrides() {
let native = NativeOverrides::new()
.with_working_dir("/app")
.with_shell("/bin/bash")
.with_path_addition("/custom/bin");
assert_eq!(native.working_dir, Some("/app".to_string()));
assert_eq!(native.shell_or_default(), "/bin/bash");
}
#[test]
fn test_native_path_building() {
let native = NativeOverrides::new()
.with_path_addition("/usr/local/bin")
.with_path_addition("/opt/bin");
let path = native.build_path(Some("/usr/bin"));
assert!(path.contains("/usr/local/bin"));
assert!(path.contains("/opt/bin"));
assert!(path.contains("/usr/bin"));
}
#[test]
fn test_runtime_overrides_serialization() {
let overrides = RuntimeOverrides::new()
.with_wasm(WasmOverrides::new().with_fuel_limit(1000))
.with_docker(DockerOverrides::new().with_image("python:3.11"));
let json = serde_json::to_string(&overrides).unwrap();
let deserialized: RuntimeOverrides = serde_json::from_str(&json).unwrap();
assert!(deserialized.wasm.is_some());
assert!(deserialized.docker.is_some());
}
}