Skip to content

Commit 461c6ab

Browse files
hack3richack3ric
committed
bpf: add branch hints
Basically 649b78a, but without [[clang::{,un}likely]] attributes.
1 parent a98b0db commit 461c6ab

File tree

6 files changed

+142
-128
lines changed

6 files changed

+142
-128
lines changed

‎bpf/egress.c

+26-23
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,11 @@ static inline int mangle_data(struct __sk_buff* skb, __u16 offset, __be32* csum_
1515
try_shot(bpf_skb_change_tail(skb, skb->len + TCP_UDP_HEADER_DIFF, 0));
1616
__u8 buf[TCP_UDP_HEADER_DIFF + 4] = {};
1717
__u32 copy_len = min(data_len, TCP_UDP_HEADER_DIFF);
18-
if (copy_len > 0) {
18+
if (likely(copy_len > 0)) {
1919
// HACK: make verifier happy
2020
// Probably related:
2121
// https://lore.kernel.org/bpf/[email protected]/T/
22-
if (copy_len < 2) copy_len = 1;
22+
if (unlikely(copy_len < 2)) copy_len = 1;
2323

2424
try_shot(bpf_skb_load_bytes(skb, offset, buf + 1, copy_len));
2525
try_shot(bpf_skb_store_bytes(skb, skb->len - copy_len, buf + 1, copy_len, 0));
@@ -41,7 +41,7 @@ static inline void update_tcp_header(struct tcphdr* tcp, __u16 payload_len, __u3
4141
tcp->doff = 5;
4242
tcp->window = htons(cwnd >> CWND_SCALE);
4343
tcp->ack = true;
44-
if (payload_len == 0) tcp->psh = true;
44+
tcp->psh = payload_len == 0;
4545
tcp->urg_ptr = 0;
4646
}
4747

@@ -54,22 +54,25 @@ int egress_handler(struct __sk_buff* skb) {
5454
struct ipv6hdr* ipv6 = NULL;
5555
__u32 ip_end, nexthdr = 0;
5656

57-
if (eth_proto == ETH_P_IP) {
58-
redecl_shot(struct iphdr, ipv4, ETH_HLEN, skb);
59-
ip_end = ETH_HLEN + (ipv4->ihl << 2);
60-
} else if (eth_proto == ETH_P_IPV6) {
61-
redecl_shot(struct ipv6hdr, ipv6, ETH_HLEN, skb);
62-
nexthdr = ipv6->nexthdr;
63-
ip_end = ETH_HLEN + sizeof(*ipv6);
64-
struct ipv6_opt_hdr* opt = NULL;
65-
for (int i = 0; i < 8; i++) {
66-
if (!ipv6_is_ext(nexthdr)) break;
67-
redecl_drop(struct ipv6_opt_hdr, opt, ip_end, skb);
68-
nexthdr = opt->nexthdr;
69-
ip_end += (opt->hdrlen + 1) << 3;
70-
}
71-
} else {
72-
return TC_ACT_OK;
57+
switch (eth_proto) {
58+
case ETH_P_IP:
59+
redecl_shot(struct iphdr, ipv4, ETH_HLEN, skb);
60+
ip_end = ETH_HLEN + (ipv4->ihl << 2);
61+
break;
62+
case ETH_P_IPV6:
63+
redecl_shot(struct ipv6hdr, ipv6, ETH_HLEN, skb);
64+
nexthdr = ipv6->nexthdr;
65+
ip_end = ETH_HLEN + sizeof(*ipv6);
66+
struct ipv6_opt_hdr* opt = NULL;
67+
for (int i = 0; i < 8; i++) {
68+
if (!ipv6_is_ext(nexthdr)) break;
69+
redecl_drop(struct ipv6_opt_hdr, opt, ip_end, skb);
70+
nexthdr = opt->nexthdr;
71+
ip_end += (opt->hdrlen + 1) << 3;
72+
}
73+
break;
74+
default:
75+
return TC_ACT_OK;
7376
}
7477

7578
__u8 ip_proto = ipv4 ? ipv4->protocol : ipv6 ? nexthdr : 0;
@@ -82,7 +85,7 @@ int egress_handler(struct __sk_buff* skb) {
8285
if (!settings) return TC_ACT_OK;
8386
struct conn_tuple conn_key = gen_conn_key(QUARTET_UDP);
8487
struct connection* conn = bpf_map_lookup_elem(&mimic_conns, &conn_key);
85-
if (!conn) {
88+
if (unlikely(!conn)) {
8689
if (settings->hi == 0) return TC_ACT_STOLEN; // passive mode
8790
struct connection conn_value = conn_init(settings, tstamp);
8891
try_shot(bpf_map_update_elem(&mimic_conns, &conn_key, &conn_value, BPF_ANY));
@@ -98,7 +101,7 @@ int egress_handler(struct __sk_buff* skb) {
98101
__u32 random = bpf_get_prandom_u32();
99102

100103
bpf_spin_lock(&conn->lock);
101-
if (conn->state == CONN_ESTABLISHED) {
104+
if (likely(conn->state == CONN_ESTABLISHED)) {
102105
seq = conn->seq;
103106
ack_seq = conn->ack_seq;
104107
conn->seq += payload_len;
@@ -115,7 +118,7 @@ int egress_handler(struct __sk_buff* skb) {
115118
ack_seq = conn->ack_seq = 0;
116119
conn->retry_tstamp = conn->reset_tstamp = tstamp;
117120
bpf_spin_unlock(&conn->lock);
118-
log_conn(LOG_INFO, LOG_CONN_INIT, &conn_key);
121+
log_conn(LOG_CONN_INIT, &conn_key);
119122
send_ctrl_packet(&conn_key, TCP_FLAG_SYN, seq, ack_seq, 0xffff);
120123
} else {
121124
bpf_spin_unlock(&conn->lock);
@@ -165,7 +168,7 @@ int egress_handler(struct __sk_buff* skb) {
165168

166169
__u32 csum_off = ip_end + offsetof(struct tcphdr, check);
167170
redecl_shot(struct tcphdr, tcp, ip_end, skb);
168-
log_tcp(LOG_TRACE, false, &conn_key, tcp, payload_len);
171+
log_tcp(false, &conn_key, tcp, payload_len);
169172

170173
tcp->check = 0;
171174
csum_diff =

‎bpf/ingress.c

+45-42
Original file line numberDiff line numberDiff line change
@@ -49,30 +49,30 @@ static inline int read_tcp_options(struct xdp_md* xdp, struct tcphdr* tcp, __u32
4949
struct tcp_options* opt) {
5050
__u8 opt_buf[80] = {};
5151
__u32 len = (tcp->doff << 2) - sizeof(*tcp);
52-
if (len > 80) { // TCP options too large
52+
if (unlikely(len > 80)) { // TCP options too large
5353
return XDP_DROP;
5454
} else if (len == 0) { // prevent zero-sized read
5555
return XDP_PASS;
5656
} else {
57-
if (len < 2) len = 1;
57+
if (unlikely(len < 2)) len = 1;
5858
try_drop(bpf_xdp_load_bytes(xdp, ip_end + sizeof(*tcp), opt_buf, len));
5959
}
6060

6161
for (__u32 i = 0; i < len; i++) {
62-
if (i > 80 - 1) return XDP_DROP;
62+
if (unlikely(i > 80 - 1)) return XDP_DROP;
6363
switch (opt_buf[i]) {
6464
case 0: // end of option list
6565
case 1: // no-op
6666
break;
6767
case 2: // MSS
68-
if (i > 80 - 4 || opt_buf[i + 1] != 4) return XDP_DROP;
68+
if (unlikely(i > 80 - 4 || opt_buf[i + 1] != 4)) return XDP_DROP;
6969
opt->mss = (opt_buf[i + 2] << 8) + opt_buf[i + 3];
7070
i += 3;
7171
break;
7272
default:
7373
// HACK: `80 - 2` -> `80 - 3`
7474
// mimic.bpf.o compiled with LLVM 18 failed eBPF verifier in Linux 6.6 or lower.
75-
if (i > 80 - 3) return XDP_DROP;
75+
if (unlikely(i > 80 - 3)) return XDP_DROP;
7676
__u8 l = opt_buf[i + 1];
7777
if (l < 2 || i + l > len) return XDP_DROP;
7878
i += l - 1;
@@ -92,24 +92,27 @@ int ingress_handler(struct xdp_md* xdp) {
9292
struct ipv6hdr* ipv6 = NULL;
9393
__u32 ip_end, ip_payload_len, nexthdr = 0;
9494

95-
if (eth_proto == ETH_P_IP) {
96-
redecl_drop(struct iphdr, ipv4, ETH_HLEN, xdp);
97-
ip_end = ETH_HLEN + (ipv4->ihl << 2);
98-
ip_payload_len = ntohs(ipv4->tot_len) - (ipv4->ihl << 2);
99-
} else if (eth_proto == ETH_P_IPV6) {
100-
redecl_drop(struct ipv6hdr, ipv6, ETH_HLEN, xdp);
101-
nexthdr = ipv6->nexthdr;
102-
ip_end = ETH_HLEN + sizeof(*ipv6);
103-
struct ipv6_opt_hdr* opt = NULL;
104-
for (int i = 0; i < 8; i++) {
105-
if (!ipv6_is_ext(nexthdr)) break;
106-
redecl_drop(struct ipv6_opt_hdr, opt, ip_end, xdp);
107-
nexthdr = opt->nexthdr;
108-
ip_end += (opt->hdrlen + 1) << 3;
109-
}
110-
ip_payload_len = ntohs(ipv6->payload_len);
111-
} else {
112-
return XDP_PASS;
95+
switch (eth_proto) {
96+
case ETH_P_IP:
97+
redecl_drop(struct iphdr, ipv4, ETH_HLEN, xdp);
98+
ip_end = ETH_HLEN + (ipv4->ihl << 2);
99+
ip_payload_len = ntohs(ipv4->tot_len) - (ipv4->ihl << 2);
100+
break;
101+
case ETH_P_IPV6:
102+
redecl_drop(struct ipv6hdr, ipv6, ETH_HLEN, xdp);
103+
nexthdr = ipv6->nexthdr;
104+
ip_end = ETH_HLEN + sizeof(*ipv6);
105+
struct ipv6_opt_hdr* opt = NULL;
106+
for (int i = 0; i < 8; i++) {
107+
if (!ipv6_is_ext(nexthdr)) break;
108+
redecl_drop(struct ipv6_opt_hdr, opt, ip_end, xdp);
109+
nexthdr = opt->nexthdr;
110+
ip_end += (opt->hdrlen + 1) << 3;
111+
}
112+
ip_payload_len = ntohs(ipv6->payload_len);
113+
break;
114+
default:
115+
return XDP_PASS;
113116
}
114117

115118
__u8 ip_proto = ipv4 ? ipv4->protocol : ipv6 ? nexthdr : 0;
@@ -121,7 +124,7 @@ int ingress_handler(struct xdp_md* xdp) {
121124
struct conn_tuple conn_key = gen_conn_key(QUARTET_TCP);
122125
__u32 payload_len = ip_payload_len - (tcp->doff << 2);
123126

124-
log_tcp(LOG_TRACE, true, &conn_key, tcp, payload_len);
127+
log_tcp(true, &conn_key, tcp, payload_len);
125128
struct connection* conn = bpf_map_lookup_elem(&mimic_conns, &conn_key);
126129

127130
struct tcp_options opt = {};
@@ -136,7 +139,7 @@ int ingress_handler(struct xdp_md* xdp) {
136139
__u64 tstamp = bpf_ktime_get_boot_ns();
137140

138141
// Quick path for RST and FIN
139-
if (tcp->rst || tcp->fin) {
142+
if (unlikely(tcp->rst || tcp->fin)) {
140143
if (conn) {
141144
__u32 cooldown;
142145
bpf_spin_lock(&conn->lock);
@@ -146,16 +149,16 @@ int ingress_handler(struct xdp_md* xdp) {
146149
bpf_spin_unlock(&conn->lock);
147150
use_pktbuf(RB_ITEM_FREE_PKTBUF, pktbuf);
148151
if (tcp->rst) {
149-
log_destroy(LOG_WARN, &conn_key, DESTROY_RECV_RST, cooldown);
152+
log_destroy(&conn_key, DESTROY_RECV_RST, cooldown);
150153
} else {
151154
send_ctrl_packet(&conn_key, TCP_FLAG_RST, htonl(tcp->ack_seq), 0, 0);
152-
log_destroy(LOG_WARN, &conn_key, DESTROY_RECV_FIN, cooldown);
155+
log_destroy(&conn_key, DESTROY_RECV_FIN, cooldown);
153156
}
154157
}
155158
return XDP_DROP;
156159
}
157160

158-
if (!conn) {
161+
if (unlikely(!conn)) {
159162
if (!tcp->syn || tcp->ack) {
160163
send_ctrl_packet(&conn_key, TCP_FLAG_RST, htonl(tcp->ack_seq), 0, 0);
161164
return XDP_DROP;
@@ -180,7 +183,7 @@ int ingress_handler(struct xdp_md* xdp) {
180183

181184
switch (conn->state) {
182185
case CONN_IDLE:
183-
if (tcp->syn && !tcp->ack) {
186+
if (likely(tcp->syn && !tcp->ack)) {
184187
conn->state = CONN_SYN_RECV;
185188
flags |= TCP_FLAG_SYN | TCP_FLAG_ACK;
186189
seq = conn->seq = random;
@@ -193,14 +196,14 @@ int ingress_handler(struct xdp_md* xdp) {
193196
break;
194197

195198
case CONN_SYN_RECV:
196-
if (tcp->syn && !tcp->ack) {
199+
if (likely(tcp->syn && !tcp->ack)) {
197200
__u32 new_ack_seq = next_ack_seq(tcp, payload_len);
198-
if (new_ack_seq != conn->ack_seq) goto fsm_error;
201+
if (unlikely(new_ack_seq != conn->ack_seq)) goto fsm_error;
199202
flags |= TCP_FLAG_SYN | TCP_FLAG_ACK;
200203
seq = conn->seq++;
201204
ack_seq = new_ack_seq;
202205
conn->peer_mss = opt.mss;
203-
} else if (!tcp->syn && tcp->ack) {
206+
} else if (likely(!tcp->syn && tcp->ack)) {
204207
will_send_ctrl_packet = false;
205208
conn->state = CONN_ESTABLISHED;
206209
conn->ack_seq = next_ack_seq(tcp, payload_len);
@@ -213,9 +216,9 @@ int ingress_handler(struct xdp_md* xdp) {
213216
break;
214217

215218
case CONN_SYN_SENT:
216-
if (tcp->syn) {
219+
if (likely(tcp->syn)) {
217220
flags |= TCP_FLAG_ACK;
218-
if (tcp->ack) {
221+
if (likely(tcp->ack)) {
219222
// 3-way handshake
220223
conn->state = CONN_ESTABLISHED;
221224
conn->cooldown_mul = 0;
@@ -234,8 +237,8 @@ int ingress_handler(struct xdp_md* xdp) {
234237
}
235238
break;
236239

237-
case CONN_ESTABLISHED:
238-
if (tcp->syn) {
240+
case CONN_ESTABLISHED: // TODO: likely
241+
if (unlikely(tcp->syn)) {
239242
goto fsm_error;
240243
} else if (ntohl(tcp->seq) == conn->ack_seq - 1) {
241244
// Received keepalive; send keepalive ACK
@@ -252,7 +255,7 @@ int ingress_handler(struct xdp_md* xdp) {
252255
} else if (!tcp->psh && payload_len == 0) {
253256
// Empty segment without PSH will be treated as control packet
254257
will_send_ctrl_packet = false;
255-
} else {
258+
} else { // TODO: likely
256259
will_send_ctrl_packet = will_drop = false;
257260
conn->ack_seq += payload_len;
258261
__u32 peer_mss = conn->peer_mss ?: 1460;
@@ -269,7 +272,7 @@ int ingress_handler(struct xdp_md* xdp) {
269272
}
270273
break;
271274

272-
default:
275+
default: // TODO: unlikely
273276
fsm_error:
274277
flags |= TCP_FLAG_RST;
275278
swap(pktbuf, conn->pktbuf);
@@ -282,14 +285,14 @@ int ingress_handler(struct xdp_md* xdp) {
282285

283286
bpf_spin_unlock(&conn->lock);
284287

285-
if (flags & TCP_FLAG_SYN && flags & TCP_FLAG_ACK) log_conn(LOG_INFO, LOG_CONN_ACCEPT, &conn_key);
288+
if (flags & TCP_FLAG_SYN && flags & TCP_FLAG_ACK) log_conn(LOG_CONN_ACCEPT, &conn_key);
286289
if (will_send_ctrl_packet)
287290
send_ctrl_packet(&conn_key, flags, seq, ack_seq, flags & TCP_FLAG_RST ? 0 : cwnd);
288-
if (flags & TCP_FLAG_RST) {
289-
log_destroy(LOG_WARN, &conn_key, DESTROY_INVALID, cooldown);
291+
if (unlikely(flags & TCP_FLAG_RST)) {
292+
log_destroy(&conn_key, DESTROY_INVALID, cooldown);
290293
use_pktbuf(RB_ITEM_FREE_PKTBUF, pktbuf);
291294
} else if (newly_estab) {
292-
log_conn(LOG_INFO, LOG_CONN_ESTABLISH, &conn_key);
295+
log_conn(LOG_CONN_ESTABLISH, &conn_key);
293296
use_pktbuf(RB_ITEM_CONSUME_PKTBUF, pktbuf);
294297
}
295298
if (will_drop) return XDP_DROP;

‎bpf/main.c

+6-5
Original file line numberDiff line numberDiff line change
@@ -31,16 +31,17 @@ int send_ctrl_packet(struct conn_tuple* conn, __be32 flags, __u32 seq, __u32 ack
3131
int store_packet(struct __sk_buff* skb, __u32 pkt_off, struct conn_tuple* key, int ip_summed) {
3232
int retcode;
3333
__u32 data_len = skb->len - pkt_off;
34-
if (!key || data_len > MAX_PACKET_SIZE) return TC_ACT_SHOT;
34+
if (unlikely(!key || data_len > MAX_PACKET_SIZE)) return TC_ACT_SHOT;
3535

3636
bool has_remainder = data_len % SEGMENT_SIZE;
3737
__u32 segments = data_len / SEGMENT_SIZE + has_remainder;
3838
__u32 alloc_size = sizeof(struct rb_item) + segments * SEGMENT_SIZE;
3939
struct bpf_dynptr ptr = {};
40-
if (bpf_ringbuf_reserve_dynptr(&mimic_rb, alloc_size, 0, &ptr) < 0) cleanup(TC_ACT_SHOT);
40+
if (unlikely(bpf_ringbuf_reserve_dynptr(&mimic_rb, alloc_size, 0, &ptr) < 0))
41+
cleanup(TC_ACT_SHOT);
4142

4243
struct rb_item* item = bpf_dynptr_data(&ptr, 0, sizeof(*item));
43-
if (!item) cleanup(TC_ACT_SHOT);
44+
if (unlikely(!item)) cleanup(TC_ACT_SHOT);
4445
item->type = RB_ITEM_STORE_PACKET;
4546
item->store_packet.conn_key = *key;
4647
item->store_packet.len = data_len;
@@ -77,10 +78,10 @@ int store_packet(struct __sk_buff* skb, __u32 pkt_off, struct conn_tuple* key, i
7778

7879
// Need to manually clear conn.pktbuf in eBPF
7980
int use_pktbuf(enum rb_item_type type, __u64 buf) {
80-
if (type != RB_ITEM_CONSUME_PKTBUF && type != RB_ITEM_FREE_PKTBUF) return -1;
81+
if (unlikely(type != RB_ITEM_CONSUME_PKTBUF && type != RB_ITEM_FREE_PKTBUF)) return -1;
8182
if (!buf) return 0;
8283
struct rb_item* item = bpf_ringbuf_reserve(&mimic_rb, sizeof(*item), 0);
83-
if (!item) return -1;
84+
if (unlikely(!item)) return -1;
8485
item->type = type;
8586
item->pktbuf = buf;
8687
bpf_ringbuf_submit(item, 0);

0 commit comments

Comments
 (0)