fix: protect port forward receive chan map from concurrent read/write

Signed-off-by: Alf-Rune Siqveland <alf.rune@northern.tech>

Author: alfrunes

cmd/portforward.go

@@ -21,6 +21,7 @@ import (
 	"os/signal"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/mendersoftware/go-lib-micro/ws"
@@ -94,7 +95,8 @@ type PortForwardCmd struct {
 	sessionID    string
 	bindingHost  string
 	portMappings []portMapping
-	recvChans    map[string]chan *ws.ProtoMsg
+	recvChanMu   sync.RWMutex
+	recvChans    map[string]chan<- *ws.ProtoMsg
 	running      bool
 	stop         chan struct{}
 	err          error
@@ -184,7 +186,7 @@ func NewPortForwardCmd(cmd *cobra.Command, args []string) (*PortForwardCmd, erro
 		deviceID:     args[0],
 		bindingHost:  bindingHost,
 		portMappings: portMappings,
-		recvChans:    make(map[string]chan *ws.ProtoMsg),
+		recvChans:    make(map[string]chan<- *ws.ProtoMsg),
 		stop:         make(chan struct{}),
 	}, nil
 }
@@ -198,6 +200,12 @@ func (c *PortForwardCmd) Run() error {
 	}
 }
 
+func (c *PortForwardCmd) registerRecvChan(connectionID string, recvChan chan<- *ws.ProtoMsg) {
+	c.recvChanMu.Lock()
+	defer c.recvChanMu.Unlock()
+	c.recvChans[connectionID] = recvChan
+}
+
 func (c *PortForwardCmd) run() error {
 	ctx, cancelContext := context.WithCancel(context.Background())
 	defer cancelContext()
@@ -239,14 +247,14 @@ func (c *PortForwardCmd) run() error {
 			if err != nil {
 				return err
 			}
-			go forwarder.Run(ctx, c.sessionID, msgChan, c.recvChans)
+			go forwarder.Run(ctx, c.sessionID, msgChan, c.registerRecvChan)
 		case protocolUDP:
 			forwarder, err := NewUDPPortForwarder(c.bindingHost, portMapping.LocalPort,
 				portMapping.RemoteHost, portMapping.RemotePort)
 			if err != nil {
 				return err
 			}
-			go forwarder.Run(ctx, c.sessionID, msgChan, c.recvChans)
+			go forwarder.Run(ctx, c.sessionID, msgChan, c.registerRecvChan)
 		default:
 			return errors.New("unknown protocol: " + portMapping.Protocol)
 		}
@@ -411,9 +419,11 @@ func (c *PortForwardCmd) processIncomingMessages(
 				m.Header.MsgType == wspf.MessageTypePortForwardStop) {
 			connectionID, _ := m.Header.Properties[wspf.PropertyConnectionID].(string)
 			if connectionID != "" {
+				c.recvChanMu.RLock()
 				if recvChan, ok := c.recvChans[connectionID]; ok {
 					recvChan <- m
 				}
+				c.recvChanMu.RUnlock()
 			}
 		}
 	}

cmd/portforward_tcp.go

@@ -60,7 +60,7 @@ func (p *TCPPortForwarder) Run(
 	ctx context.Context,
 	sessionID string,
 	msgChan chan *ws.ProtoMsg,
-	recvChans map[string]chan *ws.ProtoMsg,
+	registerRecvChan func(string, chan<- *ws.ProtoMsg),
 ) {
 	// listen for new connections
 	defer p.listen.Close()
@@ -101,7 +101,7 @@ func (p *TCPPortForwarder) Run(
 			connectionUUID, _ := uuid.NewUUID()
 			connectionID := connectionUUID.String()
 			recvChan := make(chan *ws.ProtoMsg, portForwardTCPChannelSize)
-			recvChans[connectionID] = recvChan
+			registerRecvChan(connectionID, recvChan)
 			go p.handleRequest(ctx, conn, sessionID, connectionID, recvChan, msgChan)
 		case <-ctx.Done():
 			return
@@ -135,8 +135,15 @@ func (p *TCPPortForwarder) handleInboundMessages(
 	}()
 
 	for {
+		var (
+			m    *ws.ProtoMsg
+			open bool
+		)
 		select {
-		case m := <-recvChan:
+		case m, open = <-recvChan:
+			if !open {
+				return
+			}
 			if m.Header.Proto == ws.ProtoTypePortForward &&
 				m.Header.MsgType == wspf.MessageTypePortForwardStop {
 				sendStopMessage = false
@@ -145,12 +152,13 @@ func (p *TCPPortForwarder) handleInboundMessages(
 				m.Header.MsgType == wspf.MessageTypePortForward {
 				_, err := conn.Write(m.Body)
 				if err != nil {
-					if errors.Unwrap(err) != net.ErrClosed {
+					if !errors.Is(err, net.ErrClosed) {
 						fmt.Fprintf(os.Stderr, "error: %v\n", err.Error())
 					}
+					return
 				} else {
 					// send the ack
-					m := &ws.ProtoMsg{
+					m = &ws.ProtoMsg{
 						Header: ws.ProtoHdr{
 							Proto:     ws.ProtoTypePortForward,
 							MsgType:   wspf.MessageTypePortForwardAck,
@@ -164,7 +172,10 @@ func (p *TCPPortForwarder) handleInboundMessages(
 				}
 			} else if m.Header.Proto == ws.ProtoTypePortForward &&
 				m.Header.MsgType == wspf.MessageTypePortForwardAck {
-				<-ackChan
+				_, open = <-ackChan
+				if !open {
+					return
+				}
 			}
 		case <-ctx.Done():
 			return

cmd/portforward_udp.go

@@ -73,15 +73,15 @@ func (p *UDPPortForwarder) Run(
 	ctx context.Context,
 	sessionID string,
 	msgChan chan *ws.ProtoMsg,
-	recvChans map[string]chan *ws.ProtoMsg,
+	registerRecvChan func(string, chan<- *ws.ProtoMsg),
 ) {
 	// listen for new connections
 	defer p.conn.Close()
 
 	connectionUUID, _ := uuid.NewUUID()
 	connectionID := connectionUUID.String()
 	recvChan := make(chan *ws.ProtoMsg, portForwardUDPChannelSize)
-	recvChans[connectionID] = recvChan
+	registerRecvChan(connectionID, recvChan)
 
 	protocol := portforward.PortForwardProtocol(wspf.PortForwardProtocolUDP)
 	portforwardNew := &wspf.PortForwardNew{