Skip to content

Commit a30f6fc

Browse files
adgaultieradgaultier
committed
add all ports support
denoted as "*"
1 parent 48a4078 commit a30f6fc

File tree

2 files changed

+143
-63
lines changed

2 files changed

+143
-63
lines changed

‎oryx-tui/src/ebpf.rs

Lines changed: 111 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ use oryx_common::{protocols::Protocol, RawPacket};
1818
use crate::{
1919
event::Event,
2020
notification::{Notification, NotificationLevel},
21-
section::firewall::FirewallRule,
21+
section::firewall::{BlockedPort, FirewallRule},
2222
};
2323
use mio::{event::Source, unix::SourceFd, Events, Interest, Poll, Registry, Token};
2424

@@ -66,51 +66,70 @@ impl Source for RingBuffer<'_> {
6666
fn update_ipv4_blocklist(
6767
ipv4_firewall: &mut HashMap<MapData, u32, [u16; 32]>,
6868
addr: Ipv4Addr,
69-
port: u16,
69+
port: BlockedPort,
7070
enabled: bool,
7171
) {
72+
// hashmap entry exists
7273
if let Ok(mut blocked_ports) = ipv4_firewall.get(&addr.to_bits(), 0) {
73-
if enabled {
74-
// add port to blocklist
75-
if let Some(first_zero) = blocked_ports.iter().enumerate().find(|&x| *x.1 == 0) {
76-
blocked_ports[first_zero.0] = port;
77-
// dbg!("UPSERTING");
78-
// dbg!(blocked_ports[0], blocked_ports[1]);
79-
ipv4_firewall
80-
.insert(addr.to_bits(), blocked_ports, 0)
81-
.unwrap();
82-
} else {
83-
todo!(); // list is full
84-
}
85-
} else {
86-
// remove port from blocklist
87-
// on veut rebuild une blocklist avec les ports restants non nuls
88-
// par example là [8888,0,80,0,..]
89-
// hashmap = key:[0,0,0]
90-
// => [8888,80,0 ....]
91-
let non_null_ports = blocked_ports
92-
.into_iter()
93-
.filter(|p| (*p != 0 && *p != port))
94-
.collect::<Vec<u16>>();
95-
let mut blocked_ports = [0; 32];
96-
for (idx, p) in non_null_ports.iter().enumerate() {
97-
blocked_ports[idx] = *p;
74+
match port {
75+
// single port update
76+
BlockedPort::Single(port) => {
77+
if enabled {
78+
// add port to blocklist
79+
if let Some(first_zero) = blocked_ports.iter().enumerate().find(|&x| *x.1 == 0)
80+
{
81+
blocked_ports[first_zero.0] = port;
82+
// dbg!("UPSERTING");
83+
// dbg!(blocked_ports[0], blocked_ports[1]);
84+
ipv4_firewall
85+
.insert(addr.to_bits(), blocked_ports, 0)
86+
.unwrap();
87+
} else {
88+
todo!(); // list is full
89+
}
90+
} else {
91+
// remove port from blocklist
92+
// eg: remove port 53 [8888,53,80,0,..] => [8888,0,80,0,..] => [8888,80,0 ....]
93+
let non_null_ports = blocked_ports
94+
.into_iter()
95+
.filter(|p| (*p != 0 && *p != port))
96+
.collect::<Vec<u16>>();
97+
let mut blocked_ports = [0; 32];
98+
for (idx, p) in non_null_ports.iter().enumerate() {
99+
blocked_ports[idx] = *p;
100+
}
101+
if blocked_ports.iter().sum::<u16>() == 0 {
102+
//if block_list is now empty, we need to delete key
103+
ipv4_firewall.remove(&addr.to_bits()).unwrap();
104+
} else {
105+
ipv4_firewall
106+
.insert(addr.to_bits(), blocked_ports, 0)
107+
.unwrap();
108+
}
109+
}
98110
}
99-
if blocked_ports.iter().sum::<u16>() == 0 {
100-
//now block_list is empty, we need to delete key
101-
ipv4_firewall.remove(&addr.to_bits()).unwrap();
102-
} else {
103-
ipv4_firewall
104-
.insert(addr.to_bits(), blocked_ports, 0)
105-
.unwrap();
111+
BlockedPort::All => {
112+
if enabled {
113+
ipv4_firewall.insert(addr.to_bits(), [0; 32], 0).unwrap();
114+
} else {
115+
ipv4_firewall.remove(&addr.to_bits()).unwrap();
116+
}
106117
}
107118
}
108-
} else {
119+
}
120+
// no hashmap entry, create new blocklist
121+
else {
109122
// shouldn't be disabling if blocklist is empty
110123
assert!(enabled);
111-
// create new blocklist with port as first element
124+
112125
let mut blocked_ports: [u16; 32] = [0; 32];
113-
blocked_ports[0] = port;
126+
match port {
127+
BlockedPort::Single(port) => {
128+
blocked_ports[0] = port;
129+
}
130+
BlockedPort::All => {}
131+
}
132+
114133
ipv4_firewall
115134
.insert(addr.to_bits(), blocked_ports, 0)
116135
.unwrap();
@@ -120,41 +139,75 @@ fn update_ipv4_blocklist(
120139
fn update_ipv6_blocklist(
121140
ipv6_firewall: &mut HashMap<MapData, u128, [u16; 32]>,
122141
addr: Ipv6Addr,
123-
port: u16,
142+
port: BlockedPort,
124143
enabled: bool,
125144
) {
145+
// hashmap entry exists
126146
if let Ok(mut blocked_ports) = ipv6_firewall.get(&addr.to_bits(), 0) {
127-
if enabled {
128-
// add port to blocklist
129-
if let Some(first_zero) = blocked_ports.iter().enumerate().find(|&x| *x.1 == 0) {
130-
blocked_ports[first_zero.0] = port;
131-
ipv6_firewall
132-
.insert(addr.to_bits(), blocked_ports, 0)
133-
.unwrap();
134-
} else {
135-
todo!(); // list is full
147+
match port {
148+
// single port update
149+
BlockedPort::Single(port) => {
150+
if enabled {
151+
// add port to blocklist
152+
if let Some(first_zero) = blocked_ports.iter().enumerate().find(|&x| *x.1 == 0)
153+
{
154+
blocked_ports[first_zero.0] = port;
155+
// dbg!("UPSERTING");
156+
// dbg!(blocked_ports[0], blocked_ports[1]);
157+
ipv6_firewall
158+
.insert(addr.to_bits(), blocked_ports, 0)
159+
.unwrap();
160+
} else {
161+
todo!(); // list is full
162+
}
163+
} else {
164+
// remove port from blocklist
165+
// eg: remove port 53 [8888,53,80,0,..] => [8888,0,80,0,..] => [8888,80,0 ....]
166+
let non_null_ports = blocked_ports
167+
.into_iter()
168+
.filter(|p| (*p != 0 && *p != port))
169+
.collect::<Vec<u16>>();
170+
let mut blocked_ports = [0; 32];
171+
for (idx, p) in non_null_ports.iter().enumerate() {
172+
blocked_ports[idx] = *p;
173+
}
174+
if blocked_ports.iter().sum::<u16>() == 0 {
175+
//if block_list is now empty, we need to delete key
176+
ipv6_firewall.remove(&addr.to_bits()).unwrap();
177+
} else {
178+
ipv6_firewall
179+
.insert(addr.to_bits(), blocked_ports, 0)
180+
.unwrap();
181+
}
182+
}
136183
}
137-
} else {
138-
// remove port from blocklist
139-
if let Some(matching_port) = blocked_ports.iter().enumerate().find(|&x| *x.1 == port) {
140-
blocked_ports[matching_port.0] = 0;
141-
ipv6_firewall
142-
.insert(addr.to_bits(), blocked_ports, 0)
143-
.unwrap();
184+
BlockedPort::All => {
185+
if enabled {
186+
ipv6_firewall.insert(addr.to_bits(), [0; 32], 0).unwrap();
187+
} else {
188+
ipv6_firewall.remove(&addr.to_bits()).unwrap();
189+
}
144190
}
145191
}
146-
} else {
192+
}
193+
// no hashmap entry, create new blocklist
194+
else {
147195
// shouldn't be disabling if blocklist is empty
148196
assert!(enabled);
149-
//create new blocklist with port as first element
197+
150198
let mut blocked_ports: [u16; 32] = [0; 32];
151-
blocked_ports[0] = port;
199+
match port {
200+
BlockedPort::Single(port) => {
201+
blocked_ports[0] = port;
202+
}
203+
BlockedPort::All => {}
204+
}
205+
152206
ipv6_firewall
153207
.insert(addr.to_bits(), blocked_ports, 0)
154208
.unwrap();
155209
}
156210
}
157-
158211
impl Ebpf {
159212
pub fn load_ingress(
160213
iface: String,

‎oryx-tui/src/section/firewall.rs

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use ratatui::{
77
widgets::{Block, Borders, Cell, Clear, HighlightSpacing, Padding, Row, Table, TableState},
88
Frame,
99
};
10-
use std::{net::IpAddr, str::FromStr};
10+
use std::{net::IpAddr, num::ParseIntError, str::FromStr};
1111
use tui_input::{backend::crossterm::EventHandler, Input};
1212
use uuid;
1313

@@ -19,7 +19,33 @@ pub struct FirewallRule {
1919
name: String,
2020
pub enabled: bool,
2121
pub ip: IpAddr,
22-
pub port: u16,
22+
pub port: BlockedPort,
23+
}
24+
25+
#[derive(Debug, Clone)]
26+
pub enum BlockedPort {
27+
Single(u16),
28+
All,
29+
}
30+
31+
impl Display for BlockedPort {
32+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33+
match self {
34+
BlockedPort::Single(p) => write!(f, "{}", p),
35+
BlockedPort::All => write!(f, "*"),
36+
}
37+
}
38+
}
39+
40+
impl FromStr for BlockedPort {
41+
type Err = ParseIntError;
42+
fn from_str(s: &str) -> Result<Self, Self::Err> {
43+
if s == "*" {
44+
return Ok(BlockedPort::All);
45+
} else {
46+
return Ok(BlockedPort::Single(u16::from_str(s)?));
47+
}
48+
}
2349
}
2450

2551
impl Display for FirewallRule {
@@ -80,7 +106,7 @@ impl UserInput {
80106
self.port.error = None;
81107
if self.port.field.value().is_empty() {
82108
self.port.error = Some("Required field.".to_string());
83-
} else if u16::from_str(self.port.field.value()).is_err() {
109+
} else if BlockedPort::from_str(self.port.field.value()).is_err() {
84110
self.port.error = Some("Invalid Port number.".to_string());
85111
}
86112
}
@@ -280,13 +306,14 @@ impl Firewall {
280306
// update rule with user input
281307
rule.name = user_input.name.field.to_string();
282308
rule.ip = IpAddr::from_str(user_input.ip.field.value()).unwrap();
283-
rule.port = u16::from_str(user_input.port.field.value()).unwrap();
309+
rule.port =
310+
BlockedPort::from_str(user_input.port.field.value()).unwrap();
284311
} else {
285312
let rule = FirewallRule {
286313
id: uuid::Uuid::new_v4(),
287314
name: user_input.name.field.to_string(),
288315
ip: IpAddr::from_str(user_input.ip.field.value()).unwrap(),
289-
port: u16::from_str(user_input.port.field.value()).unwrap(),
316+
port: BlockedPort::from_str(user_input.port.field.value()).unwrap(),
290317
enabled: false,
291318
};
292319
self.rules.push(rule);

0 commit comments

Comments
 (0)