1use crate::{CliError, CliErrorKind};
2
3#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4pub struct ArgLimits {
5 pub max_words: usize,
6 pub max_word_bytes: usize,
7 pub max_total_bytes: usize,
8}
9
10#[derive(Clone, Debug, PartialEq, Eq)]
11pub struct RawArgs {
12 words: Vec<String>,
13}
14
15impl Default for ArgLimits {
16 fn default() -> Self {
17 Self {
18 max_words: 1024,
19 max_word_bytes: 64 * 1024,
20 max_total_bytes: 1024 * 1024,
21 }
22 }
23}
24
25impl RawArgs {
26 pub fn from_env() -> Result<Self, CliError> {
27 let mut words = Vec::new();
28 for (index, word) in std::env::args_os().skip(1).enumerate() {
29 let word = word.into_string().map_err(|_| {
30 CliError::new(
31 CliErrorKind::InvalidInput,
32 format!("argument {} is not valid UTF-8", index + 1),
33 )
34 })?;
35 words.push(word);
36 }
37 Self::parse_with_limits(words, ArgLimits::default())
38 }
39
40 pub fn parse<I>(args: I) -> Result<Self, CliError>
41 where
42 I: IntoIterator,
43 I::Item: Into<String>,
44 {
45 Self::parse_with_limits(args, ArgLimits::default())
46 }
47
48 pub fn parse_with_limits<I>(args: I, limits: ArgLimits) -> Result<Self, CliError>
49 where
50 I: IntoIterator,
51 I::Item: Into<String>,
52 {
53 let mut words = Vec::new();
54 let mut total_bytes = 0usize;
55 for arg in args {
56 let value = arg.into();
57 validate_raw_word(&value, words.len(), total_bytes, limits)?;
58 total_bytes = total_bytes.saturating_add(value.len());
59 words.push(value);
60 }
61 Ok(Self { words })
62 }
63
64 pub fn words(&self) -> &[String] {
65 &self.words
66 }
67}
68
69fn validate_raw_word(
70 value: &str,
71 count: usize,
72 total: usize,
73 limits: ArgLimits,
74) -> Result<(), CliError> {
75 if count >= limits.max_words {
76 return Err(CliError::new(
77 CliErrorKind::InvalidInput,
78 format!("argument count exceeds limit of {}", limits.max_words),
79 ));
80 }
81 if value.len() > limits.max_word_bytes {
82 return Err(CliError::new(
83 CliErrorKind::InvalidInput,
84 format!(
85 "argument {} exceeds byte limit of {}",
86 count + 1,
87 limits.max_word_bytes
88 ),
89 ));
90 }
91 if total.saturating_add(value.len()) > limits.max_total_bytes {
92 return Err(CliError::new(
93 CliErrorKind::InvalidInput,
94 format!(
95 "total argument bytes exceed limit of {}",
96 limits.max_total_bytes
97 ),
98 ));
99 }
100 if value.contains('\0') {
101 return Err(CliError::new(
102 CliErrorKind::InvalidInput,
103 format!("argument {} contains a NUL byte", count + 1),
104 ));
105 }
106 Ok(())
107}