Merge pull request #111 from martinling/save-hitl-captures

Optionally save captures from hardware-in-the-loop testing

Author: martinling

src/capture.rs

@@ -1628,7 +1628,7 @@ mod tests {
     use std::io::{BufReader, BufWriter, BufRead, Write};
     use std::path::PathBuf;
     use crate::decoder::Decoder;
-    use crate::loader::Loader;
+    use crate::pcap::Loader;
     use itertools::Itertools;
 
     fn summarize_item(cap: &mut CaptureReader, item: &TrafficItem, depth: usize)

src/main.rs

@@ -20,8 +20,8 @@ mod decoder;
 mod expander;
 mod id;
 mod index_stream;
-mod loader;
 mod model;
+mod pcap;
 mod rcu;
 mod row_data;
 mod stream;
@@ -47,9 +47,14 @@ use ui::{
     stop_cynthion
 };
 
+fn have_argument(name: &str) -> bool {
+    std::env::args().any(|arg| arg == name)
+}
+
 fn main() {
-    if std::env::args().any(|arg| arg == "--test-cynthion") {
-        test_cynthion::run_test();
+    if have_argument("--test-cynthion") {
+        let save_captures = have_argument("--save-captures");
+        test_cynthion::run_test(save_captures);
     } else {
         let application = gtk::Application::new(
             Some("com.greatscottgadgets.packetry"),

src/loader.rs src/pcap.rssrc/loader.rs renamed to src/pcap.rs

@@ -1,14 +1,16 @@
+use std::borrow::Cow;
 use std::fs::File;
-use std::io::BufReader;
+use std::io::{BufReader, BufWriter, Write};
 use std::mem::size_of;
 use std::path::PathBuf;
 
 use pcap_file::{
-    pcap::{PcapReader, PcapHeader, RawPcapPacket},
+    pcap::{PcapReader, PcapWriter, PcapHeader, RawPcapPacket},
+    DataLink,
     TsResolution,
 };
 
-use anyhow::Error;
+use anyhow::{Context, Error};
 
 pub struct Loader {
     pcap: PcapReader<BufReader<File>>,
@@ -18,6 +20,10 @@ pub struct Loader {
     start_time: Option<u64>,
 }
 
+pub struct Writer {
+    pcap: PcapWriter<BufWriter<File>>,
+}
+
 impl Loader {
     pub fn open(path: PathBuf) -> Result<Loader, Error> {
         let file = File::open(path)?;
@@ -55,3 +61,37 @@ impl Loader {
         }
     }
 }
+
+impl Writer {
+    pub fn open(path: PathBuf) -> Result<Writer, Error> {
+        let file = File::create(path)?;
+        let writer = BufWriter::new(file);
+        let header = PcapHeader {
+            datalink: DataLink::USB_2_0,
+            ts_resolution: TsResolution::NanoSecond,
+            .. PcapHeader::default()
+        };
+        Ok(Writer{pcap: PcapWriter::with_header(writer, header)?})
+    }
+
+    pub fn add_packet(&mut self, bytes: &[u8], timestamp_ns: u64) -> Result<(), Error> {
+        let length: u32 = bytes
+            .len()
+            .try_into()
+            .context("Packet too large for pcap file")?;
+        let packet = RawPcapPacket {
+            ts_sec: (timestamp_ns / 1_000_000_000) as u32,
+            ts_frac: (timestamp_ns % 1_000_000_000) as u32,
+            incl_len: length,
+            orig_len: length,
+            data: Cow::from(bytes)
+        };
+        self.pcap.write_raw_packet(&packet)?;
+        Ok(())
+    }
+
+    pub fn close(self) -> Result<(), Error> {
+        self.pcap.into_writer().flush()?;
+        Ok(())
+    }
+}

src/test_cynthion.rs

@@ -8,25 +8,29 @@ use crate::capture::{
     PacketId,
 };
 use crate::decoder::Decoder;
+use crate::pcap::Writer;
 
 use anyhow::{Context, Error};
 use futures_lite::future::block_on;
 use nusb::transfer::RequestBuffer;
 
+use std::path::PathBuf;
 use std::thread::sleep;
 use std::time::Duration;
 
-pub fn run_test() {
-    for (speed, ep_addr, length, sof) in [
-        (Speed::High, 0x81, 4096, Some((124500,  125500, 500))),
-        (Speed::Full, 0x82,  512, Some((995000, 1005000,  50))),
-        (Speed::Low,  0x83,   64, None)]
+pub fn run_test(save_captures: bool) {
+    for (name, speed, ep_addr, length, sof) in [
+        ("HS", Speed::High, 0x81, 4096, Some((124500,  125500, 500))),
+        ("FS", Speed::Full, 0x82,  512, Some((995000, 1005000,  50))),
+        ("LS", Speed::Low,  0x83,   64, None)]
     {
-        test(speed, ep_addr, length, sof).unwrap();
+        test(save_captures, name, speed, ep_addr, length, sof).unwrap();
     }
 }
 
-fn test(speed: Speed,
+fn test(save_capture: bool,
+        name: &str,
+        speed: Speed,
         ep_addr: u8,
         length: usize,
         sof: Option<(u64, u64, u64)>)
@@ -96,6 +100,19 @@ fn test(speed: Speed,
             .context("Error decoding packet")?;
     }
 
+    if save_capture {
+        // Write the capture to a file.
+        let path = PathBuf::from(format!("./HITL-{name}.pcap"));
+        let mut writer = Writer::open(path)?;
+        for i in 0..reader.packet_index.len() {
+            let packet_id = PacketId::from(i);
+            let packet = reader.packet(packet_id)?;
+            let timestamp_ns = reader.packet_time(packet_id)?;
+            writer.add_packet(&packet, timestamp_ns)?;
+        }
+        writer.close()?;
+    }
+
     // Look for the test device in the capture.
     let device_id = DeviceId::from(1);
     let device_data = reader.device_data(&device_id)?;

src/test_replay.rs

@@ -8,7 +8,7 @@ use itertools::assert_equal;
 use serde_json::Deserializer;
 
 use crate::decoder::Decoder;
-use crate::loader::Loader;
+use crate::pcap::Loader;
 use crate::model::GenericModel;
 use crate::row_data::{GenericRowData, TrafficRowData, DeviceRowData};
 use crate::record_ui::UiAction;

src/ui.rs

@@ -1,7 +1,4 @@
-use std::borrow::Cow;
 use std::cell::RefCell;
-use std::fs::File;
-use std::io::{BufWriter, Write};
 use std::path::PathBuf;
 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::time::Duration;
@@ -46,12 +43,6 @@ use gtk::{
     ButtonsType,
 };
 
-use pcap_file::{
-    DataLink,
-    TsResolution,
-    pcap::{PcapWriter, PcapHeader, RawPcapPacket},
-};
-
 use crate::backend::cynthion::{
     CynthionDevice,
     CynthionHandle,
@@ -70,7 +61,7 @@ use crate::capture::{
 };
 use crate::decoder::Decoder;
 use crate::expander::ExpanderWrapper;
-use crate::loader::Loader;
+use crate::pcap::{Loader, Writer};
 use crate::model::{GenericModel, TrafficModel, DeviceModel};
 use crate::row_data::{
     GenericRowData,
@@ -861,36 +852,18 @@ fn start_pcap(action: FileAction, path: PathBuf) -> Result<(), Error> {
                 let packet_count = capture.packet_index.len();
                 TOTAL.store(packet_count, Ordering::Relaxed);
                 CURRENT.store(0, Ordering::Relaxed);
-                let file = File::create(path)?;
-                let writer = BufWriter::new(file);
-                let header = PcapHeader {
-                    datalink: DataLink::USB_2_0,
-                    ts_resolution: TsResolution::NanoSecond,
-                    .. PcapHeader::default()
-                };
-                let mut pcap = PcapWriter::with_header(writer, header)?;
+                let mut writer = Writer::open(path)?;
                 for i in 0..packet_count {
                     let packet_id = PacketId::from(i);
-                    let bytes = capture.packet(packet_id)?;
+                    let packet = capture.packet(packet_id)?;
                     let timestamp_ns = capture.packet_time(packet_id)?;
-                    let length: u32 = bytes
-                        .len()
-                        .try_into()
-                        .context("Packet too large for pcap file")?;
-                    let packet = RawPcapPacket {
-                        ts_sec: (timestamp_ns / 1_000_000_000) as u32,
-                        ts_frac: (timestamp_ns % 1_000_000_000) as u32,
-                        incl_len: length,
-                        orig_len: length,
-                        data: Cow::from(bytes)
-                    };
-                    pcap.write_raw_packet(&packet)?;
+                    writer.add_packet(&packet, timestamp_ns)?;
                     CURRENT.store(i + 1, Ordering::Relaxed);
                     if STOP.load(Ordering::Relaxed) {
                         break;
                     }
                 }
-                pcap.into_writer().flush()?;
+                writer.close()?;
                 Ok(())
             },
         };