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. 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. 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 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 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] 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] 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] 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