1use crate::{CliError, CliErrorKind};
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16pub struct ArgLimits {
17 pub max_words: usize,
18 pub max_word_bytes: usize,
19 pub max_total_bytes: usize,
20}
21
22#[derive(Clone, Debug, PartialEq, Eq)]
23pub struct RawArgs {
24 words: Vec<String>,
25}
26
27impl Default for ArgLimits {
28 fn default() -> Self {
29 Self {
30 max_words: 1024,
31 max_word_bytes: 64 * 1024,
32 max_total_bytes: 1024 * 1024,
33 }
34 }
35}
36
37impl RawArgs {
38 pub fn from_env() -> Result<Self, CliError> {
41 Self::from_os_words(std::env::args_os().skip(1), ArgLimits::default())
42 }
43
44 fn from_os_words(
45 args: impl IntoIterator<Item = std::ffi::OsString>,
46 limits: ArgLimits,
47 ) -> Result<Self, CliError> {
48 let mut words = Vec::new();
49 let mut total_bytes = 0usize;
50 for (index, word) in args.into_iter().enumerate() {
51 let word = word.into_string().map_err(|_| {
52 CliError::new(
53 CliErrorKind::InvalidInput,
54 format!("argument {} is not valid UTF-8", index + 1),
55 )
56 })?;
57 validate_raw_word(&word, words.len(), total_bytes, limits)?;
58 total_bytes = total_bytes.saturating_add(word.len());
59 words.push(word);
60 }
61 Ok(Self { words })
62 }
63
64 pub fn parse<I>(args: I) -> Result<Self, CliError>
65 where
66 I: IntoIterator,
67 I::Item: Into<String>,
68 {
69 Self::parse_with_limits(args, ArgLimits::default())
70 }
71
72 pub fn parse_with_limits<I>(args: I, limits: ArgLimits) -> Result<Self, CliError>
73 where
74 I: IntoIterator,
75 I::Item: Into<String>,
76 {
77 let mut words = Vec::new();
78 let mut total_bytes = 0usize;
79 for arg in args {
80 let value = arg.into();
81 validate_raw_word(&value, words.len(), total_bytes, limits)?;
82 total_bytes = total_bytes.saturating_add(value.len());
83 words.push(value);
84 }
85 Ok(Self { words })
86 }
87
88 pub fn words(&self) -> &[String] {
89 &self.words
90 }
91}
92
93fn validate_raw_word(
94 value: &str,
95 count: usize,
96 total: usize,
97 limits: ArgLimits,
98) -> Result<(), CliError> {
99 if count >= limits.max_words {
100 return Err(CliError::new(
101 CliErrorKind::InvalidInput,
102 format!("argument count exceeds limit of {}", limits.max_words),
103 ));
104 }
105 if value.len() > limits.max_word_bytes {
106 return Err(CliError::new(
107 CliErrorKind::InvalidInput,
108 format!(
109 "argument {} exceeds byte limit of {}",
110 count + 1,
111 limits.max_word_bytes
112 ),
113 ));
114 }
115 if total.saturating_add(value.len()) > limits.max_total_bytes {
116 return Err(CliError::new(
117 CliErrorKind::InvalidInput,
118 format!(
119 "total argument bytes exceed limit of {}",
120 limits.max_total_bytes
121 ),
122 ));
123 }
124 if value.contains('\0') {
125 return Err(CliError::new(
126 CliErrorKind::InvalidInput,
127 format!("argument {} contains a NUL byte", count + 1),
128 ));
129 }
130 Ok(())
131}
132
133#[cfg(test)]
134mod tests {
135 use super::*;
136 use std::ffi::OsString;
137
138 #[test]
139 fn environment_ingestion_stops_consuming_at_first_limit_failure() {
140 let mut consumed = 0;
141 let args = ["first", "second", "never-read"].into_iter().map(|word| {
142 consumed += 1;
143 OsString::from(word)
144 });
145 let limits = ArgLimits {
146 max_words: 1,
147 ..ArgLimits::default()
148 };
149 assert!(RawArgs::from_os_words(args, limits).is_err());
150 assert_eq!(consumed, 2);
151 }
152
153 #[test]
154 fn environment_ingestion_enforces_exact_aggregate_utf8_bytes() {
155 let limits = ArgLimits {
156 max_words: 2,
157 max_word_bytes: 2,
158 max_total_bytes: 3,
159 };
160 let accepted =
161 RawArgs::from_os_words([OsString::from("é"), OsString::from("a")], limits).unwrap();
162 assert_eq!(accepted.words(), &["é", "a"]);
163 assert!(
164 RawArgs::from_os_words([OsString::from("é"), OsString::from("ab")], limits).is_err()
165 );
166 }
167
168 #[cfg(unix)]
169 #[test]
170 fn invalid_environment_encoding_keeps_the_argument_index() {
171 use std::os::unix::ffi::OsStringExt;
172 let error = RawArgs::from_os_words(
173 [OsString::from("valid"), OsString::from_vec(vec![0xff])],
174 ArgLimits::default(),
175 )
176 .unwrap_err();
177 assert_eq!(error.to_string(), "argument 2 is not valid UTF-8");
178 }
179}