Skip to main content

libinject/
network.rs

1use crate::drcore::log;
2use spin::Mutex;
3use std::{
4    env,
5    fmt::Display,
6    fs::{File, metadata},
7    io::Write,
8    net::SocketAddr,
9    sync::LazyLock,
10};
11
12static PCAP: LazyLock<Mutex<Vec<Packet>>> = LazyLock::new(|| Mutex::new(Vec::with_capacity(100)));
13
14#[derive(Clone, Debug)]
15pub struct PacketMeta {
16    addr: SocketAddr,
17    payload: Vec<u8>,
18}
19
20#[derive(Clone, Debug)]
21pub enum Packet {
22    Inbound(PacketMeta),
23    Outbound(PacketMeta),
24}
25
26impl Packet {
27    pub fn to_csv(&self) -> String {
28        match self {
29            Packet::Inbound(PacketMeta { addr, payload }) => format!("{addr},self,{:?}\n", payload),
30            Packet::Outbound(PacketMeta { addr, payload }) => {
31                format!("self,{addr},{:?}\n", payload)
32            }
33        }
34    }
35}
36
37impl Display for Packet {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        match self {
40            Packet::Inbound(PacketMeta { addr, payload }) => {
41                write!(f, "{addr} -> self {:?}\n", payload)?
42            }
43            Packet::Outbound(PacketMeta { addr, payload }) => {
44                write!(f, "self -> {addr} {:?}\n", payload)?
45            }
46        }
47        Ok(())
48    }
49}
50
51impl Packet {
52    pub fn inbound(src: SocketAddr, payload: Vec<u8>) -> Packet {
53        Packet::Inbound(PacketMeta { addr: src, payload })
54    }
55    pub fn outbound(dst: SocketAddr, payload: Vec<u8>) -> Packet {
56        Packet::Outbound(PacketMeta { addr: dst, payload })
57    }
58}
59
60pub fn push(pkt: Packet) {
61    PCAP.lock().push(pkt);
62}
63
64#[unsafe(no_mangle)]
65pub extern "C" fn dump_pcap() -> i32 {
66    match dump(DumpFormat::CSV) {
67        Ok(()) => 0,
68        Err(err_msg) => {
69            log(&format!("Error dumping pcap as csv: {}\n", &err_msg));
70            1
71        }
72    }
73}
74
75pub enum DumpFormat {
76    CSV,
77    HumanReadable,
78}
79
80pub fn dump(dump_fmt: DumpFormat) -> Result<(), String> {
81    let out_file_path = match env::var("PCAP") {
82        Ok(file_path) => file_path,
83        Err(_) => return Ok(()),
84    };
85
86    let pcap = PCAP.lock();
87
88    let file_exists = metadata(&out_file_path).is_ok();
89
90    let mut out_file = File::options()
91        .create(true)
92        .append(true)
93        .open(&out_file_path)
94        .map_err(|e| format!("Failed to open file at path {out_file_path}: {e}\n"))?;
95
96    if !file_exists {
97        if let DumpFormat::CSV = dump_fmt {
98            out_file
99                .write("src,dst,payload\n".as_bytes())
100                .map_err(|e| format!("Failed writing csv header to file {out_file_path}: {e}\n"))?;
101        }
102    }
103
104    for pkt in pcap.iter() {
105        let serialise = match dump_fmt {
106            DumpFormat::HumanReadable => Packet::to_string,
107            DumpFormat::CSV => Packet::to_csv,
108        };
109        out_file
110            .write_all(serialise(pkt).as_bytes())
111            .map_err(|e| format!("Failed writing captured network packets to file: {e}\n"))?;
112    }
113    Ok(())
114}