Skip to content

Commit 7632f3a

Browse files
nbdd0121nbdd0121
committed
Move argument parsing to cli module
1 parent 5d05a35 commit 7632f3a

File tree

4 files changed

+162
-159
lines changed

4 files changed

+162
-159
lines changed

‎src/cli/mod.rs

Lines changed: 75 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,82 @@
11
pub mod device;
22
pub mod logfmt;
33
pub mod symlink;
4-
pub mod timeout;
4+
5+
use std::time::Duration;
6+
7+
use clap::{Parser, Subcommand};
8+
use clap_verbosity_flag::{InfoLevel, Verbosity};
59

610
pub use device::Device;
711
pub use logfmt::LogFormat;
812
pub use symlink::Symlink;
9-
pub use timeout::Timeout;
13+
14+
fn parse_timeout(s: &str) -> Result<Option<Duration>, humantime::DurationError> {
15+
Ok(match s {
16+
"inf" | "infinite" | "none" | "forever" => None,
17+
_ => Some(humantime::parse_duration(s)?),
18+
})
19+
}
20+
21+
#[derive(Parser)]
22+
pub struct Args {
23+
#[command(flatten)]
24+
pub verbosity: Verbosity<InfoLevel>,
25+
26+
#[arg(
27+
short = 'L',
28+
long,
29+
default_value = "+l-pmt",
30+
id = "FORMAT",
31+
global = true
32+
)]
33+
/// Log mesage format: [[+-][ltmp]*]* {n}
34+
/// +/-: enable/disable {n}
35+
/// l: level {n}
36+
/// t: timestamp {n}
37+
/// m/p: module name/path {n}
38+
///
39+
pub log_format: LogFormat,
40+
41+
#[command(subcommand)]
42+
pub action: Action,
43+
}
44+
45+
#[derive(Subcommand)]
46+
#[command(max_term_width = 180)]
47+
pub enum Action {
48+
/// Wraps a call to `docker run` to allow hot-plugging devices into a
49+
/// container as they are plugged
50+
Run(Run),
51+
}
52+
53+
#[derive(clap::Args)]
54+
pub struct Run {
55+
#[arg(short = 'd', long, id = "DEVICE")]
56+
/// Root hotplug device: [[parent-of:]*]<PREFIX>:<DEVICE> {n}
57+
/// PREFIX can be: {n}
58+
/// - usb: A USB device identified as <VID>[:<PID>[:<SERIAL>]] {n}
59+
/// - syspath: A directory path in /sys/** {n}
60+
/// - devnode: A device path in /dev/** {n}
61+
/// e.g., parent-of:usb:2b3e:c310
62+
pub root_device: Device,
63+
64+
#[arg(short = 'l', long, id = "SYMLINK")]
65+
/// Create a symlink for a device: <PREFIX>:<DEVICE>=<PATH> {n}
66+
/// PREFIX can be: {n}
67+
/// - usb: A USB device identified as <VID>:<PID>:<INTERFACE> {n}
68+
/// e.g., usb:2b3e:c310:1=/dev/ttyACM_CW310_0
69+
pub symlink: Vec<Symlink>,
70+
71+
#[arg(short = 'u', long, default_value = "5", id = "CODE")]
72+
/// Exit code to return when the root device is unplugged
73+
pub root_unplugged_exit_code: u8,
74+
75+
#[arg(short = 't', long, default_value = "20s", id = "TIMEOUT", value_parser = parse_timeout)]
76+
/// Timeout when waiting for the container to be removed
77+
pub remove_timeout: core::option::Option<Duration>, // needs to be `core::option::Option` because `Option` is treated specially by clap.
78+
79+
#[arg(trailing_var_arg = true, id = "ARGS")]
80+
/// Arguments to pass to `docker run`
81+
pub docker_args: Vec<String>,
82+
}

‎src/cli/timeout.rs

Lines changed: 0 additions & 17 deletions
This file was deleted.

‎src/docker/container.rs

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use crate::cli::Timeout;
1+
use std::time::Duration;
22

33
use super::{IoStream, IoStreamSource};
44

@@ -22,7 +22,7 @@ impl Container {
2222
&self.id
2323
}
2424

25-
pub async fn remove(&self, timeout: Timeout) -> Result<()> {
25+
pub async fn remove(&self, timeout: Option<Duration>) -> Result<()> {
2626
log::info!("Removing container {}", self.id);
2727

2828
// Since we passed "--rm" flag, docker will automatically start removing the container.
@@ -40,10 +40,15 @@ impl Container {
4040
}
4141
.await;
4242

43-
if let Timeout::Some(duration) = timeout {
44-
let _ = tokio::time::timeout(duration, self.remove_event.clone()).await;
43+
if let Some(duration) = timeout {
44+
tokio::time::timeout(duration, self.remove_event.clone())
45+
.await?
46+
.context("no destroy event")?;
4547
} else {
46-
self.remove_event.clone().await;
48+
self.remove_event
49+
.clone()
50+
.await
51+
.context("no destroy event")?;
4752
}
4853
Ok(())
4954
}

‎src/main.rs

Lines changed: 77 additions & 135 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ mod docker;
33
mod hotplug;
44
mod util;
55

6-
use cli::{Device, LogFormat, Symlink, Timeout};
6+
use cli::{Action, Device, Symlink};
77
use docker::{Container, Docker};
88
use hotplug::{Event as HotPlugEvent, HotPlug, PluggedDevice};
99

@@ -12,66 +12,12 @@ use std::{fmt::Display, path::Path};
1212
use tokio_stream::StreamExt;
1313

1414
use anyhow::{bail, Context, Result};
15-
use clap::{Parser, Subcommand};
15+
use clap::Parser;
1616
use clap_verbosity_flag::{InfoLevel, LogLevel, Verbosity};
1717
use log::info;
1818

1919
use crate::hotplug::PluggableDevice;
2020

21-
#[derive(Parser)]
22-
#[command(max_term_width = 180)]
23-
struct Args {
24-
#[command(subcommand)]
25-
action: Action,
26-
}
27-
28-
#[derive(Subcommand)]
29-
enum Action {
30-
/// Wraps a call to `docker run` to allow hot-plugging devices into a
31-
/// container as they are plugged
32-
Run {
33-
#[arg(short = 'd', long, id = "DEVICE")]
34-
/// Root hotplug device: [[parent-of:]*]<PREFIX>:<DEVICE> {n}
35-
/// PREFIX can be: {n}
36-
/// - usb: A USB device identified as <VID>[:<PID>[:<SERIAL>]] {n}
37-
/// - syspath: A directory path in /sys/** {n}
38-
/// - devnode: A device path in /dev/** {n}
39-
/// e.g., parent-of:usb:2b3e:c310
40-
root_device: Device,
41-
42-
#[arg(short = 'l', long, id = "SYMLINK")]
43-
/// Create a symlink for a device: <PREFIX>:<DEVICE>=<PATH> {n}
44-
/// PREFIX can be: {n}
45-
/// - usb: A USB device identified as <VID>:<PID>:<INTERFACE> {n}
46-
/// e.g., usb:2b3e:c310:1=/dev/ttyACM_CW310_0
47-
symlink: Vec<Symlink>,
48-
49-
#[arg(short = 'u', long, default_value = "5", id = "CODE")]
50-
/// Exit code to return when the root device is unplugged
51-
root_unplugged_exit_code: u8,
52-
53-
#[arg(short = 't', long, default_value = "20s", id = "TIMEOUT")]
54-
/// Timeout when waiting for the container to be removed
55-
remove_timeout: Timeout,
56-
57-
#[command(flatten)]
58-
verbosity: Verbosity<InfoLevel>,
59-
60-
#[arg(short = 'L', long, default_value = "+l-pmt", id = "FORMAT")]
61-
/// Log mesage format: [[+-][ltmp]*]* {n}
62-
/// +/-: enable/disable {n}
63-
/// l: level {n}
64-
/// t: timestamp {n}
65-
/// m/p: module name/path {n}
66-
///
67-
log_format: LogFormat,
68-
69-
#[arg(trailing_var_arg = true, id = "ARGS")]
70-
/// Arguments to pass to `docker run`
71-
docker_args: Vec<String>,
72-
},
73-
}
74-
7521
#[derive(Clone)]
7622
enum Event {
7723
Add(PluggedDevice),
@@ -149,101 +95,97 @@ fn run_hotplug(
14995
}
15096
}
15197

152-
async fn hotplug_main() -> Result<u8> {
153-
let args = Args::parse();
98+
async fn run(param: cli::Run, verbosity: Verbosity<InfoLevel>) -> Result<u8> {
15499
let mut status = 0;
155100

156-
match args.action {
157-
Action::Run {
158-
verbosity,
159-
log_format,
160-
remove_timeout,
161-
root_unplugged_exit_code,
162-
root_device,
163-
symlink,
164-
docker_args,
165-
} => {
166-
let log_env = env_logger::Env::default()
167-
.filter_or("LOG", "off")
168-
.write_style_or("LOG_STYLE", "auto");
169-
170-
if !Path::new("/sys/fs/cgroup/devices/").is_dir() {
171-
bail!("Could not find cgroup v1");
172-
}
173-
174-
env_logger::Builder::from_env(log_env)
175-
.filter_module("container_hotplug", verbosity.log_level_filter())
176-
.format_timestamp(if log_format.timestamp {
177-
Some(Default::default())
178-
} else {
179-
None
180-
})
181-
.format_module_path(log_format.path)
182-
.format_target(log_format.module)
183-
.format_level(log_format.level)
184-
.init();
101+
if !Path::new("/sys/fs/cgroup/devices/").is_dir() {
102+
bail!("Could not find cgroup v1");
103+
}
185104

186-
let docker = Docker::connect_with_defaults()?;
187-
let container = docker.run(docker_args).await?;
188-
let _ = container.pipe_signals();
105+
let docker = Docker::connect_with_defaults()?;
106+
let container = docker.run(param.docker_args).await?;
107+
let _ = container.pipe_signals();
108+
109+
let hub_path = param.root_device.hub()?.syspath().to_owned();
110+
let hotplug_stream = run_hotplug(
111+
param.root_device,
112+
param.symlink,
113+
container.clone(),
114+
verbosity,
115+
);
116+
let container_stream = {
117+
let container = container.clone();
118+
async_stream::try_stream! {
119+
let status = container.wait().await?;
120+
yield Event::Stopped(container.clone(), status)
121+
}
122+
};
189123

190-
let hub_path = root_device.hub()?.syspath().to_owned();
191-
let hotplug_stream = run_hotplug(root_device, symlink, container.clone(), verbosity);
192-
let container_stream = {
193-
let container = container.clone();
194-
async_stream::try_stream! {
195-
let status = container.wait().await?;
196-
yield Event::Stopped(container.clone(), status)
124+
let stream = pin!(tokio_stream::empty()
125+
.merge(hotplug_stream)
126+
.merge(container_stream));
127+
128+
let result: Result<()> = async {
129+
tokio::pin!(stream);
130+
while let Some(event) = stream.next().await {
131+
let event = event?;
132+
info!("{}", event);
133+
match event {
134+
Event::Remove(dev) if dev.syspath() == hub_path => {
135+
info!("Hub device detached. Stopping container.");
136+
status = param.root_unplugged_exit_code;
137+
container.kill(15).await?;
138+
break;
197139
}
198-
};
199-
200-
let stream = pin!(tokio_stream::empty()
201-
.merge(hotplug_stream)
202-
.merge(container_stream));
203-
204-
let result: Result<()> = async {
205-
tokio::pin!(stream);
206-
while let Some(event) = stream.next().await {
207-
let event = event?;
208-
info!("{}", event);
209-
match event {
210-
Event::Remove(dev) if dev.syspath() == hub_path => {
211-
info!("Hub device detached. Stopping container.");
212-
status = root_unplugged_exit_code;
213-
container.kill(15).await?;
214-
break;
140+
Event::Stopped(_, code) => {
141+
status = 1;
142+
if let Ok(code) = u8::try_from(code) {
143+
// Use the container exit code, but only if it won't be confused
144+
// with the pre-defined root_unplugged_exit_code.
145+
if code != param.root_unplugged_exit_code {
146+
status = code;
215147
}
216-
Event::Stopped(_, code) => {
217-
status = 1;
218-
if let Ok(code) = u8::try_from(code) {
219-
// Use the container exit code, but only if it won't be confused
220-
// with the pre-defined root_unplugged_exit_code.
221-
if code != root_unplugged_exit_code {
222-
status = code;
223-
}
224-
} else {
225-
status = 1;
226-
}
227-
break;
228-
}
229-
_ => {}
148+
} else {
149+
status = 1;
230150
}
151+
break;
231152
}
232-
Ok(())
153+
_ => {}
233154
}
234-
.await;
235-
236-
let _ = container.remove(remove_timeout).await;
237-
result?
238155
}
239-
};
156+
Ok(())
157+
}
158+
.await;
240159

160+
let _ = container.remove(param.remove_timeout).await;
161+
result?;
241162
Ok(status)
242163
}
243164

244165
#[tokio::main]
245166
async fn main() {
246-
let code = match hotplug_main().await {
167+
let args = cli::Args::parse();
168+
169+
let log_env = env_logger::Env::default()
170+
.filter_or("LOG", "off")
171+
.write_style_or("LOG_STYLE", "auto");
172+
173+
env_logger::Builder::from_env(log_env)
174+
.filter_module("container_hotplug", args.verbosity.log_level_filter())
175+
.format_timestamp(if args.log_format.timestamp {
176+
Some(Default::default())
177+
} else {
178+
None
179+
})
180+
.format_module_path(args.log_format.path)
181+
.format_target(args.log_format.module)
182+
.format_level(args.log_format.level)
183+
.init();
184+
185+
let result = match args.action {
186+
Action::Run(param) => run(param, args.verbosity).await,
187+
};
188+
let code = match result {
247189
Ok(code) => code,
248190
Err(err) => {
249191
eprintln!("Error: {err:?}");

0 commit comments

Comments
 (0)