commands.rs (17428B)
1 /* This Source Code Form is subject to the terms of the Mozilla Public 2 * License, v. 2.0. If a copy of the MPL was not distributed with this 3 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ 4 5 use serde::{Deserialize, Serialize}; 6 use std::io::{self, Read, Write}; 7 use std::path::PathBuf; 8 use std::process::Command; 9 use url::Url; 10 11 #[cfg(target_os = "windows")] 12 const OS_NAME: &str = "windows"; 13 14 #[cfg(target_os = "macos")] 15 const OS_NAME: &str = "macos"; 16 17 #[derive(Serialize, Deserialize)] 18 #[serde(tag = "command", content = "data")] 19 // { 20 // "command": "LaunchFirefox", 21 // "data": {"url": "https://example.com"}, 22 // } 23 pub enum FirefoxCommand { 24 LaunchFirefox { url: String }, 25 LaunchFirefoxPrivate { url: String }, 26 GetVersion {}, 27 GetInstallId {}, 28 } 29 #[derive(Serialize, Deserialize)] 30 // { 31 // "message": "Successful launch", 32 // "result_code": 1, 33 // } 34 pub struct Response { 35 pub message: String, 36 pub result_code: u32, 37 } 38 39 #[derive(Serialize, Deserialize)] 40 // { 41 // "installation_id": "123ABC456", 42 // } 43 pub struct InstallationId { 44 pub installation_id: String, 45 } 46 47 #[repr(u32)] 48 pub enum ResultCode { 49 Success = 0, 50 Error = 1, 51 } 52 impl From<ResultCode> for u32 { 53 fn from(m: ResultCode) -> u32 { 54 m as u32 55 } 56 } 57 58 trait CommandRunner { 59 fn new() -> Self 60 where 61 Self: Sized; 62 fn arg(&mut self, arg: &str) -> &mut Self; 63 fn args(&mut self, args: &[&str]) -> &mut Self; 64 fn spawn(&mut self) -> std::io::Result<()>; 65 fn to_string(&mut self) -> std::io::Result<String>; 66 } 67 68 impl CommandRunner for Command { 69 fn new() -> Self { 70 #[cfg(target_os = "macos")] 71 { 72 Command::new("open") 73 } 74 #[cfg(target_os = "windows")] 75 { 76 use mozbuild::config::MOZ_APP_NAME; 77 use std::env; 78 use std::path::Path; 79 // Get the current executable's path, we know Firefox is in the 80 // same folder is nmhproxy.exe so we can use that. 81 let nmh_exe_path = env::current_exe().unwrap(); 82 let nmh_exe_folder = nmh_exe_path.parent().unwrap_or_else(|| Path::new("")); 83 let moz_exe_path = nmh_exe_folder.join(format!("{}.exe", MOZ_APP_NAME)); 84 Command::new(moz_exe_path) 85 } 86 } 87 fn arg(&mut self, arg: &str) -> &mut Self { 88 self.arg(arg) 89 } 90 fn args(&mut self, args: &[&str]) -> &mut Self { 91 self.args(args) 92 } 93 fn spawn(&mut self) -> std::io::Result<()> { 94 self.spawn().map(|_| ()) 95 } 96 fn to_string(&mut self) -> std::io::Result<String> { 97 Ok("".to_string()) 98 } 99 } 100 101 // The message length is a 32-bit integer in native byte order 102 pub fn read_message_length<R: Read>(mut reader: R) -> std::io::Result<u32> { 103 let mut buffer = [0u8; 4]; 104 reader.read_exact(&mut buffer)?; 105 let length: u32 = u32::from_ne_bytes(buffer); 106 if (length > 0) && (length < 100 * 1024) { 107 Ok(length) 108 } else { 109 Err(io::Error::new( 110 io::ErrorKind::InvalidData, 111 "Invalid message length", 112 )) 113 } 114 } 115 116 pub fn read_message_string<R: Read>(mut reader: R, length: u32) -> io::Result<String> { 117 let mut buffer = vec![0u8; length.try_into().unwrap()]; 118 reader.read_exact(&mut buffer)?; 119 let message = 120 String::from_utf8(buffer).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; 121 Ok(message) 122 } 123 124 pub fn process_command(command: &FirefoxCommand) -> std::io::Result<bool> { 125 match &command { 126 FirefoxCommand::LaunchFirefox { url } => { 127 launch_firefox::<Command>(url.to_owned(), false, OS_NAME)?; 128 Ok(true) 129 } 130 FirefoxCommand::LaunchFirefoxPrivate { url } => { 131 launch_firefox::<Command>(url.to_owned(), true, OS_NAME)?; 132 Ok(true) 133 } 134 FirefoxCommand::GetVersion {} => generate_response("1", ResultCode::Success.into()), 135 FirefoxCommand::GetInstallId {} => { 136 // config_dir() evaluates to ~/Library/Application Support on macOS 137 // and %RoamingAppData% on Windows. 138 let mut json_path = match dirs::config_dir() { 139 Some(path) => path, 140 None => { 141 return generate_response( 142 "Config dir could not be found", 143 ResultCode::Error.into(), 144 ) 145 } 146 }; 147 #[cfg(target_os = "windows")] 148 json_path.push("Mozilla\\Firefox"); 149 #[cfg(target_os = "macos")] 150 json_path.push("Firefox"); 151 152 json_path.push("install_id"); 153 json_path.set_extension("json"); 154 let mut install_id = String::new(); 155 get_install_id(&mut json_path, &mut install_id) 156 } 157 } 158 } 159 160 pub fn generate_response(message: &str, result_code: u32) -> std::io::Result<bool> { 161 let response_struct = Response { 162 message: message.to_string(), 163 result_code, 164 }; 165 let response_str = serde_json::to_string(&response_struct) 166 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; 167 let response_len_bytes: [u8; 4] = (response_str.len() as u32).to_ne_bytes(); 168 std::io::stdout().write_all(&response_len_bytes)?; 169 std::io::stdout().write_all(response_str.as_bytes())?; 170 std::io::stdout().flush()?; 171 Ok(true) 172 } 173 174 fn validate_url(url: String) -> std::io::Result<String> { 175 let parsed_url = Url::parse(url.as_str()) 176 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; 177 match parsed_url.scheme() { 178 "http" | "https" | "file" => Ok(parsed_url.to_string()), 179 _ => Err(std::io::Error::new( 180 std::io::ErrorKind::InvalidInput, 181 "Invalid URL scheme", 182 )), 183 } 184 } 185 186 fn launch_firefox<C: CommandRunner>( 187 url: String, 188 private: bool, 189 os: &str, 190 ) -> std::io::Result<String> { 191 let validated_url: String = validate_url(url)?; 192 let mut command = C::new(); 193 if os == "macos" { 194 use mozbuild::config::MOZ_MACBUNDLE_ID; 195 let mut args: [&str; 2] = ["--args", "-url"]; 196 if private { 197 args[1] = "-private-window"; 198 } 199 command 200 .arg("-n") 201 .arg("-b") 202 .arg(MOZ_MACBUNDLE_ID) 203 .args(&args) 204 .arg(validated_url.as_str()); 205 } else if os == "windows" { 206 let mut args: [&str; 2] = ["-osint", "-url"]; 207 if private { 208 args[1] = "-private-window"; 209 } 210 command.args(&args).arg(validated_url.as_str()); 211 } 212 match command.spawn() { 213 Ok(_) => generate_response( 214 if private { 215 "Successful private launch" 216 } else { 217 "Sucessful launch" 218 }, 219 ResultCode::Success.into(), 220 )?, 221 Err(_) => generate_response( 222 if private { 223 "Failed private launch" 224 } else { 225 "Failed launch" 226 }, 227 ResultCode::Error.into(), 228 )?, 229 }; 230 command.to_string() 231 } 232 233 fn get_install_id(json_path: &mut PathBuf, install_id: &mut String) -> std::io::Result<bool> { 234 if !json_path.exists() { 235 return Err(std::io::Error::new( 236 std::io::ErrorKind::NotFound, 237 "Install ID file does not exist", 238 )); 239 } 240 let json_size = std::fs::metadata(&json_path) 241 .map_err(|e| std::io::Error::new(std::io::ErrorKind::NotFound, e))? 242 .len(); 243 // Set a 1 KB limit for the file size. 244 if json_size <= 0 || json_size > 1024 { 245 return Err(std::io::Error::new( 246 std::io::ErrorKind::InvalidData, 247 "Install ID file has invalid size", 248 )); 249 } 250 let mut file = 251 std::fs::File::open(json_path).or_else(|_| -> std::io::Result<std::fs::File> { 252 return Err(std::io::Error::new( 253 std::io::ErrorKind::NotFound, 254 "Failed to open file", 255 )); 256 })?; 257 let mut contents = String::new(); 258 match file.read_to_string(&mut contents) { 259 Ok(_) => match serde_json::from_str::<InstallationId>(&contents) { 260 Ok(id) => { 261 *install_id = id.installation_id.clone(); 262 generate_response(&id.installation_id, ResultCode::Success.into()) 263 } 264 Err(_) => { 265 return Err(std::io::Error::new( 266 std::io::ErrorKind::InvalidData, 267 "Failed to read installation ID", 268 )) 269 } 270 }, 271 Err(_) => generate_response("Failed to read file", ResultCode::Error.into()), 272 }?; 273 Ok(true) 274 } 275 276 #[cfg(test)] 277 mod tests { 278 use super::*; 279 use std::io::Cursor; 280 use tempfile::NamedTempFile; 281 282 struct MockCommand { 283 command_line: String, 284 } 285 286 impl CommandRunner for MockCommand { 287 fn new() -> Self { 288 MockCommand { 289 command_line: String::new(), 290 } 291 } 292 fn arg(&mut self, arg: &str) -> &mut Self { 293 self.command_line.push_str(arg); 294 self.command_line.push(' '); 295 self 296 } 297 fn args(&mut self, args: &[&str]) -> &mut Self { 298 for arg in args { 299 self.command_line.push_str(arg); 300 self.command_line.push(' '); 301 } 302 self 303 } 304 fn spawn(&mut self) -> std::io::Result<()> { 305 Ok(()) 306 } 307 fn to_string(&mut self) -> std::io::Result<String> { 308 Ok(self.command_line.clone()) 309 } 310 } 311 312 #[test] 313 fn test_validate_url() { 314 let valid_test_cases = vec![ 315 "https://example.com/".to_string(), 316 "http://example.com/".to_string(), 317 "file:///path/to/file".to_string(), 318 "https://test.example.com/".to_string(), 319 ]; 320 321 for input in valid_test_cases { 322 let result = validate_url(input.clone()); 323 assert!(result.is_ok(), "Expected Ok, got Err"); 324 // Safe to unwrap because we know the result is Ok 325 let ok_value = result.unwrap(); 326 assert_eq!(ok_value, input); 327 } 328 329 assert!(matches!( 330 validate_url("fakeprotocol://test.example.com/".to_string()).map_err(|e| e.kind()), 331 Err(std::io::ErrorKind::InvalidInput) 332 )); 333 334 assert!(matches!( 335 validate_url("invalidURL".to_string()).map_err(|e| e.kind()), 336 Err(std::io::ErrorKind::InvalidData) 337 )); 338 } 339 340 #[test] 341 fn test_read_message_length_valid() { 342 let input: [u8; 4] = 256u32.to_ne_bytes(); 343 let mut cursor = Cursor::new(input); 344 let length = read_message_length(&mut cursor); 345 assert!(length.is_ok(), "Expected Ok, got Err"); 346 assert_eq!(length.unwrap(), 256); 347 } 348 349 #[test] 350 fn test_read_message_length_invalid_too_large() { 351 let input: [u8; 4] = 1_000_000u32.to_ne_bytes(); 352 let mut cursor = Cursor::new(input); 353 let result = read_message_length(&mut cursor); 354 assert!(result.is_err()); 355 let error = result.err().unwrap(); 356 assert_eq!(error.kind(), io::ErrorKind::InvalidData); 357 } 358 359 #[test] 360 fn test_read_message_length_invalid_zero() { 361 let input: [u8; 4] = 0u32.to_ne_bytes(); 362 let mut cursor = Cursor::new(input); 363 let result = read_message_length(&mut cursor); 364 assert!(result.is_err()); 365 let error = result.err().unwrap(); 366 assert_eq!(error.kind(), io::ErrorKind::InvalidData); 367 } 368 369 #[test] 370 fn test_read_message_string_valid() { 371 let input_data = b"Valid UTF8 string!"; 372 let input_length = input_data.len() as u32; 373 let message = read_message_string(&input_data[..], input_length); 374 assert!(message.is_ok(), "Expected Ok, got Err"); 375 assert_eq!(message.unwrap(), "Valid UTF8 string!"); 376 } 377 378 #[test] 379 fn test_read_message_string_invalid() { 380 let input_data: [u8; 3] = [0xff, 0xfe, 0xfd]; 381 let input_length = input_data.len() as u32; 382 let result = read_message_string(&input_data[..], input_length); 383 assert!(result.is_err()); 384 let error = result.err().unwrap(); 385 assert_eq!(error.kind(), io::ErrorKind::InvalidData); 386 } 387 388 #[test] 389 fn test_launch_regular_command_macos() { 390 let url = "https://example.com"; 391 let result = launch_firefox::<MockCommand>(url.to_string(), false, "macos"); 392 assert!(result.is_ok()); 393 let command_line = result.unwrap(); 394 let correct_url_format = format!("-url {}", url); 395 assert!(command_line.contains(correct_url_format.as_str())); 396 } 397 398 #[test] 399 fn test_launch_regular_command_windows() { 400 let url = "https://example.com"; 401 let result = launch_firefox::<MockCommand>(url.to_string(), false, "windows"); 402 assert!(result.is_ok()); 403 let command_line = result.unwrap(); 404 let correct_url_format = format!("-osint -url {}", url); 405 assert!(command_line.contains(correct_url_format.as_str())); 406 } 407 408 #[test] 409 fn test_launch_private_command_macos() { 410 let url = "https://example.com"; 411 let result = launch_firefox::<MockCommand>(url.to_string(), true, "macos"); 412 assert!(result.is_ok()); 413 let command_line = result.unwrap(); 414 let correct_url_format = format!("-private-window {}", url); 415 assert!(command_line.contains(correct_url_format.as_str())); 416 } 417 418 #[test] 419 fn test_launch_private_command_windows() { 420 let url = "https://example.com"; 421 let result = launch_firefox::<MockCommand>(url.to_string(), true, "windows"); 422 assert!(result.is_ok()); 423 let command_line = result.unwrap(); 424 let correct_url_format = format!("-osint -private-window {}", url); 425 assert!(command_line.contains(correct_url_format.as_str())); 426 } 427 428 #[test] 429 fn test_get_install_id_valid() -> std::io::Result<()> { 430 let mut tempfile = NamedTempFile::new().unwrap(); 431 let installation_id = InstallationId { 432 installation_id: "123ABC456".to_string(), 433 }; 434 let json_string = serde_json::to_string(&installation_id); 435 let _ = tempfile.write_all(json_string?.as_bytes()); 436 let mut install_id = String::new(); 437 let result = get_install_id(&mut tempfile.path().to_path_buf(), &mut install_id); 438 assert!(result.is_ok()); 439 assert_eq!(install_id, "123ABC456"); 440 Ok(()) 441 } 442 443 #[test] 444 fn test_get_install_id_incorrect_var() -> std::io::Result<()> { 445 #[derive(Serialize, Deserialize)] 446 pub struct IncorrectJSON { 447 pub incorrect_var: String, 448 } 449 let mut tempfile = NamedTempFile::new().unwrap(); 450 let incorrect_json = IncorrectJSON { 451 incorrect_var: "incorrect_val".to_string(), 452 }; 453 let json_string = serde_json::to_string(&incorrect_json); 454 let _ = tempfile.write_all(json_string?.as_bytes()); 455 let mut install_id = String::new(); 456 let result = get_install_id(&mut tempfile.path().to_path_buf(), &mut install_id); 457 assert!(result.is_err()); 458 let error = result.err().unwrap(); 459 assert_eq!(error.kind(), std::io::ErrorKind::InvalidData); 460 Ok(()) 461 } 462 463 #[test] 464 fn test_get_install_id_partially_correct_vars() -> std::io::Result<()> { 465 #[derive(Serialize, Deserialize)] 466 pub struct IncorrectJSON { 467 pub installation_id: String, 468 pub incorrect_var: String, 469 } 470 let mut tempfile = NamedTempFile::new().unwrap(); 471 let incorrect_json = IncorrectJSON { 472 installation_id: "123ABC456".to_string(), 473 incorrect_var: "incorrect_val".to_string(), 474 }; 475 let json_string = serde_json::to_string(&incorrect_json); 476 let _ = tempfile.write_all(json_string?.as_bytes()); 477 let mut install_id = String::new(); 478 let result = get_install_id(&mut tempfile.path().to_path_buf(), &mut install_id); 479 // This still succeeds as the installation_id field is present 480 assert!(result.is_ok()); 481 Ok(()) 482 } 483 484 #[test] 485 fn test_get_install_id_file_does_not_exist() -> std::io::Result<()> { 486 let tempfile = NamedTempFile::new().unwrap(); 487 let mut path = tempfile.path().to_path_buf(); 488 tempfile.close()?; 489 let mut install_id = String::new(); 490 let result = get_install_id(&mut path, &mut install_id); 491 assert!(result.is_err()); 492 let error = result.err().unwrap(); 493 assert_eq!(error.kind(), std::io::ErrorKind::NotFound); 494 Ok(()) 495 } 496 497 #[test] 498 fn test_get_install_id_file_too_large() -> std::io::Result<()> { 499 let mut tempfile = NamedTempFile::new().unwrap(); 500 let installation_id = InstallationId { 501 // Create a ~10 KB file 502 installation_id: String::from_utf8(vec![b'X'; 10000]).unwrap(), 503 }; 504 let json_string = serde_json::to_string(&installation_id); 505 let _ = tempfile.write_all(json_string?.as_bytes()); 506 let mut install_id = String::new(); 507 let result = get_install_id(&mut tempfile.path().to_path_buf(), &mut install_id); 508 assert!(result.is_err()); 509 let error = result.err().unwrap(); 510 assert_eq!(error.kind(), std::io::ErrorKind::InvalidData); 511 Ok(()) 512 } 513 }