xref: /aosp_15_r20/external/bcc/examples/tracing/nflatency.py (revision 387f9dfdfa2baef462e92476d413c7bc2470293e)
1#!/usr/bin/python3
2#
3# nflatency     Trace netfilter hook latency.
4#
5# This attaches a kprobe and kretprobe to nf_hook_slow.
6# 2020-04-03 Casey Callendrello / <[email protected]>
7
8import argparse
9import sys
10import time
11
12from bcc import BPF
13
14BPF_SRC = """
15#include <linux/ip.h>
16#include <linux/ipv6.h>
17#include <linux/tcp.h>
18#include <linux/netfilter.h>
19#include <net/ip.h>
20#include <uapi/linux/bpf.h>
21
22static struct tcphdr *skb_to_tcphdr(const struct sk_buff *skb)
23{
24    // unstable API. verify logic in tcp_hdr() -> skb_transport_header().
25    return (struct tcphdr *)(skb->head + skb->transport_header);
26}
27
28static inline struct iphdr *skb_to_iphdr(const struct sk_buff *skb)
29{
30    // unstable API. verify logic in ip_hdr() -> skb_network_header().
31    return (struct iphdr *)(skb->head + skb->network_header);
32}
33
34static inline struct ipv6hdr *skb_to_ip6hdr(const struct sk_buff *skb)
35{
36    // unstable API. verify logic in ip_hdr() -> skb_network_header().
37    return (struct ipv6hdr *)(skb->head + skb->network_header);
38}
39
40// for correlating between kprobe and kretprobe
41struct start_data {
42    u8 hook;
43    u8 pf; // netfilter protocol
44    u8 tcp_state;
45    u64 ts;
46};
47BPF_PERCPU_ARRAY(sts, struct start_data, 1);
48
49// the histogram keys
50typedef struct nf_lat_key {
51    u8 proto; // see netfilter.h
52    u8 hook;
53    u8 tcp_state;
54} nf_lat_key_t;
55
56typedef struct hist_key {
57    nf_lat_key_t key;
58    u64 slot;
59} hist_key_t;
60BPF_HISTOGRAM(dist, hist_key_t);
61
62
63int kprobe__nf_hook_slow(struct pt_regs *ctx, struct sk_buff *skb, struct nf_hook_state *state) {
64    struct start_data data = {};
65    data.ts = bpf_ktime_get_ns();
66    data.hook = state->hook;
67    data.pf = state->pf;
68
69    COND
70
71    u8 ip_proto;
72    if (skb->protocol == htons(ETH_P_IP)) {
73        struct iphdr *ip = skb_to_iphdr(skb);
74        ip_proto = ip->protocol;
75
76    } else if (skb->protocol == htons(ETH_P_IPV6)) {
77        struct ipv6hdr *ip = skb_to_ip6hdr(skb);
78        ip_proto = ip->nexthdr;
79    }
80
81    data.tcp_state = 0;
82    if (ip_proto == 0x06) { //tcp
83        struct tcphdr *tcp = skb_to_tcphdr(skb);
84        u8 tcpflags = ((u_int8_t *)tcp)[13];
85
86        // FIN or RST
87        if (((tcpflags & 1) + (tcpflags & 4)) > 0) {
88            data.tcp_state = 3;
89        }
90        // SYN / SACK
91        else if ((tcpflags & 0x02) > 0) {
92            data.tcp_state = 1;
93            if ((tcpflags & 16) > 0) { // ACK
94                data.tcp_state = 2;
95            }
96        }
97    }
98
99    u32 idx = 0;
100    sts.update(&idx, &data);
101    return 0;
102}
103
104int kretprobe__nf_hook_slow(struct pt_regs *ctx) {
105    u32 idx = 0;
106    struct start_data *s;
107    s = sts.lookup(&idx);
108    if (!s || s->ts == 0) {
109        return 0;
110    }
111
112    s->ts = bpf_ktime_get_ns() - s->ts;
113
114    hist_key_t key = {};
115    key.key.hook = s->hook;
116    key.key.proto = s->pf;
117    key.key.tcp_state = s->tcp_state;
118    key.slot = bpf_log2l(s->ts / FACTOR );
119    dist.increment(key);
120
121    s->ts = 0;
122    sts.update(&idx, s);
123
124    return 0;
125}
126"""
127
128# constants from netfilter.h
129NF_HOOKS = {
130    0: "PRE_ROUTING",
131    1: "LOCAL_IN",
132    2: "FORWARD",
133    3: "LOCAL_OUT",
134    4: "POST_ROUTING",
135}
136
137NF_PROTOS = {
138    0: "UNSPEC",
139    1: "INET",
140    2: "IPV4",
141    3: "ARP",
142    5: "NETDEV",
143    7: "BRIDGE",
144    10: "IPV6",
145    12: "DECNET",
146}
147
148TCP_FLAGS = {
149    0: "other",
150    1: "SYN",
151    2: "SACK",
152    3: "FIN",
153}
154
155EXAMPLES = """examples:
156    nflatency                   # print netfilter latency histograms, 1 second refresh
157    nflatency -p IPV4 -p IPV6   # print only for ipv4 and ipv6
158    nflatency -k PRE_ROUTING    # only record the PRE_ROUTING hook
159    nflatency -i 5 -d 10        # run for 10 seconds, printing every 5
160"""
161
162
163parser = argparse.ArgumentParser(
164    description="Track latency added by netfilter hooks. Where possible, interesting TCP flags are included",
165    formatter_class=argparse.RawDescriptionHelpFormatter,
166    epilog=EXAMPLES)
167parser.add_argument("-p", "--proto",
168                    action='append',
169                    help="record this protocol only (multiple parameters allowed)",
170                    choices=NF_PROTOS.values())
171parser.add_argument("-k", "--hook",
172                    action='append',
173                    help="record this netfilter hook only (multiple parameters allowed)",
174                    choices=NF_HOOKS.values())
175parser.add_argument("-i", "--interval", type=int,
176                    help="summary interval, in seconds. Default is 10, unless --duration is supplied")
177parser.add_argument("-d", "--duration", type=int,
178                    help="total duration of trace, in seconds")
179parser.add_argument("--nano", action="store_true",
180                    help="bucket by nanoseconds instead of milliseconds")
181
182def main():
183    args = parser.parse_args()
184
185    src = build_src(args)
186    b = BPF(text=src)
187    dist = b.get_table("dist")
188
189    seconds = 0
190    interval = 0
191    if not args.interval:
192        interval = 1
193        if args.duration:
194            interval = args.duration
195    else:
196        interval = args.interval
197
198    sys.stderr.write("Tracing netfilter hooks... Hit Ctrl-C to end.\n")
199    while 1:
200        try:
201            dist.print_log2_hist(
202                section_header="Bucket",
203                bucket_fn=lambda k: (k.proto, k.hook, k.tcp_state),
204                section_print_fn=bucket_desc)
205            if args.duration and seconds >= args.duration:
206                sys.exit(0)
207            seconds += interval
208            time.sleep(interval)
209        except KeyboardInterrupt:
210            sys.exit(1)
211
212
213def build_src(args):
214    cond_src = ""
215    if args.proto:
216        predicate = " || ".join(map(lambda x: "data.pf == NFPROTO_%s" % x, args.proto))
217        cond_src = "if (!(%s)) { return 0; }\n" % predicate
218    if args.hook:
219        predicate = " || ".join(map(lambda x: "data.hook == NF_INET_%s" % x, args.hook))
220        cond_src = "%s    if (!(%s)) { return 0;}\n" % (cond_src, predicate)
221
222    factor = "1000"
223    if args.nano:
224        factor = "1"
225
226    return BPF_SRC.replace('COND', cond_src).replace('FACTOR', factor)
227
228
229def bucket_desc(bucket):
230    return "%s %s (tcp %s)" % (
231        NF_PROTOS[bucket[0]],
232        NF_HOOKS[bucket[1]],
233        TCP_FLAGS[bucket[2]])
234
235
236if __name__ == "__main__":
237    main()
238