1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
//! Tools for handling the Control Pipe with `async` (`tokio`).
//!
//! There are three main classes provided in this module:
//!
//! * [`ExtcapControlSender`] – Implements the sender side for sending control
//!   packets from the extcap program you are implementing to Wireshark.
//! * [`ExtcapControlReader`] – Implements the reader side that receives control
//!   packets sent from Wireshark.
//! * [`ChannelExtcapControlReader`] – A wrapper around `ExtcapControlReader`
//!   that provides simpler, but less flexible, handling of the communication
//!   using a Tokio channel.
//!
//! See Wireshark's [Adding Capture Interfaces And Log Sources Using
//! Extcap](https://www.wireshark.org/docs/wsdg_html_chunked/ChCaptureExtcap.html#_messages)
//! section 8.2.3.2.1 for a description of the protocol format.

use async_trait::async_trait;
use log::debug;
use nom_derive::Parse;
use std::path::{Path, PathBuf};
use thiserror::Error;
#[cfg(target_os = "windows")]
use tokio::fs::File;
#[cfg(not(target_os = "windows"))]
use tokio::net::unix::pipe::{Receiver, Sender};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    sync::{
        mpsc::{self, error::SendError},
        Mutex,
    },
    task::JoinHandle,
};

pub mod util;
use util::AsyncReadExt as _;

use crate::controls::{ControlCommand, ControlPacket};

/// Error type returned for control packet read operations.
#[derive(Debug, Error)]
pub enum ReadControlError {
    /// Error reading the incoming control pipe.
    #[error(transparent)]
    IoError(#[from] tokio::io::Error),

    /// Error parsing the incoming data into the [`ControlPacket`] format.
    #[error("Error parsing control packet: {0}")]
    ParseError(String),
}

/// Error associated with [`ChannelExtcapControlReader`].
#[derive(Debug, Error)]
pub enum ControlChannelError {
    /// Error returned when the control packet cannot be read. See
    /// the docs on [`ReadControlError`].
    #[error(transparent)]
    ReadControl(#[from] ReadControlError),

    /// Error returned when the control packet cannot be sent on the channel.
    /// This is caused by an underlying [`SendError`].
    #[error("Cannot send control packet to channel")]
    CannotSend,
}

impl<T> From<SendError<T>> for ControlChannelError {
    fn from(_: SendError<T>) -> Self {
        ControlChannelError::CannotSend
    }
}

/// A reader for an Extcap Control using a [`Channel`][mpsc::channel]. This is
/// the easier to use, but higher overhead way to read control packets. When the
/// reader is spawned, a thread is spawned to continuously read messages and
/// writes them into a bounded `channel`. This allows reading the control
/// messages without worrying about spawning async tasks, by calling
/// [`try_read_packet`][Self::try_read_packet] every once in a while.
///
/// Assuming the extcap `capture` implementation uses a loop to read or generate
/// the packets, it can repeatedly call `try_read_packet` to read and handle the
/// control packets until there are no more buffered messages before starting
/// the main capturing logic.
///
/// For example:
/// ```ignore
/// fn capture(reader: &ChannelExtcapControlReader) -> Result<()> {
///     let pcap_header = ...;
///     let mut pcap_writer = PcapWriter::with_header(fifo, pcap_header)?;
///     loop {
///         while let Some(packet) = reader.try_read_packet().await {
///             // Handle the control packet
///         }
///         pcap_writer.write_packet(...)?;
///     }
///     Ok(())
/// }
pub struct ChannelExtcapControlReader {
    /// The join handle for the spawned thread. In most cases there is no need
    /// to use this, as the control fifo is expected to run for the whole
    /// duration of the capture.
    pub join_handle: JoinHandle<Result<(), ControlChannelError>>,
    /// The channel to receive control packets from.
    pub read_channel: mpsc::Receiver<ControlPacket<'static>>,
}

impl ChannelExtcapControlReader {
    /// Create a `ChannelExtcapControlReader` and spawns the underlying thread
    /// it uses to start reading the control packets from the pipe given in
    /// `in_path`.
    pub fn spawn(in_path: PathBuf) -> Self {
        let (tx, rx) = mpsc::channel::<ControlPacket<'static>>(10);
        let join_handle = tokio::task::spawn(async move {
            let mut reader = ExtcapControlReader::new(&in_path).await;
            loop {
                tx.send(reader.read_control_packet().await?).await?;
            }
        });
        Self {
            join_handle,
            read_channel: rx,
        }
    }

    /// Try to read a buffered control packet, or return `None` if there are no
    /// incoming control packets.
    pub async fn try_read_packet(&mut self) -> Option<ControlPacket<'static>> {
        self.read_channel.try_recv().ok()
    }

    /// Reads a control packet. If the incoming channel is empty, this will
    /// block and wait until an incoming packet comes in. This is typically used
    /// when the extcap capture starts to wait for the `Initialized` packet from
    /// the control channel.
    ///
    /// If you are only using this method and not using `try_read_packet`,
    /// consider whether you can use [`ExtcapControlReader`] directly for lower
    /// overhead.
    pub async fn read_packet(&mut self) -> Option<ControlPacket<'static>> {
        self.read_channel.recv().await
    }
}

/// A reader for the Extcap control pipe.
pub struct ExtcapControlReader {
    /// The file to read the control packets from. This is the fifo passed with
    /// the `--extcap-control-in` flag.
    #[cfg(not(target_os = "windows"))]
    in_file: Receiver,
    /// The file to read the control packets from. This is the fifo passed with
    /// the `--extcap-control-in` flag.
    #[cfg(target_os = "windows")]
    in_file: File,
}

impl ExtcapControlReader {
    /// Creates a new instance of [`ExtcapControlReader`].
    ///
    /// * `in_path`: The path of the extcap control pipe passed with
    ///   `--extcap-control-in`.
    #[cfg(not(target_os = "windows"))]
    pub async fn new(in_path: &Path) -> Self {
        Self {
            in_file: tokio::net::unix::pipe::OpenOptions::new()
                .open_receiver(in_path)
                .unwrap(),
        }
    }

    /// Creates a new instance of [`ExtcapControlReader`].
    ///
    /// * `in_path`: The path of the extcap control pipe passed with
    ///   `--extcap-control-in`.
    #[cfg(target_os = "windows")]
    pub async fn new(in_path: &Path) -> Self {
        Self {
            in_file: File::open(in_path).await.unwrap(),
        }
    }

    /// Read one control packet, awaiting until the packet arrives. Since the
    /// control packet pipe is expected to stay open for the entire duration of
    /// the extcap program, if the pipe is closed prematurely in this function
    /// here, `UnexpectedEof` will be returned.
    pub async fn read_control_packet(
        &mut self,
    ) -> Result<ControlPacket<'static>, ReadControlError> {
        let header_bytes = self
            .in_file
            .try_read_exact::<6>()
            .await?
            .ok_or_else(|| std::io::Error::from(std::io::ErrorKind::UnexpectedEof))?;
        debug!(
            "Read header bytes from incoming control message, now parsing... {:?}",
            header_bytes
        );
        let (_rem, packet) = match ControlPacket::parse(&header_bytes) {
            Ok((rem, packet)) => (rem, packet.into_owned()),
            Err(nom::Err::Incomplete(nom::Needed::Size(size))) => {
                let mut payload_bytes = vec![0_u8; size.get()];
                self.in_file.read_exact(&mut payload_bytes).await?;
                let all_bytes = [header_bytes.as_slice(), payload_bytes.as_slice()].concat();
                ControlPacket::parse(&all_bytes)
                    .map(|(_, packet)| (&[][..], packet.into_owned()))
                    .unwrap_or_else(|e| panic!("Unable to parse header packet: {e}"))
            }
            Err(e) => Err(ReadControlError::ParseError(e.to_string()))?,
        };
        debug!("Parsed incoming control message: {packet:?}");
        Ok(packet)
    }
}

const UNUSED_CONTROL_NUMBER: u8 = 255;

/// Sender for extcap control packets. These control packets controls the UI
/// generated by Wireshark. This trait also provides convenience functions for
/// sending control packets formatted for particular usages like `info_message`
/// and `status_message`. For other functions controlling various toolbar
/// controls, see the methods in the [`control`][crate::controls] module instead.
#[async_trait]
pub trait ExtcapControlSenderTrait: Send + Sync + Sized {
    /// Sends the given `packet` by writing it to the given output file (or
    /// fifo).
    async fn send(self, packet: ControlPacket<'_>) -> Result<(), tokio::io::Error>;

    /// Shows a message in an information dialog popup. The message will show on
    /// the screen until the user dismisses the popup.
    async fn info_message(self, message: &str) -> Result<(), tokio::io::Error> {
        self.send(ControlPacket::new_with_payload(
            UNUSED_CONTROL_NUMBER,
            ControlCommand::InformationMessage,
            message.as_bytes(),
        ))
        .await
    }

    /// Shows a message in a warning dialog popup. The message will show on the
    /// screen until the user dismisses the popup.
    async fn warning_message(self, message: &str) -> Result<(), tokio::io::Error> {
        self.send(ControlPacket::new_with_payload(
            UNUSED_CONTROL_NUMBER,
            ControlCommand::WarningMessage,
            message.as_bytes(),
        ))
        .await
    }

    /// Shows a message in an error dialog popup. The message will show on the
    /// screen until the user dismisses the popup.
    async fn error_message(self, message: &str) -> Result<(), tokio::io::Error> {
        self.send(ControlPacket::new_with_payload(
            UNUSED_CONTROL_NUMBER,
            ControlCommand::ErrorMessage,
            message.as_bytes(),
        ))
        .await
    }

    /// Shows a message in the status bar at the bottom of the Wireshark window.
    /// When the message is shown, the status bar will also flash yellow to
    /// bring it to the user's attention. The message will stay on the status
    /// bar for a few seconds, or until another message overwrites it.
    async fn status_message(self, message: &str) -> Result<(), tokio::io::Error> {
        self.send(ControlPacket::new_with_payload(
            UNUSED_CONTROL_NUMBER,
            ControlCommand::StatusbarMessage,
            message.as_bytes(),
        ))
        .await
    }
}

/// A sender for the extcap control packets. `out_file` should be the file given
/// by the `--extcap-control-out` flag.
pub struct ExtcapControlSender {
    #[cfg(not(target_os = "windows"))]
    out_file: Sender,
    #[cfg(target_os = "windows")]
    out_file: File,
}

impl ExtcapControlSender {
    #[cfg(not(target_os = "windows"))]
    /// Creates a new instance of [`ExtcapControlSender`].
    ///
    /// * `out_path`: The path specified by the `--extcap-control-out` flag.
    pub async fn new(out_path: &Path) -> Self {
        use std::time::Duration;

        for i in 0..50 {
            match tokio::net::unix::pipe::OpenOptions::new().open_sender(out_path) {
                Ok(out_file) => return Self { out_file },
                Err(e) => {
                    if let Some(libc::ENXIO) = e.raw_os_error() {
                        // This seems sketchy, but the docs for pipe::Sender says "The file is a
                        // FIFO, but no process has it open for reading. Sleep for a while and try
                        // again"
                        // https://docs.rs/tokio/latest/tokio/net/unix/pipe/struct.Sender.html
                        tokio::time::sleep(Duration::from_millis(i * 100)).await;
                    } else {
                        panic!("{e:?}");
                    }
                }
            };
        }
        panic!("Failed waiting for extcap-control-out to be opened");
    }

    #[cfg(target_os = "windows")]
    /// Creates a new instance of [`ExtcapControlSender`].
    ///
    /// * `out_path`: The path specified by the `--extcap-control-out` flag.
    pub async fn new(out_path: &Path) -> Self {
        Self {
            out_file: File::create(out_path).await.unwrap(),
        }
    }
}

#[async_trait]
impl<'a> ExtcapControlSenderTrait for &'a mut ExtcapControlSender {
    async fn send(self, packet: ControlPacket<'_>) -> Result<(), tokio::io::Error> {
        debug!("Sending extcap control message: {packet:#?}");
        self.out_file.write_all(&packet.to_header_bytes()).await?;
        self.out_file.write_all(&packet.payload).await?;
        self.out_file.flush().await?;
        Ok(())
    }
}

/// An implementation of ExtcapControlSenderTrait that is no-op when the
/// `Option` is `None`. Since Wireshark may not include the
/// `--extcap-control-out` flag (e.g. when no controls are returned during
/// `--extcap-interfaces`, or when running in tshark), this allows an easier but
/// less efficient way to say `option_extcap_sender.status_message(...)` without
/// constantly checking for the option.
#[async_trait]
impl<T> ExtcapControlSenderTrait for &mut Option<T>
where
    T: Send + Sync,
    for<'a> &'a mut T: ExtcapControlSenderTrait,
{
    async fn send(self, packet: ControlPacket<'_>) -> Result<(), tokio::io::Error> {
        if let Some(s) = self {
            s.send(packet).await
        } else {
            Ok(())
        }
    }
}

/// Just for syntactic niceness when working with a control sender behind a
/// mutex. This usage allows the sender to be locked only for the duration of
/// that one control packet, so it can be interleaved in between other async
/// function calls.
#[async_trait]
impl<T> ExtcapControlSenderTrait for &Mutex<T>
where
    T: Send,
    for<'a> &'a mut T: ExtcapControlSenderTrait,
{
    /// Sends a control message to Wireshark.
    async fn send(self, packet: ControlPacket<'_>) -> Result<(), tokio::io::Error> {
        self.lock().await.send(packet).await
    }
}