Skip to content

Commit c286bb5

Browse files
adgaultieradgaultier
committed
refacto?
1 parent 33fc05d commit c286bb5

File tree

8 files changed

+313
-268
lines changed

8 files changed

+313
-268
lines changed

‎oryx-ebpf/src/main.rs

+68-36
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@ static TRANSPORT_FILTERS: Array<u32> = Array::with_max_entries(8, 0);
3333
#[map]
3434
static LINK_FILTERS: Array<u32> = Array::with_max_entries(8, 0);
3535

36+
#[map]
37+
static TRAFFIC_DIRECTION_FILTERS: Array<u32> = Array::with_max_entries(2, 0);
38+
3639
#[map]
3740
static BLOCKLIST_IPV6: HashMap<u128, [u16; MAX_RULES_PORT]> =
3841
HashMap::<u128, [u16; MAX_RULES_PORT]>::with_max_entries(MAX_FIREWALL_RULES, 0);
@@ -56,6 +59,7 @@ fn submit(packet: RawPacket) {
5659
buf.submit(0);
5760
}
5861
}
62+
5963
#[inline]
6064
fn ptr_at<T>(ctx: &TcContext, offset: usize) -> Result<*const T, ()> {
6165
let start = ctx.data();
@@ -70,12 +74,26 @@ fn ptr_at<T>(ctx: &TcContext, offset: usize) -> Result<*const T, ()> {
7074
}
7175

7276
#[inline]
73-
fn filter_for_ipv4_address(
74-
addr: u32,
75-
port: u16,
76-
blocked_ports_map: &HashMap<u32, [u16; 32]>,
77-
) -> bool {
78-
if let Some(blocked_ports) = unsafe { blocked_ports_map.get(&addr) } {
77+
fn filter_direction() -> bool {
78+
// 0(default) -> false(send to tui), 1 -> true(filter)
79+
if let Some(v) = TRAFFIC_DIRECTION_FILTERS.get(0) {
80+
return *v != 0;
81+
}
82+
false
83+
}
84+
85+
#[inline]
86+
fn is_ingress() -> bool {
87+
// 0(default) -> true(is ingress), 1 -> false (is egress)
88+
if let Some(v) = TRAFFIC_DIRECTION_FILTERS.get(1) {
89+
return *v == 0;
90+
}
91+
true
92+
}
93+
94+
#[inline]
95+
fn block_ipv4(addr: u32, port: u16) -> bool {
96+
if let Some(blocked_ports) = unsafe { BLOCKLIST_IPV4.get(&addr) } {
7997
for (idx, blocked_port) in blocked_ports.iter().enumerate() {
8098
if *blocked_port == 0 {
8199
if idx == 0 {
@@ -92,12 +110,8 @@ fn filter_for_ipv4_address(
92110
}
93111

94112
#[inline]
95-
fn filter_for_ipv6_address(
96-
addr: u128,
97-
port: u16,
98-
blocked_ports_map: &HashMap<u128, [u16; 32]>,
99-
) -> bool {
100-
if let Some(blocked_ports) = unsafe { blocked_ports_map.get(&addr) } {
113+
fn block_ipv6(addr: u128, port: u16) -> bool {
114+
if let Some(blocked_ports) = unsafe { BLOCKLIST_IPV6.get(&addr) } {
101115
for (idx, blocked_port) in blocked_ports.iter().enumerate() {
102116
if *blocked_port == 0 {
103117
if idx == 0 {
@@ -142,44 +156,53 @@ fn process(ctx: TcContext) -> Result<i32, ()> {
142156
match ethhdr.ether_type {
143157
EtherType::Ipv4 => {
144158
let header: Ipv4Hdr = ctx.load(EthHdr::LEN).map_err(|_| ())?;
145-
let src_addr = u32::from_be(header.src_addr);
146-
let dst_addr = u32::from_be(header.dst_addr);
159+
160+
let addr = if is_ingress() {
161+
u32::from_be(header.src_addr)
162+
} else {
163+
u32::from_be(header.dst_addr)
164+
};
147165

148166
match header.proto {
149167
IpProto::Tcp => {
150168
let tcphdr: *const TcpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
151-
let src_port = u16::from_be(unsafe { (*tcphdr).source });
152-
let dst_port = u16::from_be(unsafe { (*tcphdr).dest });
169+
let port = if is_ingress() {
170+
u16::from_be(unsafe { (*tcphdr).source })
171+
} else {
172+
u16::from_be(unsafe { (*tcphdr).dest })
173+
};
153174

154-
if filter_for_ipv4_address(src_addr, src_port, &BLOCKLIST_IPV4)
155-
|| filter_for_ipv4_address(dst_addr, dst_port, &BLOCKLIST_IPV4)
156-
{
175+
if block_ipv4(addr, port) {
157176
return Ok(TC_ACT_SHOT);
158177
}
159178

160179
if filter_packet(Protocol::Network(NetworkProtocol::Ipv4))
161180
|| filter_packet(Protocol::Transport(TransportProtocol::TCP))
181+
|| filter_direction()
162182
{
163183
return Ok(TC_ACT_PIPE); //DONT FWD PACKET TO TUI
164184
}
185+
165186
submit(RawPacket::Ip(
166187
IpHdr::V4(header),
167188
ProtoHdr::Tcp(unsafe { *tcphdr }),
168189
));
169190
}
170191
IpProto::Udp => {
171192
let udphdr: *const UdpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
172-
let src_port = u16::from_be(unsafe { (*udphdr).source });
173-
let dst_port = u16::from_be(unsafe { (*udphdr).dest });
193+
let port = if is_ingress() {
194+
u16::from_be(unsafe { (*udphdr).source })
195+
} else {
196+
u16::from_be(unsafe { (*udphdr).dest })
197+
};
174198

175-
if filter_for_ipv4_address(src_addr, src_port, &BLOCKLIST_IPV4)
176-
|| filter_for_ipv4_address(dst_addr, dst_port, &BLOCKLIST_IPV4)
177-
{
199+
if block_ipv4(addr, port) {
178200
return Ok(TC_ACT_SHOT);
179201
}
180202

181203
if filter_packet(Protocol::Network(NetworkProtocol::Ipv4))
182204
|| filter_packet(Protocol::Transport(TransportProtocol::UDP))
205+
|| filter_direction()
183206
{
184207
return Ok(TC_ACT_PIPE);
185208
}
@@ -204,22 +227,28 @@ fn process(ctx: TcContext) -> Result<i32, ()> {
204227
}
205228
EtherType::Ipv6 => {
206229
let header: Ipv6Hdr = ctx.load(EthHdr::LEN).map_err(|_| ())?;
207-
let src_addr = header.src_addr().to_bits();
208-
let dst_addr = header.dst_addr().to_bits();
230+
let addr = if is_ingress() {
231+
header.src_addr().to_bits()
232+
} else {
233+
header.dst_addr().to_bits()
234+
};
209235

210236
match header.next_hdr {
211237
IpProto::Tcp => {
212238
let tcphdr: *const TcpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv6Hdr::LEN)?;
213-
let src_port = u16::from_be(unsafe { (*tcphdr).source });
214-
let dst_port = u16::from_be(unsafe { (*tcphdr).dest });
239+
let port = if is_ingress() {
240+
u16::from_be(unsafe { (*tcphdr).source })
241+
} else {
242+
u16::from_be(unsafe { (*tcphdr).dest })
243+
};
215244

216-
if filter_for_ipv6_address(src_addr, src_port, &BLOCKLIST_IPV6)
217-
|| filter_for_ipv6_address(dst_addr, dst_port, &BLOCKLIST_IPV6)
218-
{
245+
if block_ipv6(addr, port) {
219246
return Ok(TC_ACT_SHOT);
220247
}
248+
221249
if filter_packet(Protocol::Network(NetworkProtocol::Ipv6))
222250
|| filter_packet(Protocol::Transport(TransportProtocol::TCP))
251+
|| filter_direction()
223252
{
224253
return Ok(TC_ACT_PIPE); //DONT FWD PACKET TO TUI
225254
}
@@ -230,16 +259,19 @@ fn process(ctx: TcContext) -> Result<i32, ()> {
230259
}
231260
IpProto::Udp => {
232261
let udphdr: *const UdpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv6Hdr::LEN)?;
233-
let src_port = u16::from_be(unsafe { (*udphdr).source });
234-
let dst_port = u16::from_be(unsafe { (*udphdr).dest });
262+
let port = if is_ingress() {
263+
u16::from_be(unsafe { (*udphdr).source })
264+
} else {
265+
u16::from_be(unsafe { (*udphdr).dest })
266+
};
235267

236-
if filter_for_ipv6_address(src_addr, src_port, &BLOCKLIST_IPV6)
237-
|| filter_for_ipv6_address(dst_addr, dst_port, &BLOCKLIST_IPV6)
238-
{
268+
if block_ipv6(addr, port) {
239269
return Ok(TC_ACT_SHOT);
240270
}
271+
241272
if filter_packet(Protocol::Network(NetworkProtocol::Ipv6))
242273
|| filter_packet(Protocol::Transport(TransportProtocol::UDP))
274+
|| filter_direction()
243275
{
244276
return Ok(TC_ACT_PIPE); //DONT FWD PACKET TO TUI
245277
}

‎oryx-tui/src/app.rs

+7-16
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,11 @@ use std::{
88
error,
99
sync::{Arc, Mutex},
1010
thread,
11+
time::Duration,
1112
};
1213

13-
use crate::notification::Notification;
1414
use crate::{filter::Filter, help::Help};
15+
use crate::{filter::IoChans, notification::Notification};
1516
use crate::{packet::AppPacket, section::Section};
1617

1718
pub type AppResult<T> = std::result::Result<T, Box<dyn error::Error>>;
@@ -58,9 +59,7 @@ impl App {
5859

5960
let (sender, receiver) = kanal::unbounded();
6061

61-
let (firewall_ingress_sender, firewall_ingress_receiver) = kanal::unbounded();
62-
let (firewall_egress_sender, firewall_egress_receiver) = kanal::unbounded();
63-
62+
let firewall_chans = IoChans::new();
6463
thread::spawn({
6564
let packets = packets.clone();
6665
move || loop {
@@ -78,20 +77,11 @@ impl App {
7877
Self {
7978
running: true,
8079
help: Help::new(),
81-
filter: Filter::new(
82-
firewall_ingress_sender.clone(),
83-
firewall_ingress_receiver,
84-
firewall_egress_sender.clone(),
85-
firewall_egress_receiver,
86-
),
80+
filter: Filter::new(firewall_chans.clone()),
8781
start_sniffing: false,
8882
packets: packets.clone(),
8983
notifications: Vec::new(),
90-
section: Section::new(
91-
packets.clone(),
92-
firewall_ingress_sender,
93-
firewall_egress_sender,
94-
),
84+
section: Section::new(packets.clone(), firewall_chans.clone()),
9585
data_channel_sender: sender,
9686
is_editing: false,
9787
active_popup: None,
@@ -138,7 +128,8 @@ impl App {
138128
if let Err(e) = self.section.firewall.save_rules() {
139129
error!("{}", e)
140130
}
141-
131+
self.filter.terminate();
132+
thread::sleep(Duration::from_millis(110));
142133
self.running = false;
143134
}
144135
}

‎oryx-tui/src/ebpf.rs

+22-3
Original file line numberDiff line numberDiff line change
@@ -295,13 +295,18 @@ pub fn load_ingress(
295295
let mut link_filters: Array<_, u32> =
296296
Array::try_from(bpf.take_map("LINK_FILTERS").unwrap()).unwrap();
297297

298-
// firewall-ebpf interface
298+
let mut traffic_direction_filters: Array<_, u32> =
299+
Array::try_from(bpf.take_map("TRAFFIC_DIRECTION_FILTERS").unwrap()).unwrap();
300+
let _ = traffic_direction_filters.set(0, 0, 0);
301+
let _ = traffic_direction_filters.set(1, 0, 0); //setup ingress flag
302+
// firewall-ebpf interface
299303
let mut ipv4_firewall: HashMap<_, u32, [u16; MAX_RULES_PORT]> =
300304
HashMap::try_from(bpf.take_map("BLOCKLIST_IPV4").unwrap()).unwrap();
301305

302306
let mut ipv6_firewall: HashMap<_, u128, [u16; MAX_RULES_PORT]> =
303307
HashMap::try_from(bpf.take_map("BLOCKLIST_IPV6").unwrap()).unwrap();
304308

309+
// firewall thread
305310
thread::spawn(move || loop {
306311
if let Ok(signal) = firewall_ingress_receiver.recv() {
307312
match signal {
@@ -327,10 +332,11 @@ pub fn load_ingress(
327332
}
328333
});
329334

335+
// packets filters thread
330336
thread::spawn(move || loop {
331337
if let Ok(signal) = filter_channel_receiver.recv() {
332338
match signal {
333-
FilterChannelSignal::Update((filter, flag)) => match filter {
339+
FilterChannelSignal::ProtoUpdate((filter, flag)) => match filter {
334340
Protocol::Transport(p) => {
335341
let _ = transport_filters.set(p as u32, flag as u32, 0);
336342
}
@@ -341,6 +347,9 @@ pub fn load_ingress(
341347
let _ = link_filters.set(p as u32, flag as u32, 0);
342348
}
343349
},
350+
FilterChannelSignal::DirectionUpdate(flag) => {
351+
let _ = traffic_direction_filters.set(0, flag as u32, 0);
352+
}
344353
FilterChannelSignal::Kill => {
345354
break;
346355
}
@@ -482,13 +491,19 @@ pub fn load_egress(
482491
let mut link_filters: Array<_, u32> =
483492
Array::try_from(bpf.take_map("LINK_FILTERS").unwrap()).unwrap();
484493

494+
let mut traffic_direction_filters: Array<_, u32> =
495+
Array::try_from(bpf.take_map("TRAFFIC_DIRECTION_FILTERS").unwrap()).unwrap();
496+
let _ = traffic_direction_filters.set(0, 0, 0);
497+
let _ = traffic_direction_filters.set(1, 1, 0); //setup egress flag
498+
485499
// firewall-ebpf interface
486500
let mut ipv4_firewall: HashMap<_, u32, [u16; MAX_RULES_PORT]> =
487501
HashMap::try_from(bpf.take_map("BLOCKLIST_IPV4").unwrap()).unwrap();
488502

489503
let mut ipv6_firewall: HashMap<_, u128, [u16; MAX_RULES_PORT]> =
490504
HashMap::try_from(bpf.take_map("BLOCKLIST_IPV6").unwrap()).unwrap();
491505

506+
// firewall thread
492507
thread::spawn(move || loop {
493508
if let Ok(signal) = firewall_egress_receiver.recv() {
494509
match signal {
@@ -514,10 +529,11 @@ pub fn load_egress(
514529
}
515530
});
516531

532+
// packets filters thread
517533
thread::spawn(move || loop {
518534
if let Ok(signal) = filter_channel_receiver.recv() {
519535
match signal {
520-
FilterChannelSignal::Update((filter, flag)) => match filter {
536+
FilterChannelSignal::ProtoUpdate((filter, flag)) => match filter {
521537
Protocol::Transport(p) => {
522538
let _ = transport_filters.set(p as u32, flag as u32, 0);
523539
}
@@ -528,6 +544,9 @@ pub fn load_egress(
528544
let _ = link_filters.set(p as u32, flag as u32, 0);
529545
}
530546
},
547+
FilterChannelSignal::DirectionUpdate(flag) => {
548+
let _ = traffic_direction_filters.set(0, flag as u32, 0);
549+
}
531550
FilterChannelSignal::Kill => {
532551
break;
533552
}

0 commit comments

Comments
 (0)