Skip to content

Commit

Permalink
Use Map for traffic direction (#28)
Browse files Browse the repository at this point in the history 
Co-authored-by: Badr <[email protected]>
  • Loading branch information
@adgaultier @pythops
adgaultier and pythops authored Oct 14, 2024
1 parent 3c00279 commit de9d981
Show file tree
Hide file tree
Showing 12 changed files with 816 additions and 835 deletions.
118 changes: 75 additions & 43 deletions oryx-ebpf/src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,9 @@ static TRANSPORT_FILTERS: Array<u32> = Array::with_max_entries(8, 0);
#[map]
static LINK_FILTERS: Array<u32> = Array::with_max_entries(8, 0);

#[map]
static TRAFFIC_DIRECTION_FILTER: Array<u8> = Array::with_max_entries(1, 0);

#[map]
static BLOCKLIST_IPV6: HashMap<u128, [u16; MAX_RULES_PORT]> =
HashMap::<u128, [u16; MAX_RULES_PORT]>::with_max_entries(MAX_FIREWALL_RULES, 0);
Expand All @@ -41,6 +44,9 @@ static BLOCKLIST_IPV6: HashMap<u128, [u16; MAX_RULES_PORT]> =
static BLOCKLIST_IPV4: HashMap<u32, [u16; MAX_RULES_PORT]> =
HashMap::<u32, [u16; MAX_RULES_PORT]>::with_max_entries(MAX_FIREWALL_RULES, 0);

#[no_mangle]
static TRAFFIC_DIRECTION: i32 = 0;

#[classifier]
pub fn oryx(ctx: TcContext) -> i32 {
match process(ctx) {
Expand All @@ -56,6 +62,7 @@ fn submit(packet: RawPacket) {
buf.submit(0);
}
}

#[inline]
fn ptr_at<T>(ctx: &TcContext, offset: usize) -> Result<*const T, ()> {
let start = ctx.data();
Expand All @@ -70,12 +77,23 @@ fn ptr_at<T>(ctx: &TcContext, offset: usize) -> Result<*const T, ()> {
}

#[inline]
fn filter_for_ipv4_address(
addr: u32,
port: u16,
blocked_ports_map: &HashMap<u32, [u16; 32]>,
) -> bool {
if let Some(blocked_ports) = unsafe { blocked_ports_map.get(&addr) } {
fn filter_direction() -> bool {
// 0(default) -> false(send to tui), 1 -> true(filter)
if let Some(v) = TRAFFIC_DIRECTION_FILTER.get(0) {
return *v != 0;
}
false
}

#[inline]
fn is_ingress() -> bool {
let traffic_direction = unsafe { core::ptr::read_volatile(&TRAFFIC_DIRECTION) };
traffic_direction == -1
}

#[inline]
fn block_ipv4(addr: u32, port: u16) -> bool {
if let Some(blocked_ports) = unsafe { BLOCKLIST_IPV4.get(&addr) } {
for (idx, blocked_port) in blocked_ports.iter().enumerate() {
if *blocked_port == 0 {
if idx == 0 {
Expand All @@ -92,12 +110,8 @@ fn filter_for_ipv4_address(
}

#[inline]
fn filter_for_ipv6_address(
addr: u128,
port: u16,
blocked_ports_map: &HashMap<u128, [u16; 32]>,
) -> bool {
if let Some(blocked_ports) = unsafe { blocked_ports_map.get(&addr) } {
fn block_ipv6(addr: u128, port: u16) -> bool {
if let Some(blocked_ports) = unsafe { BLOCKLIST_IPV6.get(&addr) } {
for (idx, blocked_port) in blocked_ports.iter().enumerate() {
if *blocked_port == 0 {
if idx == 0 {
Expand Down Expand Up @@ -142,44 +156,53 @@ fn process(ctx: TcContext) -> Result<i32, ()> {
match ethhdr.ether_type {
EtherType::Ipv4 => {
let header: Ipv4Hdr = ctx.load(EthHdr::LEN).map_err(|_| ())?;
let src_addr = u32::from_be(header.src_addr);
let dst_addr = u32::from_be(header.dst_addr);

let addr = if is_ingress() {
u32::from_be(header.src_addr)
} else {
u32::from_be(header.dst_addr)
};

match header.proto {
IpProto::Tcp => {
let tcphdr: *const TcpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
let src_port = u16::from_be(unsafe { (*tcphdr).source });
let dst_port = u16::from_be(unsafe { (*tcphdr).dest });
let port = if is_ingress() {
u16::from_be(unsafe { (*tcphdr).source })
} else {
u16::from_be(unsafe { (*tcphdr).dest })
};

if filter_for_ipv4_address(src_addr, src_port, &BLOCKLIST_IPV4)
|| filter_for_ipv4_address(dst_addr, dst_port, &BLOCKLIST_IPV4)
{
return Ok(TC_ACT_SHOT);
if block_ipv4(addr, port) {
return Ok(TC_ACT_SHOT); //block packet
}

if filter_packet(Protocol::Network(NetworkProtocol::Ipv4))
|| filter_packet(Protocol::Transport(TransportProtocol::TCP))
|| filter_direction()
{
return Ok(TC_ACT_PIPE); //DONT FWD PACKET TO TUI
return Ok(TC_ACT_PIPE);
}

submit(RawPacket::Ip(
IpHdr::V4(header),
ProtoHdr::Tcp(unsafe { *tcphdr }),
));
}
IpProto::Udp => {
let udphdr: *const UdpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?;
let src_port = u16::from_be(unsafe { (*udphdr).source });
let dst_port = u16::from_be(unsafe { (*udphdr).dest });
let port = if is_ingress() {
u16::from_be(unsafe { (*udphdr).source })
} else {
u16::from_be(unsafe { (*udphdr).dest })
};

if filter_for_ipv4_address(src_addr, src_port, &BLOCKLIST_IPV4)
|| filter_for_ipv4_address(dst_addr, dst_port, &BLOCKLIST_IPV4)
{
return Ok(TC_ACT_SHOT);
if block_ipv4(addr, port) {
return Ok(TC_ACT_SHOT); //block packet
}

if filter_packet(Protocol::Network(NetworkProtocol::Ipv4))
|| filter_packet(Protocol::Transport(TransportProtocol::UDP))
|| filter_direction()
{
return Ok(TC_ACT_PIPE);
}
Expand All @@ -204,24 +227,30 @@ fn process(ctx: TcContext) -> Result<i32, ()> {
}
EtherType::Ipv6 => {
let header: Ipv6Hdr = ctx.load(EthHdr::LEN).map_err(|_| ())?;
let src_addr = header.src_addr().to_bits();
let dst_addr = header.dst_addr().to_bits();
let addr = if is_ingress() {
header.src_addr().to_bits()
} else {
header.dst_addr().to_bits()
};

match header.next_hdr {
IpProto::Tcp => {
let tcphdr: *const TcpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv6Hdr::LEN)?;
let src_port = u16::from_be(unsafe { (*tcphdr).source });
let dst_port = u16::from_be(unsafe { (*tcphdr).dest });
let port = if is_ingress() {
u16::from_be(unsafe { (*tcphdr).source })
} else {
u16::from_be(unsafe { (*tcphdr).dest })
};

if filter_for_ipv6_address(src_addr, src_port, &BLOCKLIST_IPV6)
|| filter_for_ipv6_address(dst_addr, dst_port, &BLOCKLIST_IPV6)
{
return Ok(TC_ACT_SHOT);
if block_ipv6(addr, port) {
return Ok(TC_ACT_SHOT); //block packet
}

if filter_packet(Protocol::Network(NetworkProtocol::Ipv6))
|| filter_packet(Protocol::Transport(TransportProtocol::TCP))
|| filter_direction()
{
return Ok(TC_ACT_PIPE); //DONT FWD PACKET TO TUI
return Ok(TC_ACT_PIPE);
}
submit(RawPacket::Ip(
IpHdr::V6(header),
Expand All @@ -230,18 +259,21 @@ fn process(ctx: TcContext) -> Result<i32, ()> {
}
IpProto::Udp => {
let udphdr: *const UdpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv6Hdr::LEN)?;
let src_port = u16::from_be(unsafe { (*udphdr).source });
let dst_port = u16::from_be(unsafe { (*udphdr).dest });
let port = if is_ingress() {
u16::from_be(unsafe { (*udphdr).source })
} else {
u16::from_be(unsafe { (*udphdr).dest })
};

if filter_for_ipv6_address(src_addr, src_port, &BLOCKLIST_IPV6)
|| filter_for_ipv6_address(dst_addr, dst_port, &BLOCKLIST_IPV6)
{
return Ok(TC_ACT_SHOT);
if block_ipv6(addr, port) {
return Ok(TC_ACT_SHOT); //block packet
}

if filter_packet(Protocol::Network(NetworkProtocol::Ipv6))
|| filter_packet(Protocol::Transport(TransportProtocol::UDP))
|| filter_direction()
{
return Ok(TC_ACT_PIPE); //DONT FWD PACKET TO TUI
return Ok(TC_ACT_PIPE);
}
submit(RawPacket::Ip(
IpHdr::V6(header),
Expand Down
22 changes: 7 additions & 15 deletions oryx-tui/src/app.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,11 @@ use std::{
error,
sync::{Arc, Mutex},
thread,
time::Duration,
};

use crate::notification::Notification;
use crate::{filter::Filter, help::Help};
use crate::{filter::IoChannels, notification::Notification};
use crate::{packet::AppPacket, section::Section};

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

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

let (firewall_ingress_sender, firewall_ingress_receiver) = kanal::unbounded();
let (firewall_egress_sender, firewall_egress_receiver) = kanal::unbounded();

let firewall_channels = IoChannels::new();
thread::spawn({
let packets = packets.clone();
move || loop {
Expand All @@ -77,20 +76,11 @@ impl App {
Self {
running: true,
help: Help::new(),
filter: Filter::new(
firewall_ingress_sender.clone(),
firewall_ingress_receiver,
firewall_egress_sender.clone(),
firewall_egress_receiver,
),
filter: Filter::new(firewall_channels.clone()),
start_sniffing: false,
packets: packets.clone(),
notifications: Vec::new(),
section: Section::new(
packets.clone(),
firewall_ingress_sender,
firewall_egress_sender,
),
section: Section::new(packets.clone(), firewall_channels.clone()),
data_channel_sender: sender,
is_editing: false,
active_popup: None,
Expand Down Expand Up @@ -133,6 +123,8 @@ impl App {
}

pub fn quit(&mut self) {
self.filter.terminate();
thread::sleep(Duration::from_millis(110));
self.running = false;
}
}
Loading

0 comments on commit de9d981

Please sign in to comment.