Skip to content

Commit ef79e4e

Browse files
marselestermarselester
committed
Draft support of -count flag in tcpconnect
1 parent 7ea6d26 commit ef79e4e

4 files changed

Lines changed: 188 additions & 117 deletions

File tree

‎cmd/tcpconnect/count.go‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
package main
2+
3+
type ipv4FlowKey struct {
4+
// SrcAddr is the source address.
5+
SrcAddr [4]byte
6+
// DstAddr is the destination address.
7+
DstAddr [4]byte
8+
// DstPort is the destination port (uint16 in C struct).
9+
// Note, network byte order is big endian.
10+
DstPort [2]byte
11+
// The struct size is 12 bytes due to 2 bytes padding: pahole -C 'ipv4_flow_key' ./cmd/tcpconnect/tcpconnect_bpfel.o
12+
_ [2]byte
13+
}
14+
15+
type ipv6FlowKey struct {
16+
// SrcAddr is the source address.
17+
SrcAddr [16]byte
18+
// DstAddr is the destination address.
19+
DstAddr [16]byte
20+
// DstPort is the destination port (uint16 in C struct).
21+
// Note, network byte order is big endian.
22+
DstPort [2]byte
23+
}

‎cmd/tcpconnect/event.go‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
package main
2+
3+
import (
4+
"bytes"
5+
"context"
6+
"encoding/binary"
7+
"fmt"
8+
"io"
9+
"log"
10+
"net"
11+
"os"
12+
13+
"github.com/cilium/ebpf"
14+
"github.com/cilium/ebpf/perf"
15+
)
16+
17+
// event represents a perf event sent to userspace from the BPF program running in the kernel.
18+
// Note, that it must match the C event struct, and both C and Go structs must be aligned the same way.
19+
type event struct {
20+
// SrcAddr is the source address.
21+
SrcAddr [16]byte
22+
// DstAddr is the destination address.
23+
DstAddr [16]byte
24+
// Comm is the process name that opened the connection.
25+
Comm [16]byte
26+
// Timestamp is the timestamp in microseconds.
27+
Timestamp uint64
28+
// AddrFam is the address family, 2 is AF_INET (IPv4), 10 is AF_INET6 (IPv6).
29+
AddrFam uint32
30+
// PID is the process ID that opened the connection.
31+
PID uint32
32+
// UID is the process user ID.
33+
UID uint32
34+
// DstPort is the destination port (uint16 in C struct).
35+
// Note, network byte order is big endian.
36+
DstPort [2]byte
37+
}
38+
39+
func printEvents(ctx context.Context, events *ebpf.Map, printTimestamp, printUID bool) {
40+
// Open a perf event reader from userspace on the PERF_EVENT_ARRAY map
41+
// defined in the BPF C program.
42+
rd, err := perf.NewReader(events, os.Getpagesize())
43+
if err != nil {
44+
log.Printf("creating perf event reader: %s", err)
45+
return
46+
}
47+
go func() {
48+
<-ctx.Done()
49+
rd.Close()
50+
}()
51+
52+
printHeader(os.Stdout, printTimestamp, printUID)
53+
54+
var startTimestamp float64
55+
for {
56+
record, err := rd.Read()
57+
if err != nil {
58+
if perf.IsClosed(err) {
59+
break
60+
}
61+
log.Printf("reading from perf event reader: %v", err)
62+
}
63+
64+
if record.LostSamples != 0 {
65+
log.Printf("ring event perf buffer is full, dropped %d samples", record.LostSamples)
66+
continue
67+
}
68+
69+
var e event
70+
err = binary.Read(
71+
bytes.NewBuffer(record.RawSample),
72+
binary.LittleEndian,
73+
&e,
74+
)
75+
if err != nil {
76+
log.Printf("failed to parse perf event: %#+v", err)
77+
continue
78+
}
79+
80+
if startTimestamp == 0 {
81+
startTimestamp = float64(e.Timestamp)
82+
}
83+
84+
printEvent(os.Stdout, &e, startTimestamp, printTimestamp, printUID)
85+
}
86+
}
87+
88+
func printHeader(w io.Writer, printTimestamp, printUID bool) {
89+
if printTimestamp {
90+
fmt.Fprintf(w, "%-9s", "TIME(s)")
91+
}
92+
if printUID {
93+
fmt.Fprintf(w, "%-6s", "UID")
94+
}
95+
fmt.Fprintf(w, "%-6s %-12s %-2s %-16s %-16s %s\n", "PID", "COMM", "IP", "SADDR", "DADDR", "DPORT")
96+
}
97+
98+
func printEvent(w io.Writer, e *event, startTimestamp float64, printTimestamp, printUID bool) {
99+
var (
100+
srcAddr, dstAddr net.IP
101+
ipVer byte
102+
)
103+
switch e.AddrFam {
104+
case 2:
105+
srcAddr = net.IP(e.SrcAddr[:4])
106+
dstAddr = net.IP(e.DstAddr[:4])
107+
ipVer = 4
108+
case 10:
109+
srcAddr = net.IP(e.SrcAddr[:])
110+
dstAddr = net.IP(e.DstAddr[:])
111+
ipVer = 6
112+
}
113+
114+
if printTimestamp {
115+
fmt.Fprintf(w, "%-9.3f", (float64(e.Timestamp)-startTimestamp)/1e6)
116+
}
117+
if printUID {
118+
fmt.Fprintf(w, "%-6d", e.UID)
119+
}
120+
121+
fmt.Fprintf(
122+
w,
123+
"%-6d %-12s %-2d %-16s %-16s %d\n",
124+
e.PID,
125+
bytes.TrimRight(e.Comm[:], "\x00"),
126+
ipVer,
127+
srcAddr,
128+
dstAddr,
129+
binary.BigEndian.Uint16(e.DstPort[:]),
130+
)
131+
}

cmd/tcpconnect/main_test.go renamed to cmd/tcpconnect/event_test.go

File renamed without changes.

‎cmd/tcpconnect/main.go‎

Lines changed: 34 additions & 117 deletions
Original file line numberDiff line numberDiff line change
@@ -20,19 +20,14 @@ This works by tracing the tcp_v4_connect() and tcp_v6_connect() kernel functions
2020
package main
2121

2222
import (
23-
"bytes"
24-
"encoding/binary"
23+
"context"
2524
"flag"
26-
"fmt"
27-
"io"
2825
"log"
29-
"net"
3026
"os"
3127
"os/signal"
3228
"syscall"
3329

3430
"github.com/cilium/ebpf/link"
35-
"github.com/cilium/ebpf/perf"
3631
"golang.org/x/sys/unix"
3732
)
3833

@@ -47,6 +42,7 @@ func main() {
4742
var (
4843
printTimestamp = flag.Bool("timestamp", false, "include the time of the connect in seconds on output, counting from the first event seen")
4944
printUID = flag.Bool("print-uid", false, "include UID on output")
45+
printCount = flag.Bool("count", false, "count connects per source IP and destination IP/port")
5046
filterUID = flag.Int("uid", -1, "trace this UID only")
5147
filterPID = flag.Int("pid", 0, "trace this PID only")
5248
)
@@ -62,13 +58,17 @@ func main() {
6258
return
6359
}
6460

65-
// Replace constants in the BPF C program to filter connections by PID or UID.
6661
spec, err := LoadTCPConnect()
6762
if err != nil {
6863
log.Printf("failed to load collection spec: %v", err)
6964
return
7065
}
66+
67+
// Replace constants in the BPF C program to filter connections by PID or UID.
7168
bpfConst := make(map[string]interface{})
69+
if *printCount {
70+
bpfConst["do_count"] = *printCount
71+
}
7272
if *filterUID > -1 {
7373
bpfConst["filter_uid"] = uint32(*filterUID)
7474
}
@@ -115,124 +115,41 @@ func main() {
115115
}
116116
defer tcpv6krp.Close()
117117

118-
// Open a perf event reader from userspace on the PERF_EVENT_ARRAY map
119-
// defined in the BPF C program.
120-
rd, err := perf.NewReader(objs.TCPConnectMaps.Events, os.Getpagesize())
121-
if err != nil {
122-
log.Printf("creating perf event reader: %s", err)
123-
return
124-
}
125-
defer rd.Close()
126-
127-
sig := make(chan os.Signal, 1)
128-
signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
129-
go func() {
130-
<-sig
131-
rd.Close()
132-
}()
133-
134-
printHeader(os.Stdout, *printTimestamp, *printUID)
135-
136-
var startTimestamp float64
137-
for {
138-
record, err := rd.Read()
139-
if err != nil {
140-
if perf.IsClosed(err) {
141-
break
142-
}
143-
log.Printf("reading from perf event reader: %v", err)
144-
}
118+
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
119+
defer stop()
145120

146-
if record.LostSamples != 0 {
147-
log.Printf("ring event perf buffer is full, dropped %d samples", record.LostSamples)
148-
continue
121+
if *printCount {
122+
var (
123+
ipv4key ipv4FlowKey
124+
ipv4value uint64
125+
)
126+
entries := objs.TCPConnectMaps.Ipv4Count.Iterate()
127+
for entries.Next(&ipv4key, &ipv4value) {
128+
log.Println(ipv4key, ipv4value)
129+
}
130+
if err := entries.Err(); err != nil {
131+
panic(err)
149132
}
150133

151-
var e event
152-
err = binary.Read(
153-
bytes.NewBuffer(record.RawSample),
154-
binary.LittleEndian,
155-
&e,
134+
var (
135+
ipv6key ipv6FlowKey
136+
ipv6value uint64
156137
)
157-
if err != nil {
158-
log.Printf("failed to parse perf event: %#+v", err)
159-
continue
138+
entries = objs.TCPConnectMaps.Ipv6Count.Iterate()
139+
for entries.Next(&ipv6key, &ipv6value) {
140+
log.Println(ipv6key, ipv6value)
160141
}
161-
162-
if startTimestamp == 0 {
163-
startTimestamp = float64(e.Timestamp)
142+
if err := entries.Err(); err != nil {
143+
panic(err)
164144
}
165-
166-
printEvent(os.Stdout, &e, startTimestamp, *printTimestamp, *printUID)
145+
// if err := objs.TCPConnectMaps.Ipv4Count.Lookup(&key, &value); err != nil {
146+
// log.Printf("failed to read IP v4 map: %s", err)
147+
// return
148+
// }
149+
} else {
150+
printEvents(ctx, objs.TCPConnectMaps.Events, *printTimestamp, *printUID)
167151
}
168152

169153
// The program terminates successfully if it received INT/TERM signal.
170154
exitCode = 0
171155
}
172-
173-
// event represents a perf event sent to userspace from the BPF program running in the kernel.
174-
// Note, that it must match the C event struct, and both C and Go structs must be aligned the same way.
175-
type event struct {
176-
// SrcAddr is the source address.
177-
SrcAddr [16]byte
178-
// DstAddr is the destination address.
179-
DstAddr [16]byte
180-
// Comm is the process name that opened the connection.
181-
Comm [16]byte
182-
// Timestamp is the timestamp in microseconds.
183-
Timestamp uint64
184-
// AddrFam is the address family, 2 is AF_INET (IPv4), 10 is AF_INET6 (IPv6).
185-
AddrFam uint32
186-
// PID is the process ID that opened the connection.
187-
PID uint32
188-
// UID is the process user ID.
189-
UID uint32
190-
// DstPort is the destination port (uint16 in C struct).
191-
// Note, network byte order is big endian.
192-
DstPort [2]byte
193-
}
194-
195-
func printHeader(w io.Writer, printTimestamp, printUID bool) {
196-
if printTimestamp {
197-
fmt.Fprintf(w, "%-9s", "TIME(s)")
198-
}
199-
if printUID {
200-
fmt.Fprintf(w, "%-6s", "UID")
201-
}
202-
fmt.Fprintf(w, "%-6s %-12s %-2s %-16s %-16s %s\n", "PID", "COMM", "IP", "SADDR", "DADDR", "DPORT")
203-
}
204-
205-
func printEvent(w io.Writer, e *event, startTimestamp float64, printTimestamp, printUID bool) {
206-
var (
207-
srcAddr, dstAddr net.IP
208-
ipVer byte
209-
)
210-
switch e.AddrFam {
211-
case 2:
212-
srcAddr = net.IP(e.SrcAddr[:4])
213-
dstAddr = net.IP(e.DstAddr[:4])
214-
ipVer = 4
215-
case 10:
216-
srcAddr = net.IP(e.SrcAddr[:])
217-
dstAddr = net.IP(e.DstAddr[:])
218-
ipVer = 6
219-
}
220-
221-
if printTimestamp {
222-
fmt.Fprintf(w, "%-9.3f", (float64(e.Timestamp)-startTimestamp)/1e6)
223-
}
224-
if printUID {
225-
fmt.Fprintf(w, "%-6d", e.UID)
226-
}
227-
228-
fmt.Fprintf(
229-
w,
230-
"%-6d %-12s %-2d %-16s %-16s %d\n",
231-
e.PID,
232-
bytes.TrimRight(e.Comm[:], "\x00"),
233-
ipVer,
234-
srcAddr,
235-
dstAddr,
236-
binary.BigEndian.Uint16(e.DstPort[:]),
237-
)
238-
}

0 commit comments

Comments
 (0)