Fix cognitive complexity

relay/testec2/main.go

@@ -11,7 +11,6 @@ import (
 
 	log "github.com/sirupsen/logrus"
 
-	"github.com/netbirdio/netbird/relay/testec2/tun"
 	"github.com/netbirdio/netbird/util"
 )
 
@@ -27,14 +26,9 @@ var (
 	relaySrvAddress string
 	turnSrvAddress  string
 	signalURL       string
-	udpListener     string
+	udpListener     string // used for TURN test
 )
 
-type TurnReceiver struct {
-	conns           []*net.UDPConn
-	clientAddresses map[string]string
-	devices         []*tun.Device
-}
 type testResult struct {
 	numOfPairs int
 	duration   time.Duration
@@ -135,56 +129,24 @@ func TRUNSenderMain() {
 
 	for _, p := range pairs {
 		log.Infof("running test with %d pairs", p)
+		turnSender := &TurnSender{}
 
-		turnConns := make(map[string]*TurnConn)
-		addresses := make([]string, 0, len(pairs))
-		for i := 0; i < p; i++ {
-			tc := AllocateTurnClient(turnSrvAddress)
-			log.Infof("allocated turn client: %s", tc.Address().String())
-			turnConns[tc.Address().String()] = tc
-			addresses = append(addresses, tc.Address().String())
-		}
+		createTurnConns(p, turnSender)
 
-		log.Infof("send addresses via signal server: %d", len(addresses))
-		clientAddresses, err := ss.SendAddress(addresses)
+		log.Infof("send addresses via signal server: %d", len(turnSender.addresses))
+		clientAddresses, err := ss.SendAddress(turnSender.addresses)
 		if err != nil {
 			log.Fatalf("failed to send address: %s", err)
 		}
 		log.Infof("received addresses: %v", clientAddresses.Address)
 
-		var i int
-		devices := make([]*tun.Device, 0, len(clientAddresses.Address))
-		for k, v := range clientAddresses.Address {
-			tc, ok := turnConns[k]
-			if !ok {
-				log.Fatalf("failed to find turn conn: %s", k)
-			}
-
-			addr, err := net.ResolveUDPAddr("udp", v)
-			if err != nil {
-				log.Fatalf("failed to resolve udp address: %s", err)
-			}
-			device := &tun.Device{
-				Name:    fmt.Sprintf("mtun-sender-%d", i),
-				IP:      fmt.Sprintf("10.0.%d.1", i),
-				PConn:   tc.relayConn,
-				DstAddr: addr,
-			}
-
-			err = device.Up()
-			if err != nil {
-				log.Fatalf("failed to bring up device: %s", err)
-			}
-
-			devices = append(devices, device)
-			i++
-		}
+		createSenderDevices(turnSender, clientAddresses)
 
 		log.Infof("waiting for tcpListeners to be ready")
 		time.Sleep(2 * time.Second)
 
-		tcpConns := make([]net.Conn, 0, len(devices))
-		for i := range devices {
+		tcpConns := make([]net.Conn, 0, len(turnSender.devices))
+		for i := range turnSender.devices {
 			addr := fmt.Sprintf("10.0.%d.2:9999", i)
 			log.Infof("dialing: %s", addr)
 			tcpConn, err := net.Dial("tcp", addr)
@@ -194,44 +156,17 @@ func TRUNSenderMain() {
 			tcpConns = append(tcpConns, tcpConn)
 		}
 
-		log.Infof("start test data transfer for %d pairs", len(devices))
+		log.Infof("start test data transfer for %d pairs", p)
 		testDataLen := len(testData)
 		wg := sync.WaitGroup{}
+		wg.Add(len(tcpConns))
 		for i, tcpConn := range tcpConns {
 			log.Infof("sending test data to device: %d", i)
-			wg.Add(1)
-			go func(i int, tcpConn net.Conn) {
-				defer wg.Done()
-				defer tcpConn.Close()
-
-				log.Infof("start to sending test data: %s", tcpConn.RemoteAddr())
-
-				si := NewStartInidication(time.Now(), testDataLen)
-				_, err = tcpConn.Write(si)
-				if err != nil {
-					log.Errorf("failed to write to tcp: %s", err)
-					return
-				}
-
-				pieceSize := 1024
-				for j := 0; j < testDataLen; j += pieceSize {
-					end := j + pieceSize
-					if end > testDataLen {
-						end = testDataLen
-					}
-					_, writeErr := tcpConn.Write(testData[j:end])
-					if writeErr != nil {
-						log.Errorf("failed to write to tcp conn: %s", writeErr)
-						return
-					}
-				}
-
-				time.Sleep(3 * time.Second)
-			}(i, tcpConn)
+			go runTurnWriting(tcpConn, testData, testDataLen, &wg)
 		}
 		wg.Wait()
 
-		for _, d := range devices {
+		for _, d := range turnSender.devices {
 			_ = d.Close()
 		}
 
@@ -293,83 +228,6 @@ func TURNReaderMain() []testResult {
 	return testResults
 }
 
-func runTurnReading(d *tun.Device, durations chan time.Duration) {
-	tcpListener, err := net.Listen("tcp", d.IP+":9999")
-	if err != nil {
-		log.Fatalf("failed to listen on tcp: %s", err)
-	}
-	defer tcpListener.Close()
-	log := log.WithField("device", tcpListener.Addr())
-
-	tcpConn, err := tcpListener.Accept()
-	if err != nil {
-		log.Fatalf("failed to accept connection: %s", err)
-	}
-	log.Infof("remote peer connected")
-
-	buf := make([]byte, 103)
-	n, err := tcpConn.Read(buf)
-	if err != nil {
-		log.Fatalf(errMsgFailedReadTCP, err)
-	}
-
-	si := DecodeStartIndication(buf[:n])
-	log.Infof("received start indication: %v, %d", si, n)
-
-	buf = make([]byte, 8192)
-	i, err := tcpConn.Read(buf)
-	if err != nil {
-		log.Fatalf(errMsgFailedReadTCP, err)
-	}
-	now := time.Now()
-	for i < si.TransferSize {
-		n, err := tcpConn.Read(buf)
-		if err != nil {
-			log.Fatalf(errMsgFailedReadTCP, err)
-		}
-		i += n
-	}
-	durations <- time.Since(now)
-}
-
-func createDevices(addresses []string, receiver *TurnReceiver) error {
-	receiver.conns = make([]*net.UDPConn, 0, len(addresses))
-	receiver.clientAddresses = make(map[string]string, len(addresses))
-	receiver.devices = make([]*tun.Device, 0, len(addresses))
-	for i, addr := range addresses {
-		localAddr, err := net.ResolveUDPAddr("udp", udpListener)
-		if err != nil {
-			return fmt.Errorf("failed to resolve UDP address: %s", err)
-		}
-
-		conn, err := net.ListenUDP("udp", localAddr)
-		if err != nil {
-			return fmt.Errorf("failed to create UDP connection: %s", err)
-		}
-
-		receiver.conns = append(receiver.conns, conn)
-		receiver.clientAddresses[addr] = conn.LocalAddr().String()
-
-		dstAddr, err := net.ResolveUDPAddr("udp", addr)
-		if err != nil {
-			return fmt.Errorf("failed to resolve address: %s", err)
-		}
-
-		device := &tun.Device{
-			Name:    fmt.Sprintf("mtun-%d", i),
-			IP:      fmt.Sprintf("10.0.%d.2", i),
-			PConn:   conn,
-			DstAddr: dstAddr,
-		}
-
-		if err = device.Up(); err != nil {
-			return fmt.Errorf("failed to bring up device: %s, %s", device.Name, err)
-		}
-		receiver.devices = append(receiver.devices, device)
-	}
-	return nil
-}
-
 func main() {
 	var mode string
 

relay/testec2/turn.go

@@ -3,79 +3,174 @@ package main
 import (
 	"fmt"
 	"net"
+	"sync"
+	"time"
+
+	"github.com/netbirdio/netbird/relay/testec2/tun"
 
-	"github.com/pion/logging"
-	"github.com/pion/turn/v3"
 	log "github.com/sirupsen/logrus"
 )
 
-type TurnConn struct {
-	conn       net.Conn
-	turnClient *turn.Client
-	relayConn  net.PacketConn
+type TurnReceiver struct {
+	conns           []*net.UDPConn
+	clientAddresses map[string]string
+	devices         []*tun.Device
 }
 
-func (tc *TurnConn) Address() net.Addr {
-	return tc.relayConn.LocalAddr()
+type TurnSender struct {
+	turnConns map[string]*TurnConn
+	addresses []string
+	devices   []*tun.Device
 }
 
-func (tc *TurnConn) Close() {
-	_ = tc.relayConn.Close()
-	tc.turnClient.Close()
-	_ = tc.conn.Close()
+func runTurnWriting(tcpConn net.Conn, testData []byte, testDataLen int, wg *sync.WaitGroup) {
+	defer wg.Done()
+	defer tcpConn.Close()
+
+	log.Infof("start to sending test data: %s", tcpConn.RemoteAddr())
+
+	si := NewStartInidication(time.Now(), testDataLen)
+	_, err := tcpConn.Write(si)
+	if err != nil {
+		log.Errorf("failed to write to tcp: %s", err)
+		return
+	}
+
+	pieceSize := 1024
+	for j := 0; j < testDataLen; j += pieceSize {
+		end := j + pieceSize
+		if end > testDataLen {
+			end = testDataLen
+		}
+		_, writeErr := tcpConn.Write(testData[j:end])
+		if writeErr != nil {
+			log.Errorf("failed to write to tcp conn: %s", writeErr)
+			return
+		}
+	}
+
+	// grant time to flush out packages
+	time.Sleep(3 * time.Second)
 }
 
-func AllocateTurnClient(serverAddr string) *TurnConn {
-	conn, err := net.Dial("tcp", serverAddr)
-	if err != nil {
-		log.Fatal(err)
-	}
+func createSenderDevices(sender *TurnSender, clientAddresses *ClientPeerAddr) {
+	var i int
+	devices := make([]*tun.Device, 0, len(clientAddresses.Address))
+	for k, v := range clientAddresses.Address {
+		tc, ok := sender.turnConns[k]
+		if !ok {
+			log.Fatalf("failed to find turn conn: %s", k)
+		}
 
-	turnClient, err := getTurnClient(serverAddr, conn)
-	if err != nil {
-		log.Fatal(err)
-	}
+		addr, err := net.ResolveUDPAddr("udp", v)
+		if err != nil {
+			log.Fatalf("failed to resolve udp address: %s", err)
+		}
+		device := &tun.Device{
+			Name:    fmt.Sprintf("mtun-sender-%d", i),
+			IP:      fmt.Sprintf("10.0.%d.1", i),
+			PConn:   tc.relayConn,
+			DstAddr: addr,
+		}
 
-	relayConn, err := turnClient.Allocate()
-	if err != nil {
-		log.Fatal(err)
-	}
+		err = device.Up()
+		if err != nil {
+			log.Fatalf("failed to bring up device: %s", err)
+		}
 
-	return &TurnConn{
-		conn:       conn,
-		turnClient: turnClient,
-		relayConn:  relayConn,
+		devices = append(devices, device)
+		i++
 	}
+	sender.devices = devices
 }
 
-func getTurnClient(address string, conn net.Conn) (*turn.Client, error) {
-	// Dial TURN Server
-	addrStr := fmt.Sprintf("%s:%d", address, 443)
-
-	fac := logging.NewDefaultLoggerFactory()
-	//fac.DefaultLogLevel = logging.LogLevelTrace
-
-	// Start a new TURN Client and wrap our net.Conn in a STUNConn
-	// This allows us to simulate datagram based communication over a net.Conn
-	cfg := &turn.ClientConfig{
-		TURNServerAddr: address,
-		Conn:           turn.NewSTUNConn(conn),
-		Username:       "test",
-		Password:       "test",
-		LoggerFactory:  fac,
+func createTurnConns(p int, sender *TurnSender) {
+	turnConns := make(map[string]*TurnConn)
+	addresses := make([]string, 0, len(pairs))
+	for i := 0; i < p; i++ {
+		tc := AllocateTurnClient(turnSrvAddress)
+		log.Infof("allocated turn client: %s", tc.Address().String())
+		turnConns[tc.Address().String()] = tc
+		addresses = append(addresses, tc.Address().String())
 	}
 
-	client, err := turn.NewClient(cfg)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create TURN client for server %s: %s", addrStr, err)
-	}
-
-	// Start listening on the conn provided.
-	err = client.Listen()
-	if err != nil {
-		client.Close()
-		return nil, fmt.Errorf("failed to listen on TURN client for server %s: %s", addrStr, err)
-	}
-
-	return client, nil
+	sender.turnConns = turnConns
+	sender.addresses = addresses
+}
+
+func runTurnReading(d *tun.Device, durations chan time.Duration) {
+	tcpListener, err := net.Listen("tcp", d.IP+":9999")
+	if err != nil {
+		log.Fatalf("failed to listen on tcp: %s", err)
+	}
+	defer tcpListener.Close()
+	log := log.WithField("device", tcpListener.Addr())
+
+	tcpConn, err := tcpListener.Accept()
+	if err != nil {
+		log.Fatalf("failed to accept connection: %s", err)
+	}
+	log.Infof("remote peer connected")
+
+	buf := make([]byte, 103)
+	n, err := tcpConn.Read(buf)
+	if err != nil {
+		log.Fatalf(errMsgFailedReadTCP, err)
+	}
+
+	si := DecodeStartIndication(buf[:n])
+	log.Infof("received start indication: %v, %d", si, n)
+
+	buf = make([]byte, 8192)
+	i, err := tcpConn.Read(buf)
+	if err != nil {
+		log.Fatalf(errMsgFailedReadTCP, err)
+	}
+	now := time.Now()
+	for i < si.TransferSize {
+		n, err := tcpConn.Read(buf)
+		if err != nil {
+			log.Fatalf(errMsgFailedReadTCP, err)
+		}
+		i += n
+	}
+	durations <- time.Since(now)
+}
+
+func createDevices(addresses []string, receiver *TurnReceiver) error {
+	receiver.conns = make([]*net.UDPConn, 0, len(addresses))
+	receiver.clientAddresses = make(map[string]string, len(addresses))
+	receiver.devices = make([]*tun.Device, 0, len(addresses))
+	for i, addr := range addresses {
+		localAddr, err := net.ResolveUDPAddr("udp", udpListener)
+		if err != nil {
+			return fmt.Errorf("failed to resolve UDP address: %s", err)
+		}
+
+		conn, err := net.ListenUDP("udp", localAddr)
+		if err != nil {
+			return fmt.Errorf("failed to create UDP connection: %s", err)
+		}
+
+		receiver.conns = append(receiver.conns, conn)
+		receiver.clientAddresses[addr] = conn.LocalAddr().String()
+
+		dstAddr, err := net.ResolveUDPAddr("udp", addr)
+		if err != nil {
+			return fmt.Errorf("failed to resolve address: %s", err)
+		}
+
+		device := &tun.Device{
+			Name:    fmt.Sprintf("mtun-%d", i),
+			IP:      fmt.Sprintf("10.0.%d.2", i),
+			PConn:   conn,
+			DstAddr: dstAddr,
+		}
+
+		if err = device.Up(); err != nil {
+			return fmt.Errorf("failed to bring up device: %s, %s", device.Name, err)
+		}
+		receiver.devices = append(receiver.devices, device)
+	}
+	return nil
 }

relay/testec2/turn_allocator.go

@@ -0,0 +1,81 @@
+package main
+
+import (
+	"fmt"
+	"net"
+
+	"github.com/pion/logging"
+	"github.com/pion/turn/v3"
+	log "github.com/sirupsen/logrus"
+)
+
+type TurnConn struct {
+	conn       net.Conn
+	turnClient *turn.Client
+	relayConn  net.PacketConn
+}
+
+func (tc *TurnConn) Address() net.Addr {
+	return tc.relayConn.LocalAddr()
+}
+
+func (tc *TurnConn) Close() {
+	_ = tc.relayConn.Close()
+	tc.turnClient.Close()
+	_ = tc.conn.Close()
+}
+
+func AllocateTurnClient(serverAddr string) *TurnConn {
+	conn, err := net.Dial("tcp", serverAddr)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	turnClient, err := getTurnClient(serverAddr, conn)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	relayConn, err := turnClient.Allocate()
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return &TurnConn{
+		conn:       conn,
+		turnClient: turnClient,
+		relayConn:  relayConn,
+	}
+}
+
+func getTurnClient(address string, conn net.Conn) (*turn.Client, error) {
+	// Dial TURN Server
+	addrStr := fmt.Sprintf("%s:%d", address, 443)
+
+	fac := logging.NewDefaultLoggerFactory()
+	//fac.DefaultLogLevel = logging.LogLevelTrace
+
+	// Start a new TURN Client and wrap our net.Conn in a STUNConn
+	// This allows us to simulate datagram based communication over a net.Conn
+	cfg := &turn.ClientConfig{
+		TURNServerAddr: address,
+		Conn:           turn.NewSTUNConn(conn),
+		Username:       "test",
+		Password:       "test",
+		LoggerFactory:  fac,
+	}
+
+	client, err := turn.NewClient(cfg)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create TURN client for server %s: %s", addrStr, err)
+	}
+
+	// Start listening on the conn provided.
+	err = client.Listen()
+	if err != nil {
+		client.Close()
+		return nil, fmt.Errorf("failed to listen on TURN client for server %s: %s", addrStr, err)
+	}
+
+	return client, nil
+}