Add per-stream ID to DTLS sessions

Client and server now include a 1-byte stream ID alongside the 16-byte session ID (17 bytes total) so multiple logical streams can be multiplexed per session. Client: propagate streamID through connection helpers and send sessionID+streamID when establishing DTLS connections; default stream 0 for the first connection and subsequent goroutines pass incremental stream IDs. Server: read 17-byte header, split into sessionID and streamID, and track connections as streamEntry {id, conn} rather than a plain conn slice. AddConn/RemoveConn signatures updated to accept stream IDs; new logic evicts existing connection with the same stream ID. backendReaderLoop uses a local round-robin index (removed atomic usage) and closes connections on write errors. Also adjusted per-stream read deadlines and cleaned up connection closing in Cleanup().

client/main.go

@@ -453,7 +453,7 @@ func dtlsFunc(ctx context.Context, conn net.PacketConn, peer *net.UDPAddr) (net.
 	return dtlsConn, nil
 }
 
-func oneDtlsConnection(ctx context.Context, peer *net.UDPAddr, listenConn net.PacketConn, connchan chan<- net.PacketConn, okchan chan<- struct{}, c chan<- error, sessionID []byte) {
+func oneDtlsConnection(ctx context.Context, peer *net.UDPAddr, listenConn net.PacketConn, connchan chan<- net.PacketConn, okchan chan<- struct{}, c chan<- error, sessionID []byte, streamID byte) {
 	var err error = nil
 	defer func() { c <- err }()
 	dtlsctx, dtlscancel := context.WithCancel(ctx)
@@ -482,14 +482,17 @@ func oneDtlsConnection(ctx context.Context, peer *net.UDPAddr, listenConn net.Pa
 		log.Printf("Closed DTLS connection\n")
 	}()
 
-	// Phase 1: Send Session ID
+	// Phase 1: Send Session ID + Stream ID (17 bytes)
 	dtlsConn.SetWriteDeadline(time.Now().Add(time.Second * 5))
-	if _, err1 = dtlsConn.Write(sessionID); err1 != nil {
+	idBuf := make([]byte, 17)
+	copy(idBuf[:16], sessionID)
+	idBuf[16] = streamID
+	if _, err1 = dtlsConn.Write(idBuf); err1 != nil {
 		err = fmt.Errorf("failed to send session ID: %s", err1)
 		return
 	}
 
-	log.Printf("Established DTLS connection and sent session ID!\n")
+	log.Printf("Established DTLS connection and sent session ID with stream %d!\n", streamID)
 	go func() {
 		for {
 			select {
@@ -766,14 +769,14 @@ func oneTurnConnection(ctx context.Context, turnParams *turnParams, peer *net.UD
 	conn2.SetDeadline(time.Time{})
 }
 
-func oneDtlsConnectionLoop(ctx context.Context, peer *net.UDPAddr, listenConnChan <-chan net.PacketConn, connchan chan<- net.PacketConn, okchan chan<- struct{}, sessionID []byte) {
+func oneDtlsConnectionLoop(ctx context.Context, peer *net.UDPAddr, listenConnChan <-chan net.PacketConn, connchan chan<- net.PacketConn, okchan chan<- struct{}, sessionID []byte, streamID byte) {
 	for {
 		select {
 		case <-ctx.Done():
 			return
 		case listenConn := <-listenConnChan:
 			c := make(chan error)
-			go oneDtlsConnection(ctx, peer, listenConn, connchan, okchan, c, sessionID)
+			go oneDtlsConnection(ctx, peer, listenConn, connchan, okchan, c, sessionID, streamID)
 			if err := <-c; err != nil {
 				log.Printf("%s", err)
 			}
@@ -913,7 +916,7 @@ func main() { //nolint:cyclop
 		wg1.Add(1)
 		go func() {
 			defer wg1.Done()
-			oneDtlsConnectionLoop(ctx, peer, listenConnChan, connchan, okchan, sessionID)
+			oneDtlsConnectionLoop(ctx, peer, listenConnChan, connchan, okchan, sessionID, 0)
 		}()
 
 		wg1.Add(1)
@@ -929,10 +932,10 @@ func main() { //nolint:cyclop
 		for i := 0; i < *n-1; i++ {
 			connchan := make(chan net.PacketConn)
 			wg1.Add(1)
-			go func() {
+			go func(streamID byte) {
 				defer wg1.Done()
-				oneDtlsConnectionLoop(ctx, peer, listenConnChan, connchan, nil, sessionID)
-			}()
+				oneDtlsConnectionLoop(ctx, peer, listenConnChan, connchan, nil, sessionID, streamID)
+			}(byte(i + 1))
 			wg1.Add(1)
 			go func() {
 				defer wg1.Done()

server/main.go

@@ -11,7 +11,6 @@ import (
 	"os"
 	"os/signal"
 	"sync"
-	"sync/atomic"
 	"syscall"
 	"time"
 
@@ -19,11 +18,15 @@ import (
 	"github.com/pion/dtls/v3/pkg/crypto/selfsign"
 )
 
+type streamEntry struct {
+	id   byte
+	conn net.Conn
+}
+
 type UserSession struct {
 	ID          string
-	Conns       []net.Conn
+	Conns       []streamEntry
 	BackendConn net.Conn
-	LastUsed    uint32
 	Lock        sync.RWMutex
 	Ctx         context.Context
 	Cancel      context.CancelFunc
@@ -51,12 +54,12 @@ func (s *SessionManager) GetOrCreate(ctx context.Context, id string, connectAddr
 	sessionCtx, cancel := context.WithCancel(ctx)
 	session := &UserSession{
 		ID:          id,
+		Conns:       make([]streamEntry, 0),
 		BackendConn: backendConn,
 		Manager:     s,
 		Ctx:         sessionCtx,
 		Cancel:      cancel,
 	}
-
 	s.Sessions[id] = session
 	go session.backendReaderLoop()
 
@@ -66,6 +69,7 @@ func (s *SessionManager) GetOrCreate(ctx context.Context, id string, connectAddr
 func (s *UserSession) backendReaderLoop() {
 	defer s.Cleanup()
 	buf := make([]byte, 1600)
+	var lastUsed uint32 = 0
 	for {
 		select {
 		case <-s.Ctx.Done():
@@ -81,43 +85,54 @@ func (s *UserSession) backendReaderLoop() {
 		}
 
 		s.Lock.RLock()
-		if len(s.Conns) == 0 {
+		nConns := uint32(len(s.Conns))
+		if nConns == 0 {
 			s.Lock.RUnlock()
 			continue
 		}
-		// Round-robin selection of DTLS connection
-		idx := atomic.AddUint32(&s.LastUsed, 1) % uint32(len(s.Conns))
-		conn := s.Conns[idx]
+
+		// Fast Round-robin selection using local variable
+		lastUsed = (lastUsed + 1) % nConns
+		conn := s.Conns[lastUsed].conn
 		s.Lock.RUnlock()
 
 		conn.SetWriteDeadline(time.Now().Add(time.Second * 10))
+
 		_, err = conn.Write(buf[:n])
 		if err != nil {
 			log.Printf("Session %s DTLS write error: %v", s.ID, err)
-			// Connection will be removed by its own reader loop
+			conn.Close()
 		}
-	}
+    }
 }
 
-func (s *UserSession) AddConn(conn net.Conn) {
+func (s *UserSession) AddConn(id byte, conn net.Conn) {
 	s.Lock.Lock()
 	defer s.Lock.Unlock()
-	s.Conns = append(s.Conns, conn)
+
+	// Evict existing connection with same ID
+	for i, entry := range s.Conns {
+		if entry.id == id {
+			//log.Printf("Session %s: Evicting old stream %d", s.ID, id)
+			entry.conn.Close()
+			s.Conns[i].conn = conn
+			return
+		}
+	}
+
+	s.Conns = append(s.Conns, streamEntry{id: id, conn: conn})
 }
 
-func (s *UserSession) RemoveConn(conn net.Conn) {
+func (s *UserSession) RemoveConn(id byte, conn net.Conn) {
 	s.Lock.Lock()
 	defer s.Lock.Unlock()
-	for i, c := range s.Conns {
-		if c == conn {
+	for i, entry := range s.Conns {
+		if entry.id == id && entry.conn == conn {
 			s.Conns = append(s.Conns[:i], s.Conns[i+1:]...)
 			break
 		}
 	}
-	// If all connections are gone, we might want to start a timer to cleanup the session
-	// but for now we'll keep it alive until backendReaderLoop fails or context is cancelled.
 }
-
 func (s *UserSession) Cleanup() {
 	s.Cancel()
 	s.BackendConn.Close()
@@ -127,13 +142,12 @@ func (s *UserSession) Cleanup() {
 	s.Manager.Lock.Unlock()
 
 	s.Lock.Lock()
-	for _, c := range s.Conns {
-		c.Close()
+	for _, entry := range s.Conns {
+		entry.conn.Close()
 	}
 	s.Conns = nil
 	s.Lock.Unlock()
 }
-
 func main() {
 	listen := flag.String("listen", "0.0.0.0:56000", "listen on ip:port")
 	connect := flag.String("connect", "", "connect to ip:port")
@@ -213,15 +227,16 @@ func main() {
 				return
 			}
 
-			// Phase 1: Read Session ID (16 bytes)
-			idBuf := make([]byte, 16)
+			// Phase 1: Read Session ID + Stream ID (17 bytes)
+			idBuf := make([]byte, 17)
 			conn.SetReadDeadline(time.Now().Add(time.Second * 5))
 			_, err := io.ReadFull(conn, idBuf)
 			if err != nil {
 				log.Println("Failed to read session ID:", err)
 				return
 			}
-			sessionID := fmt.Sprintf("%x", idBuf)
+			sessionID := fmt.Sprintf("%x", idBuf[:16])
+			streamID := idBuf[16]
 
 			session, err := manager.GetOrCreate(ctx, sessionID, *connect)
 			if err != nil {
@@ -229,15 +244,15 @@ func main() {
 				return
 			}
 
-			session.AddConn(conn)
-			defer session.RemoveConn(conn)
+			session.AddConn(streamID, conn)
+			defer session.RemoveConn(streamID, conn)
 
-			log.Printf("New stream for session %s from %s", sessionID, conn.RemoteAddr())
+			log.Printf("New stream %d for session %s from %s", streamID, sessionID, conn.RemoteAddr())
 
 			// Upstream Loop: DTLS -> Backend
 			buf := make([]byte, 1600)
 			for {
-				conn.SetReadDeadline(time.Now().Add(time.Minute * 10))
+				conn.SetReadDeadline(time.Now().Add(time.Minute * 5))
 				n, err := conn.Read(buf)
 				if err != nil {
 					log.Printf("Stream %s closed: %v", sessionID, err)