ya_relay_stack/
interface.rs

1use managed::ManagedSlice;
2use smoltcp::iface::{Config, Interface, Route, SocketHandle, SocketSet};
3use smoltcp::socket::AnySocket;
4use smoltcp::time::Instant;
5use smoltcp::wire::{EthernetAddress, HardwareAddress, IpAddress, IpCidr};
6use std::io::Write;
7
8use crate::device::CaptureDevice;
9
10pub struct CaptureInterface<'a> {
11    iface: Interface,
12    device: CaptureDevice,
13    sockets: SocketSet<'a>,
14}
15
16impl<'a> CaptureInterface<'a> {
17    pub fn new(iface: Interface, device: CaptureDevice, sockets: SocketSet<'a>) -> Self {
18        Self {
19            iface,
20            device,
21            sockets,
22        }
23    }
24
25    pub fn inner(&self) -> &Interface {
26        &self.iface
27    }
28
29    pub fn inner_mut(&mut self) -> &mut Interface {
30        &mut self.iface
31    }
32
33    pub fn device(&self) -> &CaptureDevice {
34        &self.device
35    }
36
37    pub fn device_mut(&mut self) -> &mut CaptureDevice {
38        &mut self.device
39    }
40
41    pub fn sockets(&self) -> impl Iterator<Item = (SocketHandle, &smoltcp::socket::Socket<'a>)> {
42        self.sockets.iter()
43    }
44
45    pub fn sockets_mut(
46        &mut self,
47    ) -> impl Iterator<Item = (SocketHandle, &mut smoltcp::socket::Socket<'a>)> {
48        self.sockets.iter_mut()
49    }
50
51    pub fn add_socket<T: AnySocket<'a>>(&mut self, socket: T) -> SocketHandle {
52        self.sockets.add(socket)
53    }
54
55    pub fn get_socket_and_context<T: AnySocket<'a>>(
56        &mut self,
57        handle: SocketHandle,
58    ) -> (&mut T, &mut smoltcp::iface::Context) {
59        let socket = self.sockets.get_mut(handle);
60        let ctx = self.iface.context();
61        (socket, ctx)
62    }
63
64    pub fn remove_socket(&mut self, handle: SocketHandle) -> smoltcp::socket::Socket {
65        self.sockets.remove(handle)
66    }
67
68    pub fn poll(&mut self, timestamp: Instant) -> bool {
69        let sockets = &mut self.sockets;
70        let device = &mut self.device;
71        self.iface.poll(timestamp, device, sockets)
72    }
73}
74
75fn iface_and_sockets<'a>(
76    addr: HardwareAddress,
77    device: &mut CaptureDevice,
78) -> (Interface, SocketSet<'a>) {
79    let config = Config::new(addr);
80    let now = Instant::ZERO;
81    let iface = Interface::new(config, device, now);
82    let sockets = SocketSet::new(ManagedSlice::Owned(vec![]));
83    (iface, sockets)
84}
85
86/// Creates a default TAP (Ethernet) network interface
87pub fn tap_iface<'a>(mac: HardwareAddress, mtu: usize) -> CaptureInterface<'a> {
88    let mut device = CaptureDevice::tap(mtu);
89    let (iface, sockets) = iface_and_sockets(mac, &mut device);
90    CaptureInterface::new(iface, device, sockets)
91}
92
93/// Creates a default TUN (IP) network interface
94pub fn tun_iface<'a>(mtu: usize) -> CaptureInterface<'a> {
95    let mut device = CaptureDevice::tun(mtu);
96    let (iface, sockets) = iface_and_sockets(HardwareAddress::Ip, &mut device);
97    CaptureInterface::new(iface, device, sockets)
98}
99
100/// Creates a pcap TAP (Ethernet) network interface
101pub fn pcap_tap_iface<'a, W>(mac: HardwareAddress, mtu: usize, pcap: W) -> CaptureInterface<'a>
102where
103    W: Write + 'static,
104{
105    let mut device = CaptureDevice::pcap_tap(mtu, pcap);
106    let (mut iface, sockets) = iface_and_sockets(mac, &mut device);
107    iface.set_hardware_addr(mac);
108    CaptureInterface::new(iface, device, sockets)
109}
110
111/// Creates a pcap TUN (IP) network interface
112pub fn pcap_tun_iface<'a, W>(mtu: usize, pcap: W) -> CaptureInterface<'a>
113where
114    W: Write + 'static,
115{
116    let mut device = CaptureDevice::pcap_tun(mtu, pcap);
117    let (iface, sockets) = iface_and_sockets(HardwareAddress::Ip, &mut device);
118    CaptureInterface::new(iface, device, sockets)
119}
120
121/// Assigns a new interface IP address
122pub fn add_iface_address(iface: &mut CaptureInterface, node_ip: IpCidr) {
123    iface.inner_mut().update_ip_addrs(|addrs| {
124        if !addrs.iter().any(|ip| *ip == node_ip) {
125            if let Err(err) = addrs.push(node_ip) {
126                log::error!("Failed to assign new interface IP address: {err}");
127            }
128        }
129    });
130}
131
132/// Adds a new IP route
133pub fn add_iface_route(iface: &mut CaptureInterface, net_ip: IpCidr, route: Route) {
134    iface.inner_mut().routes_mut().update(|routes| {
135        for (i, r) in routes.iter().enumerate() {
136            if r.cidr == net_ip {
137                if let Err(route) = routes.insert(i, route) {
138                    log::error!("Failed to replace route at index {i}. Route: {route:?}");
139                }
140                return;
141            }
142        }
143        if let Err(route) = routes.push(route) {
144            log::error!("Failed to add new route: {route:?}");
145        }
146    });
147}
148
149pub fn to_mac(mac: &[u8]) -> HardwareAddress {
150    let mut ethernet = if mac.len() >= 6 {
151        EthernetAddress::from_bytes(&mac[..6])
152    } else {
153        EthernetAddress::from_bytes(&rand::random::<[u8; 6]>())
154    };
155
156    if !ethernet.is_unicast() {
157        ethernet.0[0] &= !0x01;
158    }
159
160    HardwareAddress::Ethernet(ethernet)
161}
162
163pub fn ip_to_mac(ip: IpAddress) -> HardwareAddress {
164    match ip {
165        IpAddress::Ipv4(ip) => to_mac(ip.as_bytes()),
166        IpAddress::Ipv6(ip) => to_mac(ip.as_bytes()),
167    }
168}