1 // Copyright 2022 Google LLC
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 //! Command Line Interface for Netsim
16 
17 mod args;
18 mod browser;
19 mod display;
20 mod ffi;
21 mod file_handler;
22 mod requests;
23 mod response;
24 
25 use netsim_common::util::os_utils::{get_instance, get_server_address};
26 use netsim_proto::frontend::{DeleteChipRequest, ListDeviceResponse};
27 use protobuf::Message;
28 use std::env;
29 use std::fs::File;
30 use std::path::PathBuf;
31 use tracing::error;
32 
33 use crate::ffi::frontend_client_ffi::{
34     new_frontend_client, ClientResult, FrontendClient, GrpcMethod,
35 };
36 use crate::ffi::ClientResponseReader;
37 use args::{BinaryProtobuf, GetCapture, NetsimArgs};
38 use clap::Parser;
39 use cxx::{let_cxx_string, UniquePtr};
40 use file_handler::FileHandler;
41 use netsim_common::util::netsim_logger;
42 
43 // helper function to process streaming Grpc request
perform_streaming_request( client: &cxx::UniquePtr<FrontendClient>, cmd: &mut GetCapture, req: &BinaryProtobuf, filename: &str, ) -> UniquePtr<ClientResult>44 fn perform_streaming_request(
45     client: &cxx::UniquePtr<FrontendClient>,
46     cmd: &mut GetCapture,
47     req: &BinaryProtobuf,
48     filename: &str,
49 ) -> UniquePtr<ClientResult> {
50     let dir = if cmd.location.is_some() {
51         PathBuf::from(cmd.location.to_owned().unwrap())
52     } else {
53         env::current_dir().unwrap()
54     };
55     let output_file = dir.join(filename);
56     cmd.current_file = output_file.display().to_string();
57     client.get_capture(
58         req,
59         &ClientResponseReader {
60             handler: Box::new(FileHandler {
61                 file: File::create(&output_file).unwrap_or_else(|_| {
62                     panic!("Failed to create file: {}", &output_file.display())
63                 }),
64                 path: output_file,
65             }),
66         },
67     )
68 }
69 
70 /// helper function to send the Grpc request(s) and handle the response(s) per the given command
perform_command( command: &mut args::Command, client: cxx::UniquePtr<FrontendClient>, grpc_method: GrpcMethod, verbose: bool, ) -> Result<(), String>71 fn perform_command(
72     command: &mut args::Command,
73     client: cxx::UniquePtr<FrontendClient>,
74     grpc_method: GrpcMethod,
75     verbose: bool,
76 ) -> Result<(), String> {
77     // Get command's gRPC request(s)
78     let requests = match command {
79         args::Command::Capture(args::Capture::Patch(_) | args::Capture::Get(_)) => {
80             command.get_requests(&client)
81         }
82         args::Command::Beacon(args::Beacon::Remove(_)) => {
83             vec![args::Command::Devices(args::Devices { continuous: false }).get_request_bytes()]
84         }
85         _ => vec![command.get_request_bytes()],
86     };
87     let mut process_error = false;
88     // Process each request
89     for (i, req) in requests.iter().enumerate() {
90         let result = match command {
91             // Continuous option sends the gRPC call every second
92             args::Command::Devices(ref cmd) if cmd.continuous => {
93                 continuous_perform_command(command, &client, grpc_method, req, verbose)?
94             }
95             args::Command::Capture(args::Capture::List(ref cmd)) if cmd.continuous => {
96                 continuous_perform_command(command, &client, grpc_method, req, verbose)?
97             }
98             // Get Capture use streaming gRPC reader request
99             args::Command::Capture(args::Capture::Get(ref mut cmd)) => {
100                 perform_streaming_request(&client, cmd, req, &cmd.filenames[i].to_owned())
101             }
102             args::Command::Beacon(args::Beacon::Remove(ref cmd)) => {
103                 let devices = client.send_grpc(&GrpcMethod::ListDevice, req);
104                 let id = find_id_for_remove(devices.byte_vec().as_slice(), cmd)?;
105                 let req = &DeleteChipRequest { id, ..Default::default() }
106                     .write_to_bytes()
107                     .map_err(|err| format!("{err}"))?;
108 
109                 client.send_grpc(&grpc_method, req)
110             }
111             // All other commands use a single gRPC call
112             _ => client.send_grpc(&grpc_method, req),
113         };
114         if let Err(e) = process_result(command, result, verbose) {
115             error!("{}", e);
116             process_error = true;
117         };
118     }
119     if process_error {
120         return Err("Not all requests were processed successfully.".to_string());
121     }
122     Ok(())
123 }
124 
find_id_for_remove(response: &[u8], cmd: &args::BeaconRemove) -> Result<u32, String>125 fn find_id_for_remove(response: &[u8], cmd: &args::BeaconRemove) -> Result<u32, String> {
126     let devices = ListDeviceResponse::parse_from_bytes(response).unwrap().devices;
127     let id = devices
128         .iter()
129         .find(|device| device.name == cmd.device_name)
130         .and_then(|device| cmd.chip_name.as_ref().map_or(
131             (device.chips.len() == 1).then_some(&device.chips[0]),
132             |chip_name| device.chips.iter().find(|chip| &chip.name == chip_name)
133         ))
134         .ok_or(cmd.chip_name.as_ref().map_or(
135             format!("failed to delete chip: device '{}' has multiple possible candidates, please specify a chip name", cmd.device_name),
136             |chip_name| format!("failed to delete chip: could not find chip '{}' on device '{}'", chip_name, cmd.device_name))
137         )?
138         .id;
139 
140     Ok(id)
141 }
142 
143 /// Check and handle the gRPC call result
continuous_perform_command( command: &args::Command, client: &cxx::UniquePtr<FrontendClient>, grpc_method: GrpcMethod, request: &Vec<u8>, verbose: bool, ) -> Result<UniquePtr<ClientResult>, String>144 fn continuous_perform_command(
145     command: &args::Command,
146     client: &cxx::UniquePtr<FrontendClient>,
147     grpc_method: GrpcMethod,
148     request: &Vec<u8>,
149     verbose: bool,
150 ) -> Result<UniquePtr<ClientResult>, String> {
151     loop {
152         process_result(command, client.send_grpc(&grpc_method, request), verbose)?;
153         std::thread::sleep(std::time::Duration::from_secs(1));
154     }
155 }
156 /// Check and handle the gRPC call result
process_result( command: &args::Command, result: UniquePtr<ClientResult>, verbose: bool, ) -> Result<(), String>157 fn process_result(
158     command: &args::Command,
159     result: UniquePtr<ClientResult>,
160     verbose: bool,
161 ) -> Result<(), String> {
162     if result.is_ok() {
163         command.print_response(result.byte_vec().as_slice(), verbose);
164     } else {
165         return Err(format!("Grpc call error: {}", result.err()));
166     }
167     Ok(())
168 }
169 #[no_mangle]
170 /// main Rust netsim CLI function to be called by C wrapper netsim.cc
rust_main()171 pub extern "C" fn rust_main() {
172     let mut args = NetsimArgs::parse();
173     netsim_logger::init("netsim", args.verbose);
174     if matches!(args.command, args::Command::Gui) {
175         println!("Opening netsim web UI on default web browser");
176         browser::open("http://localhost:7681/");
177         return;
178     } else if matches!(args.command, args::Command::Artifact) {
179         let artifact_dir = netsim_common::system::netsimd_temp_dir();
180         println!("netsim artifact directory: {}", artifact_dir.display());
181         browser::open(artifact_dir);
182         return;
183     } else if matches!(args.command, args::Command::Bumble) {
184         println!("Opening Bumble Hive on default web browser");
185         browser::open("https://google.github.io/bumble/hive/index.html");
186         return;
187     }
188     let grpc_method = args.command.grpc_method();
189     let server = match (args.vsock, args.port) {
190         (Some(vsock), _) => format!("vsock:{vsock}"),
191         (_, Some(port)) => format!("localhost:{port}"),
192         _ => get_server_address(get_instance(args.instance)).unwrap_or_default(),
193     };
194     let_cxx_string!(server = server);
195     let client = new_frontend_client(&server);
196     if client.is_null() {
197         if !server.is_empty() {
198             error!("Unable to create frontend client. Please ensure netsimd is running and listening on {server:?}.");
199         } else {
200             error!("Unable to create frontend client. Please ensure netsimd is running.");
201         }
202         return;
203     }
204     if let Err(e) = perform_command(&mut args.command, client, grpc_method, args.verbose) {
205         error!("{e}");
206     }
207 }
208 
209 #[cfg(test)]
210 mod tests {
211     use crate::args::BeaconRemove;
212     use netsim_proto::{
213         frontend::ListDeviceResponse,
214         model::{Chip as ChipProto, Device as DeviceProto},
215     };
216     use protobuf::Message;
217 
218     use crate::find_id_for_remove;
219 
220     #[test]
test_remove_device()221     fn test_remove_device() {
222         let device_name = String::from("a-device");
223         let chip_id = 7;
224 
225         let cmd = &BeaconRemove { device_name: device_name.clone(), chip_name: None };
226 
227         let response = ListDeviceResponse {
228             devices: vec![DeviceProto {
229                 id: 0,
230                 name: device_name,
231                 chips: vec![ChipProto { id: chip_id, ..Default::default() }],
232                 ..Default::default()
233             }],
234             ..Default::default()
235         };
236 
237         let id = find_id_for_remove(response.write_to_bytes().unwrap().as_slice(), cmd);
238         assert!(id.is_ok(), "{}", id.unwrap_err());
239         let id = id.unwrap();
240 
241         assert_eq!(chip_id, id);
242     }
243 
244     #[test]
test_remove_chip()245     fn test_remove_chip() {
246         let device_name = String::from("a-device");
247         let chip_name = String::from("should-be-deleted");
248         let device_id = 4;
249         let chip_id = 2;
250 
251         let cmd =
252             &BeaconRemove { device_name: device_name.clone(), chip_name: Some(chip_name.clone()) };
253 
254         let response = ListDeviceResponse {
255             devices: vec![DeviceProto {
256                 id: device_id,
257                 name: device_name,
258                 chips: vec![
259                     ChipProto { id: chip_id, name: chip_name, ..Default::default() },
260                     ChipProto {
261                         id: chip_id + 1,
262                         name: String::from("shouldnt-be-deleted"),
263                         ..Default::default()
264                     },
265                 ],
266                 ..Default::default()
267             }],
268             ..Default::default()
269         };
270 
271         let id = find_id_for_remove(response.write_to_bytes().unwrap().as_slice(), cmd);
272         assert!(id.is_ok(), "{}", id.unwrap_err());
273         let id = id.unwrap();
274 
275         assert_eq!(chip_id, id);
276     }
277 
278     #[test]
test_remove_multiple_chips_fails()279     fn test_remove_multiple_chips_fails() {
280         let device_name = String::from("a-device");
281         let device_id = 3;
282 
283         let cmd = &BeaconRemove { device_name: device_name.clone(), chip_name: None };
284 
285         let response = ListDeviceResponse {
286             devices: vec![DeviceProto {
287                 id: device_id,
288                 name: device_name,
289                 chips: vec![
290                     ChipProto { id: 1, name: String::from("chip-1"), ..Default::default() },
291                     ChipProto { id: 2, name: String::from("chip-2"), ..Default::default() },
292                 ],
293                 ..Default::default()
294             }],
295             ..Default::default()
296         };
297 
298         let id = find_id_for_remove(response.write_to_bytes().unwrap().as_slice(), cmd);
299         assert!(id.is_err());
300     }
301 
302     #[test]
test_remove_nonexistent_chip_fails()303     fn test_remove_nonexistent_chip_fails() {
304         let device_name = String::from("a-device");
305         let device_id = 1;
306 
307         let cmd = &BeaconRemove {
308             device_name: device_name.clone(),
309             chip_name: Some(String::from("nonexistent-chip")),
310         };
311 
312         let response = ListDeviceResponse {
313             devices: vec![DeviceProto {
314                 id: device_id,
315                 name: device_name,
316                 chips: vec![ChipProto {
317                     id: 1,
318                     name: String::from("this-chip-exists"),
319                     ..Default::default()
320                 }],
321                 ..Default::default()
322             }],
323             ..Default::default()
324         };
325 
326         let id = find_id_for_remove(response.write_to_bytes().unwrap().as_slice(), cmd);
327         assert!(id.is_err());
328     }
329 }
330