Merge branch 'main' into proto-ipv6-overlay

client/internal/netflow/conntrack/conntrack.go

@@ -7,7 +7,9 @@ import (
 	"fmt"
 	"net/netip"
 	"sync"
+	"time"
 
+	"github.com/cenkalti/backoff/v4"
 	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
 	nfct "github.com/ti-mo/conntrack"
@@ -17,31 +19,64 @@ import (
 	nbnet "github.com/netbirdio/netbird/client/net"
 )
 
-const defaultChannelSize = 100
+const (
+	defaultChannelSize     = 100
+	reconnectInitInterval  = 5 * time.Second
+	reconnectMaxInterval   = 5 * time.Minute
+	reconnectRandomization = 0.5
+)
+
+// listener abstracts a netlink conntrack connection for testability.
+type listener interface {
+	Listen(evChan chan<- nfct.Event, numWorkers uint8, groups []netfilter.NetlinkGroup) (chan error, error)
+	Close() error
+}
 
 // ConnTrack manages kernel-based conntrack events
 type ConnTrack struct {
 	flowLogger nftypes.FlowLogger
 	iface      nftypes.IFaceMapper
 
-	conn *nfct.Conn
+	conn listener
 	mux  sync.Mutex
 
+	dial           func() (listener, error)
 	instanceID     uuid.UUID
 	started        bool
 	done           chan struct{}
 	sysctlModified bool
 }
 
+// DialFunc is a constructor for netlink conntrack connections.
+type DialFunc func() (listener, error)
+
+// Option configures a ConnTrack instance.
+type Option func(*ConnTrack)
+
+// WithDialer overrides the default netlink dialer, primarily for testing.
+func WithDialer(dial DialFunc) Option {
+	return func(c *ConnTrack) {
+		c.dial = dial
+	}
+}
+
+func defaultDial() (listener, error) {
+	return nfct.Dial(nil)
+}
+
 // New creates a new connection tracker that interfaces with the kernel's conntrack system
-func New(flowLogger nftypes.FlowLogger, iface nftypes.IFaceMapper) *ConnTrack {
-	return &ConnTrack{
+func New(flowLogger nftypes.FlowLogger, iface nftypes.IFaceMapper, opts ...Option) *ConnTrack {
+	ct := &ConnTrack{
 		flowLogger: flowLogger,
 		iface:      iface,
 		instanceID: uuid.New(),
-		started:    false,
+		dial:       defaultDial,
 		done:       make(chan struct{}, 1),
 	}
+	for _, opt := range opts {
+		opt(ct)
+	}
+	return ct
 }
 
 // Start begins tracking connections by listening for conntrack events. This method is idempotent.
@@ -59,8 +94,9 @@ func (c *ConnTrack) Start(enableCounters bool) error {
 		c.EnableAccounting()
 	}
 
-	conn, err := nfct.Dial(nil)
+	conn, err := c.dial()
 	if err != nil {
+		c.RestoreAccounting()
 		return fmt.Errorf("dial conntrack: %w", err)
 	}
 	c.conn = conn
@@ -76,9 +112,16 @@ func (c *ConnTrack) Start(enableCounters bool) error {
 			log.Errorf("Error closing conntrack connection: %v", err)
 		}
 		c.conn = nil
+		c.RestoreAccounting()
 		return fmt.Errorf("start conntrack listener: %w", err)
 	}
 
+	// Drain any stale stop signal from a previous cycle.
+	select {
+	case <-c.done:
+	default:
+	}
+
 	c.started = true
 
 	go c.receiverRoutine(events, errChan)
@@ -92,17 +135,98 @@ func (c *ConnTrack) receiverRoutine(events chan nfct.Event, errChan chan error)
 		case event := <-events:
 			c.handleEvent(event)
 		case err := <-errChan:
-			log.Errorf("Error from conntrack event listener: %v", err)
-			if err := c.conn.Close(); err != nil {
-				log.Errorf("Error closing conntrack connection: %v", err)
+			if events, errChan = c.handleListenerError(err); events == nil {
+				return
 			}
-			return
 		case <-c.done:
 			return
 		}
 	}
 }
 
+// handleListenerError closes the failed connection and attempts to reconnect.
+// Returns new channels on success, or nil if shutdown was requested.
+func (c *ConnTrack) handleListenerError(err error) (chan nfct.Event, chan error) {
+	log.Warnf("conntrack event listener failed: %v", err)
+	c.closeConn()
+	return c.reconnect()
+}
+
+func (c *ConnTrack) closeConn() {
+	c.mux.Lock()
+	defer c.mux.Unlock()
+
+	if c.conn != nil {
+		if err := c.conn.Close(); err != nil {
+			log.Debugf("close conntrack connection: %v", err)
+		}
+		c.conn = nil
+	}
+}
+
+// reconnect attempts to re-establish the conntrack netlink listener with exponential backoff.
+// Returns new channels on success, or nil if shutdown was requested.
+func (c *ConnTrack) reconnect() (chan nfct.Event, chan error) {
+	bo := &backoff.ExponentialBackOff{
+		InitialInterval:     reconnectInitInterval,
+		RandomizationFactor: reconnectRandomization,
+		Multiplier:          backoff.DefaultMultiplier,
+		MaxInterval:         reconnectMaxInterval,
+		MaxElapsedTime:      0, // retry indefinitely
+		Clock:               backoff.SystemClock,
+	}
+	bo.Reset()
+
+	for {
+		delay := bo.NextBackOff()
+		log.Infof("reconnecting conntrack listener in %s", delay)
+
+		select {
+		case <-c.done:
+			c.mux.Lock()
+			c.started = false
+			c.mux.Unlock()
+			return nil, nil
+		case <-time.After(delay):
+		}
+
+		conn, err := c.dial()
+		if err != nil {
+			log.Warnf("reconnect conntrack dial: %v", err)
+			continue
+		}
+
+		events := make(chan nfct.Event, defaultChannelSize)
+		errChan, err := conn.Listen(events, 1, []netfilter.NetlinkGroup{
+			netfilter.GroupCTNew,
+			netfilter.GroupCTDestroy,
+		})
+		if err != nil {
+			log.Warnf("reconnect conntrack listen: %v", err)
+			if closeErr := conn.Close(); closeErr != nil {
+				log.Debugf("close conntrack connection: %v", closeErr)
+			}
+			continue
+		}
+
+		c.mux.Lock()
+		if !c.started {
+			// Stop() ran while we were reconnecting.
+			c.mux.Unlock()
+			if closeErr := conn.Close(); closeErr != nil {
+				log.Debugf("close conntrack connection: %v", closeErr)
+			}
+			return nil, nil
+		}
+		c.conn = conn
+		c.mux.Unlock()
+
+		log.Infof("conntrack listener reconnected successfully")
+
+		return events, errChan
+	}
+}
+
 // Stop stops the connection tracking. This method is idempotent.
 func (c *ConnTrack) Stop() {
 	c.mux.Lock()
@@ -136,23 +260,27 @@ func (c *ConnTrack) Close() error {
 	c.mux.Lock()
 	defer c.mux.Unlock()
 
-	if c.started {
-		select {
-		case c.done <- struct{}{}:
-		default:
-		}
+	if !c.started {
+		return nil
 	}
 
+	select {
+	case c.done <- struct{}{}:
+	default:
+	}
+
+	c.started = false
+
+	var closeErr error
 	if c.conn != nil {
-		err := c.conn.Close()
+		closeErr = c.conn.Close()
 		c.conn = nil
-		c.started = false
+	}
 
-		c.RestoreAccounting()
+	c.RestoreAccounting()
 
-		if err != nil {
-			return fmt.Errorf("close conntrack: %w", err)
-		}
+	if closeErr != nil {
+		return fmt.Errorf("close conntrack: %w", closeErr)
 	}
 
 	return nil

client/internal/netflow/conntrack/conntrack_test.go

@@ -0,0 +1,224 @@
+//go:build linux && !android
+
+package conntrack
+
+import (
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	nfct "github.com/ti-mo/conntrack"
+	"github.com/ti-mo/netfilter"
+)
+
+type mockListener struct {
+	errChan  chan error
+	closed   atomic.Bool
+	closedCh chan struct{}
+}
+
+func newMockListener() *mockListener {
+	return &mockListener{
+		errChan:  make(chan error, 1),
+		closedCh: make(chan struct{}),
+	}
+}
+
+func (m *mockListener) Listen(evChan chan<- nfct.Event, _ uint8, _ []netfilter.NetlinkGroup) (chan error, error) {
+	return m.errChan, nil
+}
+
+func (m *mockListener) Close() error {
+	if m.closed.CompareAndSwap(false, true) {
+		close(m.closedCh)
+	}
+	return nil
+}
+
+func TestReconnectAfterError(t *testing.T) {
+	first := newMockListener()
+	second := newMockListener()
+	third := newMockListener()
+	listeners := []*mockListener{first, second, third}
+	callCount := atomic.Int32{}
+
+	ct := New(nil, nil, WithDialer(func() (listener, error) {
+		n := int(callCount.Add(1)) - 1
+		return listeners[n], nil
+	}))
+
+	err := ct.Start(false)
+	require.NoError(t, err)
+
+	// Inject an error on the first listener.
+	first.errChan <- assert.AnError
+
+	// Wait for reconnect to complete.
+	require.Eventually(t, func() bool {
+		return callCount.Load() >= 2
+	}, 15*time.Second, 100*time.Millisecond, "reconnect should dial a new connection")
+
+	// The first connection must have been closed.
+	select {
+	case <-first.closedCh:
+	case <-time.After(2 * time.Second):
+		t.Fatal("first connection was not closed")
+	}
+
+	// Verify the receiver is still running by injecting and handling a second error.
+	second.errChan <- assert.AnError
+
+	require.Eventually(t, func() bool {
+		return callCount.Load() >= 3
+	}, 15*time.Second, 100*time.Millisecond, "second reconnect should succeed")
+
+	ct.Stop()
+}
+
+func TestStopDuringReconnectBackoff(t *testing.T) {
+	mock := newMockListener()
+
+	ct := New(nil, nil, WithDialer(func() (listener, error) {
+		return mock, nil
+	}))
+
+	err := ct.Start(false)
+	require.NoError(t, err)
+
+	// Trigger an error so the receiver enters reconnect.
+	mock.errChan <- assert.AnError
+
+	// Wait for the error handler to close the old listener before calling Stop.
+	select {
+	case <-mock.closedCh:
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for reconnect to start")
+	}
+
+	// Stop while reconnecting.
+	ct.Stop()
+
+	ct.mux.Lock()
+	assert.False(t, ct.started, "started should be false after Stop")
+	assert.Nil(t, ct.conn, "conn should be nil after Stop")
+	ct.mux.Unlock()
+}
+
+func TestStopRaceWithReconnectDial(t *testing.T) {
+	first := newMockListener()
+	dialStarted := make(chan struct{})
+	dialProceed := make(chan struct{})
+	second := newMockListener()
+	callCount := atomic.Int32{}
+
+	ct := New(nil, nil, WithDialer(func() (listener, error) {
+		n := callCount.Add(1)
+		if n == 1 {
+			return first, nil
+		}
+		// Second dial: signal that we're in progress, wait for test to call Stop.
+		close(dialStarted)
+		<-dialProceed
+		return second, nil
+	}))
+
+	err := ct.Start(false)
+	require.NoError(t, err)
+
+	// Trigger error to enter reconnect.
+	first.errChan <- assert.AnError
+
+	// Wait for reconnect's second dial to begin.
+	select {
+	case <-dialStarted:
+	case <-time.After(15 * time.Second):
+		t.Fatal("timed out waiting for reconnect dial")
+	}
+
+	// Stop while dial is in progress (conn is nil at this point).
+	ct.Stop()
+
+	// Let the dial complete. reconnect should detect started==false and close the new conn.
+	close(dialProceed)
+
+	// The second connection should be closed (not leaked).
+	select {
+	case <-second.closedCh:
+	case <-time.After(2 * time.Second):
+		t.Fatal("second connection was leaked after Stop")
+	}
+
+	ct.mux.Lock()
+	assert.False(t, ct.started)
+	assert.Nil(t, ct.conn)
+	ct.mux.Unlock()
+}
+
+func TestCloseRaceWithReconnectDial(t *testing.T) {
+	first := newMockListener()
+	dialStarted := make(chan struct{})
+	dialProceed := make(chan struct{})
+	second := newMockListener()
+	callCount := atomic.Int32{}
+
+	ct := New(nil, nil, WithDialer(func() (listener, error) {
+		n := callCount.Add(1)
+		if n == 1 {
+			return first, nil
+		}
+		close(dialStarted)
+		<-dialProceed
+		return second, nil
+	}))
+
+	err := ct.Start(false)
+	require.NoError(t, err)
+
+	first.errChan <- assert.AnError
+
+	select {
+	case <-dialStarted:
+	case <-time.After(15 * time.Second):
+		t.Fatal("timed out waiting for reconnect dial")
+	}
+
+	// Close while dial is in progress (conn is nil).
+	require.NoError(t, ct.Close())
+
+	close(dialProceed)
+
+	// The second connection should be closed (not leaked).
+	select {
+	case <-second.closedCh:
+	case <-time.After(2 * time.Second):
+		t.Fatal("second connection was leaked after Close")
+	}
+
+	ct.mux.Lock()
+	assert.False(t, ct.started)
+	assert.Nil(t, ct.conn)
+	ct.mux.Unlock()
+}
+
+func TestStartIsIdempotent(t *testing.T) {
+	mock := newMockListener()
+	callCount := atomic.Int32{}
+
+	ct := New(nil, nil, WithDialer(func() (listener, error) {
+		callCount.Add(1)
+		return mock, nil
+	}))
+
+	err := ct.Start(false)
+	require.NoError(t, err)
+
+	// Second Start should be a no-op.
+	err = ct.Start(false)
+	require.NoError(t, err)
+
+	assert.Equal(t, int32(1), callCount.Load(), "dial should only be called once")
+
+	ct.Stop()
+}