「ルーター自作でわかるパケットの流れ」をRustで取り組む

use std::fs;
use std::io::Error;

pub fn write_ip_forward(turn_on: bool) -> std::io::Result<()> {
    fs::write("/proc/sys/net/ipv4/ip_forward", if turn_on { "1" } else { "0" })
}


pub fn read_ip_forward() -> std::io::Result<bool> {
    let buf = fs::read_to_string("/proc/sys/net/ipv4/ip_forward")?;
    match buf.trim_end() {
        "0" => Ok(false),
        "1" => Ok(true),
        _ => Err(Error::other("invalid value"))
    }
}


#[cfg(test)]
mod tests {
    use crate::net::ip_forward::{read_ip_forward, write_ip_forward};

    #[test]
    fn ip_forward() {
        let prev = read_ip_forward().unwrap();
        write_ip_forward(true).unwrap();
        assert!(read_ip_forward().unwrap());
        write_ip_forward(false).unwrap();
        assert!(!read_ip_forward().unwrap());
        write_ip_forward(prev).unwrap();
    }
}

use pnet::datalink::{Channel, Config, DataLinkReceiver, DataLinkSender};
use crate::eth::{ReadEthPacket, WriteEthPacket};
use crate::nic::Nic;

pub fn connect(nic: &Nic) -> (Box<dyn WriteEthPacket>, Box<dyn ReadEthPacket>) {
    let target_nic = pnet::datalink::interfaces()
        .into_iter()
        .find(|ni| ni.name == nic.name)
        .unwrap();
    let Channel::Ethernet(tx, rx) = pnet::datalink::channel(&target_nic, Config::default()).unwrap()
        else {
            panic!("unreachable")
        };
    (Box::new(PNetSocketWriter(tx)), Box::new(PNetSocketReader(rx)))
}


struct PNetSocketReader(Box<dyn DataLinkReceiver>);

impl ReadEthPacket for PNetSocketReader {
    fn next(&mut self) -> crate::error::Result<&[u8]> {
        Ok(self.0.next()?)
    }
}


struct PNetSocketWriter(Box<dyn DataLinkSender>);

impl WriteEthPacket for PNetSocketWriter {
    fn send(&mut self, packet: &[u8]) -> crate::error::Result {
        if let Some(res) = self.0.send_to(packet, None) {
            res?;
        }
        Ok(())
    }
}

use std::ffi::c_ushort;

use linux_api::net::htons;
use linux_api::net::socket::{sockaddr_ll, SocketFd};
use linux_api::net::socket::family::ProtocolFamily;
use linux_api::net::socket::protocol::Protocol;
use linux_api::net::socket::socket_type::SocketType;

use crate::eth::{ReadEthPacket, WriteEthPacket};
use crate::nic::Nic;

pub fn connect(nic: &Nic) -> crate::error::Result<(Box<dyn WriteEthPacket>, Box<dyn ReadEthPacket>)> {
    let fd = SocketFd::new(
        ProtocolFamily::AF_PACKET,
        SocketType::SOCK_RAW,
        Protocol::ETH_P_ALL,
    )?;

    let socket_addr_ll = socket_addr(nic, Protocol::ETH_P_ALL);
    fd.bind(&socket_addr_ll)?;
    Ok((Box::new(SocketWriter(fd)), Box::new(SocketReader {
        fd,
        buff: vec![0; 4096],
    })))
}

struct SocketReader {
    fd: SocketFd,
    buff: Vec<u8>,
}

impl ReadEthPacket for SocketReader {
    fn next(&mut self) -> crate::error::Result<&[u8]> {
        let buff_len = self.buff.len();
        let len = self.fd.read(self.buff.as_mut_slice(), buff_len)?;
        Ok(&self.buff[0..len as usize])
    }
}

pub struct SocketWriter(SocketFd);

impl WriteEthPacket for SocketWriter {
    fn send(&mut self, packet: &[u8]) -> crate::error::Result {
        self.0.write(packet)?;
        Ok(())
    }
}

fn socket_addr(nic: &Nic, protocol: Protocol) -> sockaddr_ll {
    sockaddr_ll {
        sll_family: ProtocolFamily::AF_PACKET.0 as c_ushort,
        sll_protocol: htons(protocol.0 as u16) as c_ushort,
        sll_ifindex: nic.index as i32,
        sll_hatype: 0,
        sll_pkttype: 0,
        sll_halen: 6,
        sll_addr: nic
            .mac
            .map(|m| {
                let m = &m.0;
                [m[0], m[1], m[2], m[3], m[4], m[5], 0, 0]
            })
            .unwrap_or([0, 0, 0, 0, 0, 0, 0, 0]),
    }
}

#[cfg(test)]
mod tests {
    use crate::eth::socket::connect;
    use crate::nic::pnet::PNetNic;
    use crate::nic::ReadAllNic;

    #[test]
    fn it() {
        let nic = PNetNic.read_all_nic();
        let (_, mut rx) = connect(&nic[1]).unwrap();
        let buff = rx.next().unwrap();
        println!("buff = {buff:?}");
    }
}

pub mod family;
pub mod protocol;
pub mod socket_type;

use crate::net::socket::family::ProtocolFamily;
use crate::net::socket::protocol::Protocol;
use crate::net::socket::socket_type::SocketType;
use std::ffi::c_void;

use libc::{c_int, sockaddr, socklen_t};
use std::io;

pub use libc::sockaddr_ll;

#[derive(Debug, Eq, PartialEq, Copy, Clone)]
pub struct SocketFd(c_int);

impl SocketFd {
    pub fn new(
        family: ProtocolFamily,
        socket_type: SocketType,
        protocol: Protocol,
    ) -> io::Result<Self> {
        let socket_fd = unsafe { libc::socket(family.0, socket_type.0, protocol.0) };
        if socket_fd != -1 {
            Ok(SocketFd(socket_fd))
        } else {
            Err(io::Error::last_os_error())
        }
    }

    pub fn bind(&self, socket_addr_ll: &sockaddr_ll) -> io::Result<()> {
        let socket_addr = (socket_addr_ll as *const _) as *const sockaddr;
        let result = unsafe {
            libc::bind(
                self.0,
                socket_addr,
                std::mem::size_of::<libc::sockaddr_ll>() as socklen_t,
            )
        };
        self.result(result)
    }

    pub fn listen(&self, backlog: c_int) -> io::Result<()> {
        self.result(unsafe { libc::listen(self.0, backlog) })
    }

    pub fn close(&self) -> io::Result<()> {
        unsafe {
            if libc::close(self.0) == -1 {
                Err(io::Error::last_os_error())
            } else {
                Ok(())
            }
        }
    }

    pub fn write(&self, buf: &[u8]) -> io::Result<()> {
        unsafe {
            let len = buf.len();
            let write_bytes = libc::write(self.0, buf as *const _ as *const c_void, len);
            if write_bytes <= 0 {
                Err(io::Error::last_os_error())
            } else {
                Ok(())
            }
        }
    }

    pub fn read(&self, buf: &mut [u8], size: usize) -> io::Result<isize> {
        let read_bytes = unsafe { libc::read(self.0, buf as *mut _ as *mut c_void, size) };
        if 0 < read_bytes {
            Ok(read_bytes)
        } else {
            Err(std::io::Error::last_os_error())
        }
    }

    fn result(&self, result: c_int) -> io::Result<()> {
        if result == 0 {
            Ok(())
        } else {
            self.close()?;
            Err(std::io::Error::last_os_error())
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::net::socket::family::ProtocolFamily;
    use crate::net::socket::protocol::Protocol;
    use crate::net::socket::socket_type::SocketType;
    use crate::net::socket::SocketFd;

    #[test]
    fn it_new_socket() {
        let socket = SocketFd::new(
            ProtocolFamily::AF_PACKET,
            SocketType::SOCK_RAW,
            Protocol::ETH_P_ALL,
        )
            .unwrap();
        println!("{socket:?}");
        socket.close().unwrap();
    }
}