Extend netbird status command to include health information (#1471)

* Adds management, signal, and relay (STUN/TURN) health probes to the status command.

* Adds a reason when the management or signal connections are disconnected.

* Adds last wireguard handshake and received/sent bytes per peer

client/internal/connect.go

@@ -27,11 +27,33 @@ import (
 
 // RunClient with main logic.
 func RunClient(ctx context.Context, config *Config, statusRecorder *peer.Status) error {
-	return runClient(ctx, config, statusRecorder, MobileDependency{})
+	return runClient(ctx, config, statusRecorder, MobileDependency{}, nil, nil, nil, nil)
+}
+
+// RunClientWithProbes runs the client's main logic with probes attached
+func RunClientWithProbes(
+	ctx context.Context,
+	config *Config,
+	statusRecorder *peer.Status,
+	mgmProbe *Probe,
+	signalProbe *Probe,
+	relayProbe *Probe,
+	wgProbe *Probe,
+) error {
+	return runClient(ctx, config, statusRecorder, MobileDependency{}, mgmProbe, signalProbe, relayProbe, wgProbe)
 }
 
 // RunClientMobile with main logic on mobile system
-func RunClientMobile(ctx context.Context, config *Config, statusRecorder *peer.Status, tunAdapter iface.TunAdapter, iFaceDiscover stdnet.ExternalIFaceDiscover, networkChangeListener listener.NetworkChangeListener, dnsAddresses []string, dnsReadyListener dns.ReadyListener) error {
+func RunClientMobile(
+	ctx context.Context,
+	config *Config,
+	statusRecorder *peer.Status,
+	tunAdapter iface.TunAdapter,
+	iFaceDiscover stdnet.ExternalIFaceDiscover,
+	networkChangeListener listener.NetworkChangeListener,
+	dnsAddresses []string,
+	dnsReadyListener dns.ReadyListener,
+) error {
 	// in case of non Android os these variables will be nil
 	mobileDependency := MobileDependency{
 		TunAdapter:            tunAdapter,
@@ -40,19 +62,35 @@ func RunClientMobile(ctx context.Context, config *Config, statusRecorder *peer.S
 		HostDNSAddresses:      dnsAddresses,
 		DnsReadyListener:      dnsReadyListener,
 	}
-	return runClient(ctx, config, statusRecorder, mobileDependency)
+	return runClient(ctx, config, statusRecorder, mobileDependency, nil, nil, nil, nil)
 }
 
-func RunClientiOS(ctx context.Context, config *Config, statusRecorder *peer.Status, fileDescriptor int32, networkChangeListener listener.NetworkChangeListener, dnsManager dns.IosDnsManager) error {
+func RunClientiOS(
+	ctx context.Context,
+	config *Config,
+	statusRecorder *peer.Status,
+	fileDescriptor int32,
+	networkChangeListener listener.NetworkChangeListener,
+	dnsManager dns.IosDnsManager,
+) error {
 	mobileDependency := MobileDependency{
 		FileDescriptor:        fileDescriptor,
 		NetworkChangeListener: networkChangeListener,
 		DnsManager:            dnsManager,
 	}
-	return runClient(ctx, config, statusRecorder, mobileDependency)
+	return runClient(ctx, config, statusRecorder, mobileDependency, nil, nil, nil, nil)
 }
 
-func runClient(ctx context.Context, config *Config, statusRecorder *peer.Status, mobileDependency MobileDependency) error {
+func runClient(
+	ctx context.Context,
+	config *Config,
+	statusRecorder *peer.Status,
+	mobileDependency MobileDependency,
+	mgmProbe *Probe,
+	signalProbe *Probe,
+	relayProbe *Probe,
+	wgProbe *Probe,
+) error {
 	log.Infof("starting NetBird client version %s", version.NetbirdVersion())
 
 	backOff := &backoff.ExponentialBackOff{
@@ -103,7 +141,7 @@ func runClient(ctx context.Context, config *Config, statusRecorder *peer.Status,
 
 		engineCtx, cancel := context.WithCancel(ctx)
 		defer func() {
-			statusRecorder.MarkManagementDisconnected()
+			statusRecorder.MarkManagementDisconnected(state.err)
 			statusRecorder.CleanLocalPeerState()
 			cancel()
 		}()
@@ -152,8 +190,10 @@ func runClient(ctx context.Context, config *Config, statusRecorder *peer.Status,
 
 		statusRecorder.UpdateSignalAddress(signalURL)
 
-		statusRecorder.MarkSignalDisconnected()
-		defer statusRecorder.MarkSignalDisconnected()
+		statusRecorder.MarkSignalDisconnected(nil)
+		defer func() {
+			statusRecorder.MarkSignalDisconnected(state.err)
+		}()
 
 		// with the global Wiretrustee config in hand connect (just a connection, no stream yet) Signal
 		signalClient, err := connectToSignal(engineCtx, loginResp.GetWiretrusteeConfig(), myPrivateKey)
@@ -181,7 +221,7 @@ func runClient(ctx context.Context, config *Config, statusRecorder *peer.Status,
 			return wrapErr(err)
 		}
 
-		engine := NewEngine(engineCtx, cancel, signalClient, mgmClient, engineConfig, mobileDependency, statusRecorder)
+		engine := NewEngineWithProbes(engineCtx, cancel, signalClient, mgmClient, engineConfig, mobileDependency, statusRecorder, mgmProbe, signalProbe, relayProbe, wgProbe)
 		err = engine.Start()
 		if err != nil {
 			log.Errorf("error while starting Netbird Connection Engine: %s", err)

client/internal/dns/server_test.go

@@ -59,6 +59,10 @@ func (w *mocWGIface) SetFilter(filter iface.PacketFilter) error {
 	return nil
 }
 
+func (w *mocWGIface) GetStats(_ string) (iface.WGStats, error) {
+	return iface.WGStats{}, nil
+}
+
 var zoneRecords = []nbdns.SimpleRecord{
 	{
 		Name:  "peera.netbird.cloud",

client/internal/dns/wgiface.go

@@ -11,4 +11,5 @@ type WGIface interface {
 	IsUserspaceBind() bool
 	GetFilter() iface.PacketFilter
 	GetDevice() *iface.DeviceWrapper
+	GetStats(peerKey string) (iface.WGStats, error)
 }

client/internal/dns/wgiface_windows.go

@@ -9,5 +9,6 @@ type WGIface interface {
 	IsUserspaceBind() bool
 	GetFilter() iface.PacketFilter
 	GetDevice() *iface.DeviceWrapper
+	GetStats(peerKey string) (iface.WGStats, error)
 	GetInterfaceGUIDString() (string, error)
 }

client/internal/engine.go

@@ -22,6 +22,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/acl"
 	"github.com/netbirdio/netbird/client/internal/dns"
 	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/client/internal/relay"
 	"github.com/netbirdio/netbird/client/internal/rosenpass"
 	"github.com/netbirdio/netbird/client/internal/routemanager"
 	"github.com/netbirdio/netbird/client/internal/wgproxy"
@@ -125,6 +126,11 @@ type Engine struct {
 	acl          acl.Manager
 
 	dnsServer dns.Server
+
+	mgmProbe    *Probe
+	signalProbe *Probe
+	relayProbe  *Probe
+	wgProbe     *Probe
 }
 
 // Peer is an instance of the Connection Peer
@@ -135,11 +141,43 @@ type Peer struct {
 
 // NewEngine creates a new Connection Engine
 func NewEngine(
-	ctx context.Context, cancel context.CancelFunc,
-	signalClient signal.Client, mgmClient mgm.Client,
-	config *EngineConfig, mobileDep MobileDependency, statusRecorder *peer.Status,
+	ctx context.Context,
+	cancel context.CancelFunc,
+	signalClient signal.Client,
+	mgmClient mgm.Client,
+	config *EngineConfig,
+	mobileDep MobileDependency,
+	statusRecorder *peer.Status,
 ) *Engine {
+	return NewEngineWithProbes(
+		ctx,
+		cancel,
+		signalClient,
+		mgmClient,
+		config,
+		mobileDep,
+		statusRecorder,
+		nil,
+		nil,
+		nil,
+		nil,
+	)
+}
 
+// NewEngineWithProbes creates a new Connection Engine with probes attached
+func NewEngineWithProbes(
+	ctx context.Context,
+	cancel context.CancelFunc,
+	signalClient signal.Client,
+	mgmClient mgm.Client,
+	config *EngineConfig,
+	mobileDep MobileDependency,
+	statusRecorder *peer.Status,
+	mgmProbe *Probe,
+	signalProbe *Probe,
+	relayProbe *Probe,
+	wgProbe *Probe,
+) *Engine {
 	return &Engine{
 		ctx:            ctx,
 		cancel:         cancel,
@@ -155,6 +193,10 @@ func NewEngine(
 		sshServerFunc:  nbssh.DefaultSSHServer,
 		statusRecorder: statusRecorder,
 		wgProxyFactory: wgproxy.NewFactory(config.WgPort),
+		mgmProbe:       mgmProbe,
+		signalProbe:    signalProbe,
+		relayProbe:     relayProbe,
+		wgProbe:        wgProbe,
 	}
 }
 
@@ -251,6 +293,7 @@ func (e *Engine) Start() error {
 
 	e.receiveSignalEvents()
 	e.receiveManagementEvents()
+	e.receiveProbeEvents()
 
 	return nil
 }
@@ -512,9 +555,7 @@ func (e *Engine) updateConfig(conf *mgmProto.PeerConfig) error {
 // E.g. when a new peer has been registered and we are allowed to connect to it.
 func (e *Engine) receiveManagementEvents() {
 	go func() {
-		err := e.mgmClient.Sync(func(update *mgmProto.SyncResponse) error {
-			return e.handleSync(update)
-		})
+		err := e.mgmClient.Sync(e.handleSync)
 		if err != nil {
 			// happens if management is unavailable for a long time.
 			// We want to cancel the operation of the whole client
@@ -1175,3 +1216,69 @@ func (e *Engine) getRosenpassAddr() string {
 	}
 	return ""
 }
+
+func (e *Engine) receiveProbeEvents() {
+	if e.signalProbe != nil {
+		go e.signalProbe.Receive(e.ctx, func() bool {
+			healthy := e.signal.IsHealthy()
+			log.Debugf("received signal probe request, healthy: %t", healthy)
+			return healthy
+		})
+	}
+
+	if e.mgmProbe != nil {
+		go e.mgmProbe.Receive(e.ctx, func() bool {
+			healthy := e.mgmClient.IsHealthy()
+			log.Debugf("received management probe request, healthy: %t", healthy)
+			return healthy
+		})
+	}
+
+	if e.relayProbe != nil {
+		go e.relayProbe.Receive(e.ctx, func() bool {
+			healthy := true
+
+			results := append(e.probeSTUNs(), e.probeTURNs()...)
+			e.statusRecorder.UpdateRelayStates(results)
+
+			// A single failed server will result in a "failed" probe
+			for _, res := range results {
+				if res.Err != nil {
+					healthy = false
+					break
+				}
+			}
+
+			log.Debugf("received relay probe request, healthy: %t", healthy)
+			return healthy
+		})
+	}
+
+	if e.wgProbe != nil {
+		go e.wgProbe.Receive(e.ctx, func() bool {
+			log.Debug("received wg probe request")
+
+			for _, peer := range e.peerConns {
+				key := peer.GetKey()
+				wgStats, err := peer.GetConf().WgConfig.WgInterface.GetStats(key)
+				if err != nil {
+					log.Debugf("failed to get wg stats for peer %s: %s", key, err)
+				}
+				// wgStats could be zero value, in which case we just reset the stats
+				if err := e.statusRecorder.UpdateWireguardPeerState(key, wgStats); err != nil {
+					log.Debugf("failed to update wg stats for peer %s: %s", key, err)
+				}
+			}
+
+			return true
+		})
+	}
+}
+
+func (e *Engine) probeSTUNs() []relay.ProbeResult {
+	return relay.ProbeAll(e.ctx, relay.ProbeSTUN, e.STUNs)
+}
+
+func (e *Engine) probeTURNs() []relay.ProbeResult {
+	return relay.ProbeAll(e.ctx, relay.ProbeTURN, e.TURNs)
+}

client/internal/peer/conn.go

@@ -408,12 +408,14 @@ func (conn *Conn) configureConnection(remoteConn net.Conn, remoteWgPort int, rem
 	conn.status = StatusConnected
 
 	peerState := State{
-		PubKey:                 conn.config.Key,
-		ConnStatus:             conn.status,
-		ConnStatusUpdate:       time.Now(),
-		LocalIceCandidateType:  pair.Local.Type().String(),
-		RemoteIceCandidateType: pair.Remote.Type().String(),
-		Direct:                 !isRelayCandidate(pair.Local),
+		PubKey:                     conn.config.Key,
+		ConnStatus:                 conn.status,
+		ConnStatusUpdate:           time.Now(),
+		LocalIceCandidateType:      pair.Local.Type().String(),
+		RemoteIceCandidateType:     pair.Remote.Type().String(),
+		LocalIceCandidateEndpoint:  fmt.Sprintf("%s:%d", pair.Local.Address(), pair.Local.Port()),
+		RemoteIceCandidateEndpoint: fmt.Sprintf("%s:%d", pair.Remote.Address(), pair.Local.Port()),
+		Direct:                     !isRelayCandidate(pair.Local),
 	}
 	if pair.Local.Type() == ice.CandidateTypeRelay || pair.Remote.Type() == ice.CandidateTypeRelay {
 		peerState.Relayed = true
@@ -500,6 +502,9 @@ func (conn *Conn) cleanup() error {
 		// todo rethink status updates
 		log.Debugf("error while updating peer's %s state, err: %v", conn.config.Key, err)
 	}
+	if err := conn.statusRecorder.UpdateWireguardPeerState(conn.config.Key, iface.WGStats{}); err != nil {
+		log.Debugf("failed to reset wireguard stats for peer %s: %s", conn.config.Key, err)
+	}
 
 	log.Debugf("cleaned up connection to peer %s", conn.config.Key)
 	if err1 != nil {

client/internal/peer/status.go

@@ -4,19 +4,27 @@ import (
 	"errors"
 	"sync"
 	"time"
+
+	"github.com/netbirdio/netbird/client/internal/relay"
+	"github.com/netbirdio/netbird/iface"
 )
 
 // State contains the latest state of a peer
 type State struct {
-	IP                     string
-	PubKey                 string
-	FQDN                   string
-	ConnStatus             ConnStatus
-	ConnStatusUpdate       time.Time
-	Relayed                bool
-	Direct                 bool
-	LocalIceCandidateType  string
-	RemoteIceCandidateType string
+	IP                         string
+	PubKey                     string
+	FQDN                       string
+	ConnStatus                 ConnStatus
+	ConnStatusUpdate           time.Time
+	Relayed                    bool
+	Direct                     bool
+	LocalIceCandidateType      string
+	RemoteIceCandidateType     string
+	LocalIceCandidateEndpoint  string
+	RemoteIceCandidateEndpoint string
+	LastWireguardHandshake     time.Time
+	BytesTx                    int64
+	BytesRx                    int64
 }
 
 // LocalPeerState contains the latest state of the local peer
@@ -31,12 +39,14 @@ type LocalPeerState struct {
 type SignalState struct {
 	URL       string
 	Connected bool
+	Error     error
 }
 
 // ManagementState contains the latest state of a management connection
 type ManagementState struct {
 	URL       string
 	Connected bool
+	Error     error
 }
 
 // FullStatus contains the full state held by the Status instance
@@ -45,15 +55,19 @@ type FullStatus struct {
 	ManagementState ManagementState
 	SignalState     SignalState
 	LocalPeerState  LocalPeerState
+	Relays          []relay.ProbeResult
 }
 
-// Status holds a state of peers, signal and management connections
+// Status holds a state of peers, signal, management connections and relays
 type Status struct {
 	mux             sync.Mutex
 	peers           map[string]State
 	changeNotify    map[string]chan struct{}
 	signalState     bool
+	signalError     error
 	managementState bool
+	managementError error
+	relayStates     []relay.ProbeResult
 	localPeer       LocalPeerState
 	offlinePeers    []State
 	mgmAddress      string
@@ -156,6 +170,8 @@ func (d *Status) UpdatePeerState(receivedState State) error {
 		peerState.Relayed = receivedState.Relayed
 		peerState.LocalIceCandidateType = receivedState.LocalIceCandidateType
 		peerState.RemoteIceCandidateType = receivedState.RemoteIceCandidateType
+		peerState.LocalIceCandidateEndpoint = receivedState.LocalIceCandidateEndpoint
+		peerState.RemoteIceCandidateEndpoint = receivedState.RemoteIceCandidateEndpoint
 	}
 
 	d.peers[receivedState.PubKey] = peerState
@@ -174,6 +190,25 @@ func (d *Status) UpdatePeerState(receivedState State) error {
 	return nil
 }
 
+// UpdateWireguardPeerState updates the wireguard bits of the peer state
+func (d *Status) UpdateWireguardPeerState(pubKey string, wgStats iface.WGStats) error {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	peerState, ok := d.peers[pubKey]
+	if !ok {
+		return errors.New("peer doesn't exist")
+	}
+
+	peerState.LastWireguardHandshake = wgStats.LastHandshake
+	peerState.BytesRx = wgStats.RxBytes
+	peerState.BytesTx = wgStats.TxBytes
+
+	d.peers[pubKey] = peerState
+
+	return nil
+}
+
 func shouldSkipNotify(received, curr State) bool {
 	switch {
 	case received.ConnStatus == StatusConnecting:
@@ -248,12 +283,13 @@ func (d *Status) CleanLocalPeerState() {
 }
 
 // MarkManagementDisconnected sets ManagementState to disconnected
-func (d *Status) MarkManagementDisconnected() {
+func (d *Status) MarkManagementDisconnected(err error) {
 	d.mux.Lock()
 	defer d.mux.Unlock()
 	defer d.onConnectionChanged()
 
 	d.managementState = false
+	d.managementError = err
 }
 
 // MarkManagementConnected sets ManagementState to connected
@@ -263,6 +299,7 @@ func (d *Status) MarkManagementConnected() {
 	defer d.onConnectionChanged()
 
 	d.managementState = true
+	d.managementError = nil
 }
 
 // UpdateSignalAddress update the address of the signal server
@@ -280,12 +317,13 @@ func (d *Status) UpdateManagementAddress(mgmAddress string) {
 }
 
 // MarkSignalDisconnected sets SignalState to disconnected
-func (d *Status) MarkSignalDisconnected() {
+func (d *Status) MarkSignalDisconnected(err error) {
 	d.mux.Lock()
 	defer d.mux.Unlock()
 	defer d.onConnectionChanged()
 
 	d.signalState = false
+	d.signalError = err
 }
 
 // MarkSignalConnected sets SignalState to connected
@@ -295,6 +333,33 @@ func (d *Status) MarkSignalConnected() {
 	defer d.onConnectionChanged()
 
 	d.signalState = true
+	d.signalError = nil
+}
+
+func (d *Status) UpdateRelayStates(relayResults []relay.ProbeResult) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+	d.relayStates = relayResults
+}
+
+func (d *Status) GetManagementState() ManagementState {
+	return ManagementState{
+		d.mgmAddress,
+		d.managementState,
+		d.managementError,
+	}
+}
+
+func (d *Status) GetSignalState() SignalState {
+	return SignalState{
+		d.signalAddress,
+		d.signalState,
+		d.signalError,
+	}
+}
+
+func (d *Status) GetRelayStates() []relay.ProbeResult {
+	return d.relayStates
 }
 
 // GetFullStatus gets full status
@@ -303,15 +368,10 @@ func (d *Status) GetFullStatus() FullStatus {
 	defer d.mux.Unlock()
 
 	fullStatus := FullStatus{
-		ManagementState: ManagementState{
-			d.mgmAddress,
-			d.managementState,
-		},
-		SignalState: SignalState{
-			d.signalAddress,
-			d.signalState,
-		},
-		LocalPeerState: d.localPeer,
+		ManagementState: d.GetManagementState(),
+		SignalState:     d.GetSignalState(),
+		LocalPeerState:  d.localPeer,
+		Relays:          d.GetRelayStates(),
 	}
 
 	for _, status := range d.peers {

client/internal/peer/status_test.go

@@ -1,6 +1,7 @@
 package peer
 
 import (
+	"errors"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -152,9 +153,10 @@ func TestUpdateSignalState(t *testing.T) {
 		name      string
 		connected bool
 		want      bool
+		err       error
 	}{
-		{"should mark as connected", true, true},
-		{"should mark as disconnected", false, false},
+		{"should mark as connected", true, true, nil},
+		{"should mark as disconnected", false, false, errors.New("test")},
 	}
 
 	status := NewRecorder("https://mgm")
@@ -165,9 +167,10 @@ func TestUpdateSignalState(t *testing.T) {
 			if test.connected {
 				status.MarkSignalConnected()
 			} else {
-				status.MarkSignalDisconnected()
+				status.MarkSignalDisconnected(test.err)
 			}
 			assert.Equal(t, test.want, status.signalState, "signal status should be equal")
+			assert.Equal(t, test.err, status.signalError)
 		})
 	}
 }
@@ -178,9 +181,10 @@ func TestUpdateManagementState(t *testing.T) {
 		name      string
 		connected bool
 		want      bool
+		err       error
 	}{
-		{"should mark as connected", true, true},
-		{"should mark as disconnected", false, false},
+		{"should mark as connected", true, true, nil},
+		{"should mark as disconnected", false, false, errors.New("test")},
 	}
 
 	status := NewRecorder(url)
@@ -190,9 +194,10 @@ func TestUpdateManagementState(t *testing.T) {
 			if test.connected {
 				status.MarkManagementConnected()
 			} else {
-				status.MarkManagementDisconnected()
+				status.MarkManagementDisconnected(test.err)
 			}
 			assert.Equal(t, test.want, status.managementState, "signalState status should be equal")
+			assert.Equal(t, test.err, status.managementError)
 		})
 	}
 }

client/internal/probe.go

@@ -0,0 +1,51 @@
+package internal
+
+import "context"
+
+// Probe allows to run on-demand callbacks from different code locations.
+// Pass the probe to a receiving and a sending end. The receiving end starts listening
+// to requests with Receive and executes a callback when the sending end requests it
+// by calling Probe.
+type Probe struct {
+	request chan struct{}
+	result  chan bool
+	ready   bool
+}
+
+// NewProbe returns a new initialized probe.
+func NewProbe() *Probe {
+	return &Probe{
+		request: make(chan struct{}),
+		result:  make(chan bool),
+	}
+}
+
+// Probe requests the callback to be run and returns a bool indicating success.
+// It always returns true as long as the receiver is not ready.
+func (p *Probe) Probe() bool {
+	if !p.ready {
+		return true
+	}
+
+	p.request <- struct{}{}
+	return <-p.result
+}
+
+// Receive starts listening for probe requests. On such a request it runs the supplied
+// callback func which must return a bool indicating success.
+// Blocks until the passed context is cancelled.
+func (p *Probe) Receive(ctx context.Context, callback func() bool) {
+	p.ready = true
+	defer func() {
+		p.ready = false
+	}()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-p.request:
+			p.result <- callback()
+		}
+	}
+}

client/internal/relay/relay.go

@@ -0,0 +1,171 @@
+package relay
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/pion/stun/v2"
+	"github.com/pion/turn/v3"
+	log "github.com/sirupsen/logrus"
+)
+
+// ProbeResult holds the info about the result of a relay probe request
+type ProbeResult struct {
+	URI  *stun.URI
+	Err  error
+	Addr string
+}
+
+// ProbeSTUN tries binding to the given STUN uri and acquiring an address
+func ProbeSTUN(ctx context.Context, uri *stun.URI) (addr string, probeErr error) {
+	defer func() {
+		if probeErr != nil {
+			log.Debugf("stun probe error from %s: %s", uri, probeErr)
+		}
+	}()
+
+	client, err := stun.DialURI(uri, &stun.DialConfig{})
+	if err != nil {
+		probeErr = fmt.Errorf("dial: %w", err)
+		return
+	}
+
+	defer func() {
+		if err := client.Close(); err != nil && probeErr == nil {
+			probeErr = fmt.Errorf("close: %w", err)
+		}
+	}()
+
+	done := make(chan struct{})
+	if err = client.Start(stun.MustBuild(stun.TransactionID, stun.BindingRequest), func(res stun.Event) {
+		if res.Error != nil {
+			probeErr = fmt.Errorf("request: %w", err)
+			return
+		}
+
+		var xorAddr stun.XORMappedAddress
+		if getErr := xorAddr.GetFrom(res.Message); getErr != nil {
+			probeErr = fmt.Errorf("get xor addr: %w", err)
+			return
+		}
+
+		log.Debugf("stun probe received address from %s: %s", uri, xorAddr)
+		addr = xorAddr.String()
+
+		done <- struct{}{}
+	}); err != nil {
+		probeErr = fmt.Errorf("client: %w", err)
+		return
+	}
+
+	select {
+	case <-ctx.Done():
+		probeErr = fmt.Errorf("stun request: %w", ctx.Err())
+		return
+	case <-done:
+	}
+
+	return addr, nil
+}
+
+// ProbeTURN tries allocating a session from the given TURN URI
+func ProbeTURN(ctx context.Context, uri *stun.URI) (addr string, probeErr error) {
+	defer func() {
+		if probeErr != nil {
+			log.Debugf("turn probe error from %s: %s", uri, probeErr)
+		}
+	}()
+
+	turnServerAddr := fmt.Sprintf("%s:%d", uri.Host, uri.Port)
+
+	var conn net.PacketConn
+	switch uri.Proto {
+	case stun.ProtoTypeUDP:
+		var err error
+		conn, err = net.ListenPacket("udp", "")
+		if err != nil {
+			probeErr = fmt.Errorf("listen: %w", err)
+			return
+		}
+	case stun.ProtoTypeTCP:
+		dialer := net.Dialer{}
+		tcpConn, err := dialer.DialContext(ctx, "tcp", turnServerAddr)
+		if err != nil {
+			probeErr = fmt.Errorf("dial: %w", err)
+			return
+		}
+		conn = turn.NewSTUNConn(tcpConn)
+	default:
+		probeErr = fmt.Errorf("conn: unknown proto: %s", uri.Proto)
+		return
+	}
+
+	defer func() {
+		if err := conn.Close(); err != nil && probeErr == nil {
+			probeErr = fmt.Errorf("conn close: %w", err)
+		}
+	}()
+
+	cfg := &turn.ClientConfig{
+		STUNServerAddr: turnServerAddr,
+		TURNServerAddr: turnServerAddr,
+		Conn:           conn,
+		Username:       uri.Username,
+		Password:       uri.Password,
+	}
+	client, err := turn.NewClient(cfg)
+	if err != nil {
+		probeErr = fmt.Errorf("create client: %w", err)
+		return
+	}
+	defer client.Close()
+
+	if err := client.Listen(); err != nil {
+		probeErr = fmt.Errorf("client listen: %w", err)
+		return
+	}
+
+	relayConn, err := client.Allocate()
+	if err != nil {
+		probeErr = fmt.Errorf("allocate: %w", err)
+		return
+	}
+	defer func() {
+		if err := relayConn.Close(); err != nil && probeErr == nil {
+			probeErr = fmt.Errorf("close relay conn: %w", err)
+		}
+	}()
+
+	log.Debugf("turn probe relay address from %s: %s", uri, relayConn.LocalAddr())
+
+	return relayConn.LocalAddr().String(), nil
+}
+
+// ProbeAll probes all given servers asynchronously and returns the results
+func ProbeAll(
+	ctx context.Context,
+	fn func(ctx context.Context, uri *stun.URI) (addr string, probeErr error),
+	relays []*stun.URI,
+) []ProbeResult {
+	results := make([]ProbeResult, len(relays))
+
+	var wg sync.WaitGroup
+	for i, uri := range relays {
+		ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
+		defer cancel()
+
+		wg.Add(1)
+		go func(res *ProbeResult, stunURI *stun.URI) {
+			defer wg.Done()
+			res.URI = stunURI
+			res.Addr, res.Err = fn(ctx, stunURI)
+		}(&results[i], uri)
+	}
+
+	wg.Wait()
+
+	return results
+}