[client] Fix/ice handshake (#4281)

In this PR, speed up the GRPC message processing, force the recreation of the ICE agent when getting a new, remote offer (do not wait for local STUN timeout).

shared/signal/client/client.go

@@ -52,7 +52,7 @@ func UnMarshalCredential(msg *proto.Message) (*Credential, error) {
 }
 
 // MarshalCredential marshal a Credential instance and returns a Message object
-func MarshalCredential(myKey wgtypes.Key, myPort int, remoteKey string, credential *Credential, t proto.Body_Type, rosenpassPubKey []byte, rosenpassAddr string, relaySrvAddress string) (*proto.Message, error) {
+func MarshalCredential(myKey wgtypes.Key, myPort int, remoteKey string, credential *Credential, t proto.Body_Type, rosenpassPubKey []byte, rosenpassAddr string, relaySrvAddress string, sessionID []byte) (*proto.Message, error) {
 	return &proto.Message{
 		Key:       myKey.PublicKey().String(),
 		RemoteKey: remoteKey,
@@ -66,6 +66,7 @@ func MarshalCredential(myKey wgtypes.Key, myPort int, remoteKey string, credenti
 				RosenpassServerAddr: rosenpassAddr,
 			},
 			RelayServerAddress: relaySrvAddress,
+			SessionId:          sessionID,
 		},
 	}, nil
 }

shared/signal/client/grpc.go

@@ -45,19 +45,10 @@ type GrpcClient struct {
 	connStateCallbackLock sync.RWMutex
 
 	onReconnectedListenerFn func()
-}
 
-func (c *GrpcClient) StreamConnected() bool {
-	return c.status == StreamConnected
-}
-
-func (c *GrpcClient) GetStatus() Status {
-	return c.status
-}
-
-// Close Closes underlying connections to the Signal Exchange
-func (c *GrpcClient) Close() error {
-	return c.signalConn.Close()
+	decryptionWorker       *Worker
+	decryptionWorkerCancel context.CancelFunc
+	decryptionWg           sync.WaitGroup
 }
 
 // NewClient creates a new Signal client
@@ -93,6 +84,25 @@ func NewClient(ctx context.Context, addr string, key wgtypes.Key, tlsEnabled boo
 	}, nil
 }
 
+func (c *GrpcClient) StreamConnected() bool {
+	return c.status == StreamConnected
+}
+
+func (c *GrpcClient) GetStatus() Status {
+	return c.status
+}
+
+// Close Closes underlying connections to the Signal Exchange
+func (c *GrpcClient) Close() error {
+	if c.decryptionWorkerCancel != nil {
+		c.decryptionWorkerCancel()
+	}
+	c.decryptionWg.Wait()
+	c.decryptionWorker = nil
+
+	return c.signalConn.Close()
+}
+
 // SetConnStateListener set the ConnStateNotifier
 func (c *GrpcClient) SetConnStateListener(notifier ConnStateNotifier) {
 	c.connStateCallbackLock.Lock()
@@ -148,8 +158,12 @@ func (c *GrpcClient) Receive(ctx context.Context, msgHandler func(msg *proto.Mes
 
 		log.Infof("connected to the Signal Service stream")
 		c.notifyConnected()
+
+		// Start worker pool if not already started
+		c.startEncryptionWorker(msgHandler)
+
 		// start receiving messages from the Signal stream (from other peers through signal)
-		err = c.receive(stream, msgHandler)
+		err = c.receive(stream)
 		if err != nil {
 			if s, ok := status.FromError(err); ok && s.Code() == codes.Canceled {
 				log.Debugf("signal connection context has been canceled, this usually indicates shutdown")
@@ -174,6 +188,7 @@ func (c *GrpcClient) Receive(ctx context.Context, msgHandler func(msg *proto.Mes
 
 	return nil
 }
+
 func (c *GrpcClient) notifyStreamDisconnected() {
 	c.mux.Lock()
 	defer c.mux.Unlock()
@@ -382,11 +397,11 @@ func (c *GrpcClient) Send(msg *proto.Message) error {
 }
 
 // receive receives messages from other peers coming through the Signal Exchange
-func (c *GrpcClient) receive(stream proto.SignalExchange_ConnectStreamClient,
-	msgHandler func(msg *proto.Message) error) error {
-
+// and distributes them to worker threads for processing
+func (c *GrpcClient) receive(stream proto.SignalExchange_ConnectStreamClient) error {
 	for {
 		msg, err := stream.Recv()
+		// Handle errors immediately
 		switch s, ok := status.FromError(err); {
 		case ok && s.Code() == codes.Canceled:
 			log.Debugf("stream canceled (usually indicates shutdown)")
@@ -398,24 +413,37 @@ func (c *GrpcClient) receive(stream proto.SignalExchange_ConnectStreamClient,
 			log.Debugf("Signal Service stream closed by server")
 			return err
 		case err != nil:
+			log.Errorf("Stream receive error: %v", err)
 			return err
 		}
-		log.Tracef("received a new message from Peer [fingerprint: %s]", msg.Key)
 
-		decryptedMessage, err := c.decryptMessage(msg)
-		if err != nil {
-			log.Errorf("failed decrypting message of Peer [key: %s] error: [%s]", msg.Key, err.Error())
+		if msg == nil {
+			continue
 		}
 
-		err = msgHandler(decryptedMessage)
-
-		if err != nil {
-			log.Errorf("error while handling message of Peer [key: %s] error: [%s]", msg.Key, err.Error())
-			// todo send something??
+		if err := c.decryptionWorker.AddMsg(c.ctx, msg); err != nil {
+			log.Errorf("failed to add message to decryption worker: %v", err)
 		}
 	}
 }
 
+func (c *GrpcClient) startEncryptionWorker(handler func(msg *proto.Message) error) {
+	if c.decryptionWorker != nil {
+		return
+	}
+
+	c.decryptionWorker = NewWorker(c.decryptMessage, handler)
+	workerCtx, workerCancel := context.WithCancel(context.Background())
+	c.decryptionWorkerCancel = workerCancel
+
+	c.decryptionWg.Add(1)
+	go func() {
+		defer workerCancel()
+		c.decryptionWorker.Work(workerCtx)
+		c.decryptionWg.Done()
+	}()
+}
+
 func (c *GrpcClient) notifyDisconnected(err error) {
 	c.connStateCallbackLock.RLock()
 	defer c.connStateCallbackLock.RUnlock()

shared/signal/client/worker.go

@@ -0,0 +1,55 @@
+package client
+
+import (
+	"context"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/shared/signal/proto"
+)
+
+type Worker struct {
+	decryptMessage func(msg *proto.EncryptedMessage) (*proto.Message, error)
+	handler        func(msg *proto.Message) error
+
+	encryptedMsgPool chan *proto.EncryptedMessage
+}
+
+func NewWorker(decryptFn func(msg *proto.EncryptedMessage) (*proto.Message, error), handlerFn func(msg *proto.Message) error) *Worker {
+	return &Worker{
+		decryptMessage:   decryptFn,
+		handler:          handlerFn,
+		encryptedMsgPool: make(chan *proto.EncryptedMessage, 1),
+	}
+}
+
+func (w *Worker) AddMsg(ctx context.Context, msg *proto.EncryptedMessage) error {
+	// this is blocker because do not want to drop messages here
+	select {
+	case w.encryptedMsgPool <- msg:
+	case <-ctx.Done():
+	}
+	return nil
+}
+
+func (w *Worker) Work(ctx context.Context) {
+	for {
+		select {
+		case msg := <-w.encryptedMsgPool:
+			decryptedMessage, err := w.decryptMessage(msg)
+			if err != nil {
+				log.Errorf("failed to decrypt message: %v", err)
+				continue
+			}
+
+			if err := w.handler(decryptedMessage); err != nil {
+				log.Errorf("failed to handle message: %v", err)
+				continue
+			}
+
+		case <-ctx.Done():
+			log.Infof("Message worker stopping due to context cancellation")
+			return
+		}
+	}
+}