netbird/shared/relay/client/fallback_opener_test.go

package client

import (
	"context"
	"errors"
	"net"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/netbirdio/netbird/relay/server"
)

type fakeConn struct {
	net.Conn
	closed chan struct{}
}

func newFakeConn() *fakeConn {
	return &fakeConn{closed: make(chan struct{})}
}

func (c *fakeConn) Close() error {
	close(c.closed)
	return nil
}

func newTestConnRace(t *testing.T) *connRace {
	t.Helper()
	raceCtx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	_, cancelPreferred := context.WithCancel(raceCtx)
	otherCtx, cancelOther := context.WithCancel(raceCtx)
	timer := time.NewTimer(time.Hour)
	timer.Stop()
	return &connRace{
		opener:          &FallbackOpener{},
		peerKey:         "peerKey",
		raceCtx:         raceCtx,
		otherCtx:        otherCtx,
		cancelPreferred: cancelPreferred,
		cancelOther:     cancelOther,
		results:         make(chan raceAttempt, 2),
		fallbackTimer:   timer,
	}
}

func TestHandleResult_PreferredSucceeds(t *testing.T) {
	c := newTestConnRace(t)

	conn := newFakeConn()
	o := c.handleResult(raceAttempt{conn: conn})

	require.True(t, o.done)
	require.NoError(t, o.err)
	require.Same(t, net.Conn(conn), o.conn)
	require.False(t, c.otherStarted, "fallback must not start once the preferred attempt wins")
}

func TestHandleResult_ConnAlreadyExistsIsSuccess(t *testing.T) {
	c := newTestConnRace(t)

	o := c.handleResult(raceAttempt{err: ErrConnAlreadyExists})

	require.True(t, o.done)
	require.ErrorIs(t, o.err, ErrConnAlreadyExists)
	require.False(t, c.otherStarted)
}

func TestHandleResult_PreferredFailsStartsOther(t *testing.T) {
	c := newTestConnRace(t)
	// The fallback attempt opens against a stalling listener so startOther's
	// goroutine blocks on Connect until raceCtx is cancelled by t.Cleanup.
	serverAddr, _ := stallingRelayListener(t)
	c.opener.foreignStore = NewForeignRelaysStore(c.raceCtx, hmacTokenStore, "alice", 1280, newTransportFallback(), nil, func(string) {}, keepUnusedServerTime)
	c.remoteRelayServer = RelayServer{Addr: serverAddr}
	c.preferForeign = false

	o := c.handleResult(raceAttempt{err: errors.New("boom")})

	require.False(t, o.done, "a single failure must not settle the race")
	require.True(t, c.otherStarted, "the fallback attempt must start after the preferred one fails")
	require.EqualError(t, c.lastErr, "boom")
}

func TestHandleResult_BothFailReturnsLastErr(t *testing.T) {
	c := newTestConnRace(t)
	c.otherStarted = true
	c.settled = 1
	c.lastErr = errors.New("first")

	o := c.handleResult(raceAttempt{err: errors.New("second")})

	require.True(t, o.done)
	require.EqualError(t, o.err, "second")
}

func TestOnTimeout_PrefersLastErr(t *testing.T) {
	c := newTestConnRace(t)
	c.lastErr = errors.New("dial failed")

	_, err := c.onTimeout()
	require.EqualError(t, err, "dial failed")
}

func TestOnTimeout_FallsBackToCtxErr(t *testing.T) {
	c := newTestConnRace(t)
	raceCtx, cancel := context.WithCancel(context.Background())
	cancel()
	c.raceCtx = raceCtx

	_, err := c.onTimeout()
	require.ErrorIs(t, err, context.Canceled)
}

func TestDrainLoser_ClosesLateWinner(t *testing.T) {
	r := &FallbackOpener{}
	results := make(chan raceAttempt, 2)

	loser := newFakeConn()
	results <- raceAttempt{conn: loser}

	done := make(chan struct{})
	go func() {
		defer close(done)
		r.drainLoser(results, 1, true)
	}()

	select {
	case <-loser.closed:
	case <-time.After(2 * time.Second):
		t.Fatal("losing connection was not closed")
	}
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("drainLoser did not return")
	}
}

func TestDrainLoser_NoOtherAttempt(t *testing.T) {
	r := &FallbackOpener{}
	results := make(chan raceAttempt)

	done := make(chan struct{})
	go func() {
		defer close(done)
		r.drainLoser(results, 1, false)
	}()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("drainLoser blocked with no outstanding attempt")
	}
}

func startTestRelayServer(t *testing.T, addr string) string {
	t.Helper()

	srv, err := server.NewServer(newManagerTestServerConfig(addr))
	require.NoError(t, err)

	errChan := make(chan error, 1)
	go func() {
		if err := srv.Listen(server.ListenerConfig{Address: addr}); err != nil {
			errChan <- err
		}
	}()
	t.Cleanup(func() { _ = srv.Shutdown(context.Background()) })

	require.NoError(t, waitForServerToStart(errChan))
	return addr
}