1 // SPDX-License-Identifier: MPL-2.0
2 // Note: This file is derived from the asterinas osdk,
3 // and we extend our heartfelt thanks to the developers of Asterinas!
4
5 //! This module contains utilities for manipulating common Unix command-line arguments.
6
7 use anyhow::{anyhow, Result};
8 use indexmap::{IndexMap, IndexSet};
9
10 /// Split a string of Unix arguments into an array of key-value strings or switches.
11 /// Positional arguments are not supported.
split_to_kv_array(args: &str) -> Result<Vec<String>>12 pub fn split_to_kv_array(args: &str) -> Result<Vec<String>> {
13 let target = shlex::split(args).ok_or(anyhow!("Failed to split unix arguments"))?;
14
15 // Join the key value arguments as a single element
16 let mut joined = Vec::<String>::new();
17 let mut last_has_value = false;
18 for elem in target {
19 if !elem.starts_with('-') && !last_has_value {
20 if let Some(last) = joined.last_mut() {
21 last.push(' ');
22 last.push_str(&elem);
23 last_has_value = true;
24 continue;
25 }
26 }
27 joined.push(elem);
28 last_has_value = false;
29 }
30
31 Ok(joined)
32 }
33
34 /// Apply key-value pairs to an array of strings.
35 ///
36 /// The provided arguments will be appended to the array if the key is not already present or if the key is a multi-value key.
37 /// Otherwise, the value will be updated.
apply_kv_array( array: &mut Vec<String>, args: &Vec<String>, separator: &str, multi_value_keys: &[&str], single_value_keys: &[&str], ) -> Result<()>38 pub fn apply_kv_array(
39 array: &mut Vec<String>,
40 args: &Vec<String>,
41 separator: &str,
42 multi_value_keys: &[&str],
43 single_value_keys: &[&str],
44 ) -> Result<()> {
45 let multi_value_keys = {
46 let mut inferred_keys = infer_multi_value_keys(array, separator);
47 for key in multi_value_keys {
48 inferred_keys.insert(key.to_string());
49 }
50 inferred_keys
51 };
52
53 log::debug!("apply_kv_array: multi value keys: {:?}", multi_value_keys);
54
55 // We use IndexMap to keep key orders
56 let mut key_strings = IndexMap::new();
57 let mut multi_value_key_strings: IndexMap<String, Vec<String>> = IndexMap::new();
58 for item in array.drain(..) {
59 // Each key-value string has two patterns:
60 // 1. Separated by separator: key value / key=value
61 if let Some(key) = get_key(&item, separator) {
62 if multi_value_keys.contains(&key) {
63 if let Some(v) = multi_value_key_strings.get_mut(&key) {
64 v.push(item);
65 } else {
66 let v = vec![item];
67 multi_value_key_strings.insert(key, v);
68 }
69 continue;
70 }
71
72 key_strings.insert(key, item);
73 continue;
74 }
75 // 2. Only key, no value
76 key_strings.insert(item.clone(), item);
77 }
78
79 for arg in args {
80 if let Some(key) = get_key(arg, separator) {
81 if multi_value_keys.contains(&key) {
82 if let Some(v) = multi_value_key_strings.get_mut(&key) {
83 v.push(arg.to_owned());
84 } else {
85 let v = vec![arg.to_owned()];
86 multi_value_key_strings.insert(key, v);
87 }
88 continue;
89 }
90
91 key_strings.insert(key, arg.to_owned());
92 continue;
93 } else {
94 if single_value_keys.contains(&arg.as_str()) || multi_value_keys.contains(arg) {
95 return Err(anyhow!("Invalid argument: {}", arg));
96 } else {
97 key_strings.insert(arg.to_owned(), arg.to_owned());
98 }
99 }
100 }
101
102 *array = key_strings.into_iter().map(|(_, value)| value).collect();
103
104 for (_, mut values) in multi_value_key_strings {
105 array.append(&mut values);
106 }
107 Ok(())
108 }
109
infer_multi_value_keys(array: &Vec<String>, separator: &str) -> IndexSet<String>110 fn infer_multi_value_keys(array: &Vec<String>, separator: &str) -> IndexSet<String> {
111 let mut multi_val_keys = IndexSet::new();
112
113 let mut occurred_keys = IndexSet::new();
114 for item in array {
115 let Some(key) = get_key(item, separator) else {
116 continue;
117 };
118
119 if occurred_keys.contains(&key) {
120 multi_val_keys.insert(key);
121 } else {
122 occurred_keys.insert(key);
123 }
124 }
125
126 multi_val_keys
127 }
128
get_key(item: &str, separator: &str) -> Option<String>129 pub fn get_key(item: &str, separator: &str) -> Option<String> {
130 let split = item.split(separator).collect::<Vec<_>>();
131 let len = split.len();
132 if len > 2 || len == 0 {
133 log::error!("{} is an invalid argument.", item);
134 return None;
135 }
136
137 if len == 1 {
138 return None;
139 }
140
141 let key = split.first().unwrap();
142
143 Some(key.to_string())
144 }
145
146 #[cfg(test)]
147 pub mod test {
148 use super::*;
149
150 #[test]
test_get_key()151 fn test_get_key() {
152 let string1 = "init=/bin/init";
153 let key = get_key(string1, "=").unwrap();
154 assert_eq!(key.as_str(), "init");
155
156 let string2 = "-m 8G";
157 let key = get_key(string2, " ").unwrap();
158 assert_eq!(key.as_str(), "-m");
159
160 let string3 = "-device virtio-keyboard-pci,disable-legacy=on,disable-modern=off";
161 let key = get_key(string3, " ").unwrap();
162 assert_eq!(key.as_str(), "-device");
163
164 let string4 = "-device";
165 assert!(get_key(string4, " ").is_none());
166
167 let string5 = "-m 8G a";
168 assert!(get_key(string5, " ").is_none());
169 let string6 = "-m=8G";
170 assert!(get_key(string6, " ").is_none());
171 }
172
173 #[test]
test_apply_kv_array()174 fn test_apply_kv_array() {
175 let qemu_args = &[
176 "-enable-kvm",
177 "-m 8G",
178 "-device virtio-blk-pci,bus=pcie.0,addr=0x6,drive=x0,disable-legacy=on,disable-modern=off",
179 "-device virtio-keyboard-pci,disable-legacy=on,disable-modern=off",
180 ];
181
182 let args = &["-m 100G", "-device ioh3420,id=pcie.0,chassis=1"];
183
184 let expected = &[
185 "-enable-kvm",
186 "-m 100G",
187 "-device virtio-blk-pci,bus=pcie.0,addr=0x6,drive=x0,disable-legacy=on,disable-modern=off",
188 "-device virtio-keyboard-pci,disable-legacy=on,disable-modern=off",
189 "-device ioh3420,id=pcie.0,chassis=1",
190 ];
191
192 let mut array = qemu_args.iter().map(ToString::to_string).collect();
193 let args = args.iter().map(ToString::to_string).collect();
194 assert!(apply_kv_array(&mut array, &args, " ", &["-device"], &[]).is_ok());
195
196 let expected: Vec<_> = expected.iter().map(ToString::to_string).collect();
197 assert_eq!(expected, array);
198 }
199
200 #[test]
test_apply_kv_array_insert_multi_value_key()201 fn test_apply_kv_array_insert_multi_value_key() {
202 let mut array = vec!["key1=value1".to_string(), "key2=value2".to_string()];
203 let args = vec!["key3=value3".to_string()];
204 let separator = "=";
205 let multi_value_keys = vec!["key3"];
206 let single_value_keys: Vec<&str> = vec![];
207
208 let result = apply_kv_array(
209 &mut array,
210 &args,
211 separator,
212 &multi_value_keys,
213 &single_value_keys,
214 );
215
216 assert!(result.is_ok());
217 assert_eq!(
218 array,
219 vec![
220 "key1=value1".to_string(),
221 "key2=value2".to_string(),
222 "key3=value3".to_string(),
223 ]
224 );
225 }
226 }
227