diff --git a/probe/endpoint/ebpf.go b/probe/endpoint/ebpf.go index f9d57dcec..f7bbd6861 100644 --- a/probe/endpoint/ebpf.go +++ b/probe/endpoint/ebpf.go @@ -28,13 +28,14 @@ type ebpfConnection struct { // Closed connections are kept in the `closedConnections` slice for one iteration of `walkConnections`. type EbpfTracker struct { sync.Mutex - tracer *tracer.Tracer - readyToHandleConnections bool - dead bool - lastTimestampV4 uint64 + tracer *tracer.Tracer + ready bool + dead bool + lastTimestampV4 uint64 openConnections map[fourTuple]ebpfConnection closedConnections []ebpfConnection + closedDuringInit map[fourTuple]struct{} } var releaseRegex = regexp.MustCompile(`^(\d+)\.(\d+).*$`) @@ -77,7 +78,8 @@ func newEbpfTracker() (*EbpfTracker, error) { } tracker := &EbpfTracker{ - openConnections: map[fourTuple]ebpfConnection{}, + openConnections: map[fourTuple]ebpfConnection{}, + closedDuringInit: map[fourTuple]struct{}{}, } tracer, err := tracer.NewTracer(tracker.tcpEventCbV4, tracker.tcpEventCbV6, tracker.lostCb) @@ -194,10 +196,6 @@ func (t *EbpfTracker) handleConnection(ev tracer.EventType, tuple fourTuple, pid t.Lock() defer t.Unlock() - if !t.isReadyToHandleConnections() { - return - } - log.Debugf("handleConnection(%v, [%v:%v --> %v:%v], pid=%v, netNS=%v)", ev, tuple.fromAddr, tuple.fromPort, tuple.toAddr, tuple.toPort, pid, networkNamespace) @@ -217,6 +215,9 @@ func (t *EbpfTracker) handleConnection(ev tracer.EventType, tuple fourTuple, pid networkNamespace: networkNamespace, } case tracer.EventClose: + if !t.ready { + t.closedDuringInit[tuple] = struct{}{} + } if deadConn, ok := t.openConnections[tuple]; ok { delete(t.openConnections, tuple) t.closedConnections = append(t.closedConnections, deadConn) @@ -247,14 +248,19 @@ func (t *EbpfTracker) feedInitialConnections(conns procspy.ConnIter, seenTuples t.Lock() for conn := conns.Next(); conn != nil; conn = conns.Next() { tuple, namespaceID, incoming := connectionTuple(conn, seenTuples) - t.openConnections[tuple] = ebpfConnection{ - incoming: incoming, - tuple: tuple, - pid: int(conn.Proc.PID), - networkNamespace: namespaceID, + if _, ok := t.closedDuringInit[tuple]; !ok { + if _, ok := t.openConnections[tuple]; !ok { + t.openConnections[tuple] = ebpfConnection{ + incoming: incoming, + tuple: tuple, + pid: int(conn.Proc.PID), + networkNamespace: namespaceID, + } + } } } - t.readyToHandleConnections = true + t.closedDuringInit = nil + t.ready = true t.Unlock() for _, p := range processesWaitingInAccept { @@ -263,10 +269,6 @@ func (t *EbpfTracker) feedInitialConnections(conns procspy.ConnIter, seenTuples } } -func (t *EbpfTracker) isReadyToHandleConnections() bool { - return t.readyToHandleConnections -} - func (t *EbpfTracker) isDead() bool { t.Lock() defer t.Unlock() diff --git a/probe/endpoint/ebpf_test.go b/probe/endpoint/ebpf_test.go index 24d96aab5..6586cf9d8 100644 --- a/probe/endpoint/ebpf_test.go +++ b/probe/endpoint/ebpf_test.go @@ -9,6 +9,15 @@ import ( "github.com/weaveworks/tcptracer-bpf/pkg/tracer" ) +func newMockEbpfTracker() *EbpfTracker { + return &EbpfTracker{ + ready: true, + dead: false, + + openConnections: map[fourTuple]ebpfConnection{}, + } +} + func TestHandleConnection(t *testing.T) { var ( ServerPid uint32 = 42 @@ -92,13 +101,7 @@ func TestHandleConnection(t *testing.T) { } ) - mockEbpfTracker := &EbpfTracker{ - readyToHandleConnections: true, - dead: false, - - openConnections: map[fourTuple]ebpfConnection{}, - closedConnections: []ebpfConnection{}, - } + mockEbpfTracker := newMockEbpfTracker() tuple := fourTuple{IPv4ConnectEvent.SAddr.String(), IPv4ConnectEvent.DAddr.String(), uint16(IPv4ConnectEvent.SPort), uint16(IPv4ConnectEvent.DPort)} mockEbpfTracker.handleConnection(IPv4ConnectEvent.Type, tuple, int(IPv4ConnectEvent.Pid), strconv.FormatUint(uint64(IPv4ConnectEvent.NetNS), 10)) @@ -114,13 +117,7 @@ func TestHandleConnection(t *testing.T) { mockEbpfTracker.openConnections[tuple]) } - mockEbpfTracker = &EbpfTracker{ - readyToHandleConnections: true, - dead: false, - - openConnections: map[fourTuple]ebpfConnection{}, - closedConnections: []ebpfConnection{}, - } + mockEbpfTracker = newMockEbpfTracker() tuple = fourTuple{IPv4AcceptEvent.SAddr.String(), IPv4AcceptEvent.DAddr.String(), uint16(IPv4AcceptEvent.SPort), uint16(IPv4AcceptEvent.DPort)} mockEbpfTracker.handleConnection(IPv4AcceptEvent.Type, tuple, int(IPv4AcceptEvent.Pid), strconv.FormatUint(uint64(IPv4AcceptEvent.NetNS), 10)) @@ -155,26 +152,20 @@ func TestWalkConnections(t *testing.T) { toPort: 0, } ) - mockEbpfTracker := &EbpfTracker{ - readyToHandleConnections: true, - dead: false, - openConnections: map[fourTuple]ebpfConnection{ - activeTuple: { - tuple: activeTuple, - networkNamespace: "12345", - incoming: true, - pid: 0, - }, - }, - closedConnections: []ebpfConnection{ - { - tuple: inactiveTuple, - networkNamespace: "12345", - incoming: false, - pid: 0, - }, - }, + mockEbpfTracker := newMockEbpfTracker() + mockEbpfTracker.openConnections[activeTuple] = ebpfConnection{ + tuple: activeTuple, + networkNamespace: "12345", + incoming: true, + pid: 0, } + mockEbpfTracker.closedConnections = append(mockEbpfTracker.closedConnections, + ebpfConnection{ + tuple: inactiveTuple, + networkNamespace: "12345", + incoming: false, + pid: 0, + }) mockEbpfTracker.walkConnections(func(e ebpfConnection) { cnt++ }) @@ -204,11 +195,7 @@ func TestInvalidTimeStampDead(t *testing.T) { NetNS: NetNS, } ) - mockEbpfTracker := &EbpfTracker{ - readyToHandleConnections: true, - dead: false, - openConnections: map[fourTuple]ebpfConnection{}, - } + mockEbpfTracker := newMockEbpfTracker() event.Timestamp = 0 mockEbpfTracker.tcpEventCbV4(event) event2 := event