293 lines
8.5 KiB
Rust
293 lines
8.5 KiB
Rust
use std::{collections::HashSet, fmt};
|
|
|
|
use clap::ValueEnum;
|
|
use irtt_client::{
|
|
managed::{ManagedPacing, ManagedTargetConfig, TargetAuth, TargetId},
|
|
Authentication, HmacKey,
|
|
};
|
|
|
|
/// One raw positional target captured by Clap.
|
|
///
|
|
/// Parsing happens during run preparation so invalid target syntax never makes
|
|
/// Clap include the original argument (which may carry an HMAC key) in an
|
|
/// error diagnostic.
|
|
#[derive(Clone, PartialEq, Eq)]
|
|
pub struct TargetArg {
|
|
input: String,
|
|
}
|
|
|
|
impl TargetArg {
|
|
pub fn new(input: impl Into<String>) -> Self {
|
|
Self {
|
|
input: input.into(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl fmt::Debug for TargetArg {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
f.write_str("TargetArg(<redacted>)")
|
|
}
|
|
}
|
|
|
|
/// Capture one positional target without exposing it through Clap diagnostics.
|
|
pub fn parse_target(input: &str) -> Result<TargetArg, String> {
|
|
Ok(TargetArg::new(input))
|
|
}
|
|
|
|
#[derive(Clone, PartialEq, Eq)]
|
|
pub struct TargetSpec {
|
|
pub label: String,
|
|
pub addr: String,
|
|
pub auth: TargetAuth,
|
|
}
|
|
|
|
impl fmt::Debug for TargetSpec {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
f.debug_struct("TargetSpec")
|
|
.field("label", &self.label)
|
|
.field("addr", &self.addr)
|
|
.field("auth", &self.auth)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct PreparedTarget {
|
|
pub label: String,
|
|
pub managed: ManagedTargetConfig,
|
|
}
|
|
|
|
impl fmt::Debug for PreparedTarget {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
f.debug_struct("PreparedTarget")
|
|
.field("label", &self.label)
|
|
.field("server_addr", &self.managed.server_addr)
|
|
.field("auth", &self.managed.auth)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
|
|
pub enum GroupPacingArg {
|
|
Staggered,
|
|
Burst,
|
|
}
|
|
|
|
impl From<GroupPacingArg> for ManagedPacing {
|
|
fn from(value: GroupPacingArg) -> Self {
|
|
match value {
|
|
GroupPacingArg::Staggered => Self::Staggered,
|
|
GroupPacingArg::Burst => Self::Burst,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
enum TargetParseError {
|
|
EmptyLabel,
|
|
EmptyAddress,
|
|
InvalidEscape,
|
|
EmptyTarget,
|
|
}
|
|
|
|
impl fmt::Display for TargetParseError {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
let message = match self {
|
|
Self::EmptyLabel => "target label must not be empty",
|
|
Self::EmptyAddress => "target address must not be empty",
|
|
Self::InvalidEscape => "invalid target escape sequence",
|
|
Self::EmptyTarget => "target must not be empty",
|
|
};
|
|
f.write_str(message)
|
|
}
|
|
}
|
|
|
|
fn split_unescaped(input: &str, delimiter: char) -> Result<Vec<&str>, TargetParseError> {
|
|
let mut parts = Vec::new();
|
|
let mut start = 0;
|
|
let mut escaped = false;
|
|
for (index, ch) in input.char_indices() {
|
|
if escaped {
|
|
escaped = false;
|
|
} else if ch == '\\' {
|
|
escaped = true;
|
|
} else if ch == delimiter {
|
|
parts.push(&input[start..index]);
|
|
start = index + ch.len_utf8();
|
|
}
|
|
}
|
|
if escaped {
|
|
return Err(TargetParseError::InvalidEscape);
|
|
}
|
|
parts.push(&input[start..]);
|
|
Ok(parts)
|
|
}
|
|
|
|
fn first_unescaped(input: &str, delimiter: char) -> Result<Option<usize>, TargetParseError> {
|
|
let mut escaped = false;
|
|
for (index, ch) in input.char_indices() {
|
|
if escaped {
|
|
escaped = false;
|
|
} else if ch == '\\' {
|
|
escaped = true;
|
|
} else if ch == delimiter {
|
|
return Ok(Some(index));
|
|
}
|
|
}
|
|
if escaped {
|
|
return Err(TargetParseError::InvalidEscape);
|
|
}
|
|
Ok(None)
|
|
}
|
|
|
|
fn unescape(input: &str) -> Result<String, TargetParseError> {
|
|
let mut output = String::with_capacity(input.len());
|
|
let mut escaped = false;
|
|
for ch in input.chars() {
|
|
if escaped {
|
|
if !matches!(ch, '\\' | '=' | ';' | ',' | '@') {
|
|
return Err(TargetParseError::InvalidEscape);
|
|
}
|
|
output.push(ch);
|
|
escaped = false;
|
|
} else if ch == '\\' {
|
|
escaped = true;
|
|
} else {
|
|
output.push(ch);
|
|
}
|
|
}
|
|
if escaped {
|
|
return Err(TargetParseError::InvalidEscape);
|
|
}
|
|
Ok(output)
|
|
}
|
|
|
|
fn hmac_modifier(input: &str) -> Result<Option<usize>, TargetParseError> {
|
|
let mut escaped = false;
|
|
for (index, ch) in input.char_indices() {
|
|
if escaped {
|
|
escaped = false;
|
|
} else if ch == '\\' {
|
|
escaped = true;
|
|
} else if input[index..].starts_with("@hmac=") {
|
|
return Ok(Some(index));
|
|
}
|
|
}
|
|
if escaped {
|
|
return Err(TargetParseError::InvalidEscape);
|
|
}
|
|
Ok(None)
|
|
}
|
|
|
|
fn parse_target_syntax(
|
|
input: &str,
|
|
) -> Result<(Option<String>, String, TargetAuth), TargetParseError> {
|
|
if input.is_empty() {
|
|
return Err(TargetParseError::EmptyTarget);
|
|
}
|
|
let (main, auth) = match hmac_modifier(input)? {
|
|
Some(index) => {
|
|
let value = unescape(&input[index + "@hmac=".len()..])?;
|
|
let auth = if value.is_empty() {
|
|
TargetAuth::Override(Authentication::Unauthenticated)
|
|
} else {
|
|
TargetAuth::Override(Authentication::Hmac(HmacKey::new(value.into_bytes())))
|
|
};
|
|
(&input[..index], auth)
|
|
}
|
|
None => (input, TargetAuth::Inherit),
|
|
};
|
|
let (label, address) = match first_unescaped(main, '=')? {
|
|
None => (None, unescape(main)?),
|
|
Some(index) => {
|
|
let label = unescape(&main[..index])?;
|
|
let address = unescape(&main[index + '='.len_utf8()..])?;
|
|
(Some(label), address)
|
|
}
|
|
};
|
|
if label.as_deref().is_some_and(str::is_empty) {
|
|
return Err(TargetParseError::EmptyLabel);
|
|
}
|
|
if address.is_empty() {
|
|
return Err(TargetParseError::EmptyAddress);
|
|
}
|
|
Ok((label, address, auth))
|
|
}
|
|
|
|
pub(crate) fn target_specs_with_empty(
|
|
targets: &[TargetArg],
|
|
allow_empty: bool,
|
|
) -> Result<Vec<TargetSpec>, String> {
|
|
let mut specs = Vec::with_capacity(targets.len());
|
|
let mut unlabeled_counts = std::collections::HashMap::<String, usize>::new();
|
|
for (index, target) in targets.iter().enumerate() {
|
|
let (explicit_label, addr, auth) = parse_target_syntax(&target.input)
|
|
.map_err(|error| format!("invalid target {}: {error}", index + 1))?;
|
|
let label = match explicit_label {
|
|
Some(label) => label,
|
|
None => {
|
|
let count = unlabeled_counts.entry(addr.clone()).or_default();
|
|
*count += 1;
|
|
if *count == 1 {
|
|
addr.clone()
|
|
} else {
|
|
format!("{}#{}", addr, *count)
|
|
}
|
|
}
|
|
};
|
|
specs.push(TargetSpec { label, addr, auth });
|
|
}
|
|
|
|
if specs.is_empty() && !allow_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("duplicate target label".to_owned());
|
|
}
|
|
}
|
|
|
|
Ok(specs)
|
|
}
|
|
|
|
pub fn target_specs(targets: &[TargetArg]) -> Result<Vec<TargetSpec>, String> {
|
|
target_specs_with_empty(targets, false)
|
|
}
|
|
|
|
pub fn prepare_managed_targets(specs: Vec<TargetSpec>) -> Result<Vec<PreparedTarget>, String> {
|
|
let mut targets = Vec::with_capacity(specs.len());
|
|
for spec in specs {
|
|
let mut managed = ManagedTargetConfig::new(TargetId::from(spec.label.clone()), spec.addr);
|
|
managed.auth = spec.auth;
|
|
targets.push(PreparedTarget {
|
|
label: spec.label,
|
|
managed,
|
|
});
|
|
}
|
|
Ok(targets)
|
|
}
|
|
|
|
/// Parse one complete stdin target set after its line terminator was removed.
|
|
///
|
|
/// Commas frame stdin elements only; positional target arguments do not use
|
|
/// this framing and may therefore contain raw commas.
|
|
pub fn parse_stdin_target_set(
|
|
record: &str,
|
|
maximum_targets: usize,
|
|
) -> Result<Vec<PreparedTarget>, String> {
|
|
if record == "[]" {
|
|
return Ok(Vec::new());
|
|
}
|
|
let elements = split_unescaped(record, ',').map_err(|error| error.to_string())?;
|
|
if elements.len() > maximum_targets {
|
|
return Err(format!(
|
|
"target set exceeds the {maximum_targets}-target limit"
|
|
));
|
|
}
|
|
let args = elements.into_iter().map(TargetArg::new).collect::<Vec<_>>();
|
|
prepare_managed_targets(target_specs_with_empty(&args, true)?)
|
|
}
|