Feature/embedded STUN (#5062)

stun/server.go

@@ -0,0 +1,170 @@
+// Package stun provides an embedded STUN server for NAT traversal discovery.
+package stun
+
+import (
+	"errors"
+	"fmt"
+	"net"
+	"sync"
+
+	"github.com/hashicorp/go-multierror"
+	nberrors "github.com/netbirdio/netbird/client/errors"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/formatter"
+	"github.com/pion/stun/v3"
+)
+
+// ErrServerClosed is returned by Listen when the server is shut down gracefully.
+var ErrServerClosed = errors.New("stun: server closed")
+
+// ErrNoListeners is returned by Listen when no UDP connections were provided.
+var ErrNoListeners = errors.New("stun: no listeners configured")
+
+// Server implements a STUN server that responds to binding requests
+// with the client's reflexive transport address.
+type Server struct {
+	conns    []*net.UDPConn
+	logger   *log.Entry
+	logLevel log.Level
+
+	wg sync.WaitGroup
+}
+
+// NewServer creates a new STUN server with the given UDP listeners.
+// The caller is responsible for creating and providing the listeners.
+// logLevel can be: panic, fatal, error, warn, info, debug, trace
+func NewServer(conns []*net.UDPConn, logLevel string) *Server {
+	level, err := log.ParseLevel(logLevel)
+	if err != nil {
+		level = log.InfoLevel
+	}
+
+	// Create a separate logger with its own level setting
+	// This allows --stun-log-level to work independently of --log-level
+	stunLogger := log.New()
+	stunLogger.SetOutput(log.StandardLogger().Out)
+	stunLogger.SetLevel(level)
+	// Use the formatter package to set up formatter, ReportCaller, and context hook
+	formatter.SetTextFormatter(stunLogger)
+
+	logger := stunLogger.WithField("component", "stun-server")
+	logger.Infof("STUN server log level set to: %s", level.String())
+
+	return &Server{
+		conns:    conns,
+		logger:   logger,
+		logLevel: level,
+	}
+}
+
+// Listen starts the STUN server and blocks until the server is shut down.
+// Returns ErrServerClosed when shut down gracefully via Shutdown.
+// Returns ErrNoListeners if no UDP connections were provided.
+func (s *Server) Listen() error {
+	if len(s.conns) == 0 {
+		return ErrNoListeners
+	}
+
+	// Start a read loop for each listener
+	for _, conn := range s.conns {
+		s.logger.Infof("STUN server listening on %s", conn.LocalAddr())
+		s.wg.Add(1)
+		go s.readLoop(conn)
+	}
+
+	s.wg.Wait()
+	return ErrServerClosed
+}
+
+// readLoop continuously reads UDP packets and handles STUN requests.
+func (s *Server) readLoop(conn *net.UDPConn) {
+	defer s.wg.Done()
+	buf := make([]byte, 1500) // Standard MTU size
+	for {
+		n, remoteAddr, err := conn.ReadFromUDP(buf)
+
+		if err != nil {
+			// Check if the connection was closed externally
+			if errors.Is(err, net.ErrClosed) {
+				s.logger.Info("UDP connection closed, stopping read loop")
+				return
+			}
+			s.logger.Warnf("failed to read UDP packet: %v", err)
+			continue
+		}
+
+		// Handle packet in the same goroutine to avoid complexity
+		// STUN responses are small and fast
+		s.handlePacket(conn, buf[:n], remoteAddr)
+	}
+}
+
+// handlePacket processes a STUN request and sends a response.
+func (s *Server) handlePacket(conn *net.UDPConn, data []byte, addr *net.UDPAddr) {
+	localPort := conn.LocalAddr().(*net.UDPAddr).Port
+
+	s.logger.Debugf("[port:%d] received %d bytes from %s", localPort, len(data), addr)
+
+	// Check if it's a STUN message
+	if !stun.IsMessage(data) {
+		s.logger.Debugf("[port:%d] not a STUN message (first bytes: %x)", localPort, data[:min(len(data), 8)])
+		return
+	}
+
+	// Parse the STUN message
+	msg := &stun.Message{Raw: data}
+	if err := msg.Decode(); err != nil {
+		s.logger.Warnf("[port:%d] failed to decode STUN message from %s: %v", localPort, addr, err)
+		return
+	}
+
+	s.logger.Debugf("[port:%d] received STUN %s from %s (tx=%x)", localPort, msg.Type, addr, msg.TransactionID[:8])
+
+	// Only handle binding requests
+	if msg.Type != stun.BindingRequest {
+		s.logger.Debugf("[port:%d] ignoring non-binding request: %s", localPort, msg.Type)
+		return
+	}
+
+	// Build the response
+	response, err := stun.Build(
+		stun.NewTransactionIDSetter(msg.TransactionID),
+		stun.BindingSuccess,
+		&stun.XORMappedAddress{
+			IP:   addr.IP,
+			Port: addr.Port,
+		},
+		stun.Fingerprint,
+	)
+	if err != nil {
+		s.logger.Errorf("[port:%d] failed to build STUN response: %v", localPort, err)
+		return
+	}
+
+	// Send the response on the same connection it was received on
+	n, err := conn.WriteToUDP(response.Raw, addr)
+	if err != nil {
+		s.logger.Errorf("[port:%d] failed to send STUN response to %s: %v", localPort, addr, err)
+		return
+	}
+
+	s.logger.Debugf("[port:%d] sent STUN BindingSuccess to %s (%d bytes) with XORMappedAddress %s:%d", localPort, addr, n, addr.IP, addr.Port)
+}
+
+// Shutdown gracefully stops the STUN server.
+func (s *Server) Shutdown() error {
+	s.logger.Info("shutting down STUN server")
+
+	var merr *multierror.Error
+
+	for _, conn := range s.conns {
+		if err := conn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
+			merr = multierror.Append(merr, fmt.Errorf("close STUN UDP connection: %w", err))
+		}
+	}
+
+	// Wait for all readLoops to finish
+	s.wg.Wait()
+	return nberrors.FormatErrorOrNil(merr)
+}