libinject/
connections.rs

1use crate::{
2    network::{self, Packet},
3    socket::{Socket, SocketData},
4};
5use itertools::Itertools;
6use std::{
7    collections::HashMap,
8    net::SocketAddr,
9    sync::{Arc, LazyLock, RwLock},
10};
11
12/// A spoofed network connection with the most recent request (sent by the binary over "send").
13#[derive(Debug)]
14pub struct Connection {
15    pub addr: SocketAddr,
16    pub socket: Socket,
17}
18
19impl ToString for Connection {
20    fn to_string(&self) -> String {
21        self.addr.to_string()
22    }
23}
24
25/// HashMap of currently open network connections.
26static CONNECTION_MAP: LazyLock<RwLock<HashMap<usize, Arc<Connection>>>> =
27    LazyLock::new(|| RwLock::new(HashMap::new()));
28
29pub fn insert(socket_id: usize, socket_addr: SocketAddr, socket_data: SocketData) {
30    CONNECTION_MAP
31        .write()
32        .expect("RwLock should not be poisoned")
33        .insert(
34            socket_id,
35            Arc::new(Connection {
36                addr: socket_addr,
37                socket: Socket {
38                    id: socket_id,
39                    data: socket_data,
40                },
41            }),
42        );
43}
44
45/// Record a reqest on this connection with the network module.
46pub fn record_request(socket_id: usize, request: Vec<u8>) {
47    let connection =
48        self::get(socket_id).expect("tried recording a request on an unconnected socket");
49    network::push(Packet::outbound(connection.addr, request));
50}
51
52/// Record a response on this connection with the network module.
53pub fn record_response(socket_id: usize, response: Vec<u8>) {
54    let connection =
55        self::get(socket_id).expect("tried recording a response on an unconnected socket");
56    network::push(Packet::inbound(connection.addr, response));
57}
58
59pub fn get(socket_id: usize) -> Option<Arc<Connection>> {
60    CONNECTION_MAP
61        .read()
62        .expect("RwLock should not be poisoned")
63        .get(&socket_id)
64        .cloned()
65}
66
67pub fn exactly_one() -> Result<Arc<Connection>, String> {
68    CONNECTION_MAP
69        .read()
70        .expect("RwLock should not be poisoned")
71        .iter()
72        .exactly_one()
73        .map_err(|e| format!("Not exactly one connection: {e}"))
74        .map(|(_socket_id, conn)| conn.clone())
75}
libinject/connections.rs

libinject/
connections.rs
