• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/python
2# @lint-avoid-python-3-compatibility-imports
3#
4# Copyright (c) 2021 Chenyue Zhou
5
6from __future__ import print_function
7import os
8import sys
9import time
10import atexit
11import argparse
12
13from bcc import BPF, BPFAttachType, lib
14
15
16examples = """examples:
17    ./sockmap.py -c /root/cgroup # attach to /root/cgroup
18"""
19parser = argparse.ArgumentParser(
20        description="pipe data across multiple sockets",
21        formatter_class=argparse.RawDescriptionHelpFormatter,
22        epilog=examples)
23parser.add_argument("-c", "--cgroup", required=True,
24        help="Specify the cgroup address. Note. must be cgroup2")
25
26bpf_text = '''
27#include <net/sock.h>
28
29#define MAX_SOCK_OPS_MAP_ENTRIES 65535
30
31struct sock_key {
32    u32 remote_ip4;
33    u32 local_ip4;
34    u32 remote_port;
35    u32 local_port;
36    u32 family;
37};
38
39BPF_SOCKHASH(sock_hash, struct sock_key, MAX_SOCK_OPS_MAP_ENTRIES);
40
41static __always_inline void bpf_sock_ops_ipv4(struct bpf_sock_ops *skops) {
42    struct sock_key skk = {
43        .remote_ip4 = skops->remote_ip4,
44        .local_ip4  = skops->local_ip4,
45        .local_port = skops->local_port,
46        .remote_port  = bpf_ntohl(skops->remote_port),
47        .family = skops->family,
48    };
49    int ret;
50
51    bpf_trace_printk("remote-port: %d, local-port: %d\\n", skk.remote_port,
52                     skk.local_port);
53    ret = sock_hash.sock_hash_update(skops, &skk, BPF_NOEXIST);
54    if (ret) {
55        bpf_trace_printk("bpf_sock_hash_update() failed. %d\\n", -ret);
56        return;
57    }
58
59    bpf_trace_printk("Sockhash op: %d, port %d --> %d\\n", skops->op,
60                     skk.local_port, skk.remote_port);
61}
62
63int bpf_sockhash(struct bpf_sock_ops *skops) {
64    u32 op = skops->op;
65
66    /* ipv4 only */
67    if (skops->family != AF_INET)
68	return 0;
69
70    switch (op) {
71        case BPF_SOCK_OPS_PASSIVE_ESTABLISHED_CB:
72        case BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB:
73            bpf_sock_ops_ipv4(skops);
74            break;
75        default:
76            break;
77    }
78
79    return 0;
80}
81
82int bpf_redir(struct sk_msg_md *msg) {
83    if (msg->family != AF_INET)
84        return SK_PASS;
85
86    if (msg->remote_ip4 != msg->local_ip4)
87        return SK_PASS;
88
89    struct sock_key skk = {
90        .remote_ip4 = msg->local_ip4,
91        .local_ip4  = msg->remote_ip4,
92        .local_port = bpf_ntohl(msg->remote_port),
93        .remote_port = msg->local_port,
94        .family = msg->family,
95    };
96    int ret = 0;
97
98    ret = sock_hash.msg_redirect_hash(msg, &skk, BPF_F_INGRESS);
99    bpf_trace_printk("try redirect port %d --> %d\\n", msg->local_port,
100                     bpf_ntohl(msg->remote_port));
101    if (ret != SK_PASS)
102        bpf_trace_printk("redirect port %d --> %d failed\\n", msg->local_port,
103                         bpf_ntohl(msg->remote_port));
104
105    return ret;
106}
107'''
108args = parser.parse_args()
109bpf = BPF(text=bpf_text)
110func_sock_ops = bpf.load_func("bpf_sockhash", bpf.SOCK_OPS)
111func_sock_redir = bpf.load_func("bpf_redir", bpf.SK_MSG)
112# raise if error
113fd = os.open(args.cgroup, os.O_RDONLY)
114map_fd = lib.bpf_table_fd(bpf.module, b"sock_hash")
115bpf.attach_func(func_sock_ops, fd, BPFAttachType.CGROUP_SOCK_OPS)
116bpf.attach_func(func_sock_redir, map_fd, BPFAttachType.SK_MSG_VERDICT)
117
118def detach_all():
119    bpf.detach_func(func_sock_ops, fd, BPFAttachType.CGROUP_SOCK_OPS)
120    bpf.detach_func(func_sock_redir, map_fd, BPFAttachType.SK_MSG_VERDICT)
121    print("Detaching...")
122
123atexit.register(detach_all)
124
125while True:
126    try:
127        bpf.trace_print()
128        sleep(1)
129    except KeyboardInterrupt:
130        sys.exit(0)
131