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 ChipId 19 //! and FacadeId to CaptureInfo. 20 21 use std::collections::btree_map::{Iter, Values}; 22 use std::collections::{BTreeMap, HashMap}; 23 use std::fs::{File, OpenOptions}; 24 use std::io::Result; 25 use std::sync::{Arc, Mutex}; 26 use std::time::{SystemTime, UNIX_EPOCH}; 27 28 use frontend_proto::{ 29 common::ChipKind, 30 model::{Capture as ProtoCapture, State}, 31 }; 32 use protobuf::well_known_types::timestamp::Timestamp; 33 34 use crate::ffi::get_facade_id; 35 36 use super::pcap_util::write_pcap_header; 37 38 pub type ChipId = i32; 39 pub type FacadeId = i32; 40 41 pub struct CaptureInfo { 42 facade_id: FacadeId, 43 pub file: Option<File>, 44 // Following items will be returned as ProtoCapture. (state: file.is_some()) 45 id: ChipId, 46 pub chip_kind: ChipKind, 47 pub device_name: String, 48 pub size: usize, 49 pub records: i32, 50 pub seconds: i64, 51 pub nanos: i32, 52 pub valid: bool, 53 } 54 55 // Captures contains a recent copy of all chips and their ChipKind, chip_id, 56 // and owning device name. Information for any recent or ongoing captures is 57 // also stored in the ProtoCapture. 58 // facade_key_to_capture allows for fast lookups when handle_request, handle_response 59 // is invoked from packet_hub. 60 pub struct Captures { 61 pub facade_key_to_capture: HashMap<(ChipKind, FacadeId), Arc<Mutex<CaptureInfo>>>, 62 // BTreeMap is used for chip_id_to_capture, so that the CaptureInfo can always be 63 // ordered by ChipId. ListCaptureResponse will produce a ordered list of CaptureInfos. 64 pub chip_id_to_capture: BTreeMap<ChipId, Arc<Mutex<CaptureInfo>>>, 65 } 66 67 impl CaptureInfo { new(chip_kind: ChipKind, chip_id: ChipId, device_name: String) -> Self68 pub fn new(chip_kind: ChipKind, chip_id: ChipId, device_name: String) -> Self { 69 CaptureInfo { 70 facade_id: get_facade_id(chip_id), 71 id: chip_id, 72 chip_kind, 73 device_name, 74 size: 0, 75 records: 0, 76 seconds: 0, 77 nanos: 0, 78 valid: true, 79 file: None, 80 } 81 } 82 83 // Creates a pcap file with headers and store it under temp directory 84 // The lifecycle of the file is NOT tied to the lifecycle of the struct 85 // Format: /tmp/netsim-pcaps/{chip_id}-{device_name}-{chip_kind}.pcap start_capture(&mut self) -> Result<()>86 pub fn start_capture(&mut self) -> Result<()> { 87 if self.file.is_some() { 88 return Ok(()); 89 } 90 let mut filename = std::env::temp_dir(); 91 filename.push("netsim-pcaps"); 92 std::fs::create_dir_all(&filename)?; 93 filename.push(format!("{:?}-{:}-{:?}.pcap", self.id, self.device_name, self.chip_kind)); 94 let mut file = OpenOptions::new().write(true).truncate(true).create(true).open(filename)?; 95 let size = write_pcap_header(&mut file)?; 96 let timestamp = SystemTime::now().duration_since(UNIX_EPOCH).expect("Time went backwards"); 97 self.size = size; 98 self.records = 0; 99 self.seconds = timestamp.as_secs() as i64; 100 self.nanos = timestamp.subsec_nanos() as i32; 101 self.file = Some(file); 102 Ok(()) 103 } 104 105 // Closes file by removing ownership of self.file 106 // Capture info will still retain the size and record count 107 // So it can be downloaded easily when GetCapture is invoked. stop_capture(&mut self)108 pub fn stop_capture(&mut self) { 109 self.file = None; 110 } 111 new_facade_key(kind: ChipKind, facade_id: FacadeId) -> (ChipKind, FacadeId)112 pub fn new_facade_key(kind: ChipKind, facade_id: FacadeId) -> (ChipKind, FacadeId) { 113 (kind, facade_id) 114 } 115 get_facade_key(&self) -> (ChipKind, FacadeId)116 pub fn get_facade_key(&self) -> (ChipKind, FacadeId) { 117 CaptureInfo::new_facade_key(self.chip_kind, self.facade_id) 118 } 119 get_capture_proto(&self) -> ProtoCapture120 pub fn get_capture_proto(&self) -> ProtoCapture { 121 let timestamp = 122 Timestamp { seconds: self.seconds, nanos: self.nanos, ..Default::default() }; 123 ProtoCapture { 124 id: self.id, 125 chip_kind: self.chip_kind.into(), 126 device_name: self.device_name.clone(), 127 state: match self.file.is_some() { 128 true => State::ON.into(), 129 false => State::OFF.into(), 130 }, 131 size: self.size as i32, 132 records: self.records, 133 timestamp: Some(timestamp).into(), 134 valid: self.valid, 135 ..Default::default() 136 } 137 } 138 } 139 140 impl Captures { new() -> Self141 pub fn new() -> Self { 142 Captures { 143 facade_key_to_capture: HashMap::<(ChipKind, FacadeId), Arc<Mutex<CaptureInfo>>>::new(), 144 chip_id_to_capture: BTreeMap::<ChipId, Arc<Mutex<CaptureInfo>>>::new(), 145 } 146 } 147 contains(&self, key: ChipId) -> bool148 pub fn contains(&self, key: ChipId) -> bool { 149 self.chip_id_to_capture.contains_key(&key) 150 } 151 get(&mut self, key: ChipId) -> Option<&mut Arc<Mutex<CaptureInfo>>>152 pub fn get(&mut self, key: ChipId) -> Option<&mut Arc<Mutex<CaptureInfo>>> { 153 self.chip_id_to_capture.get_mut(&key) 154 } 155 insert(&mut self, capture: CaptureInfo)156 pub fn insert(&mut self, capture: CaptureInfo) { 157 let chip_id = capture.id; 158 let facade_key = capture.get_facade_key(); 159 let arc_capture = Arc::new(Mutex::new(capture)); 160 self.chip_id_to_capture.insert(chip_id, arc_capture.clone()); 161 self.facade_key_to_capture.insert(facade_key, arc_capture); 162 } 163 is_empty(&self) -> bool164 pub fn is_empty(&self) -> bool { 165 self.chip_id_to_capture.is_empty() 166 } 167 iter(&self) -> Iter<ChipId, Arc<Mutex<CaptureInfo>>>168 pub fn iter(&self) -> Iter<ChipId, Arc<Mutex<CaptureInfo>>> { 169 self.chip_id_to_capture.iter() 170 } 171 172 // When Capture is removed, remove from each map and also invoke closing of files. remove(&mut self, key: &ChipId)173 pub fn remove(&mut self, key: &ChipId) { 174 if let Some(arc_capture) = self.chip_id_to_capture.get(key) { 175 if let Ok(mut capture) = arc_capture.lock() { 176 self.facade_key_to_capture.remove(&capture.get_facade_key()); 177 capture.stop_capture(); 178 } 179 } else { 180 println!("key does not exist in Captures"); 181 return; 182 } 183 self.chip_id_to_capture.remove(key); 184 } 185 values(&self) -> Values<ChipId, Arc<Mutex<CaptureInfo>>>186 pub fn values(&self) -> Values<ChipId, Arc<Mutex<CaptureInfo>>> { 187 self.chip_id_to_capture.values() 188 } 189 } 190