sparse/pcap-sys/src/lib.rs

// Copyright (C) 2023 Andrew Rioux
// 
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
// 
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
// 
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use std::{
    ffi::{CStr, CString},
    os::fd::{AsRawFd, RawFd},
    pin::Pin,
    ptr, slice,
    task::{self, Poll}, collections::HashMap,
};

pub mod error;
mod ffi;
pub mod packets;

pub mod consts {
    pub use super::ffi::{
        BUFSIZ, DLT_ARCNET, DLT_AX25, DLT_CHAOS, DLT_EN10MB, DLT_EN3MB, DLT_FDDI, DLT_IEEE802,
        DLT_NULL, DLT_PPP, DLT_PRONET, DLT_SLIP, PCAP_BUF_SIZE, PCAP_CHAR_ENC_LOCAL,
        PCAP_CHAR_ENC_UTF_8, PCAP_ERRBUF_SIZE, PCAP_ERROR, PCAP_ERROR_ACTIVATED, PCAP_ERROR_BREAK,
        PCAP_ERROR_CANTSET_TSTAMP_TYPE, PCAP_ERROR_IFACE_NOT_UP, PCAP_ERROR_NOT_ACTIVATED,
        PCAP_ERROR_NOT_RFMON, PCAP_ERROR_NO_SUCH_DEVICE, PCAP_ERROR_PERM_DENIED,
        PCAP_ERROR_PROMISC_PERM_DENIED, PCAP_ERROR_RFMON_NOTSUP,
        PCAP_ERROR_TSTAMP_PRECISION_NOTSUP, PCAP_IF_CONNECTION_STATUS,
        PCAP_IF_CONNECTION_STATUS_CONNECTED, PCAP_IF_CONNECTION_STATUS_DISCONNECTED,
        PCAP_IF_CONNECTION_STATUS_NOT_APPLICABLE, PCAP_IF_CONNECTION_STATUS_UNKNOWN,
        PCAP_IF_LOOPBACK, PCAP_IF_RUNNING, PCAP_IF_UP, PCAP_IF_WIRELESS, PCAP_NETMASK_UNKNOWN,
        PCAP_OPENFLAG_DATATX_UDP, PCAP_OPENFLAG_MAX_RESPONSIVENESS, PCAP_OPENFLAG_NOCAPTURE_LOCAL,
        PCAP_OPENFLAG_NOCAPTURE_RPCAP, PCAP_OPENFLAG_PROMISCUOUS, PCAP_SAMP_1_EVERY_N,
        PCAP_SAMP_FIRST_AFTER_N_MS, PCAP_SAMP_NOSAMP, PCAP_SRC_FILE, PCAP_SRC_FILE_STRING,
        PCAP_SRC_IFLOCAL, PCAP_SRC_IFREMOTE, PCAP_SRC_IF_STRING, PCAP_TSTAMP_ADAPTER,
        PCAP_TSTAMP_ADAPTER_UNSYNCED, PCAP_TSTAMP_HOST, PCAP_TSTAMP_HOST_HIPREC,
        PCAP_TSTAMP_HOST_HIPREC_UNSYNCED, PCAP_TSTAMP_HOST_LOWPREC, PCAP_TSTAMP_PRECISION_MICRO,
        PCAP_TSTAMP_PRECISION_NANO, PCAP_VERSION_MAJOR, PCAP_VERSION_MINOR, PCAP_WARNING,
        PCAP_WARNING_PROMISC_NOTSUP, PCAP_WARNING_TSTAMP_TYPE_NOTSUP, RPCAP_HOSTLIST_SIZE,
        RPCAP_RMTAUTH_NULL, RPCAP_RMTAUTH_PWD,
    };
}

use ffi::PcapDevIf;
use futures::{ready, Stream, StreamExt};
use tokio::io::unix::AsyncFd;
use tokio_stream::StreamMap;

pub struct PcapDevIterator {
    dev_if: *const PcapDevIf,
    current_dev: *const PcapDevIf,
}

impl PcapDevIterator {
    pub fn new() -> error::Result<PcapDevIterator> {
        let mut errbuf = [0i8; ffi::PCAP_ERRBUF_SIZE];
        let mut dev_if = ptr::null();

        if unsafe { ffi::pcap_findalldevs(&mut dev_if, errbuf.as_mut_ptr()) } != 0 {
            Err(&errbuf)?;
        }

        Ok(PcapDevIterator {
            dev_if,
            current_dev: dev_if,
        })
    }
}

impl Drop for PcapDevIterator {
    fn drop(&mut self) {
        unsafe {
            ffi::pcap_freealldevs(self.dev_if);
        }
    }
}

impl std::iter::Iterator for PcapDevIterator {
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        unsafe {
            if self.current_dev.is_null() {
                return None;
            }

            let current_dev = self.current_dev;

            let dev_name = CStr::from_ptr((*current_dev).name)
                .to_str()
                .unwrap()
                .to_string();

            self.current_dev = (*current_dev).next;

            Some(dev_name)
        }
    }
}

pub trait State {}
pub trait Activated: State {}
pub trait NotListening: Activated {}
pub trait Listening: Activated {}
pub trait Disabled: State {}

pub enum DevActivated {}
impl State for DevActivated {}
impl Activated for DevActivated {}
impl NotListening for DevActivated {}

pub enum DevDisabled {}
impl State for DevDisabled {}
impl Disabled for DevDisabled {}

pub enum DevListening {}
impl State for DevListening {}
impl Activated for DevListening {}
impl Listening for DevListening {}

pub struct BpfProgram {}

pub struct Interface<T: State> {
    dev_name: CString,
    dev: *mut ffi::PcapDev,
    marker: std::marker::PhantomData<T>,
    absorbed: bool,
    nonblocking: bool,
}

impl<T: State> Drop for Interface<T> {
    fn drop(&mut self) {
        if !self.absorbed {
            unsafe { ffi::pcap_close(self.dev) };
        }
    }
}

unsafe impl<T: State> Send for Interface<T> {}

impl<T: State> Interface<T> {
    pub fn new(name: &str) -> error::Result<Interface<DevDisabled>> {
        let mut errbuf = [0i8; ffi::PCAP_ERRBUF_SIZE];

        let dev_name = CString::new(name)?;
        let dev = unsafe { ffi::pcap_create(dev_name.as_ptr(), errbuf.as_mut_ptr()) };

        if dev.is_null() {
            Err(&errbuf)?;
        }

        Ok(Interface::<DevDisabled> {
            dev_name,
            dev,
            marker: std::marker::PhantomData,
            absorbed: false,
            nonblocking: false,
        })
    }

    pub fn lookupnet(&self) -> error::Result<(u32, u32)> {
        let mut errbuf = [0i8; ffi::PCAP_ERRBUF_SIZE];
        let mut net: u32 = 0;
        let mut mask: u32 = 0;

        if unsafe {
            ffi::pcap_lookupnet(
                self.dev_name.as_ptr(),
                &mut net as *mut u32,
                &mut mask as *mut u32,
                errbuf.as_mut_ptr(),
            )
        } == -1
        {
            Err(error::Error::from(&errbuf))?;
        }

        Ok((net, mask))
    }

    pub fn set_non_blocking(&mut self, nonblocking: bool) -> error::Result<()> {
        self.nonblocking = nonblocking;

        let mut errbuf = [0i8; ffi::PCAP_ERRBUF_SIZE];

        if unsafe { ffi::pcap_setnonblock(self.dev, i32::from(nonblocking), errbuf.as_mut_ptr()) }
            == -1
        {
            Err(error::Error::from(&errbuf))?;
        }

        Ok(())
    }
}

impl<T: Disabled> Interface<T> {
    pub fn set_promisc(&mut self, promisc: bool) -> error::Result<()> {
        if unsafe { ffi::pcap_set_promisc(self.dev, i32::from(promisc)) } != 0 {
            Err(unsafe { ffi::pcap_geterr(self.dev) })?;
        }

        Ok(())
    }

    pub fn set_buffer_size(&mut self, bufsize: i32) -> error::Result<()> {
        if unsafe { ffi::pcap_set_buffer_size(self.dev, bufsize) } != 0 {
            Err(unsafe { ffi::pcap_geterr(self.dev) })?;
        }

        Ok(())
    }

    pub fn set_timeout(&mut self, timeout: i32) -> error::Result<()> {
        if unsafe { ffi::pcap_set_timeout(self.dev, timeout) } != 0 {
            Err(unsafe { ffi::pcap_geterr(self.dev) })?;
        }

        Ok(())
    }

    pub fn activate(mut self) -> error::Result<Interface<DevActivated>> {
        if unsafe { ffi::pcap_activate(self.dev) } != 0 {
            Err(unsafe { ffi::pcap_geterr(self.dev) })?;
        }

        self.absorbed = true;

        Ok(Interface::<DevActivated> {
            dev_name: self.dev_name.clone(),
            dev: self.dev,
            marker: std::marker::PhantomData,
            absorbed: false,
            nonblocking: self.nonblocking,
        })
    }
}

impl<T: Activated> Interface<T> {
    pub fn datalink(&self) -> i32 {
        unsafe { ffi::pcap_datalink(self.dev) }
    }

    pub fn set_filter(
        &mut self,
        filter: &str,
        optimize: bool,
        mask: Option<u32>,
    ) -> error::Result<Box<ffi::BpfProgram>> {
        let mut bpf = ffi::BpfProgram {
            bf_len: 0,
            bpf_insn: ptr::null(),
        };

        let mask = if let Some(m) = mask {
            m
        } else {
            let (_, m) = self.lookupnet()?;
            m
        };

        let filter = CString::new(filter)?;

        if unsafe {
            ffi::pcap_compile(
                self.dev,
                &mut bpf as *mut ffi::BpfProgram,
                filter.as_ptr(),
                i32::from(optimize),
                mask,
            )
        } == -1
        {
            Err(unsafe { ffi::pcap_geterr(self.dev) })?;
        }

        if unsafe { ffi::pcap_setfilter(self.dev, &bpf as *const ffi::BpfProgram) } == -1 {
            Err(unsafe { ffi::pcap_geterr(self.dev) })?;
        }

        Ok(Box::new(bpf))
    }

    pub fn sendpacket(&self, packet: packets::EthernetPkt) -> error::Result<()> {
        if unsafe {
            ffi::pcap_sendpacket(
                self.dev,
                packet.data.as_ptr(),
                packet.data.len().try_into().unwrap(),
            )
        } == ffi::PCAP_ERROR
        {
            Err(unsafe { ffi::pcap_geterr(self.dev) })?;
        }

        Ok(())
    }

    pub fn next_packet(&mut self) -> error::Result<packets::EthernetPacket> {
        let mut header: *mut ffi::PktHeader = ptr::null_mut();
        let mut data: *mut libc::c_char = ptr::null_mut();

        if unsafe { ffi::pcap_next_ex(self.dev, &mut header as *mut _, &mut data as *mut _) < 1 } {
            return unsafe { Err(ffi::pcap_geterr(self.dev))? };
        }

        let rdata = unsafe { slice::from_raw_parts(data as *mut u8, (*header).caplen as usize) };
        if rdata.len() < 24 {
            return Err(error::Error::PacketLengthInvalid);
        }

        Ok(packets::EthernetPkt { data: rdata }.to_owned())
    }
}

struct ListenHandler<'a, F>
where
    F: Fn(&Interface<DevListening>, packets::EthernetPkt) -> error::Result<bool>,
{
    packet_handler: F,
    break_on_fail: bool,
    fail_error: Option<error::Error>,
    interface: &'a Interface<DevListening>,
}

impl<T: NotListening> Interface<T> {
    pub fn listen<F>(
        &self,
        packet_handler: F,
        break_on_fail: bool,
        packet_count: i32,
    ) -> (Option<error::Error>, i32)
    where
        F: Fn(&Interface<DevListening>, packets::EthernetPkt) -> error::Result<bool>,
    {
        unsafe extern "C" fn cback<F>(
            user: *mut libc::c_void,
            header: *const ffi::PktHeader,
            data: *const u8,
        ) where
            F: Fn(&Interface<DevListening>, packets::EthernetPkt) -> error::Result<bool>,
        {
            let info = &mut *(user as *mut ListenHandler<F>);

            let data = slice::from_raw_parts(data, (*header).caplen as usize);

            if data.len() < 14 {
                eprintln!(
                    " * Failed to get full packet, captured {} bytes",
                    data.len()
                );
            }

            let result = (info.packet_handler)(info.interface, packets::EthernetPkt { data });

            match result {
                Err(e) => {
                    eprintln!(" * Packet handle error: {:?}", e);
                    if info.break_on_fail {
                        info.fail_error = Some(e);
                        ffi::pcap_breakloop(info.interface.dev);
                    }
                }
                Ok(b) => {
                    if b {
                        ffi::pcap_breakloop(info.interface.dev);
                    }
                }
            }
        }

        let interface = Interface {
            dev_name: self.dev_name.clone(),
            dev: self.dev,
            marker: std::marker::PhantomData,
            absorbed: true,
            nonblocking: self.nonblocking,
        };

        let mut info = ListenHandler::<F> {
            packet_handler,
            break_on_fail,
            fail_error: None,
            interface: &interface,
        };

        let count = unsafe {
            ffi::pcap_loop(
                self.dev,
                packet_count,
                cback::<F>,
                &mut info as *mut ListenHandler<F> as *mut libc::c_void,
            )
        };

        (info.fail_error, count)
    }

    pub fn stream(mut self) -> error::Result<InterfaceStream<DevActivated>> {
        self.set_non_blocking(true)?;

        Ok(InterfaceStream {
            inner: AsyncFd::with_interest(
                InternalInterfaceStream::<DevActivated>::new(unsafe { std::mem::transmute(self) })?,
                tokio::io::Interest::READABLE,
            )?,
        })
    }
}

struct InternalInterfaceStream<T: Activated> {
    interface: Interface<T>,
    fd: RawFd,
}

impl<T: Activated> InternalInterfaceStream<T> {
    fn new(interface: Interface<T>) -> error::Result<InternalInterfaceStream<T>> {
        let fd = unsafe { ffi::pcap_get_selectable_fd(interface.dev) };
        if fd == -1 {
            return Err(error::Error::InvalidPcapFd);
        }
        Ok(Self { interface, fd })
    }
}

impl<T: Activated> AsRawFd for InternalInterfaceStream<T> {
    fn as_raw_fd(&self) -> RawFd {
        self.fd
    }
}

pub struct InterfaceStream<T: Activated> {
    inner: AsyncFd<InternalInterfaceStream<T>>,
}

impl<T: Activated> InterfaceStream<T> {
    pub fn sendpacket(&mut self, packet: packets::EthernetPkt) -> error::Result<()> {
        self.inner.get_mut().interface.sendpacket(packet)
    }
}

impl<T: Activated> Unpin for InterfaceStream<T> {}

impl<T: Activated> futures::Stream for InterfaceStream<T> {
    type Item = error::Result<packets::EthernetPacket>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
        let stream = Pin::into_inner(self);

        loop {
            let mut guard = ready!(stream.inner.poll_read_ready_mut(cx))?;

            match guard.try_io(|inner| match inner.get_mut().interface.next_packet() {
                Ok(p) => Ok(Ok(p)),
                Err(e) => Ok(Err(e)),
            }) {
                Ok(result) => {
                    return Poll::Ready(Some(result?));
                }
                Err(_would_block) => continue,
            }
        }
    }
}

pub fn new_aggregate_interface_filtered<F>(crash: bool, mut f: F) -> error::Result<AggregateInterface<DevDisabled>>
where
    F: FnMut(&str) -> bool
{
    let interfaces = if crash {
        PcapDevIterator::new()?
            .filter(|s| (f)(s))
            .map(|if_name| {
                let new_name = if_name.clone();
                Interface::<DevDisabled>::new(&if_name)
                    .map(|interface| (if_name, interface))
                    .map_err(|e| e.add_ifname(&new_name))
            })
            .collect::<error::Result<HashMap<_, _>>>()?
    } else {
        PcapDevIterator::new()?
            .filter(|s| (f)(s))
            .filter_map(|if_name| {
                let new_name = if_name.clone();
                Interface::<DevDisabled>::new(&if_name)
                    .map(|interface| (if_name, interface))
                    .ok()
            })
            .collect::<HashMap<_, _>>()
    };

    Ok(AggregateInterface {
        interfaces,
        crash
    })
}

pub fn new_aggregate_interface(crash: bool) -> error::Result<AggregateInterface<DevDisabled>> {
    new_aggregate_interface_filtered(crash, |_| true)
}

pub struct AggregateInterface<T: State> {
    interfaces: HashMap<String, Interface<T>>,
    crash: bool
}

impl<T: State> AggregateInterface<T> {
    pub fn set_non_blocking(&mut self, nonblocking: bool) -> error::Result<()> {
        for (n, i) in self.interfaces.iter_mut() {
            i
                .set_non_blocking(nonblocking)
                .map_err(|e| e.add_ifname(n))?;
        }

        Ok(())
    }

    pub fn lookupnets(&self) -> error::Result<HashMap<&str, (u32, u32)>> {
        self.interfaces
            .iter()
            .map(|(name, interface)| {
                interface
                    .lookupnet()
                    .map(|net| (&**name, net))
                    .map_err(|e| e.add_ifname(&name))
            })
            .collect::<error::Result<_>>()
    }
}

impl<T: Disabled> AggregateInterface<T> {
    pub fn set_promisc(&mut self, promisc: bool) -> error::Result<()> {
        for (n, i) in self.interfaces.iter_mut() {
            i.set_promisc(promisc).map_err(|e| e.add_ifname(n))?;
        }

        Ok(())
    }

    pub fn set_buffer_size(&mut self, bufsize: i32) -> error::Result<()> {
        for (n, i) in self.interfaces.iter_mut() {
            i.set_buffer_size(bufsize).map_err(|e| e.add_ifname(n))?;
        }

        Ok(())
    }

    pub fn set_timeout(&mut self, timeout: i32) -> error::Result<()> {
        for (n, i) in self.interfaces.iter_mut() {
            i.set_timeout(timeout).map_err(|e| e.add_ifname(n))?;
        }

        Ok(())
    }

    pub fn activate(self) -> error::Result<AggregateInterface<DevActivated>> {
        Ok(AggregateInterface {
            interfaces: self.interfaces
                .into_iter()
                .map(|(name, interface)| {
                    let new_name = name.clone();
                    interface
                        .activate()
                        .map(|interface| (name, interface))
                        .map_err(|e| e.add_ifname(&new_name))
                })
                .collect::<error::Result<_>>()?
        })
    }
}

impl<T: Activated> AggregateInterface<T> {
    pub fn datalinks(&self) -> HashMap<&str, i32> {
        self.interfaces
            .iter()
            .map(|(name, interface)| {
                (&**name, interface.datalink())
            })
            .collect::<_>()
    }

    pub fn prune<F>(&mut self, mut f: F)
    where
        F: FnMut(&str, &mut Interface<T>) -> bool
    {
        let to_prune = self.interfaces
            .iter_mut()
            .filter_map(|(k,v)| {
                if (f)(k, v) {
                    Some(k.clone())
                } else {
                    None
                }
            })
            .collect::<Vec<_>>();

        for name in to_prune {
            self.interfaces.remove(&name);
        }
    }

    pub fn set_filter(
        &mut self,
        filter: &str,
        optimize: bool,
        mask: Option<u32>
    ) -> error::Result<HashMap<&str, Box<ffi::BpfProgram>>> {
        self.interfaces
            .iter_mut()
            .map(|(name, interface)| {
                interface.set_filter(filter, optimize, mask)
                    .map(|bpf| (&**name, bpf))
                    .map_err(|e| e.add_ifname(&name))
            })
            .collect::<error::Result<_>>()
    }

    pub fn sendpacket(&self, ifname: &str, packet: packets::EthernetPkt) -> error::Result<()> {
        if let Some(interface) = self.interfaces.get(ifname) {
            interface.sendpacket(packet).map_err(|e| e.add_ifname(ifname))?;
        }

        Ok(())
    }
}

impl<T: NotListening> AggregateInterface<T> {
    pub fn stream(self) -> error::Result<AggregateInterfaceStream<DevActivated>> {
        Ok(AggregateInterfaceStream {
            streams: self.interfaces
                .into_iter()
                .map(|(ifname, interface)| {
                    let new_name = ifname.clone();
                    interface
                        .stream()
                        .map(|stream| (ifname, stream))
                        .map_err(|e| e.add_ifname(&new_name))
                })
                .collect::<error::Result<_>>()?
        })
    }
}

pub struct AggregateInterfaceStream<T: Activated> {
    streams: StreamMap<String, InterfaceStream<T>>
}

impl<T: Activated> AggregateInterfaceStream<T> {
    pub fn sendpacket(&mut self, ifname: &str, packet: packets::EthernetPkt) -> error::Result<()> {
        if let Some(interface) = self.streams
            .values_mut()
            .find(|interface| {
                let dev_name = interface.inner
                        .get_ref()
                        .interface
                        .dev_name
                        .clone();
                
                CString::new(ifname).map(|ifname|
                    ifname == dev_name)
                    .unwrap_or(false)
            }) {
            interface.sendpacket(packet)?;
        }

        Ok(())
    }
}

impl<T: Activated> futures::Stream for AggregateInterfaceStream<T> {
    type Item = (String, error::Result<packets::EthernetPacket>);

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> {
        self.streams.poll_next_unpin(cx)
    }
}