1 // Copyright 2023 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 //! The internal structure of CaptureInfo and CaptureMaps
16 //!
17 //! CaptureInfo is the internal structure of any Capture that includes
18 //! the protobuf structure. CaptureMaps contains mappings of
19 //! ChipIdentifier and <FacadeIdentifier, Kind> to CaptureInfo.
20
21 use bytes::Bytes;
22
23 use std::collections::btree_map::{Iter, Values};
24 use std::collections::BTreeMap;
25 use std::fs::{File, OpenOptions};
26 use std::io::{Error, ErrorKind, Result};
27 use std::sync::mpsc::channel;
28 use std::sync::mpsc::{Receiver, Sender};
29 use std::sync::{Arc, Mutex};
30 use std::thread;
31 use std::time::{SystemTime, UNIX_EPOCH};
32
33 use super::pcap_util::{write_pcap_header, write_pcapng_header, LinkType};
34 use log::{info, warn};
35
36 use netsim_proto::{common::ChipKind, model::Capture as ProtoCapture};
37 use protobuf::well_known_types::timestamp::Timestamp;
38
39 use crate::events::{ChipAdded, ChipRemoved, Event};
40 use crate::resource::clone_captures;
41
42 use crate::captures::captures_handler::handle_packet;
43 use crate::captures::pcap_util::PacketDirection;
44 use crate::devices::chip::ChipIdentifier;
45
46 /// Internal Capture struct
47 pub struct CaptureInfo {
48 /// Some(File) if the file is opened and capture is actively happening.
49 /// None if the file is not opened.
50 pub file: Option<File>,
51 // Following items will be returned as ProtoCapture. (state: file.is_some())
52 id: ChipIdentifier,
53 /// ChipKind (BLUETOOTH, WIFI, or UWB)
54 pub chip_kind: ChipKind,
55 /// Device name
56 pub device_name: String,
57 /// Size of pcap file
58 pub size: usize,
59 /// Number of packet records
60 pub records: i32,
61 /// Timestamp as seconds
62 pub seconds: i64,
63 /// Timestamp as sub-nanoseconds
64 pub nanos: i32,
65 /// Boolean status of whether the device is connected to netsim
66 pub valid: bool,
67 /// Extension (pcap or pcapng)
68 pub extension: String,
69 }
70
71 /// Captures contains a recent copy of all chips and their ChipKind, chip_id,
72 /// and owning device name.
73 ///
74 /// Information for any recent or ongoing captures is also stored in the ProtoCapture.
75 pub struct Captures {
76 /// A mapping of chip id to CaptureInfo.
77 ///
78 /// BTreeMap is used for chip_id_to_capture, so that the CaptureInfo can always be
79 /// ordered by ChipId. ListCaptureResponse will produce a ordered list of CaptureInfos.
80 pub chip_id_to_capture: BTreeMap<ChipIdentifier, Arc<Mutex<CaptureInfo>>>,
81 sender: Sender<CapturePacket>,
82 }
83
84 impl CaptureInfo {
85 /// Create an instance of CaptureInfo
new(chip_kind: ChipKind, chip_id: ChipIdentifier, device_name: String) -> Self86 pub fn new(chip_kind: ChipKind, chip_id: ChipIdentifier, device_name: String) -> Self {
87 let extension = match chip_kind {
88 ChipKind::UWB => "pcapng".to_string(),
89 _ => "pcap".to_string(),
90 };
91 CaptureInfo {
92 id: chip_id,
93 chip_kind,
94 device_name,
95 size: 0,
96 records: 0,
97 seconds: 0,
98 nanos: 0,
99 valid: true,
100 file: None,
101 extension,
102 }
103 }
104
105 /// Creates a pcap file with headers and store it under temp directory.
106 ///
107 /// The lifecycle of the file is NOT tied to the lifecycle of the struct
108 /// Format: /tmp/netsimd/$USER/pcaps/netsim-{chip_id}-{device_name}-{chip_kind}.{extension}
start_capture(&mut self) -> Result<()>109 pub fn start_capture(&mut self) -> Result<()> {
110 if self.file.is_some() {
111 return Ok(());
112 }
113 let mut filename = netsim_common::system::netsimd_temp_dir();
114 filename.push("pcaps");
115 std::fs::create_dir_all(&filename)?;
116 filename.push(format!(
117 "netsim-{:?}-{:}-{:?}.{}",
118 self.id, self.device_name, self.chip_kind, self.extension
119 ));
120 let mut file = OpenOptions::new().write(true).truncate(true).create(true).open(filename)?;
121 let link_type = match self.chip_kind {
122 ChipKind::BLUETOOTH => LinkType::BluetoothHciH4WithPhdr,
123 ChipKind::BLUETOOTH_BEACON => LinkType::BluetoothHciH4WithPhdr,
124 ChipKind::WIFI => LinkType::Ieee80211RadioTap,
125 ChipKind::UWB => LinkType::FiraUci,
126 _ => return Err(Error::new(ErrorKind::Other, "Unsupported link type")),
127 };
128 let size = match self.extension.as_str() {
129 "pcap" => write_pcap_header(link_type, &mut file)?,
130 "pcapng" => write_pcapng_header(link_type, &mut file)?,
131 _ => return Err(Error::new(ErrorKind::Other, "Incorrect Extension for file")),
132 };
133 let timestamp = SystemTime::now().duration_since(UNIX_EPOCH).expect("Time went backwards");
134 self.size = size;
135 self.records = 0;
136 self.seconds = timestamp.as_secs() as i64;
137 self.nanos = timestamp.subsec_nanos() as i32;
138 self.file = Some(file);
139 Ok(())
140 }
141
142 /// Closes file by removing ownership of self.file.
143 ///
144 /// Capture info will still retain the size and record count
145 /// So it can be downloaded easily when GetCapture is invoked.
stop_capture(&mut self)146 pub fn stop_capture(&mut self) {
147 self.file = None;
148 }
149
150 /// Returns a Capture protobuf from CaptureInfo
get_capture_proto(&self) -> ProtoCapture151 pub fn get_capture_proto(&self) -> ProtoCapture {
152 let timestamp =
153 Timestamp { seconds: self.seconds, nanos: self.nanos, ..Default::default() };
154 ProtoCapture {
155 id: self.id.0,
156 chip_kind: self.chip_kind.into(),
157 device_name: self.device_name.clone(),
158 state: Some(self.file.is_some()),
159 size: self.size as i32,
160 records: self.records,
161 timestamp: Some(timestamp).into(),
162 valid: self.valid,
163 ..Default::default()
164 }
165 }
166 }
167
168 struct CapturePacket {
169 chip_id: ChipIdentifier,
170 packet: Bytes,
171 packet_type: u32,
172 direction: PacketDirection,
173 }
174
175 impl Captures {
176 /// Create an instance of Captures, which includes 2 empty hashmaps
new() -> Self177 pub fn new() -> Self {
178 let (sender, rx) = channel::<CapturePacket>();
179 let _ =
180 thread::Builder::new().name("capture_packet_handler".to_string()).spawn(move || {
181 while let Ok(CapturePacket { chip_id, packet, packet_type, direction }) = rx.recv()
182 {
183 handle_packet(chip_id, &packet, packet_type, direction);
184 }
185 });
186 Captures {
187 chip_id_to_capture: BTreeMap::<ChipIdentifier, Arc<Mutex<CaptureInfo>>>::new(),
188 sender,
189 }
190 }
191
192 /// Sends a captured packet to the output processing thread.
send( &self, chip_id: ChipIdentifier, packet: &Bytes, packet_type: u32, direction: PacketDirection, )193 pub fn send(
194 &self,
195 chip_id: ChipIdentifier,
196 packet: &Bytes,
197 packet_type: u32,
198 direction: PacketDirection,
199 ) {
200 let _ = self.sender.send(CapturePacket {
201 chip_id,
202 packet: packet.clone(),
203 packet_type,
204 direction,
205 });
206 }
207
208 /// Returns true if key exists in Captures.chip_id_to_capture
contains(&self, key: ChipIdentifier) -> bool209 pub fn contains(&self, key: ChipIdentifier) -> bool {
210 self.chip_id_to_capture.contains_key(&key)
211 }
212
213 /// Returns an Option of lockable and mutable CaptureInfo with given key
get(&mut self, key: ChipIdentifier) -> Option<&mut Arc<Mutex<CaptureInfo>>>214 pub fn get(&mut self, key: ChipIdentifier) -> Option<&mut Arc<Mutex<CaptureInfo>>> {
215 self.chip_id_to_capture.get_mut(&key)
216 }
217
218 /// Inserts the given CatpureInfo into Captures hashmaps
insert(&mut self, capture: CaptureInfo)219 pub fn insert(&mut self, capture: CaptureInfo) {
220 let chip_id = capture.id;
221 let arc_capture = Arc::new(Mutex::new(capture));
222 self.chip_id_to_capture.insert(chip_id, arc_capture.clone());
223 }
224
225 /// Returns true if chip_id_to_capture is empty
is_empty(&self) -> bool226 pub fn is_empty(&self) -> bool {
227 self.chip_id_to_capture.is_empty()
228 }
229
230 /// Returns an iterable object of chip_id_to_capture hashmap
iter(&self) -> Iter<ChipIdentifier, Arc<Mutex<CaptureInfo>>>231 pub fn iter(&self) -> Iter<ChipIdentifier, Arc<Mutex<CaptureInfo>>> {
232 self.chip_id_to_capture.iter()
233 }
234
235 /// Removes a CaptureInfo with given key from Captures
236 ///
237 /// When Capture is removed, remove from each map and also invoke closing of files.
remove(&mut self, key: &ChipIdentifier)238 pub fn remove(&mut self, key: &ChipIdentifier) {
239 if let Some(arc_capture) = self.chip_id_to_capture.get(key) {
240 let mut capture = arc_capture.lock().expect("Failed to acquire lock on CaptureInfo");
241 // Valid is marked false when chip is disconnected from netsim
242 capture.valid = false;
243 capture.stop_capture();
244 } else {
245 info!("key does not exist in Captures");
246 }
247 // CaptureInfo is not removed even after chip is removed
248 }
249
250 /// Returns Values of chip_id_to_capture hashmap values
values(&self) -> Values<ChipIdentifier, Arc<Mutex<CaptureInfo>>>251 pub fn values(&self) -> Values<ChipIdentifier, Arc<Mutex<CaptureInfo>>> {
252 self.chip_id_to_capture.values()
253 }
254 }
255
256 impl Default for Captures {
default() -> Self257 fn default() -> Self {
258 Self::new()
259 }
260 }
261
262 /// Create a thread to process events that matter to the Capture resource.
263 ///
264 /// We maintain a CaptureInfo for each chip that has been
265 /// connected to the simulation. This procedure monitors ChipAdded
266 /// and ChipRemoved events and updates the collection of CaptureInfo.
267 ///
spawn_capture_event_subscriber(event_rx: Receiver<Event>, capture: bool)268 pub fn spawn_capture_event_subscriber(event_rx: Receiver<Event>, capture: bool) {
269 let _ =
270 thread::Builder::new().name("capture_event_subscriber".to_string()).spawn(move || loop {
271 match event_rx.recv() {
272 Ok(Event::ChipAdded(ChipAdded { chip_id, chip_kind, device_name, .. })) => {
273 let mut capture_info =
274 CaptureInfo::new(chip_kind, chip_id, device_name.clone());
275 if capture {
276 if let Err(err) = capture_info.start_capture() {
277 warn!("{err:?}");
278 }
279 }
280 clone_captures().write().unwrap().insert(capture_info);
281 info!("Capture event: ChipAdded chip_id: {chip_id} device_name: {device_name}");
282 }
283 Ok(Event::ChipRemoved(ChipRemoved { chip_id, .. })) => {
284 clone_captures().write().unwrap().remove(&chip_id);
285 info!("Capture event: ChipRemoved chip_id: {chip_id}");
286 }
287 _ => {}
288 }
289 });
290 }
291