Skip to content

Commit de9d981

Browse files
adgaultierpythops
adgaultier
and
pythops
authored
Use Map for traffic direction (#28)
Co-authored-by: Badr <[email protected]>
1 parent 3c00279 commit de9d981

File tree

12 files changed

+816
-835
lines changed

12 files changed

+816
-835
lines changed

‎oryx-ebpf/src/main.rs

+75-43
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_FILTER: Array<u8> = Array::with_max_entries(1, 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);
@@ -41,6 +44,9 @@ static BLOCKLIST_IPV6: HashMap<u128, [u16; MAX_RULES_PORT]> =
4144
static BLOCKLIST_IPV4: HashMap<u32, [u16; MAX_RULES_PORT]> =
4245
HashMap::<u32, [u16; MAX_RULES_PORT]>::with_max_entries(MAX_FIREWALL_RULES, 0);
4346

47+
#[no_mangle]
48+
static TRAFFIC_DIRECTION: i32 = 0;
49+
4450
#[classifier]
4551
pub fn oryx(ctx: TcContext) -> i32 {
4652
match process(ctx) {
@@ -56,6 +62,7 @@ fn submit(packet: RawPacket) {
5662
buf.submit(0);
5763
}
5864
}
65+
5966
#[inline]
6067
fn ptr_at<T>(ctx: &TcContext, offset: usize) -> Result<*const T, ()> {
6168
let start = ctx.data();
@@ -70,12 +77,23 @@ fn ptr_at<T>(ctx: &TcContext, offset: usize) -> Result<*const T, ()> {
7077
}
7178

7279
#[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) } {
80+
fn filter_direction() -> bool {
81+
// 0(default) -> false(send to tui), 1 -> true(filter)
82+
if let Some(v) = TRAFFIC_DIRECTION_FILTER.get(0) {
83+
return *v != 0;
84+
}
85+
false
86+
}
87+
88+
#[inline]
89+
fn is_ingress() -> bool {
90+
let traffic_direction = unsafe { core::ptr::read_volatile(&TRAFFIC_DIRECTION) };
91+
traffic_direction == -1
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-
{
157-
return Ok(TC_ACT_SHOT);
175+
if block_ipv4(addr, port) {
176+
return Ok(TC_ACT_SHOT); //block packet
158177
}
159178

160179
if filter_packet(Protocol::Network(NetworkProtocol::Ipv4))
161180
|| filter_packet(Protocol::Transport(TransportProtocol::TCP))
181+
|| filter_direction()
162182
{
163-
return Ok(TC_ACT_PIPE); //DONT FWD PACKET TO TUI
183+
return Ok(TC_ACT_PIPE);
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-
{
178-
return Ok(TC_ACT_SHOT);
199+
if block_ipv4(addr, port) {
200+
return Ok(TC_ACT_SHOT); //block packet
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,24 +227,30 @@ 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-
{
219-
return Ok(TC_ACT_SHOT);
245+
if block_ipv6(addr, port) {
246+
return Ok(TC_ACT_SHOT); //block packet
220247
}
248+
221249
if filter_packet(Protocol::Network(NetworkProtocol::Ipv6))
222250
|| filter_packet(Protocol::Transport(TransportProtocol::TCP))
251+
|| filter_direction()
223252
{
224-
return Ok(TC_ACT_PIPE); //DONT FWD PACKET TO TUI
253+
return Ok(TC_ACT_PIPE);
225254
}
226255
submit(RawPacket::Ip(
227256
IpHdr::V6(header),
@@ -230,18 +259,21 @@ 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-
{
239-
return Ok(TC_ACT_SHOT);
268+
if block_ipv6(addr, port) {
269+
return Ok(TC_ACT_SHOT); //block packet
240270
}
271+
241272
if filter_packet(Protocol::Network(NetworkProtocol::Ipv6))
242273
|| filter_packet(Protocol::Transport(TransportProtocol::UDP))
274+
|| filter_direction()
243275
{
244-
return Ok(TC_ACT_PIPE); //DONT FWD PACKET TO TUI
276+
return Ok(TC_ACT_PIPE);
245277
}
246278
submit(RawPacket::Ip(
247279
IpHdr::V6(header),

‎oryx-tui/src/app.rs

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

12-
use crate::notification::Notification;
1313
use crate::{filter::Filter, help::Help};
14+
use crate::{filter::IoChannels, notification::Notification};
1415
use crate::{packet::AppPacket, section::Section};
1516

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

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

60-
let (firewall_ingress_sender, firewall_ingress_receiver) = kanal::unbounded();
61-
let (firewall_egress_sender, firewall_egress_receiver) = kanal::unbounded();
62-
61+
let firewall_channels = IoChannels::new();
6362
thread::spawn({
6463
let packets = packets.clone();
6564
move || loop {
@@ -77,20 +76,11 @@ impl App {
7776
Self {
7877
running: true,
7978
help: Help::new(),
80-
filter: Filter::new(
81-
firewall_ingress_sender.clone(),
82-
firewall_ingress_receiver,
83-
firewall_egress_sender.clone(),
84-
firewall_egress_receiver,
85-
),
79+
filter: Filter::new(firewall_channels.clone()),
8680
start_sniffing: false,
8781
packets: packets.clone(),
8882
notifications: Vec::new(),
89-
section: Section::new(
90-
packets.clone(),
91-
firewall_ingress_sender,
92-
firewall_egress_sender,
93-
),
83+
section: Section::new(packets.clone(), firewall_channels.clone()),
9484
data_channel_sender: sender,
9585
is_editing: false,
9686
active_popup: None,
@@ -133,6 +123,8 @@ impl App {
133123
}
134124

135125
pub fn quit(&mut self) {
126+
self.filter.terminate();
127+
thread::sleep(Duration::from_millis(110));
136128
self.running = false;
137129
}
138130
}

0 commit comments

Comments
 (0)