forked from iovisor/bcc
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
bcc-python: support attach_func() and detach_func() (iovisor#3479)
- support attach_func() and detach_func(). - add an sockmap issue to demonstrate using these two functions.
- Loading branch information
1 parent
ce7e523
commit 5934161
Showing
4 changed files
with
230 additions
and
1 deletion.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,130 @@ | ||
#!/usr/bin/python | ||
# @lint-avoid-python-3-compatibility-imports | ||
# | ||
# Copyright (c) 2021 Chenyue Zhou | ||
|
||
from __future__ import print_function | ||
import os | ||
import sys | ||
import time | ||
import atexit | ||
import argparse | ||
|
||
from bcc import BPF, lib | ||
|
||
|
||
examples = """examples: | ||
./sockmap.py -c /root/cgroup # attach to /root/cgroup | ||
""" | ||
parser = argparse.ArgumentParser( | ||
description="pipe data across multiple sockets", | ||
formatter_class=argparse.RawDescriptionHelpFormatter, | ||
epilog=examples) | ||
parser.add_argument("-c", "--cgroup", required=True, | ||
help="Specify the cgroup address. Note. must be cgroup2") | ||
|
||
bpf_text = ''' | ||
#include <net/sock.h> | ||
#define MAX_SOCK_OPS_MAP_ENTRIES 65535 | ||
struct sock_key { | ||
u32 remote_ip4; | ||
u32 local_ip4; | ||
u32 remote_port; | ||
u32 local_port; | ||
u32 family; | ||
}; | ||
BPF_SOCKHASH(sock_hash, struct sock_key, MAX_SOCK_OPS_MAP_ENTRIES); | ||
static __always_inline void bpf_sock_ops_ipv4(struct bpf_sock_ops *skops) { | ||
struct sock_key skk = { | ||
.remote_ip4 = skops->remote_ip4, | ||
.local_ip4 = skops->local_ip4, | ||
.local_port = skops->local_port, | ||
.remote_port = bpf_ntohl(skops->remote_port), | ||
.family = skops->family, | ||
}; | ||
int ret; | ||
bpf_trace_printk("remote-port: %d, local-port: %d\\n", skk.remote_port, | ||
skk.local_port); | ||
ret = sock_hash.sock_hash_update(skops, &skk, BPF_NOEXIST); | ||
if (ret) { | ||
bpf_trace_printk("bpf_sock_hash_update() failed. %d\\n", -ret); | ||
return; | ||
} | ||
bpf_trace_printk("Sockhash op: %d, port %d --> %d\\n", skops->op, | ||
skk.local_port, skk.remote_port); | ||
} | ||
int bpf_sockhash(struct bpf_sock_ops *skops) { | ||
u32 op = skops->op; | ||
/* ipv4 only */ | ||
if (skops->family != AF_INET) | ||
return 0; | ||
switch (op) { | ||
case BPF_SOCK_OPS_PASSIVE_ESTABLISHED_CB: | ||
case BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB: | ||
bpf_sock_ops_ipv4(skops); | ||
break; | ||
default: | ||
break; | ||
} | ||
return 0; | ||
} | ||
int bpf_redir(struct sk_msg_md *msg) { | ||
if (msg->family != AF_INET) | ||
return SK_PASS; | ||
if (msg->remote_ip4 != msg->local_ip4) | ||
return SK_PASS; | ||
struct sock_key skk = { | ||
.remote_ip4 = msg->local_ip4, | ||
.local_ip4 = msg->remote_ip4, | ||
.local_port = bpf_ntohl(msg->remote_port), | ||
.remote_port = msg->local_port, | ||
.family = msg->family, | ||
}; | ||
int ret = 0; | ||
ret = sock_hash.msg_redirect_hash(msg, &skk, BPF_F_INGRESS); | ||
bpf_trace_printk("try redirect port %d --> %d\\n", msg->local_port, | ||
bpf_ntohl(msg->remote_port)); | ||
if (ret != SK_PASS) | ||
bpf_trace_printk("redirect port %d --> %d failed\\n", msg->local_port, | ||
bpf_ntohl(msg->remote_port)); | ||
return ret; | ||
} | ||
''' | ||
args = parser.parse_args() | ||
bpf = BPF(text=bpf_text) | ||
func_sock_ops = bpf.load_func("bpf_sockhash", bpf.SOCK_OPS) | ||
func_sock_redir = bpf.load_func("bpf_redir", bpf.SK_MSG) | ||
# raise if error | ||
fd = os.open(args.cgroup, os.O_RDONLY) | ||
map_fd = lib.bpf_table_fd(bpf.module, b"sock_hash") | ||
bpf.attach_func(func_sock_ops, fd, bpf.CGROUP_SOCK_OPS) | ||
bpf.attach_func(func_sock_redir, map_fd, bpf.SK_MSG_VERDICT) | ||
|
||
def detach_all(): | ||
bpf.detach_func(func_sock_ops, fd, bpf.CGROUP_SOCK_OPS) | ||
bpf.detach_func(func_sock_redir, map_fd, bpf.SK_MSG_VERDICT) | ||
print("Detaching...") | ||
|
||
atexit.register(detach_all) | ||
|
||
while True: | ||
try: | ||
bpf.trace_print() | ||
sleep(1) | ||
except KeyboardInterrupt: | ||
sys.exit(0) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters