Add global configuration file support

src/client.rs

@@ -21,15 +21,245 @@ use std::io::{Read, Write};
 use std::net::{TcpStream, Shutdown};
 use std::path::Path;
 use std::process::exit;
-use std::str;
+use std::{fs, str};
 use json::object;
 use log::*;
 
-use crate::constants::MESSAGE_ACK;
+use crate::cli::{DEFAULT_ADDRESS, DEFAULT_PORT};
+use crate::constants::*;
 use crate::pipeline::streamer::streamer;
+use std::{collections::HashMap, env, path::PathBuf};
+
+use config::Config;
 
 const BUFFER_SIZE: usize = 32768 * 4;
 
+
+pub struct ClientConfig {
+    log_file: String,
+    address: String,
+    port: String,
+    id: String,
+    dependencies: String,
+}
+
+impl ClientConfig {
+    fn new(log_file: String, address: String, port: String, id: String, dependencies: String) -> Self {
+        ClientConfig {
+            log_file,
+            address,
+            port,
+            id,
+            dependencies,
+        }
+    }
+
+    pub fn get_log_file(&self) -> &str {
+        &self.log_file
+    }
+
+    pub fn get_address(&self) -> &str {
+        &self.address
+    }
+
+    pub fn get_port(&self) -> &str {
+        &self.port
+    }
+
+    pub fn get_id(&self) -> &str {
+        &self.id
+    }
+
+    pub fn get_dependencies(&self) -> &str {
+        &self.dependencies
+    }
+}
+
+const CONFIG_KEY_ID: &str = "id";
+const CONFIG_KEY_DEPS: &str = "dependencies";
+const CONFIG_KEY_ADDR: &str = "address";
+const CONFIG_KEY_PORT: &str = "port";
+const CONFIG_KEY_LOG: &str = "log-file";
+
+pub fn load_config_file<P: AsRef<Path>>(images_dir: P, action: &str) -> ClientConfig {
+    let images_dir = images_dir.as_ref();
+    let local_config_file = images_dir.join(Path::new(CONFIG_FILE));
+
+    // Handle per-process configuration workflow
+    if local_config_file.is_file() {
+        // Example of per-process config file:
+        // {
+        //    "id": "A",
+        //    "dependencies": "B:C",
+        //    "address": "127.0.0.1",
+        //    "port": "8080",
+        //    "log-file": "/var/log/criu-coordinator.log"
+        // }
+        let settings = Config::builder().add_source(config::File::from(local_config_file)).build().unwrap();
+        let settings_map = settings.try_deserialize::<HashMap<String, String>>().unwrap();
+
+        return ClientConfig::new(
+            settings_map.get(CONFIG_KEY_LOG).cloned().unwrap_or_else(|| "-".to_string()),
+            settings_map.get(CONFIG_KEY_ADDR).cloned().unwrap_or_else(|| DEFAULT_ADDRESS.to_string()),
+            settings_map.get(CONFIG_KEY_PORT).cloned().unwrap_or_else(|| DEFAULT_PORT.to_string()),
+            settings_map.get(CONFIG_KEY_ID).unwrap().clone(),
+            settings_map.get(CONFIG_KEY_DEPS).cloned().unwrap_or_default(),
+        );
+    }
+
+    // The following allows us to load global config files from /etc/criu.
+    // This is useful for example when we want to use the same config file
+    // for multiple containers.
+    // Example of global config file:
+    // {
+    //    "address": "127.0.0.1",
+    //    "port": "8080",
+    //    "log-file": "/var/log/criu-coordinator.log",
+    //    "dependencies": {
+    //        "A": ["B", "C"],
+    //        "B": ["C", "A"],
+    //        "C": ["A"]
+    //    }
+    // }
+    // Where dependencies is a map of IDs (e.g: container IDs) to a list of dependencies.
+    let global_config_file = PathBuf::from("/etc/criu").join(Path::new(CONFIG_FILE));
+
+    if !global_config_file.is_file() {
+        panic!("Global config file {:?} is not found", global_config_file);
+    }
+
+    let global_settings = Config::builder().add_source(config::File::from(global_config_file)).build().unwrap();
+    let global_map = global_settings.try_deserialize::<HashMap<String, config::Value>>().unwrap();
+
+    let address = global_map.get(CONFIG_KEY_ADDR).map(|v| v.clone().into_string().unwrap()).unwrap_or_else(|| DEFAULT_ADDRESS.to_string());
+    let port = global_map.get(CONFIG_KEY_PORT).map(|v| v.clone().into_string().unwrap()).unwrap_or_else(|| DEFAULT_PORT.to_string());
+    let log_file = global_map.get(CONFIG_KEY_LOG).map(|v| v.clone().into_string().unwrap()).unwrap_or_else(|| "-".to_string());
+
+    if is_dump_action(action) {
+        let pid_str = env::var(ENV_INIT_PID)
+            .unwrap_or_else(|_| panic!("{} not set", ENV_INIT_PID));
+        let pid: u32 = pid_str.parse().expect("Invalid PID");
+
+
+        let deps_map: HashMap<String, Vec<String>> = global_map
+            .get(CONFIG_KEY_DEPS)
+            .unwrap_or_else(|| panic!("'{}' map is missing in global config", CONFIG_KEY_DEPS))
+            .clone().into_table().unwrap()
+            .into_iter().map(|(k, v)| {
+                let deps = v.into_array().unwrap().into_iter().map(|val| val.into_string().unwrap()).collect();
+                (k, deps)
+            }).collect();
+
+        // We first try to find a container ID.
+        let id = match find_container_id_from_pid(pid) {
+            Ok(container_id) => container_id,
+            Err(_) => {
+                // If the PID is not in a container cgroup, we consider it a regular process.
+                // We identify it by its process name from /proc/<pid>/comm.
+                let process_name_path = format!("/proc/{pid}/comm");
+                if let Ok(name) = fs::read_to_string(process_name_path) {
+                    name.trim().to_string()
+                } else {
+                    // Fallback to using the PID as the ID if comm is unreadable
+                    pid.to_string()
+                }
+            }
+        };
+
+        let dependencies = find_dependencies_in_global_config(&deps_map, &id).unwrap();
+
+        // Write the local config for each container during dump
+        if action == ACTION_PRE_DUMP || action == ACTION_PRE_STREAM {
+            write_checkpoint_config(images_dir, &id, &dependencies);
+        }
+
+        ClientConfig::new(
+            log_file,
+            address,
+            port,
+            id,
+            dependencies,
+        )
+    } else { // Restore action
+        if !local_config_file.is_file() {
+            panic!("Restore action initiated, but no {CONFIG_FILE} found in the image directory {:?}", images_dir);
+        }
+
+        let local_settings = Config::builder().add_source(config::File::from(local_config_file)).build().unwrap();
+        let local_map = local_settings.try_deserialize::<HashMap<String, String>>().unwrap();
+
+        ClientConfig::new(
+            log_file,
+            address,
+            port,
+            local_map.get(CONFIG_KEY_ID).unwrap().clone(),
+            local_map.get(CONFIG_KEY_DEPS).cloned().unwrap_or_default(),
+        )
+    }
+}
+
+ /// Find containers dependencies by matching the discovered ID as a prefix of a key in the map
+fn find_dependencies_in_global_config(
+    deps_map: &HashMap<String, Vec<String>>,
+    id: &str,
+) -> Result<String, String> {
+    let deps = deps_map
+        .iter()
+        .find(|(key, _)| id.starts_with(*key))
+        .map(|(_, deps)| deps.join(":"))
+        .ok_or_else(|| {
+            format!("No dependency entry found for container ID matching '{id}'")
+        })?;
+    Ok(deps)
+}
+
+/// Find a container ID from the host PID by inspecting the process's cgroup.
+fn find_container_id_from_pid(pid: u32) -> Result<String, String> {
+    let cgroup_path = format!("/proc/{pid}/cgroup");
+    let cgroup_content = fs::read_to_string(&cgroup_path)
+        .map_err(|e| format!("Failed to read {cgroup_path}: {e}"))?;
+
+    let mut container_id: Option<String> = None;
+    for line in cgroup_content.lines() {
+        if line.len() < 64 {
+            continue;
+        }
+        for i in 0..=(line.len() - 64) {
+            let potential_id = &line[i..i + 64];
+            if potential_id.chars().all(|c| c.is_ascii_hexdigit()) {
+                let is_start = i == 0 || !line.chars().nth(i - 1).unwrap().is_ascii_hexdigit();
+                let is_end = (i + 64 == line.len())
+                    || !line.chars().nth(i + 64).unwrap().is_ascii_hexdigit();
+                if is_start && is_end {
+                    container_id = Some(potential_id.to_string());
+                }
+            }
+        }
+    }
+
+    container_id.ok_or_else(|| {
+        format!("Could not find container ID from cgroup file for PID {pid}")
+    })
+}
+
+/// Write per-checkpoint configuration file into the checkpoint images directory.
+fn write_checkpoint_config(img_dir: &Path, id: &str, dependencies: &str) {
+    let config_path = img_dir.join(CONFIG_FILE);
+    let content = format!("{{\n   \"id\": \"{id}\",\n   \"dependencies\": \"{dependencies}\"\n}}",);
+
+    fs::write(&config_path, content)
+        .unwrap_or_else(|_| panic!("Failed to write checkpoint config file to {:?}", config_path))
+}
+
+
+pub fn is_dump_action(action: &str) -> bool {
+    matches!(action, ACTION_PRE_DUMP | ACTION_NETWORK_LOCK | ACTION_POST_DUMP | ACTION_PRE_STREAM)
+}
+
+pub fn is_restore_action(action: &str) -> bool {
+    matches!(action, ACTION_PRE_RESTORE | ACTION_POST_RESTORE | ACTION_NETWORK_UNLOCK | ACTION_POST_RESUME)
+}
+
 pub fn run_client(address: &str, port: u16, id: &str, deps: &str, action: &str, images_dir: &Path, enable_streaming: bool) {
     let server_address = format!("{address}:{port}");
 

src/constants.rs

@@ -32,6 +32,8 @@ pub const ACTION_ADD_DEPENDENCIES: &str = "add-dependencies";
 pub const ENV_ACTION: &str = "CRTOOLS_SCRIPT_ACTION";
 /// ENV_IMAGE_DIR specifies path as used a base directory for CRIU images.
 pub const ENV_IMAGE_DIR: &str = "CRTOOLS_IMAGE_DIR";
+/// ENV_INIT_PID specifies the PID of the process to be checkpointed by CRIU.
+pub const ENV_INIT_PID: &str = "CRTOOLS_INIT_PID";
 
 /// Unix socket used for "criu dump".
 pub const IMG_STREAMER_CAPTURE_SOCKET_NAME: &str = "streamer-capture.sock";

src/main.rs

@@ -25,92 +25,31 @@ mod pipeline;
 mod logger;
 
 use constants::*;
-use config::Config;
-use std::collections::HashMap;
+
 use std::{env, path::PathBuf, process::exit, fs, os::unix::prelude::FileTypeExt};
-use std::path::Path;
 
 use clap::{CommandFactory, Parser};
 use clap_complete::{generate, Shell};
 use std::io;
 
-use cli::{Opts, Mode, DEFAULT_ADDRESS, DEFAULT_PORT};
+use cli::{Opts, Mode};
 use client::run_client;
 use server::run_server;
 use logger::init_logger;
 
-struct ClientConfig {
-    log_file: String,
-    address: String,
-    port: String,
-    id: String,
-    dependencies: String,
-}
-
-const CONFIG_KEY_ID: &str = "id";
-const CONFIG_KEY_DEPS: &str = "dependencies";
-const CONFIG_KEY_ADDR: &str = "address";
-const CONFIG_KEY_PORT: &str = "port";
-const CONFIG_KEY_LOG: &str = "log-file";
-
-fn load_config_file<P: AsRef<Path>>(images_dir: P) -> ClientConfig {
-    let images_dir = images_dir.as_ref();
-    let mut config_file = images_dir.join(Path::new(CONFIG_FILE));
-    if !config_file.is_file() {
-        // The following allows us to load global config files from /etc/criu.
-        // This is useful for example when we want to use the same config file
-        // for multiple containers.
-        let config_dir = PathBuf::from("/etc/criu");
-        config_file = config_dir.join(Path::new(CONFIG_FILE));
-        if !config_file.is_file() {
-            panic!("config file does not exist")
-        }
-    }
-
-    let settings = Config::builder().add_source(config::File::from(config_file)).build().unwrap();
-    let settings_map = settings.try_deserialize::<HashMap<String, String>>().unwrap();
+use crate::client::{load_config_file, is_dump_action, is_restore_action};
 
-    if !settings_map.contains_key(CONFIG_KEY_ID) {
-        panic!("id missing in config file")
-    }
-    let id = settings_map.get(CONFIG_KEY_ID).unwrap();
-
-    let mut dependencies = String::new();
-    if settings_map.contains_key(CONFIG_KEY_DEPS) {
-        dependencies = settings_map.get(CONFIG_KEY_DEPS).unwrap().to_string();
-    }
-
-    let mut address = DEFAULT_ADDRESS;
-    if settings_map.contains_key(CONFIG_KEY_ADDR) {
-        address = settings_map.get(CONFIG_KEY_ADDR).unwrap();
-    }
-
-    let mut port = DEFAULT_PORT;
-    if settings_map.contains_key(CONFIG_KEY_PORT) {
-        port = settings_map.get(CONFIG_KEY_PORT).unwrap();
-    }
-
-    let mut log_file = "-";
-    if settings_map.contains_key(CONFIG_KEY_LOG) {
-        log_file = settings_map.get(CONFIG_KEY_LOG).unwrap();
-    }
-
-    ClientConfig {
-        log_file: log_file.to_string(),
-        address: address.to_string(),
-        port: port.to_string(),
-        id: id.to_string(),
-        dependencies,
-    }
-}
 
 fn main() {
     if let Ok(action) = env::var(ENV_ACTION) {
+        if !is_dump_action(&action) && !is_restore_action(&action) {
+            exit(0)
+        }
 
         let images_dir = PathBuf::from(env::var(ENV_IMAGE_DIR)
             .unwrap_or_else(|_| panic!("Missing {} environment variable", ENV_IMAGE_DIR)));
 
-        let client_config = load_config_file(&images_dir);
+        let client_config = load_config_file(&images_dir, &action);
 
         // Ignore all action hooks other than "pre-stream", "pre-dump" and "pre-restore".
         let enable_streaming = match action.as_str() {
@@ -127,20 +66,16 @@ fn main() {
                     Err(_) => false
                 }
             },
-            ACTION_PRE_RESTORE => false,
-            ACTION_POST_DUMP => false,
-            ACTION_NETWORK_LOCK => false,
-            ACTION_NETWORK_UNLOCK => false,
-            _ => exit(0)
+            _ => false,
         };
 
-        init_logger(Some(&images_dir), client_config.log_file);
+        init_logger(Some(&images_dir), client_config.get_log_file().to_string());
 
         run_client(
-            &client_config.address,
-            client_config.port.parse().unwrap(),
-            &client_config.id,
-            &client_config.dependencies,
+            client_config.get_address(),
+            client_config.get_port().parse().unwrap(),
+            client_config.get_id(),
+            client_config.get_dependencies(),
             &action,
             &images_dir,
             enable_streaming