Эх сурвалжийг харах

phy: simplify PcapSink trait

Dario Nieuwenhuis 3 жил өмнө
parent
commit
dfe8265772

+ 3 - 5
examples/utils.rs

@@ -5,12 +5,10 @@ use env_logger::Builder;
 use getopts::{Matches, Options};
 #[cfg(feature = "log")]
 use log::{trace, Level, LevelFilter};
-use std::cell::RefCell;
 use std::env;
 use std::fs::File;
 use std::io::{self, Write};
 use std::process;
-use std::rc::Rc;
 use std::str::{self, FromStr};
 use std::time::{SystemTime, UNIX_EPOCH};
 
@@ -18,7 +16,7 @@ use smoltcp::phy::RawSocket;
 #[cfg(feature = "phy-tuntap_interface")]
 use smoltcp::phy::TunTapInterface;
 use smoltcp::phy::{Device, FaultInjector, Medium, Tracer};
-use smoltcp::phy::{PcapMode, PcapSink, PcapWriter};
+use smoltcp::phy::{PcapMode, PcapWriter};
 use smoltcp::time::{Duration, Instant};
 
 #[cfg(feature = "log")]
@@ -165,7 +163,7 @@ pub fn parse_middleware_options<D>(
     matches: &mut Matches,
     device: D,
     loopback: bool,
-) -> FaultInjector<Tracer<PcapWriter<D, Rc<dyn PcapSink>>>>
+) -> FaultInjector<Tracer<PcapWriter<D, Box<dyn io::Write>>>>
 where
     D: for<'a> Device<'a>,
 {
@@ -208,7 +206,7 @@ where
 
     let device = PcapWriter::new(
         device,
-        Rc::new(RefCell::new(pcap_writer)) as Rc<dyn PcapSink>,
+        pcap_writer,
         if loopback {
             PcapMode::TxOnly
         } else {

+ 1 - 1
src/phy/fault_injector.rs

@@ -289,7 +289,7 @@ impl<'a, Rx: phy::RxToken> phy::RxToken for RxToken<'a, Rx> {
                 let mut corrupt = &mut corrupt[..buffer.len()];
                 corrupt.copy_from_slice(buffer);
                 state.borrow_mut().corrupt(&mut corrupt);
-                f(&mut corrupt)
+                f(corrupt)
             } else {
                 f(buffer)
             }

+ 2 - 2
src/phy/fuzz_injector.rs

@@ -122,8 +122,8 @@ impl<'a, Tx: phy::TxToken, FTx: Fuzzer> phy::TxToken for TxToken<'a, Tx, FTx> {
         F: FnOnce(&mut [u8]) -> Result<R>,
     {
         let Self { fuzzer, token } = self;
-        token.consume(timestamp, len, |mut buf| {
-            fuzzer.fuzz_packet(&mut buf);
+        token.consume(timestamp, len, |buf| {
+            fuzzer.fuzz_packet(buf);
             f(buf)
         })
     }

+ 49 - 58
src/phy/pcap_writer.rs

@@ -1,8 +1,7 @@
 use byteorder::{ByteOrder, NativeEndian};
+use core::cell::RefCell;
 use phy::Medium;
 #[cfg(feature = "std")]
-use std::cell::RefCell;
-#[cfg(feature = "std")]
 use std::io::Write;
 
 use crate::phy::{self, Device, DeviceCapabilities};
@@ -34,17 +33,20 @@ pub enum PcapMode {
 /// A packet capture sink.
 pub trait PcapSink {
     /// Write data into the sink.
-    fn write(&self, data: &[u8]);
+    fn write(&mut self, data: &[u8]);
+
+    /// Flush data written into the sync.
+    fn flush(&mut self) {}
 
     /// Write an `u16` into the sink, in native byte order.
-    fn write_u16(&self, value: u16) {
+    fn write_u16(&mut self, value: u16) {
         let mut bytes = [0u8; 2];
         NativeEndian::write_u16(&mut bytes, value);
         self.write(&bytes[..])
     }
 
     /// Write an `u32` into the sink, in native byte order.
-    fn write_u32(&self, value: u32) {
+    fn write_u32(&mut self, value: u32) {
         let mut bytes = [0u8; 4];
         NativeEndian::write_u32(&mut bytes, value);
         self.write(&bytes[..])
@@ -53,7 +55,7 @@ pub trait PcapSink {
     /// Write the libpcap global header into the sink.
     ///
     /// This method may be overridden e.g. if special synchronization is necessary.
-    fn global_header(&self, link_type: PcapLinkType) {
+    fn global_header(&mut self, link_type: PcapLinkType) {
         self.write_u32(0xa1b2c3d4); // magic number
         self.write_u16(2); // major version
         self.write_u16(4); // minor version
@@ -69,7 +71,7 @@ pub trait PcapSink {
     ///
     /// # Panics
     /// This function panics if `length` is greater than 65535.
-    fn packet_header(&self, timestamp: Instant, length: usize) {
+    fn packet_header(&mut self, timestamp: Instant, length: usize) {
         assert!(length <= 65535);
 
         self.write_u32(timestamp.secs() as u32); // timestamp seconds
@@ -81,28 +83,21 @@ pub trait PcapSink {
     /// Write the libpcap packet header followed by packet data into the sink.
     ///
     /// See also the note for [global_header](#method.global_header).
-    fn packet(&self, timestamp: Instant, packet: &[u8]) {
+    fn packet(&mut self, timestamp: Instant, packet: &[u8]) {
         self.packet_header(timestamp, packet.len());
-        self.write(packet)
-    }
-}
-
-impl<T: AsRef<dyn PcapSink>> PcapSink for T {
-    fn write(&self, data: &[u8]) {
-        self.as_ref().write(data)
+        self.write(packet);
+        self.flush();
     }
 }
 
 #[cfg(feature = "std")]
-impl<T: Write> PcapSink for RefCell<T> {
-    fn write(&self, data: &[u8]) {
-        self.borrow_mut().write_all(data).expect("cannot write")
+impl<T: Write> PcapSink for T {
+    fn write(&mut self, data: &[u8]) {
+        T::write_all(self, data).expect("cannot write")
     }
 
-    fn packet(&self, timestamp: Instant, packet: &[u8]) {
-        self.packet_header(timestamp, packet.len());
-        PcapSink::write(self, packet);
-        self.borrow_mut().flush().expect("cannot flush")
+    fn flush(&mut self) {
+        T::flush(self).expect("cannot flush")
     }
 }
 
@@ -119,20 +114,19 @@ impl<T: Write> PcapSink for RefCell<T> {
 /// [libpcap]: https://wiki.wireshark.org/Development/LibpcapFileFormat
 /// [sink]: trait.PcapSink.html
 #[derive(Debug)]
-#[cfg_attr(feature = "defmt", derive(defmt::Format))]
 pub struct PcapWriter<D, S>
 where
     D: for<'a> Device<'a>,
-    S: PcapSink + Clone,
+    S: PcapSink,
 {
     lower: D,
-    sink: S,
+    sink: RefCell<S>,
     mode: PcapMode,
 }
 
-impl<D: for<'a> Device<'a>, S: PcapSink + Clone> PcapWriter<D, S> {
+impl<D: for<'a> Device<'a>, S: PcapSink> PcapWriter<D, S> {
     /// Creates a packet capture writer.
-    pub fn new(lower: D, sink: S, mode: PcapMode) -> PcapWriter<D, S> {
+    pub fn new(lower: D, mut sink: S, mode: PcapMode) -> PcapWriter<D, S> {
         let medium = lower.capabilities().medium;
         let link_type = match medium {
             #[cfg(feature = "medium-ip")]
@@ -141,7 +135,11 @@ impl<D: for<'a> Device<'a>, S: PcapSink + Clone> PcapWriter<D, S> {
             Medium::Ethernet => PcapLinkType::Ethernet,
         };
         sink.global_header(link_type);
-        PcapWriter { lower, sink, mode }
+        PcapWriter {
+            lower,
+            sink: RefCell::new(sink),
+            mode,
+        }
     }
 
     /// Get a reference to the underlying device.
@@ -163,31 +161,27 @@ impl<D: for<'a> Device<'a>, S: PcapSink + Clone> PcapWriter<D, S> {
 impl<'a, D, S> Device<'a> for PcapWriter<D, S>
 where
     D: for<'b> Device<'b>,
-    S: PcapSink + Clone + 'a,
+    S: PcapSink + 'a,
 {
-    type RxToken = RxToken<<D as Device<'a>>::RxToken, S>;
-    type TxToken = TxToken<<D as Device<'a>>::TxToken, S>;
+    type RxToken = RxToken<'a, <D as Device<'a>>::RxToken, S>;
+    type TxToken = TxToken<'a, <D as Device<'a>>::TxToken, S>;
 
     fn capabilities(&self) -> DeviceCapabilities {
         self.lower.capabilities()
     }
 
     fn receive(&'a mut self) -> Option<(Self::RxToken, Self::TxToken)> {
-        let &mut Self {
-            ref mut lower,
-            ref sink,
-            mode,
-            ..
-        } = self;
-        lower.receive().map(|(rx_token, tx_token)| {
+        let sink = &self.sink;
+        let mode = self.mode;
+        self.lower.receive().map(move |(rx_token, tx_token)| {
             let rx = RxToken {
                 token: rx_token,
-                sink: sink.clone(),
+                sink,
                 mode,
             };
             let tx = TxToken {
                 token: tx_token,
-                sink: sink.clone(),
+                sink,
                 mode,
             };
             (rx, tx)
@@ -195,32 +189,29 @@ where
     }
 
     fn transmit(&'a mut self) -> Option<Self::TxToken> {
-        let &mut Self {
-            ref mut lower,
-            ref sink,
-            mode,
-        } = self;
-        lower.transmit().map(|token| TxToken {
-            token,
-            sink: sink.clone(),
-            mode,
-        })
+        let sink = &self.sink;
+        let mode = self.mode;
+        self.lower
+            .transmit()
+            .map(move |token| TxToken { token, sink, mode })
     }
 }
 
 #[doc(hidden)]
-pub struct RxToken<Rx: phy::RxToken, S: PcapSink> {
+pub struct RxToken<'a, Rx: phy::RxToken, S: PcapSink> {
     token: Rx,
-    sink: S,
+    sink: &'a RefCell<S>,
     mode: PcapMode,
 }
 
-impl<Rx: phy::RxToken, S: PcapSink> phy::RxToken for RxToken<Rx, S> {
+impl<'a, Rx: phy::RxToken, S: PcapSink> phy::RxToken for RxToken<'a, Rx, S> {
     fn consume<R, F: FnOnce(&mut [u8]) -> Result<R>>(self, timestamp: Instant, f: F) -> Result<R> {
         let Self { token, sink, mode } = self;
         token.consume(timestamp, |buffer| {
             match mode {
-                PcapMode::Both | PcapMode::RxOnly => sink.packet(timestamp, buffer.as_ref()),
+                PcapMode::Both | PcapMode::RxOnly => {
+                    sink.borrow_mut().packet(timestamp, buffer.as_ref())
+                }
                 PcapMode::TxOnly => (),
             }
             f(buffer)
@@ -229,13 +220,13 @@ impl<Rx: phy::RxToken, S: PcapSink> phy::RxToken for RxToken<Rx, S> {
 }
 
 #[doc(hidden)]
-pub struct TxToken<Tx: phy::TxToken, S: PcapSink> {
+pub struct TxToken<'a, Tx: phy::TxToken, S: PcapSink> {
     token: Tx,
-    sink: S,
+    sink: &'a RefCell<S>,
     mode: PcapMode,
 }
 
-impl<Tx: phy::TxToken, S: PcapSink> phy::TxToken for TxToken<Tx, S> {
+impl<'a, Tx: phy::TxToken, S: PcapSink> phy::TxToken for TxToken<'a, Tx, S> {
     fn consume<R, F>(self, timestamp: Instant, len: usize, f: F) -> Result<R>
     where
         F: FnOnce(&mut [u8]) -> Result<R>,
@@ -244,7 +235,7 @@ impl<Tx: phy::TxToken, S: PcapSink> phy::TxToken for TxToken<Tx, S> {
         token.consume(timestamp, len, |buffer| {
             let result = f(buffer);
             match mode {
-                PcapMode::Both | PcapMode::TxOnly => sink.packet(timestamp, buffer),
+                PcapMode::Both | PcapMode::TxOnly => sink.borrow_mut().packet(timestamp, buffer),
                 PcapMode::RxOnly => (),
             };
             result

+ 3 - 7
utils/packet2pcap.rs

@@ -1,10 +1,9 @@
 use getopts::Options;
 use smoltcp::phy::{PcapLinkType, PcapSink};
 use smoltcp::time::Instant;
-use std::cell::RefCell;
 use std::env;
 use std::fs::File;
-use std::io::{self, Read, Write};
+use std::io::{self, Read};
 use std::path::Path;
 use std::process::exit;
 
@@ -17,12 +16,9 @@ fn convert(
     let mut packet = Vec::new();
     packet_file.read_to_end(&mut packet)?;
 
-    let pcap = RefCell::new(Vec::new());
-    PcapSink::global_header(&pcap, link_type);
-    PcapSink::packet(&pcap, Instant::from_millis(0), &packet[..]);
-
     let mut pcap_file = File::create(pcap_filename)?;
-    pcap_file.write_all(&pcap.borrow()[..])?;
+    PcapSink::global_header(&mut pcap_file, link_type);
+    PcapSink::packet(&mut pcap_file, Instant::from_millis(0), &packet[..]);
 
     Ok(())
 }