netbird/shared/relay/client/fallback_opener.go

package client

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

	log "github.com/sirupsen/logrus"
)

const (
	raceTotalTimeout  = 40 * time.Second
	raceFallbackDelay = 10 * time.Second
)

type raceAttempt struct {
	conn net.Conn
	err  error
}

type raceOutcome struct {
	conn net.Conn
	err  error
	done bool
}

type connRace struct {
	opener            *FallbackOpener
	peerKey           string
	remoteRelayServer RelayServer
	preferForeign     bool

	raceCtx         context.Context
	otherCtx        context.Context
	cancelPreferred context.CancelFunc
	cancelOther     context.CancelFunc
	results         chan raceAttempt
	fallbackTimer   *time.Timer

	otherStarted bool
	settled      int
	lastErr      error
}

type FallbackOpener struct {
	home         *Client
	foreignStore *ForeignRelaysStore

	fallbackDelay time.Duration
	totalTimeout  time.Duration
	// openFn performs a single attempt. It is overridable in tests; when nil the
	// real home/foreign dispatch in open is used.
	openFn func(ctx context.Context, peerKey string, remoteRelayServer RelayServer, foreign bool) raceAttempt
}

func NewFallbackOpener(home *Client, foreignStore *ForeignRelaysStore) *FallbackOpener {
	return &FallbackOpener{
		home:          home,
		foreignStore:  foreignStore,
		fallbackDelay: raceFallbackDelay,
		totalTimeout:  raceTotalTimeout,
	}
}

func (r *FallbackOpener) Run(ctx context.Context, peerKey string, remoteRelayServer RelayServer, preferForeign bool) (net.Conn, error) {
	raceCtx, cancel := context.WithTimeout(ctx, r.totalTimeout)
	defer cancel()

	preferredCtx, cancelPreferred := context.WithCancel(raceCtx)
	otherCtx, cancelOther := context.WithCancel(raceCtx)

	race := &connRace{
		opener:            r,
		peerKey:           peerKey,
		remoteRelayServer: remoteRelayServer,
		preferForeign:     preferForeign,
		raceCtx:           raceCtx,
		otherCtx:          otherCtx,
		cancelPreferred:   cancelPreferred,
		cancelOther:       cancelOther,
		results:           make(chan raceAttempt, 2),
		fallbackTimer:     time.NewTimer(r.fallbackDelay),
	}
	defer race.fallbackTimer.Stop()

	go func() {
		race.results <- r.open(preferredCtx, peerKey, remoteRelayServer, preferForeign)
	}()

	for {
		select {
		case <-race.fallbackTimer.C:
			race.startOther()
		case res := <-race.results:
			if o := race.handleResult(res); o.done {
				return o.conn, o.err
			}
		case <-raceCtx.Done():
			return race.onTimeout()
		}
	}
}

func (c *connRace) startOther() {
	if c.otherStarted {
		return
	}
	c.otherStarted = true
	c.fallbackTimer.Stop()
	go func() {
		c.results <- c.opener.open(c.otherCtx, c.peerKey, c.remoteRelayServer, !c.preferForeign)
	}()
}

func (c *connRace) handleResult(res raceAttempt) raceOutcome {
	if (res.err == nil && res.conn != nil) || errors.Is(res.err, ErrConnAlreadyExists) {
		c.settled++
		c.stop()
		return raceOutcome{conn: res.conn, err: res.err, done: true}
	}

	c.lastErr = res.err
	c.settled++
	if !c.otherStarted {
		c.startOther()
		return raceOutcome{}
	}
	if c.settled == 2 {
		c.cancelPreferred()
		c.cancelOther()
		return raceOutcome{err: c.lastErr, done: true}
	}
	return raceOutcome{}
}

func (c *connRace) onTimeout() (net.Conn, error) {
	c.stop()
	if c.lastErr != nil {
		return nil, c.lastErr
	}
	return nil, c.raceCtx.Err()
}

func (c *connRace) stop() {
	c.cancelPreferred()
	c.cancelOther()
	go c.opener.drainLoser(c.results, c.settled, c.otherStarted)
}

func (r *FallbackOpener) open(ctx context.Context, peerKey string, remoteRelayServer RelayServer, foreign bool) raceAttempt {
	if r.openFn != nil {
		return r.openFn(ctx, peerKey, remoteRelayServer, foreign)
	}
	if foreign {
		conn, err := r.foreignStore.OpenConn(ctx, peerKey, remoteRelayServer)
		return raceAttempt{conn: conn, err: err}
	}
	conn, err := r.home.OpenConn(ctx, peerKey)
	return raceAttempt{conn: conn, err: err}
}

func (r *FallbackOpener) drainLoser(results chan raceAttempt, settled int, otherStarted bool) {
	started := 1
	if otherStarted {
		started = 2
	}
	for i := settled; i < started; i++ {
		res := <-results
		if res.conn != nil {
			if err := res.conn.Close(); err != nil {
				log.Debugf("failed to close losing relay connection: %v", err)
			}
		}
	}
}