add bfp filter to pushack sockets

Author: Arian8j2

forwarder/src/poll/pushack.rs

@@ -8,7 +8,10 @@ use crate::{
 };
 use parking_lot::RwLock;
 use smoltcp::wire::{IPV4_HEADER_LEN, TCP_HEADER_LEN};
-use std::{net::SocketAddr, sync::Arc};
+use std::{
+    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
+    sync::Arc,
+};
 
 #[derive(Debug)]
 pub struct PushackPoll {
@@ -25,11 +28,21 @@ impl Poll for PushackPoll {
         peers: Arc<RwLock<PeerManager>>,
         on_peer_recv: Box<dyn Fn(&Peer, &mut [u8])>,
     ) -> anyhow::Result<()> {
-        let socket = PushackSocket::bind(&"0.0.0.0:0".parse().unwrap())?;
+        let addr = SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0);
+        let socket = PushackSocket::inner_bind(addr.into())?;
         let buffer = create_socket_buffer!(MAX_PACKET_SIZE);
 
+        #[cfg(target_os = "linux")]
+        {
+            use crate::socket::pushack::{create_bfp_filter, TcpBpfFilter};
+            let filter = create_bfp_filter(TcpBpfFilter::SrcPort(self.remote_addr.port()));
+            if let Err(error) = socket.attach_filter(&filter) {
+                log::warn!("couldn't attach bpf filter: {error:?}");
+            }
+        }
+
         loop {
-            let Ok(size) = socket.inner_socket().recv(cast_maybe_uninit(buffer)) else {
+            let Ok(size) = socket.recv(cast_maybe_uninit(buffer)) else {
                 continue;
             };
             let Some(tcp_packet) = parse_pushack_packet(&buffer[..size]) else {

forwarder/src/socket/pushack.rs

@@ -23,8 +23,15 @@ impl PushackSocket {
     pub fn bind(addr: &SocketAddr) -> io::Result<Self> {
         let _udp_socket = UdpSocket::bind(addr)?;
         let addr = _udp_socket.local_addr()?;
-
         let socket = Self::inner_bind(addr)?;
+
+        #[cfg(target_os = "linux")]
+        if let Err(error) =
+            socket.attach_filter(&create_bfp_filter(TcpBpfFilter::DstPort(addr.port())))
+        {
+            log::warn!("couldn't attach bpf filter: {error:?}");
+        }
+
         Ok(Self {
             _udp_socket,
             addr,
@@ -33,11 +40,7 @@ impl PushackSocket {
     }
 
     pub fn inner_bind(addr: SocketAddr) -> io::Result<socket2::Socket> {
-        let socket = if addr.is_ipv4() {
-            socket2::Socket::new(Domain::IPV4, Type::RAW, Some(Protocol::TCP))
-        } else {
-            socket2::Socket::new(Domain::IPV6, Type::RAW, Some(Protocol::TCP))
-        }?;
+        let socket = socket2::Socket::new(Domain::IPV4, Type::RAW, Some(Protocol::TCP))?;
         socket.bind(&addr.into())?;
         Ok(socket)
     }
@@ -166,3 +169,26 @@ pub fn parse_pushack_packet(packet: &[u8]) -> Option<TcpPacketPort> {
     };
     Some(ports)
 }
+
+#[cfg(target_os = "linux")]
+#[derive(Debug)]
+pub enum TcpBpfFilter {
+    SrcPort(u16),
+    DstPort(u16),
+}
+
+#[cfg(target_os = "linux")]
+pub fn create_bfp_filter(filter: TcpBpfFilter) -> [libc::sock_filter; 4] {
+    let (filter_offset, value) = match filter {
+        TcpBpfFilter::SrcPort(port) => (0, port),
+        TcpBpfFilter::DstPort(port) => (2, port),
+    };
+    let total_offset = IPV4_HEADER_LEN + filter_offset;
+    [
+        (0x28, 0, 0, total_offset as u32), // ldh [offset]
+        (0x15, 0, 1, value as u32),        // jne val, drop
+        (0x06, 0, 0, 0xffffffff),          // ret #-1
+        (0x06, 0, 0, 0000000000),          // drop: ret #0
+    ]
+    .map(|(code, jt, jf, k)| libc::sock_filter { code, jt, jf, k })
+}