Add udp listener and did some change for debug purpose.

relay/server/listener/udp/listener.go

@@ -0,0 +1,96 @@
+package udp
+
+import (
+	"net"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/relay/server/listener"
+)
+
+type Listener struct {
+	address string
+
+	onAcceptFn func(conn net.Conn)
+
+	conns    map[string]*UDPConn
+	wg       sync.WaitGroup
+	quit     chan struct{}
+	lock     sync.Mutex
+	listener *net.UDPConn
+}
+
+func NewListener(address string) listener.Listener {
+	return &Listener{
+		address: address,
+		conns:   make(map[string]*UDPConn),
+	}
+}
+
+func (l *Listener) Listen(onAcceptFn func(conn net.Conn)) error {
+	l.lock.Lock()
+
+	l.onAcceptFn = onAcceptFn
+	l.quit = make(chan struct{})
+
+	addr := &net.UDPAddr{
+		Port: 1234,
+		IP:   net.ParseIP("0.0.0.0"),
+	}
+	li, err := net.ListenUDP("udp", addr)
+	if err != nil {
+		log.Errorf("%s", err)
+		l.lock.Unlock()
+		return err
+	}
+	log.Debugf("udp server is listening on address: %s", l.address)
+	l.listener = li
+	l.wg.Add(1)
+	go l.readLoop()
+
+	l.lock.Unlock()
+	<-l.quit
+	return nil
+}
+
+// Close todo: prevent multiple call (do not close two times the channel)
+func (l *Listener) Close() error {
+	l.lock.Lock()
+	defer l.lock.Unlock()
+
+	close(l.quit)
+	err := l.listener.Close()
+	l.wg.Wait()
+	return err
+}
+
+func (l *Listener) readLoop() {
+	defer l.wg.Done()
+
+	for {
+		buf := make([]byte, 1500)
+		n, addr, err := l.listener.ReadFromUDP(buf)
+		if err != nil {
+			select {
+			case <-l.quit:
+				return
+			default:
+				log.Errorf("failed to accept connection: %s", err)
+				continue
+			}
+		}
+
+		pConn, ok := l.conns[addr.String()]
+		if ok {
+			pConn.onNewMsg(buf[:n])
+			continue
+		}
+
+		pConn = NewConn(l.listener, addr)
+		l.conns[addr.String()] = pConn
+		go l.onAcceptFn(pConn)
+		pConn.onNewMsg(buf[:n])
+
+	}
+}

relay/server/listener/udp/udp_conn.go

@@ -0,0 +1,68 @@
+package udp
+
+import (
+	"io"
+	"net"
+	"time"
+)
+
+type UDPConn struct {
+	*net.UDPConn
+	addr       *net.UDPAddr
+	msgChannel chan []byte
+}
+
+func NewConn(conn *net.UDPConn, addr *net.UDPAddr) *UDPConn {
+	return &UDPConn{
+		UDPConn:    conn,
+		addr:       addr,
+		msgChannel: make(chan []byte),
+	}
+}
+
+func (u *UDPConn) Read(b []byte) (n int, err error) {
+	msg, ok := <-u.msgChannel
+	if !ok {
+		return 0, io.EOF
+	}
+
+	n = copy(b, msg)
+	return n, nil
+}
+
+func (u *UDPConn) Write(b []byte) (n int, err error) {
+	return u.UDPConn.WriteTo(b, u.addr)
+}
+
+func (u *UDPConn) Close() error {
+	//TODO implement me
+	//panic("implement me")
+	return nil
+}
+
+func (u *UDPConn) LocalAddr() net.Addr {
+	return u.UDPConn.LocalAddr()
+}
+
+func (u *UDPConn) RemoteAddr() net.Addr {
+	return u.addr
+}
+
+func (u *UDPConn) SetDeadline(t time.Time) error {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (u *UDPConn) SetReadDeadline(t time.Time) error {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (u *UDPConn) SetWriteDeadline(t time.Time) error {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (u *UDPConn) onNewMsg(b []byte) {
+	u.msgChannel <- b
+}

relay/server/listener/ws/server_conn.go

@@ -2,6 +2,7 @@ package ws
 
 import (
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -10,11 +11,13 @@ import (
 
 type Conn struct {
 	*websocket.Conn
+
+	mu sync.Mutex
 }
 
 func NewConn(wsConn *websocket.Conn) *Conn {
 	return &Conn{
-		wsConn,
+		Conn: wsConn,
 	}
 }
 
@@ -33,7 +36,9 @@ func (c *Conn) Read(b []byte) (n int, err error) {
 }
 
 func (c *Conn) Write(b []byte) (int, error) {
+	c.mu.Lock()
 	err := c.WriteMessage(websocket.BinaryMessage, b)
+	c.mu.Unlock()
 	return len(b), err
 }
 

relay/server/server.go

@@ -9,7 +9,7 @@ import (
 
 	"github.com/netbirdio/netbird/relay/messages"
 	"github.com/netbirdio/netbird/relay/server/listener"
-	"github.com/netbirdio/netbird/relay/server/listener/ws"
+	"github.com/netbirdio/netbird/relay/server/listener/udp"
 )
 
 // Server
@@ -30,7 +30,7 @@ func NewServer() *Server {
 }
 
 func (r *Server) Listen(address string) error {
-	r.listener = ws.NewListener(address)
+	r.listener = udp.NewListener(address)
 	return r.listener.Listen(r.accept)
 }
 
@@ -51,7 +51,7 @@ func (r *Server) accept(conn net.Conn) {
 		}
 		return
 	}
-	peer.Log.Debugf("on new connection: %s", conn.RemoteAddr())
+	peer.Log.Debugf("peer connected from: %s", conn.RemoteAddr())
 
 	r.store.AddPeer(peer)
 	defer func() {
@@ -59,8 +59,8 @@ func (r *Server) accept(conn net.Conn) {
 		r.store.DeletePeer(peer)
 	}()
 
-	buf := make([]byte, 65535) // todo: optimize buffer size
 	for {
+		buf := make([]byte, 1500) // todo: optimize buffer size
 		n, err := conn.Read(buf)
 		if err != nil {
 			if err != io.EOF {
@@ -98,18 +98,18 @@ func (r *Server) accept(conn net.Conn) {
 				peer.Log.Errorf("failed to unmarshal transport message: %s", err)
 				continue
 			}
+			go func() {
+				foreignChannelID, remoteConn, err := peer.ConnByChannelID(channelId)
+				if err != nil {
+					peer.Log.Errorf("failed to transport message from peer '%s' to '%d': %s", peer.ID(), channelId, err)
+					return
+				}
 
-			foreignChannelID, remoteConn, err := peer.ConnByChannelID(channelId)
-			if err != nil {
-				peer.Log.Errorf("failed to transport message from peer '%s' to '%d': %s", peer.ID(), channelId, err)
-				continue
-			}
-
-			err = transportTo(remoteConn, foreignChannelID, msg)
-			if err != nil {
-				peer.Log.Errorf("failed to transport message from peer '%s' to '%d': %s", peer.ID(), channelId, err)
-				continue
-			}
+				err = transportTo(remoteConn, foreignChannelID, msg)
+				if err != nil {
+					peer.Log.Errorf("failed to transport message from peer '%s' to '%d': %s", peer.ID(), channelId, err)
+				}
+			}()
 		}
 	}
 }