Skip to content

Commit 01b9f1c

Browse files
WaujitoWaujito
committed
Reorganize code
1 parent f04bedc commit 01b9f1c

3 files changed

Lines changed: 83 additions & 304 deletions

File tree

‎src/main.rs‎

Lines changed: 6 additions & 191 deletions
Original file line numberDiff line numberDiff line change
@@ -1,42 +1,21 @@
11
#![allow(clippy::let_unit_value)]
22

3-
use std::cell::Cell;
4-
use std::cell::RefCell;
5-
use std::mem;
63
use std::mem::MaybeUninit;
7-
use std::os::fd::BorrowedFd;
84
use std::os::unix::io::AsFd as _;
95

10-
use anyhow::anyhow;
11-
use anyhow::Context;
12-
use anyhow::Result;
6+
use anyhow::{Context, Result};
7+
8+
use ebpf_dpit::{
9+
ebpf_prog, sni_logging_handle, DpitSkelLib, init_skel,
10+
TcController, XdpController
11+
};
1312

1413
use ebpf_prog::types::sni_action;
15-
use ebpf_prog::DpitSkel;
16-
use libbpf_rs::MapCore;
17-
use libbpf_rs::MapFlags;
18-
use libbpf_rs::OpenObject;
19-
use libbpf_rs::XdpFlags;
2014
use sni_logging_handle::init_sni_logging;
2115
use tokio::signal;
2216

2317
use clap::Parser;
2418

25-
use libbpf_rs::skel::OpenSkel;
26-
use libbpf_rs::skel::SkelBuilder;
27-
use libbpf_rs::TcHook;
28-
use libbpf_rs::Xdp;
29-
use libbpf_rs::TC_EGRESS;
30-
31-
use nix::net::if_::if_nametoindex;
32-
33-
mod ebpf_prog {
34-
include!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/dpit.skel.rs"));
35-
}
36-
37-
mod sni_logging_handle;
38-
mod utils;
39-
4019
#[derive(Debug, Parser)]
4120
struct Command {
4221
/// interface to attach to
@@ -54,170 +33,6 @@ struct Command {
5433
block_domains: String,
5534
}
5635

57-
fn init_skel<'obj>(
58-
open_object: &'obj mut MaybeUninit<OpenObject>,
59-
) -> Result<ebpf_prog::DpitSkel<'obj>> {
60-
let builder = ebpf_prog::DpitSkelBuilder::default();
61-
let open = builder.open(open_object)?;
62-
let skel = open.load()?;
63-
64-
Ok(skel)
65-
}
66-
67-
struct TcController<'obj> {
68-
pub _fd: BorrowedFd<'obj>,
69-
pub tc_eggress: RefCell<TcHook>,
70-
created: Cell<bool>,
71-
}
72-
73-
impl<'obj> TcController<'obj> {
74-
fn new(fd: BorrowedFd<'obj>, iface: &str) -> Result<Self> {
75-
let ifidx = if_nametoindex(iface)? as i32;
76-
77-
let mut tc_eggress = TcHook::new(fd);
78-
tc_eggress
79-
.ifindex(ifidx)
80-
.replace(true)
81-
.handle(1)
82-
.priority(1)
83-
.attach_point(TC_EGRESS);
84-
85-
Ok(Self {
86-
_fd: fd,
87-
tc_eggress: RefCell::new(tc_eggress),
88-
created: Cell::new(false),
89-
})
90-
}
91-
92-
fn create(&self) -> Result<()> {
93-
if let Err(err) = self.tc_eggress.borrow_mut().create() {
94-
eprintln!("Failed to create eggress hook: {err}");
95-
eprintln!("Attempting to destroy the eggress hook");
96-
97-
self.created.set(true);
98-
let _ = self.tc_eggress.borrow_mut().destroy();
99-
100-
return Err(anyhow!(err));
101-
}
102-
103-
self.created.set(true);
104-
105-
Ok(())
106-
}
107-
108-
pub fn attach(&self) -> Result<()> {
109-
if !self.created.get() {
110-
self.create()?;
111-
}
112-
113-
if let Err(err) = self.tc_eggress.borrow_mut().attach() {
114-
eprintln!("Failed to attach egress hook {err}");
115-
return Err(anyhow!(err));
116-
}
117-
118-
Ok(())
119-
}
120-
121-
pub fn destroy(&self) {
122-
if !self.created.get() {
123-
return;
124-
}
125-
126-
if let Err(e) = self.tc_eggress.borrow_mut().detach() {
127-
println!("Failed to detach egress hook {e}");
128-
}
129-
130-
if let Err(e) = self.tc_eggress.borrow_mut().destroy() {
131-
println!("Failed to destroy eggress hook {e}");
132-
}
133-
134-
self.created.set(false);
135-
}
136-
}
137-
138-
impl<'obj> Drop for TcController<'obj> {
139-
fn drop(&mut self) {
140-
self.destroy();
141-
}
142-
}
143-
144-
struct XdpController<'obj> {
145-
pub xdp: Xdp<'obj>,
146-
ifidx: i32,
147-
}
148-
149-
impl<'obj> XdpController<'obj> {
150-
fn new(fd: BorrowedFd<'obj>, iface: &str) -> Result<Self> {
151-
let ifidx = if_nametoindex(iface)? as i32;
152-
let xdp = Xdp::new(fd);
153-
154-
Ok(Self { xdp, ifidx })
155-
}
156-
157-
pub fn attach(&self) -> Result<()> {
158-
self.xdp
159-
.attach(self.ifidx, XdpFlags::NONE)
160-
.context("Failed to attach xdp")
161-
}
162-
163-
pub fn destroy(&mut self) -> Result<()> {
164-
self.xdp
165-
.detach(self.ifidx, XdpFlags::NONE)
166-
.context("Failed to detach xdp")
167-
}
168-
}
169-
170-
impl<'obj> Drop for XdpController<'obj> {
171-
fn drop(&mut self) {
172-
let _ = self.destroy();
173-
}
174-
}
175-
176-
impl<'obj> DpitSkel<'obj> {
177-
/// Adds SNI LPM entry to the TRIE map.
178-
/// For more detailed description see sni_lpm_map definition in eBPF code.
179-
/// Note, that Point-terminator is specified explicitly by user if needed.
180-
///
181-
/// This function also implicitly reverses the destination string,
182-
/// so user should pass the domain normally, like `.google.com`
183-
fn sni_lpm_add_entry(&self, domain: &str, action: sni_action) -> Result<()> {
184-
if domain.len() != domain.len() {
185-
return Err(anyhow!(
186-
"The domain MUST NOT contain Unicode. Use normal form, in ASCII"
187-
));
188-
}
189-
190-
let mut ks = ebpf_prog::types::sni_buf::default();
191-
let srev: String = domain.chars().rev().collect::<String>();
192-
let bytes_string = srev.as_bytes();
193-
// One for NULL-terminator and point-terminator
194-
if bytes_string.len() > ks.data.len() + 2 {
195-
return Err(anyhow!("Map SNI string length exceeds the limit."));
196-
};
197-
for (i, c) in bytes_string.iter().enumerate() {
198-
ks.data[i] = *c;
199-
}
200-
// NULL-terminator
201-
ks.data[bytes_string.len()] = b'\0';
202-
203-
ks.prefixlen = bytes_string.len() as u32 * 8;
204-
205-
let key = &ks as *const ebpf_prog::types::sni_buf as *const u8;
206-
let size = mem::size_of::<ebpf_prog::types::sni_buf>();
207-
let key = unsafe { std::slice::from_raw_parts(key, size) };
208-
209-
let value = &action as *const ebpf_prog::types::sni_action as *const u8;
210-
let size = mem::size_of::<ebpf_prog::types::sni_action>();
211-
let value = unsafe { std::slice::from_raw_parts(value, size) };
212-
self.maps
213-
.sni_lpm_map
214-
.update(key, value, MapFlags::ANY)
215-
.context("Error while updating the trie map")?;
216-
217-
Ok(())
218-
}
219-
}
220-
22136
#[tokio::main]
22237
async fn main() -> Result<()> {
22338
let opts = Command::parse();

‎src/sni_logging_handle.rs‎

Lines changed: 10 additions & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,14 @@
11
use anyhow::{anyhow, Result};
2-
use etherparse::{IpNumber, Ipv4Header, Ipv6Header, TcpHeader, TransportHeader};
3-
use nix::sys::socket::{sendto, MsgFlags, SockaddrIn, SockaddrIn6, SockaddrLike};
2+
use etherparse::{IpNumber, Ipv4Header, TcpHeader, TransportHeader};
43

5-
use std::{mem, sync::Arc, thread::JoinHandle, time::Duration};
4+
use std::{sync::Arc, thread::JoinHandle, time::Duration};
65

76
use crate::{
87
ebpf_prog::{
98
self,
109
types::{chlo_tls_atype, sni_action, tls_sni_signaling},
1110
},
12-
utils::{NetworkHeader, TransportHeaderParse},
11+
utils::{NetworkHeader, TransportHeaderParse, RawSocket},
1312
};
1413
use libbpf_rs::PerfBufferBuilder;
1514
use plain::Plain;
@@ -48,8 +47,7 @@ fn sni_tls_get_domain(sni_tls_data: &tls_sni_signaling) -> Result<Option<String>
4847
}
4948

5049
struct PerfProcessContext {
51-
ip_rawsocket: RawSocket,
52-
ip6_rawsocket: RawSocket,
50+
rawsocket: RawSocket,
5351
}
5452

5553
fn process_perf_signal(
@@ -108,7 +106,7 @@ fn send_bi_tcp_rst(
108106
let mut sent_packet = iph.to_bytes().to_vec();
109107
sent_packet.extend(tcph.to_bytes());
110108

111-
ctx.ip_rawsocket.send_ipv4(&iph, &sent_packet[..])?;
109+
ctx.rawsocket.send_ipv4(&iph, &sent_packet[..])?;
112110

113111
iph.destination = iph.source;
114112
iph.source = oiph.destination;
@@ -125,7 +123,7 @@ fn send_bi_tcp_rst(
125123
let mut sent_packet = iph.to_bytes().to_vec();
126124
sent_packet.extend(tcph.to_bytes());
127125

128-
ctx.ip_rawsocket.send_ipv4(&iph, &sent_packet[..])?;
126+
ctx.rawsocket.send_ipv4(&iph, &sent_packet[..])?;
129127
} else if let NetworkHeader::Ipv6(oiph) = network_hdr {
130128
let mut iph = oiph.clone();
131129
iph.hop_limit = 128;
@@ -137,7 +135,7 @@ fn send_bi_tcp_rst(
137135
let mut sent_packet = iph.to_bytes().to_vec();
138136
sent_packet.extend(tcph.to_bytes());
139137

140-
ctx.ip6_rawsocket.send_ipv6(&iph, &sent_packet[..])?;
138+
ctx.rawsocket.send_ipv6(&iph, &sent_packet[..])?;
141139

142140
iph.destination = iph.source;
143141
iph.source = oiph.destination;
@@ -153,115 +151,15 @@ fn send_bi_tcp_rst(
153151
let mut sent_packet = iph.to_bytes().to_vec();
154152
sent_packet.extend(tcph.to_bytes());
155153

156-
ctx.ip6_rawsocket.send_ipv6(&iph, &sent_packet[..])?;
154+
ctx.rawsocket.send_ipv6(&iph, &sent_packet[..])?;
157155
}
158156

159157
Ok(())
160158
}
161159

162-
#[derive(Copy, Clone, PartialEq)]
163-
enum RawSocketType {
164-
AfInet = libc::AF_INET as isize,
165-
AfInet6 = libc::AF_INET6 as isize,
166-
}
167-
struct RawSocket {
168-
fd: libc::c_int,
169-
domain_type: RawSocketType,
170-
}
171-
172-
impl RawSocket {
173-
pub fn new(skel: &ebpf_prog::DpitSkel, domain: RawSocketType) -> Result<Self> {
174-
let rawsocket =
175-
unsafe { libc::socket(domain as libc::c_int, libc::SOCK_RAW, libc::IPPROTO_RAW) };
176-
177-
if rawsocket == -1 {
178-
return Err(anyhow!("Open rawsocket"));
179-
}
180-
let mark: libc::c_int = skel.maps.rodata_data.unwrap().RAWSOCKET_MARK as libc::c_int;
181-
let ret = unsafe {
182-
libc::setsockopt(
183-
rawsocket,
184-
libc::SOL_SOCKET,
185-
libc::SO_MARK,
186-
(&mark) as *const libc::c_int as *const libc::c_void,
187-
mem::size_of_val(&mark) as libc::socklen_t,
188-
)
189-
};
190-
if ret < 0 {
191-
unsafe { libc::close(rawsocket) };
192-
return Err(anyhow!("setsockopt(SO_MARK, {}) failed", mark));
193-
}
194-
195-
Ok(Self {
196-
fd: rawsocket,
197-
domain_type: domain,
198-
})
199-
}
200-
201-
pub fn fd(&self) -> i32 {
202-
self.fd
203-
}
204-
205-
pub fn send_ipv4(&self, iph: &Ipv4Header, pkt: &[u8]) -> Result<()> {
206-
if RawSocketType::AfInet != self.domain_type {
207-
unreachable!();
208-
}
209-
210-
let daddr = SockaddrIn::new(
211-
iph.destination[0],
212-
iph.destination[1],
213-
iph.destination[2],
214-
iph.destination[3],
215-
0,
216-
);
217-
218-
let t = sendto(self.fd(), pkt, &daddr, MsgFlags::MSG_DONTWAIT)?;
219-
println!("Sent {t} bytes");
220-
221-
Ok(())
222-
}
223-
224-
pub fn send_ipv6(&self, ip6h: &Ipv6Header, pkt: &[u8]) -> Result<()> {
225-
if RawSocketType::AfInet6 != self.domain_type {
226-
unreachable!();
227-
}
228-
229-
let daddr = libc::sockaddr_in6 {
230-
sin6_family: libc::AF_INET6 as u16,
231-
/* Always 0 for raw socket */
232-
sin6_port: 0,
233-
sin6_addr: libc::in6_addr {
234-
s6_addr: ip6h.destination,
235-
},
236-
sin6_flowinfo: unsafe { mem::zeroed() },
237-
sin6_scope_id: unsafe { mem::zeroed() },
238-
};
239-
240-
let daddr = unsafe {
241-
SockaddrIn6::from_raw(
242-
&daddr as *const libc::sockaddr_in6 as *const libc::sockaddr,
243-
None,
244-
)
245-
}
246-
.ok_or(anyhow!("SockaddrIn6 from_raw"))?;
247-
248-
sendto(self.fd(), pkt, &daddr, MsgFlags::MSG_DONTWAIT)?;
249-
250-
Ok(())
251-
}
252-
}
253-
254-
impl Drop for RawSocket {
255-
fn drop(&mut self) {
256-
unsafe { libc::close(self.fd) };
257-
self.fd = 0;
258-
}
259-
}
260-
261160
pub fn init_sni_logging(skel: &ebpf_prog::DpitSkel) -> Result<JoinHandle<()>> {
262-
let ip_rawsocket = RawSocket::new(skel, RawSocketType::AfInet)?;
263-
let ip6_rawsocket = RawSocket::new(skel, RawSocketType::AfInet6)?;
264-
let perf_ctx = Arc::new(PerfProcessContext { ip_rawsocket, ip6_rawsocket });
161+
let rawsocket = RawSocket::new(skel.maps.rodata_data.unwrap().RAWSOCKET_MARK as u32)?;
162+
let perf_ctx = Arc::new(PerfProcessContext { rawsocket });
265163

266164
let handle_perf_event = move |_cpu: i32, data: &[u8]| {
267165
let mut sni_tls_data = tls_sni_signaling::default();

0 commit comments

Comments
 (0)