• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 file_handler;
21 mod grpc_client;
22 mod requests;
23 mod response;
24 
25 use netsim_common::util::ini_file::get_server_address;
26 use netsim_common::util::os_utils::get_instance;
27 use netsim_proto::frontend;
28 
29 use anyhow::{anyhow, Result};
30 use grpcio::{ChannelBuilder, EnvBuilder};
31 use std::env;
32 use std::fs::File;
33 use std::path::PathBuf;
34 use tracing::error;
35 
36 use crate::grpc_client::{ClientResponseReader, GrpcRequest, GrpcResponse};
37 use netsim_proto::frontend_grpc::FrontendServiceClient;
38 
39 use args::{GetCapture, NetsimArgs};
40 use clap::Parser;
41 use file_handler::FileHandler;
42 use netsim_common::util::netsim_logger;
43 
44 // helper function to process streaming Grpc request
perform_streaming_request( client: &FrontendServiceClient, cmd: &mut GetCapture, req: &frontend::GetCaptureRequest, filename: &str, ) -> Result<()>45 fn perform_streaming_request(
46     client: &FrontendServiceClient,
47     cmd: &mut GetCapture,
48     req: &frontend::GetCaptureRequest,
49     filename: &str,
50 ) -> Result<()> {
51     let dir = if cmd.location.is_some() {
52         PathBuf::from(cmd.location.to_owned().unwrap())
53     } else {
54         env::current_dir().unwrap()
55     };
56     let output_file = dir.join(filename);
57     cmd.current_file = output_file.display().to_string();
58     grpc_client::get_capture(
59         client,
60         req,
61         &mut ClientResponseReader {
62             handler: Box::new(FileHandler {
63                 file: File::create(&output_file).unwrap_or_else(|_| {
64                     panic!("Failed to create file: {}", &output_file.display())
65                 }),
66                 path: output_file,
67             }),
68         },
69     )
70 }
71 
72 /// helper function to send the Grpc request(s) and handle the response(s) per the given command
perform_command( command: &mut args::Command, client: FrontendServiceClient, verbose: bool, ) -> anyhow::Result<()>73 fn perform_command(
74     command: &mut args::Command,
75     client: FrontendServiceClient,
76     verbose: bool,
77 ) -> anyhow::Result<()> {
78     // Get command's gRPC request(s)
79     let requests = match command {
80         args::Command::Capture(args::Capture::Patch(_) | args::Capture::Get(_)) => {
81             command.get_requests(&client)
82         }
83         args::Command::Beacon(args::Beacon::Remove(_)) => {
84             vec![args::Command::Devices(args::Devices { continuous: false }).get_request()]
85         }
86         _ => vec![command.get_request()],
87     };
88     let mut process_error = false;
89     // Process each request
90     for (i, req) in requests.iter().enumerate() {
91         let result = match command {
92             // Continuous option sends the gRPC call every second
93             args::Command::Devices(ref cmd) if cmd.continuous => {
94                 continuous_perform_command(command, &client, req, verbose)?;
95                 panic!("Continuous command interrupted. Exiting.");
96             }
97             args::Command::Capture(args::Capture::List(ref cmd)) if cmd.continuous => {
98                 continuous_perform_command(command, &client, req, verbose)?;
99                 panic!("Continuous command interrupted. Exiting.");
100             }
101             // Get Capture use streaming gRPC reader request
102             args::Command::Capture(args::Capture::Get(ref mut cmd)) => {
103                 let GrpcRequest::GetCapture(request) = req else {
104                     panic!("Expected to find GetCaptureRequest. Got: {:?}", req);
105                 };
106                 perform_streaming_request(&client, cmd, request, &cmd.filenames[i].to_owned())?;
107                 Ok(None)
108             }
109             args::Command::Beacon(args::Beacon::Remove(ref cmd)) => {
110                 let response = grpc_client::send_grpc(&client, &GrpcRequest::ListDevice)?;
111                 let GrpcResponse::ListDevice(response) = response else {
112                     panic!("Expected to find ListDeviceResponse. Got: {:?}", response);
113                 };
114                 let id = find_id_for_remove(response, cmd)?;
115                 let res = grpc_client::send_grpc(
116                     &client,
117                     &GrpcRequest::DeleteChip(frontend::DeleteChipRequest {
118                         id,
119                         ..Default::default()
120                     }),
121                 )?;
122                 Ok(Some(res))
123             }
124             // All other commands use a single gRPC call
125             _ => {
126                 let response = grpc_client::send_grpc(&client, req)?;
127                 Ok(Some(response))
128             }
129         };
130         if let Err(e) = process_result(command, result, verbose) {
131             error!("{}", e);
132             process_error = true;
133         };
134     }
135     if process_error {
136         return Err(anyhow!("Not all requests were processed successfully."));
137     }
138     Ok(())
139 }
140 
find_id_for_remove( response: frontend::ListDeviceResponse, cmd: &args::BeaconRemove, ) -> anyhow::Result<u32>141 fn find_id_for_remove(
142     response: frontend::ListDeviceResponse,
143     cmd: &args::BeaconRemove,
144 ) -> anyhow::Result<u32> {
145     let devices = response.devices;
146     let id = devices
147         .iter()
148         .find(|device| device.name == cmd.device_name)
149         .and_then(|device| cmd.chip_name.as_ref().map_or(
150             (device.chips.len() == 1).then_some(&device.chips[0]),
151             |chip_name| device.chips.iter().find(|chip| &chip.name == chip_name)
152         ))
153         .ok_or(
154             cmd.chip_name
155                 .as_ref()
156                 .map_or(
157                     anyhow!("failed to delete chip: device '{}' has multiple possible candidates, please specify a chip name", cmd.device_name),
158                     |chip_name| {
159                         anyhow!(
160                             "failed to delete chip: could not find chip '{}' on device '{}'",
161                             chip_name, cmd.device_name
162                         )
163                     },
164                 )
165         )?
166         .id;
167 
168     Ok(id)
169 }
170 
171 /// Check and handle the gRPC call result
continuous_perform_command( command: &args::Command, client: &FrontendServiceClient, grpc_request: &GrpcRequest, verbose: bool, ) -> anyhow::Result<()>172 fn continuous_perform_command(
173     command: &args::Command,
174     client: &FrontendServiceClient,
175     grpc_request: &GrpcRequest,
176     verbose: bool,
177 ) -> anyhow::Result<()> {
178     loop {
179         let response = grpc_client::send_grpc(client, grpc_request)?;
180         process_result(command, Ok(Some(response)), verbose)?;
181         std::thread::sleep(std::time::Duration::from_secs(1));
182     }
183 }
184 /// Check and handle the gRPC call result
process_result( command: &args::Command, result: anyhow::Result<Option<GrpcResponse>>, verbose: bool, ) -> anyhow::Result<()>185 fn process_result(
186     command: &args::Command,
187     result: anyhow::Result<Option<GrpcResponse>>,
188     verbose: bool,
189 ) -> anyhow::Result<()> {
190     match result {
191         Ok(grpc_response) => {
192             let response = grpc_response.unwrap_or(GrpcResponse::Unknown);
193             command.print_response(&response, verbose);
194             Ok(())
195         }
196         Err(e) => Err(anyhow!("Grpc call error: {}", e)),
197     }
198 }
199 #[no_mangle]
200 /// main Rust netsim CLI function to be called by C wrapper netsim.cc
rust_main()201 pub extern "C" fn rust_main() {
202     let mut args = NetsimArgs::parse();
203     netsim_logger::init("netsim", args.verbose);
204     if matches!(args.command, args::Command::Gui) {
205         println!("Opening netsim web UI on default web browser");
206         browser::open("http://localhost:7681/");
207         return;
208     } else if matches!(args.command, args::Command::Artifact) {
209         let artifact_dir = netsim_common::system::netsimd_temp_dir();
210         println!("netsim artifact directory: {}", artifact_dir.display());
211         browser::open(artifact_dir);
212         return;
213     } else if matches!(args.command, args::Command::Bumble) {
214         println!("Opening Bumble Hive on default web browser");
215         browser::open("https://google.github.io/bumble/hive/index.html");
216         return;
217     }
218     let server = match (args.vsock, args.port) {
219         (Some(vsock), _) => format!("vsock:{vsock}"),
220         (_, Some(port)) => format!("localhost:{port}"),
221         _ => get_server_address(get_instance(args.instance)).unwrap_or_default(),
222     };
223     let channel =
224         ChannelBuilder::new(std::sync::Arc::new(EnvBuilder::new().build())).connect(&server);
225     let client = FrontendServiceClient::new(channel);
226     if let Err(e) = perform_command(&mut args.command, client, args.verbose) {
227         error!("{e}");
228     }
229 }
230 
231 #[cfg(test)]
232 mod tests {
233     use crate::args::BeaconRemove;
234     use netsim_proto::{
235         frontend::ListDeviceResponse,
236         model::{Chip as ChipProto, Device as DeviceProto},
237     };
238 
239     use crate::find_id_for_remove;
240 
241     #[test]
test_remove_device()242     fn test_remove_device() {
243         let device_name = String::from("a-device");
244         let chip_id = 7;
245 
246         let cmd = &BeaconRemove { device_name: device_name.clone(), chip_name: None };
247 
248         let response = ListDeviceResponse {
249             devices: vec![DeviceProto {
250                 id: 0,
251                 name: device_name,
252                 chips: vec![ChipProto { id: chip_id, ..Default::default() }],
253                 ..Default::default()
254             }],
255             ..Default::default()
256         };
257 
258         let id = find_id_for_remove(response, cmd);
259         assert!(id.is_ok(), "{}", id.unwrap_err());
260         let id = id.unwrap();
261 
262         assert_eq!(chip_id, id);
263     }
264 
265     #[test]
test_remove_chip()266     fn test_remove_chip() {
267         let device_name = String::from("a-device");
268         let chip_name = String::from("should-be-deleted");
269         let device_id = 4;
270         let chip_id = 2;
271 
272         let cmd =
273             &BeaconRemove { device_name: device_name.clone(), chip_name: Some(chip_name.clone()) };
274 
275         let response = ListDeviceResponse {
276             devices: vec![DeviceProto {
277                 id: device_id,
278                 name: device_name,
279                 chips: vec![
280                     ChipProto { id: chip_id, name: chip_name, ..Default::default() },
281                     ChipProto {
282                         id: chip_id + 1,
283                         name: String::from("shouldnt-be-deleted"),
284                         ..Default::default()
285                     },
286                 ],
287                 ..Default::default()
288             }],
289             ..Default::default()
290         };
291 
292         let id = find_id_for_remove(response, cmd);
293         assert!(id.is_ok(), "{}", id.unwrap_err());
294         let id = id.unwrap();
295 
296         assert_eq!(chip_id, id);
297     }
298 
299     #[test]
test_remove_multiple_chips_fails()300     fn test_remove_multiple_chips_fails() {
301         let device_name = String::from("a-device");
302         let device_id = 3;
303 
304         let cmd = &BeaconRemove { device_name: device_name.clone(), chip_name: None };
305 
306         let response = ListDeviceResponse {
307             devices: vec![DeviceProto {
308                 id: device_id,
309                 name: device_name,
310                 chips: vec![
311                     ChipProto { id: 1, name: String::from("chip-1"), ..Default::default() },
312                     ChipProto { id: 2, name: String::from("chip-2"), ..Default::default() },
313                 ],
314                 ..Default::default()
315             }],
316             ..Default::default()
317         };
318 
319         let id = find_id_for_remove(response, cmd);
320         assert!(id.is_err());
321     }
322 
323     #[test]
test_remove_nonexistent_chip_fails()324     fn test_remove_nonexistent_chip_fails() {
325         let device_name = String::from("a-device");
326         let device_id = 1;
327 
328         let cmd = &BeaconRemove {
329             device_name: device_name.clone(),
330             chip_name: Some(String::from("nonexistent-chip")),
331         };
332 
333         let response = ListDeviceResponse {
334             devices: vec![DeviceProto {
335                 id: device_id,
336                 name: device_name,
337                 chips: vec![ChipProto {
338                     id: 1,
339                     name: String::from("this-chip-exists"),
340                     ..Default::default()
341                 }],
342                 ..Default::default()
343             }],
344             ..Default::default()
345         };
346 
347         let id = find_id_for_remove(response, cmd);
348         assert!(id.is_err());
349     }
350 }
351