feat(tui): add multi-target support and improved dashboard visualization

Enable simultaneous probing of multiple targets within the TUI using a new target status table and multi-series RTT graph. Refactor target resolution logic into a shared module to support consistent behavior across CLI and TUI components.
This commit is contained in:
Nils Andreas Svee committed 2026-06-23 22:36:49 +07:00
1 parent 2065fe5e11
commit 836a8abbeb
7 files changed
+1242 -444

No files matched your search

+10 -190
View File
@@ -1,9 +1,12 @@
use std::{collections::HashSet, net::ToSocketAddrs, time::Duration};
use std::time::Duration;
use clap::{Parser, ValueEnum};
use irtt_client::{ClientConfig, ManagedGroupPacing, ManagedTargetConfig, TargetId};
use irtt_client::ClientConfig;
use crate::shared::client::{parse_test_duration, CommonClientArgs, TimestampArg};
use crate::shared::client::{
parse_labelled_target, parse_test_duration, resolved_managed_targets, target_specs,
CommonClientArgs, GroupPacingArg, LabelledTargetArg, ResolvedTarget, TargetSpec, TimestampArg,
};
pub const DEFAULT_CLIENT_DURATION: Duration = Duration::from_secs(10);
@@ -105,71 +108,13 @@ impl ClientArgs {
}
pub fn target_specs(&self) -> Result<Vec<CliTargetSpec>, String> {
let mut specs = Vec::new();
let mut positional_counts = std::collections::HashMap::<&str, usize>::new();
for target in &self.targets {
let count = positional_counts.entry(target.as_str()).or_default();
*count += 1;
let label = if *count == 1 {
target.clone()
} else {
format!("{target}#{}", *count)
};
specs.push(CliTargetSpec {
label,
addr: target.clone(),
});
}
for target in &self.labelled_targets {
specs.push(CliTargetSpec {
label: target.label.clone(),
addr: target.addr.clone(),
});
}
if specs.is_empty() {
return Err("at least one target is required unless --list-columns is set".to_owned());
}
let mut labels = HashSet::new();
for spec in &specs {
if !labels.insert(spec.label.clone()) {
return Err(format!("duplicate target label {:?}", spec.label));
}
}
Ok(specs)
target_specs(&self.targets, &self.labelled_targets)
}
pub fn resolved_managed_targets(&self) -> Result<Vec<ResolvedCliTarget>, String> {
let specs = self.target_specs()?;
let config = self.to_client_config();
let mut remotes = HashSet::new();
let mut targets = Vec::with_capacity(specs.len());
for spec in specs {
let remote = resolve_cli_target(&spec.addr, &config).map_err(|err| {
format!(
"failed to resolve target {} ({:?}): {err}",
spec.label, spec.addr
)
})?;
if !remotes.insert(remote) {
return Err(format!(
"duplicate resolved target address {remote} for label {}",
spec.label
));
}
targets.push(ResolvedCliTarget {
label: spec.label.clone(),
managed: ManagedTargetConfig {
id: TargetId::from(spec.label),
remote,
auth: None,
},
});
}
Ok(targets)
resolved_managed_targets(specs, &config)
}
fn primary_target_addr(&self) -> Option<&str> {
@@ -201,38 +146,8 @@ pub enum OutputFormat {
Jsonl,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LabelledTargetArg {
pub label: String,
pub addr: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CliTargetSpec {
pub label: String,
pub addr: String,
}
#[derive(Debug, Clone)]
pub struct ResolvedCliTarget {
pub label: String,
pub managed: ManagedTargetConfig,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
pub enum GroupPacingArg {
Staggered,
Burst,
}
impl From<GroupPacingArg> for ManagedGroupPacing {
fn from(value: GroupPacingArg) -> Self {
match value {
GroupPacingArg::Staggered => Self::Staggered,
GroupPacingArg::Burst => Self::Burst,
}
}
}
pub type CliTargetSpec = TargetSpec;
pub type ResolvedCliTarget = ResolvedTarget;
impl OutputFormat {
pub fn prints_summary(self) -> bool {
@@ -250,58 +165,6 @@ pub enum HeaderMode {
Never,
}
fn parse_labelled_target(input: &str) -> Result<LabelledTargetArg, String> {
let (label, addr) = input
.split_once('=')
.ok_or_else(|| "target must use LABEL=TARGET syntax".to_owned())?;
if label.is_empty() {
return Err("target label must not be empty".to_owned());
}
if addr.is_empty() {
return Err("target address must not be empty".to_owned());
}
Ok(LabelledTargetArg {
label: label.to_owned(),
addr: addr.to_owned(),
})
}
fn resolve_cli_target(addr: &str, config: &ClientConfig) -> Result<std::net::SocketAddr, String> {
let normalized = normalize_target_addr(addr);
let mut addrs = normalized
.to_socket_addrs()
.map_err(|_| format!("failed to resolve address {normalized:?}"))?;
addrs
.find(|addr| {
(!config.socket_config.ipv4_only || addr.is_ipv4())
&& (!config.socket_config.ipv6_only || addr.is_ipv6())
})
.ok_or_else(|| format!("failed to resolve address {normalized:?}"))
}
fn normalize_target_addr(addr: &str) -> String {
const DEFAULT_PORT: u16 = 2112;
if addr.parse::<std::net::SocketAddr>().is_ok() {
return addr.to_owned();
}
if addr.starts_with('[') && addr.ends_with(']') {
return format!("{addr}:{DEFAULT_PORT}");
}
if addr.starts_with('[') {
return addr.to_owned();
}
if addr.parse::<std::net::Ipv6Addr>().is_ok() {
return format!("[{addr}]:{DEFAULT_PORT}");
}
if addr
.rsplit_once(':')
.is_some_and(|(_, port)| port.parse::<u16>().is_ok())
{
return addr.to_owned();
}
format!("{addr}:{DEFAULT_PORT}")
}
#[cfg(test)]
mod tests {
use super::*;
@@ -406,49 +269,6 @@ mod tests {
assert!(err.contains("duplicate resolved target address 127.0.0.1:2112"));
}
#[test]
fn target_normalization_matches_single_target_forms() {
assert_eq!(normalize_target_addr("127.0.0.1"), "127.0.0.1:2112");
assert_eq!(normalize_target_addr("127.0.0.1:1234"), "127.0.0.1:1234");
assert_eq!(normalize_target_addr("example.com"), "example.com:2112");
assert_eq!(
normalize_target_addr("example.com:1234"),
"example.com:1234"
);
assert_eq!(normalize_target_addr("[::1]"), "[::1]:2112");
assert_eq!(normalize_target_addr("[::1]:1234"), "[::1]:1234");
assert_eq!(normalize_target_addr("::1"), "[::1]:2112");
}
#[test]
fn target_resolution_respects_family_filters() {
let config = ClientConfig::default();
assert_eq!(
resolve_cli_target("127.0.0.1", &config).unwrap(),
"127.0.0.1:2112".parse().unwrap()
);
assert_eq!(
resolve_cli_target("[::1]", &config).unwrap(),
"[::1]:2112".parse().unwrap()
);
assert_eq!(
resolve_cli_target("::1", &config).unwrap(),
"[::1]:2112".parse().unwrap()
);
let mut ipv4_only = ClientConfig::default();
ipv4_only.socket_config.ipv4_only = true;
assert!(resolve_cli_target("127.0.0.1", &ipv4_only)
.unwrap()
.is_ipv4());
assert!(resolve_cli_target("::1", &ipv4_only).is_err());
let mut ipv6_only = ClientConfig::default();
ipv6_only.socket_config.ipv6_only = true;
assert!(resolve_cli_target("::1", &ipv6_only).unwrap().is_ipv6());
assert!(resolve_cli_target("127.0.0.1", &ipv6_only).is_err());
}
#[test]
fn invalid_labelled_target_syntax_is_rejected() {
assert!(parse(&["--target", "missing-equals"]).is_err());
+147 -5
View File
@@ -5,15 +5,42 @@ use irtt_client::ClientConfig;
#[cfg(test)]
use crate::shared::client::TimestampArg;
use crate::shared::client::{parse_test_duration, CommonClientArgs};
use crate::shared::client::{
parse_labelled_target, parse_test_duration, resolved_managed_targets, target_specs,
CommonClientArgs, GroupPacingArg, LabelledTargetArg, ResolvedTarget, TargetSpec,
};
pub const DEFAULT_TUI_DURATION: Duration = Duration::ZERO;
#[derive(Debug, Clone, Parser)]
#[command(name = "irtt-tui", about = "Minimal IRTT-compatible TUI client")]
pub struct TuiArgs {
/// Server address or host, with optional port.
pub server: String,
/// Server address or host, with optional port. Repeat for multi-target mode.
#[arg(
value_name = "TARGET",
num_args = 0..,
required_unless_present = "labelled_targets",
long_help = "Server address or host, with optional port. Repeat the positional target for multi-target mode, for example: irtt-tui host-a:2112 host-b:2112."
)]
pub targets: Vec<String>,
/// Explicit labelled target, as label=host:port. Repeat for multi-target mode.
#[arg(
long = "target",
value_name = "LABEL=TARGET",
value_parser = parse_labelled_target,
long_help = "Explicit labelled target, as label=host:port. Repeat for multi-target mode. The label is used in the legend and status table and must be unique."
)]
pub labelled_targets: Vec<LabelledTargetArg>,
/// Managed group pacing for multi-target mode.
#[arg(
long,
value_enum,
default_value_t = GroupPacingArg::Staggered,
long_help = "Managed group pacing for multi-target mode.\n\nstaggered spaces active targets across the probe interval. burst sends one probe to every active target back-to-back once per interval."
)]
pub pacing: GroupPacingArg,
#[arg(
long,
@@ -30,19 +57,44 @@ pub struct TuiArgs {
impl TuiArgs {
pub fn to_client_config(&self) -> ClientConfig {
self.common.to_client_config(&self.server, self.duration)
self.common.to_client_config(
self.primary_target_addr()
.expect("at least one target is required"),
self.duration,
)
}
pub fn is_continuous(&self) -> bool {
self.duration == Duration::ZERO
}
pub fn target_specs(&self) -> Result<Vec<TuiTargetSpec>, String> {
target_specs(&self.targets, &self.labelled_targets)
}
pub fn resolved_managed_targets(&self) -> Result<Vec<ResolvedTuiTarget>, String> {
let specs = self.target_specs()?;
let config = self.to_client_config();
resolved_managed_targets(specs, &config)
}
#[cfg(test)]
pub fn timestamp_mode(&self) -> TimestampArg {
self.common.tstamp
}
fn primary_target_addr(&self) -> Option<&str> {
self.targets.first().map(String::as_str).or_else(|| {
self.labelled_targets
.first()
.map(|target| target.addr.as_str())
})
}
}
pub type TuiTargetSpec = TargetSpec;
pub type ResolvedTuiTarget = ResolvedTarget;
impl std::ops::Deref for TuiArgs {
type Target = CommonClientArgs;
@@ -77,7 +129,9 @@ mod tests {
#[test]
fn tui_parser_defaults_to_continuous_and_has_no_output_option() {
let args = parse(&["127.0.0.1:2112"]).unwrap();
assert_eq!(args.server, "127.0.0.1:2112");
assert_eq!(args.targets, ["127.0.0.1:2112"]);
assert!(args.labelled_targets.is_empty());
assert_eq!(args.pacing, GroupPacingArg::Staggered);
assert_eq!(args.duration, DEFAULT_TUI_DURATION);
assert!(args.is_continuous());
assert_eq!(args.to_client_config().duration, None);
@@ -94,6 +148,87 @@ mod tests {
assert!(!help.contains("--output"));
}
#[test]
fn multiple_positional_targets_parse() {
let args = parse(&["host-a:2112", "host-b:2112"]).unwrap();
let specs = args.target_specs().unwrap();
assert_eq!(specs[0].label, "host-a:2112");
assert_eq!(specs[0].addr, "host-a:2112");
assert_eq!(specs[1].label, "host-b:2112");
assert_eq!(specs[1].addr, "host-b:2112");
}
#[test]
fn repeated_labelled_targets_parse() {
let args = parse(&[
"--target",
"ams=ams.example.com:2112",
"--target",
"sg=sg.example.com:2112",
])
.unwrap();
let specs = args.target_specs().unwrap();
assert_eq!(specs[0].label, "ams");
assert_eq!(specs[0].addr, "ams.example.com:2112");
assert_eq!(specs[1].label, "sg");
assert_eq!(specs[1].addr, "sg.example.com:2112");
}
#[test]
fn at_least_one_target_is_required() {
assert!(parse(&[]).is_err());
}
#[test]
fn duplicate_labels_are_rejected() {
let args = parse(&["host-a:2112", "--target", "host-a:2112=host-b:2112"]).unwrap();
let err = args.target_specs().unwrap_err();
assert!(err.contains("duplicate target label"));
}
#[test]
fn duplicate_positional_target_strings_get_stable_suffixes() {
let args = parse(&["host-a:2112", "host-a:2112"]).unwrap();
let specs = args.target_specs().unwrap();
assert_eq!(specs[0].label, "host-a:2112");
assert_eq!(specs[1].label, "host-a:2112#2");
}
#[test]
fn duplicate_resolved_target_addresses_are_rejected() {
let args = parse(&["127.0.0.1:2112", "127.0.0.1"]).unwrap();
let err = args.resolved_managed_targets().unwrap_err();
assert!(err.contains("duplicate resolved target address 127.0.0.1:2112"));
}
#[test]
fn invalid_labelled_target_syntax_is_rejected() {
assert!(parse(&["--target", "missing-equals"]).is_err());
assert!(parse(&["--target", "=127.0.0.1:2112"]).is_err());
assert!(parse(&["--target", "label="]).is_err());
}
#[test]
fn pacing_option_accepts_supported_values() {
assert_eq!(
parse(&["--pacing", "staggered", "127.0.0.1:2112"])
.unwrap()
.pacing,
GroupPacingArg::Staggered
);
assert_eq!(
parse(&["--pacing", "burst", "127.0.0.1:2112"])
.unwrap()
.pacing,
GroupPacingArg::Burst
);
}
#[test]
fn shared_client_options_match_tui_config_mapping() {
let shared = [
@@ -158,4 +293,11 @@ mod tests {
assert_eq!(shared.negotiation_policy, tui.negotiation_policy);
assert_eq!(shared.socket_config.ttl, tui.socket_config.ttl);
}
#[test]
fn tui_help_lists_multi_target_options() {
let help = TuiArgs::command().render_help().to_string();
assert!(help.contains("--target <LABEL=TARGET>"));
assert!(help.contains("--pacing <PACING>"));
}
}
+166 -32
View File
@@ -1,39 +1,41 @@
use std::{
io,
sync::atomic::{AtomicBool, Ordering},
thread,
time::{Duration, Instant},
};
use crossterm::event::{self, Event, KeyCode, KeyEventKind, KeyModifiers};
use irtt_stats::StatsCollector;
use irtt_client::{
ClientEvent, EventSubscriptionError, ManagedClientGroup, ManagedClientGroupConfig,
ManagedGroupEndReason, SubscriberConfig, SubscriberOverflow,
};
use crate::{
cmd::tui::args::TuiArgs,
cmd::tui::args::{ResolvedTuiTarget, TuiArgs},
shared::client::{is_shutdown_requested, ClientSession},
};
use super::ui::{should_render, TuiConfig, TuiState, TuiStatus, TuiTerminal};
fn stats_config(continuous: bool) -> irtt_stats::StatsConfig {
if continuous {
irtt_stats::StatsConfig::continuous()
} else {
irtt_stats::StatsConfig::finite()
}
}
const RENDER_INTERVAL: Duration = Duration::from_millis(250);
const TUI_WAIT_SLICE: Duration = Duration::from_millis(20);
const IDLE_SLEEP: Duration = Duration::from_millis(5);
const GROUP_COMPLETION_GRACE: Duration = Duration::from_secs(1);
pub fn run_tui(
args: TuiArgs,
shutdown_requested: &AtomicBool,
) -> Result<(), Box<dyn std::error::Error>> {
let targets = args
.resolved_managed_targets()
.map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?;
let continuous = args.is_continuous();
let mut terminal = TuiTerminal::enter()?;
let mut state = TuiState::new(TuiConfig::from_args(&args));
let mut stats = StatsCollector::new(stats_config(continuous));
let mut state = TuiState::with_target_labels(
TuiConfig::from_args(&args),
targets.iter().map(|target| target.label.clone()),
);
let mut next_render = Instant::now();
if is_shutdown_requested(shutdown_requested) {
@@ -41,13 +43,24 @@ pub fn run_tui(
}
state.set_status(TuiStatus::Opening);
render_if_due(&mut terminal, &state, &stats, &mut next_render, true)?;
render_if_due(&mut terminal, &state, &mut next_render, true)?;
if targets.len() > 1 {
return run_group_tui(
args,
targets,
&mut terminal,
&mut state,
&mut next_render,
shutdown_requested,
);
}
let mut session = match ClientSession::connect(args.to_client_config(), continuous) {
Ok(session) => session,
Err(err) => {
state.set_error(err.to_string());
render_if_due(&mut terminal, &state, &stats, &mut next_render, true)?;
render_if_due(&mut terminal, &state, &mut next_render, true)?;
return Err(Box::new(err));
}
};
@@ -57,8 +70,8 @@ pub fn run_tui(
}
let events = session.open()?;
state.process_events(&events, &mut stats);
render_if_due(&mut terminal, &state, &stats, &mut next_render, true)?;
state.process_events(&events);
render_if_due(&mut terminal, &state, &mut next_render, true)?;
let mut interrupted = false;
// Keep this in lockstep with run_stream: send due probes, drain available
@@ -66,7 +79,7 @@ pub fn run_tui(
// then perform the same final drain, timeout poll, and close sequence.
while session.should_continue(shutdown_requested) {
if handle_input(&mut state, shutdown_requested)? {
render_if_due(&mut terminal, &state, &stats, &mut next_render, true)?;
render_if_due(&mut terminal, &state, &mut next_render, true)?;
}
if state.quit_requested {
interrupted = true;
@@ -74,19 +87,18 @@ pub fn run_tui(
}
let events = session.step(shutdown_requested)?;
state.process_events(&events, &mut stats);
state.process_events(&events);
if is_shutdown_requested(shutdown_requested) {
interrupted = true;
break;
}
render_if_due(&mut terminal, &state, &stats, &mut next_render, false)?;
render_if_due(&mut terminal, &state, &mut next_render, false)?;
wait_for_tui_activity(
session.next_send_deadline(),
&mut next_render,
&mut state,
&stats,
&mut terminal,
shutdown_requested,
)?;
@@ -95,32 +107,156 @@ pub fn run_tui(
if interrupted {
state.set_status(TuiStatus::Interrupted);
render_if_due(&mut terminal, &state, &stats, &mut next_render, true)?;
render_if_due(&mut terminal, &state, &mut next_render, true)?;
}
if session.should_drain_final(interrupted) {
state.set_status(TuiStatus::Draining);
let mut drain_render = Instant::now();
session.drain_final(|events| {
state.process_events(events, &mut stats);
let _ = render_if_due(&mut terminal, &state, &stats, &mut drain_render, false);
state.process_events(events);
let _ = render_if_due(&mut terminal, &state, &mut drain_render, false);
})?;
render_if_due(&mut terminal, &state, &stats, &mut next_render, true)?;
render_if_due(&mut terminal, &state, &mut next_render, true)?;
}
let events = session.poll_timeouts()?;
state.process_events(&events, &mut stats);
state.process_events(&events);
state.set_status(TuiStatus::Closing);
render_if_due(&mut terminal, &state, &stats, &mut next_render, true)?;
render_if_due(&mut terminal, &state, &mut next_render, true)?;
let events = session.close()?;
state.process_events(&events, &mut stats);
state.process_events(&events);
state.set_status(TuiStatus::Complete);
render_if_due(&mut terminal, &state, &stats, &mut next_render, true)?;
render_if_due(&mut terminal, &state, &mut next_render, true)?;
Ok(())
}
fn run_group_tui(
args: TuiArgs,
targets: Vec<ResolvedTuiTarget>,
terminal: &mut TuiTerminal,
state: &mut TuiState,
next_render: &mut Instant,
shutdown_requested: &AtomicBool,
) -> Result<(), Box<dyn std::error::Error>> {
let mut config = args.to_client_config();
let managed_targets = targets
.iter()
.map(|target| target.managed.clone())
.collect::<Vec<_>>();
let expected_target_count = managed_targets.len();
if let Some(first) = managed_targets.first() {
config.server_addr = first.remote.to_string();
}
let group_config = ManagedClientGroupConfig {
client: config,
pacing: args.pacing.into(),
};
let (session, events) = ManagedClientGroup::start_with_subscription(
group_config,
managed_targets,
SubscriberConfig {
capacity: 16_384,
overflow: SubscriberOverflow::DropOldest,
},
)?;
let mut interrupted = false;
let mut terminal_targets = std::collections::HashSet::new();
let mut saw_target_event = false;
let mut last_event_at = Instant::now();
loop {
if is_shutdown_requested(shutdown_requested) {
interrupted = true;
session.stop();
}
if handle_input(state, shutdown_requested)? {
render_if_due(terminal, state, next_render, true)?;
}
if state.quit_requested {
interrupted = true;
session.stop();
}
match events.try_recv() {
Ok(Some(target_event)) => {
saw_target_event = true;
last_event_at = Instant::now();
if is_terminal_target_event(&target_event.event) {
terminal_targets.insert(target_event.target.as_str().to_owned());
}
state.process_target_event(&target_event);
render_if_due(terminal, state, next_render, false)?;
}
Ok(None) => {
if interrupted {
break;
}
if terminal_targets.len() >= expected_target_count {
break;
}
if should_join_group_after_idle(&args, saw_target_event, last_event_at) {
break;
}
wait_for_tui_activity(None, next_render, state, terminal, shutdown_requested)?;
thread::sleep(IDLE_SLEEP);
}
Err(EventSubscriptionError::Disconnected) => break,
}
}
if interrupted {
state.set_status(TuiStatus::Interrupted);
render_if_due(terminal, state, next_render, true)?;
}
state.set_status(TuiStatus::Closing);
render_if_due(terminal, state, next_render, true)?;
let outcome = session.join()?;
while let Ok(Some(target_event)) = events.try_recv() {
state.process_target_event(&target_event);
}
if outcome.end_reason == ManagedGroupEndReason::Cancelled && !interrupted {
return Err("managed client group was cancelled".into());
}
state.set_status(TuiStatus::Complete);
render_if_due(terminal, state, next_render, true)?;
Ok(())
}
fn is_terminal_target_event(event: &ClientEvent) -> bool {
matches!(
event,
ClientEvent::SessionClosed { .. } | ClientEvent::NoTestCompleted { .. }
)
}
fn estimated_group_completion_grace(args: &TuiArgs) -> Duration {
let open_timeout: Duration = args.to_client_config().open_timeouts.iter().sum();
open_timeout
.saturating_add(args.duration)
.saturating_add(GROUP_COMPLETION_GRACE)
}
fn should_join_group_after_idle(
args: &TuiArgs,
saw_target_event: bool,
last_event_at: Instant,
) -> bool {
if args.is_continuous() && saw_target_event {
return false;
}
last_event_at.elapsed() > estimated_group_completion_grace(args)
}
fn handle_input(state: &mut TuiState, shutdown_requested: &AtomicBool) -> io::Result<bool> {
let mut force_render = false;
while event::poll(Duration::ZERO)? {
@@ -166,7 +302,6 @@ fn handle_input(state: &mut TuiState, shutdown_requested: &AtomicBool) -> io::Re
fn render_if_due(
terminal: &mut TuiTerminal,
state: &TuiState,
stats: &StatsCollector,
next_render: &mut Instant,
force: bool,
) -> io::Result<()> {
@@ -174,7 +309,7 @@ fn render_if_due(
if !should_render(now, *next_render, state.paused, force) {
return Ok(());
}
terminal.draw(state, &stats.snapshot())?;
terminal.draw(state)?;
*next_render = now + RENDER_INTERVAL;
Ok(())
}
@@ -183,7 +318,6 @@ fn wait_for_tui_activity(
next_send_deadline: Option<Instant>,
next_render: &mut Instant,
state: &mut TuiState,
stats: &StatsCollector,
terminal: &mut TuiTerminal,
shutdown_requested: &AtomicBool,
) -> io::Result<()> {
@@ -193,7 +327,7 @@ fn wait_for_tui_activity(
}
if handle_input(state, shutdown_requested)? {
render_if_due(terminal, state, stats, next_render, true)?;
render_if_due(terminal, state, next_render, true)?;
}
Ok(())
}
File diff suppressed because it is too large. Load diff
+5
View File
@@ -1,6 +1,7 @@
pub mod args;
pub mod config;
pub mod session;
pub mod targets;
pub use args::{
parse_dscp, parse_duration, parse_length, parse_server_fill, parse_test_duration, parse_ttl,
@@ -8,3 +9,7 @@ pub use args::{
};
pub use config::{expected_probe_count, DEFAULT_RECV_TIMEOUT};
pub use session::{final_drain_duration, is_shutdown_requested, ClientSession, RECV_BUDGET};
pub use targets::{
parse_labelled_target, resolved_managed_targets, target_specs, GroupPacingArg,
LabelledTargetArg, ResolvedTarget, TargetSpec,
};
@@ -0,0 +1,240 @@
use std::{collections::HashSet, net::ToSocketAddrs};
use clap::ValueEnum;
use irtt_client::{ClientConfig, ManagedGroupPacing, ManagedTargetConfig, TargetId};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LabelledTargetArg {
pub label: String,
pub addr: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TargetSpec {
pub label: String,
pub addr: String,
}
#[derive(Debug, Clone)]
pub struct ResolvedTarget {
pub label: String,
pub managed: ManagedTargetConfig,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
pub enum GroupPacingArg {
Staggered,
Burst,
}
impl From<GroupPacingArg> for ManagedGroupPacing {
fn from(value: GroupPacingArg) -> Self {
match value {
GroupPacingArg::Staggered => Self::Staggered,
GroupPacingArg::Burst => Self::Burst,
}
}
}
pub fn parse_labelled_target(input: &str) -> Result<LabelledTargetArg, String> {
let (label, addr) = input
.split_once('=')
.ok_or_else(|| "target must use LABEL=TARGET syntax".to_owned())?;
if label.is_empty() {
return Err("target label must not be empty".to_owned());
}
if addr.is_empty() {
return Err("target address must not be empty".to_owned());
}
Ok(LabelledTargetArg {
label: label.to_owned(),
addr: addr.to_owned(),
})
}
pub fn target_specs(
positional_targets: &[String],
labelled_targets: &[LabelledTargetArg],
) -> Result<Vec<TargetSpec>, String> {
let mut specs = Vec::new();
let mut positional_counts = std::collections::HashMap::<&str, usize>::new();
for target in positional_targets {
let count = positional_counts.entry(target.as_str()).or_default();
*count += 1;
let label = if *count == 1 {
target.clone()
} else {
format!("{target}#{}", *count)
};
specs.push(TargetSpec {
label,
addr: target.clone(),
});
}
for target in labelled_targets {
specs.push(TargetSpec {
label: target.label.clone(),
addr: target.addr.clone(),
});
}
if specs.is_empty() {
return Err("at least one target is required unless --list-columns is set".to_owned());
}
let mut labels = HashSet::new();
for spec in &specs {
if !labels.insert(spec.label.clone()) {
return Err(format!("duplicate target label {:?}", spec.label));
}
}
Ok(specs)
}
pub fn resolved_managed_targets(
specs: Vec<TargetSpec>,
config: &ClientConfig,
) -> Result<Vec<ResolvedTarget>, String> {
let mut remotes = HashSet::new();
let mut targets = Vec::with_capacity(specs.len());
for spec in specs {
let remote = resolve_target(&spec.addr, config).map_err(|err| {
format!(
"failed to resolve target {} ({:?}): {err}",
spec.label, spec.addr
)
})?;
if !remotes.insert(remote) {
return Err(format!(
"duplicate resolved target address {remote} for label {}",
spec.label
));
}
targets.push(ResolvedTarget {
label: spec.label.clone(),
managed: ManagedTargetConfig {
id: TargetId::from(spec.label),
remote,
auth: None,
},
});
}
Ok(targets)
}
pub fn resolve_target(addr: &str, config: &ClientConfig) -> Result<std::net::SocketAddr, String> {
let normalized = normalize_target_addr(addr);
let mut addrs = normalized
.to_socket_addrs()
.map_err(|_| format!("failed to resolve address {normalized:?}"))?;
addrs
.find(|addr| {
(!config.socket_config.ipv4_only || addr.is_ipv4())
&& (!config.socket_config.ipv6_only || addr.is_ipv6())
})
.ok_or_else(|| format!("failed to resolve address {normalized:?}"))
}
pub fn normalize_target_addr(addr: &str) -> String {
const DEFAULT_PORT: u16 = 2112;
if addr.parse::<std::net::SocketAddr>().is_ok() {
return addr.to_owned();
}
if addr.starts_with('[') && addr.ends_with(']') {
return format!("{addr}:{DEFAULT_PORT}");
}
if addr.starts_with('[') {
return addr.to_owned();
}
if addr.parse::<std::net::Ipv6Addr>().is_ok() {
return format!("[{addr}]:{DEFAULT_PORT}");
}
if addr
.rsplit_once(':')
.is_some_and(|(_, port)| port.parse::<u16>().is_ok())
{
return addr.to_owned();
}
format!("{addr}:{DEFAULT_PORT}")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn target_normalization_matches_single_target_forms() {
assert_eq!(normalize_target_addr("127.0.0.1"), "127.0.0.1:2112");
assert_eq!(normalize_target_addr("127.0.0.1:1234"), "127.0.0.1:1234");
assert_eq!(normalize_target_addr("example.com"), "example.com:2112");
assert_eq!(
normalize_target_addr("example.com:1234"),
"example.com:1234"
);
assert_eq!(normalize_target_addr("[::1]"), "[::1]:2112");
assert_eq!(normalize_target_addr("[::1]:1234"), "[::1]:1234");
assert_eq!(normalize_target_addr("::1"), "[::1]:2112");
}
#[test]
fn target_resolution_respects_family_filters() {
let config = ClientConfig::default();
assert_eq!(
resolve_target("127.0.0.1", &config).unwrap(),
"127.0.0.1:2112".parse().unwrap()
);
assert_eq!(
resolve_target("[::1]", &config).unwrap(),
"[::1]:2112".parse().unwrap()
);
assert_eq!(
resolve_target("::1", &config).unwrap(),
"[::1]:2112".parse().unwrap()
);
let mut ipv4_only = ClientConfig::default();
ipv4_only.socket_config.ipv4_only = true;
assert!(resolve_target("127.0.0.1", &ipv4_only).unwrap().is_ipv4());
assert!(resolve_target("::1", &ipv4_only).is_err());
let mut ipv6_only = ClientConfig::default();
ipv6_only.socket_config.ipv6_only = true;
assert!(resolve_target("::1", &ipv6_only).unwrap().is_ipv6());
assert!(resolve_target("127.0.0.1", &ipv6_only).is_err());
}
#[test]
fn target_specs_suffix_repeated_positionals_and_reject_duplicate_labels() {
let positionals = vec!["host-a:2112".to_owned(), "host-a:2112".to_owned()];
let specs = target_specs(&positionals, &[]).unwrap();
assert_eq!(specs[0].label, "host-a:2112");
assert_eq!(specs[1].label, "host-a:2112#2");
let labels = vec![LabelledTargetArg {
label: "host-a:2112".to_owned(),
addr: "host-b:2112".to_owned(),
}];
let err = target_specs(&positionals[..1], &labels).unwrap_err();
assert!(err.contains("duplicate target label"));
}
#[test]
fn duplicate_resolved_target_addresses_are_rejected() {
let specs = vec![
TargetSpec {
label: "a".to_owned(),
addr: "127.0.0.1:2112".to_owned(),
},
TargetSpec {
label: "b".to_owned(),
addr: "127.0.0.1".to_owned(),
},
];
let err = resolved_managed_targets(specs, &ClientConfig::default()).unwrap_err();
assert!(err.contains("duplicate resolved target address 127.0.0.1:2112"));
}
}