// Copyright 2023 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // https://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. //! The internal structure of CaptureInfo and CaptureMaps //! //! CaptureInfo is the internal structure of any Capture that includes //! the protobuf structure. CaptureMaps contains mappings of //! ChipIdentifier and to CaptureInfo. use bytes::Bytes; use std::collections::btree_map::{Iter, Values}; use std::collections::BTreeMap; use std::fs::{File, OpenOptions}; use std::io::{Error, ErrorKind, Result}; use std::sync::mpsc::channel; use std::sync::mpsc::{Receiver, Sender}; use std::sync::{Arc, Mutex}; use std::thread; use std::time::{SystemTime, UNIX_EPOCH}; use super::pcap_util::{write_pcap_header, write_pcapng_header, LinkType}; use log::{info, warn}; use netsim_proto::{common::ChipKind, model::Capture as ProtoCapture}; use protobuf::well_known_types::timestamp::Timestamp; use crate::events::{ChipAdded, ChipRemoved, Event}; use crate::resource::clone_captures; use crate::captures::captures_handler::handle_packet; use crate::captures::pcap_util::PacketDirection; use crate::devices::chip::ChipIdentifier; /// Internal Capture struct pub struct CaptureInfo { /// Some(File) if the file is opened and capture is actively happening. /// None if the file is not opened. pub file: Option, // Following items will be returned as ProtoCapture. (state: file.is_some()) id: ChipIdentifier, /// ChipKind (BLUETOOTH, WIFI, or UWB) pub chip_kind: ChipKind, /// Device name pub device_name: String, /// Size of pcap file pub size: usize, /// Number of packet records pub records: i32, /// Timestamp as seconds pub seconds: i64, /// Timestamp as sub-nanoseconds pub nanos: i32, /// Boolean status of whether the device is connected to netsim pub valid: bool, /// Extension (pcap or pcapng) pub extension: String, } /// Captures contains a recent copy of all chips and their ChipKind, chip_id, /// and owning device name. /// /// Information for any recent or ongoing captures is also stored in the ProtoCapture. pub struct Captures { /// A mapping of chip id to CaptureInfo. /// /// BTreeMap is used for chip_id_to_capture, so that the CaptureInfo can always be /// ordered by ChipId. ListCaptureResponse will produce a ordered list of CaptureInfos. pub chip_id_to_capture: BTreeMap>>, sender: Sender, } impl CaptureInfo { /// Create an instance of CaptureInfo pub fn new(chip_kind: ChipKind, chip_id: ChipIdentifier, device_name: String) -> Self { let extension = match chip_kind { ChipKind::UWB => "pcapng".to_string(), _ => "pcap".to_string(), }; CaptureInfo { id: chip_id, chip_kind, device_name, size: 0, records: 0, seconds: 0, nanos: 0, valid: true, file: None, extension, } } /// Creates a pcap file with headers and store it under temp directory. /// /// The lifecycle of the file is NOT tied to the lifecycle of the struct /// Format: /tmp/netsimd/$USER/pcaps/netsim-{chip_id}-{device_name}-{chip_kind}.{extension} pub fn start_capture(&mut self) -> Result<()> { if self.file.is_some() { return Ok(()); } let mut filename = netsim_common::system::netsimd_temp_dir(); filename.push("pcaps"); std::fs::create_dir_all(&filename)?; filename.push(format!( "netsim-{:?}-{:}-{:?}.{}", self.id, self.device_name, self.chip_kind, self.extension )); let mut file = OpenOptions::new().write(true).truncate(true).create(true).open(filename)?; let link_type = match self.chip_kind { ChipKind::BLUETOOTH => LinkType::BluetoothHciH4WithPhdr, ChipKind::BLUETOOTH_BEACON => LinkType::BluetoothHciH4WithPhdr, ChipKind::WIFI => LinkType::Ieee80211RadioTap, ChipKind::UWB => LinkType::FiraUci, _ => return Err(Error::new(ErrorKind::Other, "Unsupported link type")), }; let size = match self.extension.as_str() { "pcap" => write_pcap_header(link_type, &mut file)?, "pcapng" => write_pcapng_header(link_type, &mut file)?, _ => return Err(Error::new(ErrorKind::Other, "Incorrect Extension for file")), }; let timestamp = SystemTime::now().duration_since(UNIX_EPOCH).expect("Time went backwards"); self.size = size; self.records = 0; self.seconds = timestamp.as_secs() as i64; self.nanos = timestamp.subsec_nanos() as i32; self.file = Some(file); Ok(()) } /// Closes file by removing ownership of self.file. /// /// Capture info will still retain the size and record count /// So it can be downloaded easily when GetCapture is invoked. pub fn stop_capture(&mut self) { self.file = None; } /// Returns a Capture protobuf from CaptureInfo pub fn get_capture_proto(&self) -> ProtoCapture { let timestamp = Timestamp { seconds: self.seconds, nanos: self.nanos, ..Default::default() }; ProtoCapture { id: self.id.0, chip_kind: self.chip_kind.into(), device_name: self.device_name.clone(), state: Some(self.file.is_some()), size: self.size as i32, records: self.records, timestamp: Some(timestamp).into(), valid: self.valid, ..Default::default() } } } struct CapturePacket { chip_id: ChipIdentifier, packet: Bytes, packet_type: u32, direction: PacketDirection, } impl Captures { /// Create an instance of Captures, which includes 2 empty hashmaps pub fn new() -> Self { let (sender, rx) = channel::(); let _ = thread::Builder::new().name("capture_packet_handler".to_string()).spawn(move || { while let Ok(CapturePacket { chip_id, packet, packet_type, direction }) = rx.recv() { handle_packet(chip_id, &packet, packet_type, direction); } }); Captures { chip_id_to_capture: BTreeMap::>>::new(), sender, } } /// Sends a captured packet to the output processing thread. pub fn send( &self, chip_id: ChipIdentifier, packet: &Bytes, packet_type: u32, direction: PacketDirection, ) { let _ = self.sender.send(CapturePacket { chip_id, packet: packet.clone(), packet_type, direction, }); } /// Returns true if key exists in Captures.chip_id_to_capture pub fn contains(&self, key: ChipIdentifier) -> bool { self.chip_id_to_capture.contains_key(&key) } /// Returns an Option of lockable and mutable CaptureInfo with given key pub fn get(&mut self, key: ChipIdentifier) -> Option<&mut Arc>> { self.chip_id_to_capture.get_mut(&key) } /// Inserts the given CatpureInfo into Captures hashmaps pub fn insert(&mut self, capture: CaptureInfo) { let chip_id = capture.id; let arc_capture = Arc::new(Mutex::new(capture)); self.chip_id_to_capture.insert(chip_id, arc_capture.clone()); } /// Returns true if chip_id_to_capture is empty pub fn is_empty(&self) -> bool { self.chip_id_to_capture.is_empty() } /// Returns an iterable object of chip_id_to_capture hashmap pub fn iter(&self) -> Iter>> { self.chip_id_to_capture.iter() } /// Removes a CaptureInfo with given key from Captures /// /// When Capture is removed, remove from each map and also invoke closing of files. pub fn remove(&mut self, key: &ChipIdentifier) { if let Some(arc_capture) = self.chip_id_to_capture.get(key) { let mut capture = arc_capture.lock().expect("Failed to acquire lock on CaptureInfo"); // Valid is marked false when chip is disconnected from netsim capture.valid = false; capture.stop_capture(); } else { info!("key does not exist in Captures"); } // CaptureInfo is not removed even after chip is removed } /// Returns Values of chip_id_to_capture hashmap values pub fn values(&self) -> Values>> { self.chip_id_to_capture.values() } } impl Default for Captures { fn default() -> Self { Self::new() } } /// Create a thread to process events that matter to the Capture resource. /// /// We maintain a CaptureInfo for each chip that has been /// connected to the simulation. This procedure monitors ChipAdded /// and ChipRemoved events and updates the collection of CaptureInfo. /// pub fn spawn_capture_event_subscriber(event_rx: Receiver, capture: bool) { let _ = thread::Builder::new().name("capture_event_subscriber".to_string()).spawn(move || loop { match event_rx.recv() { Ok(Event::ChipAdded(ChipAdded { chip_id, chip_kind, device_name, .. })) => { let mut capture_info = CaptureInfo::new(chip_kind, chip_id, device_name.clone()); if capture { if let Err(err) = capture_info.start_capture() { warn!("{err:?}"); } } clone_captures().write().unwrap().insert(capture_info); info!("Capture event: ChipAdded chip_id: {chip_id} device_name: {device_name}"); } Ok(Event::ChipRemoved(ChipRemoved { chip_id, .. })) => { clone_captures().write().unwrap().remove(&chip_id); info!("Capture event: ChipRemoved chip_id: {chip_id}"); } _ => {} } }); }