From fd83e8521af57af6f171a16205516c6ded63db22 Mon Sep 17 00:00:00 2001 From: Nicolas Viennot Date: Tue, 17 Aug 2021 23:57:58 +0000 Subject: [PATCH] Add --tcp-listen-remap CLI option Signed-off-by: Nicolas Viennot --- README.md | 5 ++++- src/extract.rs | 1 + src/main.rs | 46 +++++++++++++++++++++++++++++++++++++++++++--- tests/tests.rs | 2 +- 4 files changed, 49 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 3c2b4543..5d06509d 100644 --- a/README.md +++ b/README.md @@ -68,7 +68,10 @@ OPTIONS: Multiple external files may be passed as a comma separated list. -p, --progress-fd File descriptor where to report progress. Defaults to 2. - + --tcp-listen-remap ... When serving the image, remap on the fly the TCP listen socket + ports. Format is old_port:new_port. May only be used with the + serve operation. Multiple tcp port remaps may be passed as a comma + separated list. SUBCOMMANDS: capture Capture a CRIU image serve Serve a captured CRIU image to CRIU diff --git a/src/extract.rs b/src/extract.rs index d847dad4..e6d6f711 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -356,6 +356,7 @@ pub fn serve(images_dir: &Path, mut progress_pipe: fs::File, shard_pipes: Vec, ext_file_pipes: Vec<(String, UnixPipe)>, + tcp_listen_remaps: Vec<(u16, u16)>, ) -> Result<()> { create_dir_all(images_dir)?; diff --git a/src/main.rs b/src/main.rs index 5c204980..0b830fd9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -37,8 +37,8 @@ use anyhow::{Result, Context}; fn parse_ext_fd(s: &str) -> Result<(String, i32)> { let mut parts = s.split(':'); - Ok(match (parts.next(), parts.next()) { - (Some(filename), Some(fd)) => { + Ok(match (parts.next(), parts.next(), parts.next()) { + (Some(filename), Some(fd), None) => { let filename = filename.to_string(); let fd = fd.parse().context("Provided ext fd is not an integer")?; (filename, fd) @@ -47,6 +47,18 @@ fn parse_ext_fd(s: &str) -> Result<(String, i32)> { }) } +fn parse_port_remap(s: &str) -> Result<(u16, u16)> { + let mut parts = s.split(':'); + Ok(match (parts.next(), parts.next(), parts.next()) { + (Some(old_port), Some(new_port), None) => { + let old_port = old_port.parse().context("Provided old_port is not a u16 integer")?; + let new_port = new_port.parse().context("Provided new_port is not a u16 integer")?; + (old_port, new_port) + }, + _ => bail!("Format is old_port:new_port") + }) +} + #[derive(StructOpt, PartialEq, Debug)] #[structopt(about, // When showing --help, we want to keep the order of arguments defined @@ -81,6 +93,12 @@ struct Opts { #[structopt(short, long)] progress_fd: Option, + /// When serving the image, remap on the fly the TCP listen socket ports. + /// Format is old_port:new_port. May only be used with the serve operation. + /// Multiple tcp port remaps may be passed as a comma separated list. + #[structopt(long, parse(try_from_str=parse_port_remap), require_delimiter = true)] + tcp_listen_remap: Vec<(u16, u16)>, + #[structopt(subcommand)] operation: Operation, } @@ -128,10 +146,13 @@ fn main() -> Result<()> { .map(|(filename, fd)| Ok((filename, UnixPipe::new(fd)?))) .collect::>()?; + ensure!(opts.operation == Serve || opts.tcp_listen_remap.is_empty(), + "--tcp-listen-remap is only supported when serving the image"); + match opts.operation { Capture => capture(&opts.images_dir, progress_pipe, shard_pipes, ext_file_pipes), Extract => extract(&opts.images_dir, progress_pipe, shard_pipes, ext_file_pipes), - Serve => serve(&opts.images_dir, progress_pipe, shard_pipes, ext_file_pipes), + Serve => serve(&opts.images_dir, progress_pipe, shard_pipes, ext_file_pipes, opts.tcp_listen_remap), } } @@ -147,6 +168,7 @@ mod cli_tests { images_dir: PathBuf::from("imgdir"), shard_fds: vec![], ext_file_fds: vec![], + tcp_listen_remap: vec![], progress_fd: None, operation: Operation::Capture, }) @@ -159,6 +181,7 @@ mod cli_tests { images_dir: PathBuf::from("imgdir"), shard_fds: vec![], ext_file_fds: vec![], + tcp_listen_remap: vec![], progress_fd: None, operation: Operation::Extract, }) @@ -171,6 +194,7 @@ mod cli_tests { images_dir: PathBuf::from("imgdir"), shard_fds: vec![], ext_file_fds: vec![], + tcp_listen_remap: vec![], progress_fd: None, operation: Operation::Serve, }) @@ -184,6 +208,7 @@ mod cli_tests { images_dir: PathBuf::from("imgdir"), shard_fds: vec![1,2,3], ext_file_fds: vec![], + tcp_listen_remap: vec![], progress_fd: None, operation: Operation::Capture, }) @@ -196,11 +221,25 @@ mod cli_tests { images_dir: PathBuf::from("imgdir"), shard_fds: vec![], ext_file_fds: vec![(String::from("file1"), 1), (String::from("file2"), 2)], + tcp_listen_remap: vec![], progress_fd: None, operation: Operation::Capture, }) } + #[test] + fn test_tcp_listen_remaps() { + assert_eq!(Opts::from_iter(&vec!["prog", "--images-dir", "imgdir", "--tcp-listen-remap", "2000:3000,5000:6000", "serve"]), + Opts { + images_dir: PathBuf::from("imgdir"), + shard_fds: vec![], + ext_file_fds: vec![], + tcp_listen_remap: vec![(2000,3000),(5000,6000)], + progress_fd: None, + operation: Operation::Serve, + }) + } + #[test] fn test_progess_fd() { assert_eq!(Opts::from_iter(&vec!["prog", "--images-dir", "imgdir", "--progress-fd", "3", "capture"]), @@ -208,6 +247,7 @@ mod cli_tests { images_dir: PathBuf::from("imgdir"), shard_fds: vec![], ext_file_fds: vec![], + tcp_listen_remap: vec![], progress_fd: Some(3), operation: Operation::Capture, }) diff --git a/tests/tests.rs b/tests/tests.rs index 10e94870..97ef7c70 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -104,7 +104,7 @@ trait TestImpl { thread::spawn(move || { if serve_image { - serve(&images_dir, extract_progress_w, shard_pipes_r, ext_files) + serve(&images_dir, extract_progress_w, shard_pipes_r, ext_files, vec![]) .expect("serve() failed"); } else { extract(&images_dir, extract_progress_w, shard_pipes_r, ext_files)