From ca58296e65e1da23c82ec6b096b2e4762082f4f9 Mon Sep 17 00:00:00 2001
From: Kleissner <peter@kleissner.investments>
Date: Mon, 13 Dec 2021 13:30:23 +0100
Subject: [PATCH] UDT: Fix race condition 'send on closed channel' by
 introducing a separate termination signal.

---
 udt/udtsocket.go      | 60 +++++++++++++++++++++++--------------------
 udt/udtsocket_send.go | 14 +++++-----
 2 files changed, 38 insertions(+), 36 deletions(-)

diff --git a/udt/udtsocket.go b/udt/udtsocket.go
index 49e4733..0205f5e 100644
--- a/udt/udtsocket.go
+++ b/udt/udtsocket.go
@@ -80,15 +80,16 @@ type udtSocket struct {
 	bandwidth       uint         // bandwidth reported from peer (packets/sec)
 
 	// channels
-	messageIn     chan []byte          // inbound messages. Sender is goReceiveEvent->ingestData, Receiver is client caller (Read)
-	messageOut    chan sendMessage     // outbound messages. Sender is client caller (Write), Receiver is goSendEvent. Closed when socket is closed
-	recvEvent     chan recvPktEvent    // receiver: ingest the specified packet. Sender is readPacket, receiver is goReceiveEvent
-	sendEvent     chan recvPktEvent    // sender: ingest the specified packet. Sender is readPacket, receiver is goSendEvent
-	sendPacket    chan packet.Packet   // packets to send out on the wire (once goManageConnection is running)
-	shutdownEvent chan shutdownMessage // channel signals the connection to be shutdown
-	sockClosed    chan struct{}        // closed when socket is closed
-	closeMutex    sync.Mutex
-	isClosed      bool
+	messageIn       chan []byte          // inbound messages. Sender is goReceiveEvent->ingestData, Receiver is client caller (Read)
+	messageOut      chan sendMessage     // outbound messages. Sender is client caller (Write), Receiver is goSendEvent. Closed when socket is closed
+	recvEvent       chan recvPktEvent    // receiver: ingest the specified packet. Sender is readPacket, receiver is goReceiveEvent
+	sendEvent       chan recvPktEvent    // sender: ingest the specified packet. Sender is readPacket, receiver is goSendEvent
+	sendPacket      chan packet.Packet   // packets to send out on the wire (once goManageConnection is running)
+	shutdownEvent   chan shutdownMessage // channel signals the connection to be shutdown
+	sockClosed      chan struct{}        // closed when socket is closed
+	terminateSignal chan struct{}        // termination signal
+	closeMutex      sync.Mutex
+	isClosed        bool
 
 	// timers
 	connTimeout <-chan time.Time // connecting: fires when connection attempt times out
@@ -271,6 +272,8 @@ func (s *udtSocket) Write(p []byte) (n int, err error) {
 			deadline = s.writeDeadline.C
 		}
 		select {
+		case <-s.terminateSignal:
+			return n, errors.New("terminate signal")
 		case s.messageOut <- sendMessage{content: data, tim: time.Now()}:
 			// send successful
 			return
@@ -300,7 +303,7 @@ func (s *udtSocket) Close() error {
 
 	s.isClosed = true
 
-	close(s.messageOut)
+	close(s.terminateSignal)
 	return nil
 }
 
@@ -419,24 +422,25 @@ func newSocket(m *multiplexer, config *Config, sockID uint32, isServer bool, isD
 		m:      m,
 		Config: config,
 		//raddr:          raddr,
-		created:        now,
-		sockState:      sockStateInit,
-		udtVer:         4,
-		isServer:       isServer,
-		maxPacketSize:  uint32(config.MaxPacketSize),
-		maxFlowWinSize: maxFlowWinSize,
-		isDatagram:     isDatagram,
-		sockID:         sockID,
-		initPktSeq:     packet.RandomPacketSequence(),
-		messageIn:      make(chan []byte, 256),
-		messageOut:     make(chan sendMessage, 256),
-		recvEvent:      make(chan recvPktEvent, 256),
-		sendEvent:      make(chan recvPktEvent, 256),
-		sockClosed:     make(chan struct{}, 1),
-		deliveryRate:   16,
-		bandwidth:      1,
-		sendPacket:     make(chan packet.Packet, 256),
-		shutdownEvent:  make(chan shutdownMessage, 5),
+		created:         now,
+		sockState:       sockStateInit,
+		udtVer:          4,
+		isServer:        isServer,
+		maxPacketSize:   uint32(config.MaxPacketSize),
+		maxFlowWinSize:  maxFlowWinSize,
+		isDatagram:      isDatagram,
+		sockID:          sockID,
+		initPktSeq:      packet.RandomPacketSequence(),
+		messageIn:       make(chan []byte, 256),
+		messageOut:      make(chan sendMessage, 256),
+		recvEvent:       make(chan recvPktEvent, 256),
+		sendEvent:       make(chan recvPktEvent, 256),
+		sockClosed:      make(chan struct{}, 1),
+		terminateSignal: make(chan struct{}),
+		deliveryRate:    16,
+		bandwidth:       1,
+		sendPacket:      make(chan packet.Packet, 256),
+		shutdownEvent:   make(chan shutdownMessage, 5),
 	}
 	s.cong = newUdtSocketCc(s)
 
diff --git a/udt/udtsocket_send.go b/udt/udtsocket_send.go
index 82cd226..2325cb3 100644
--- a/udt/udtsocket_send.go
+++ b/udt/udtsocket_send.go
@@ -140,14 +140,7 @@ func (s *udtSocketSend) goSendEvent() {
 
 		// wait for a channel to fire
 		select {
-		case msg, ok := <-messageOut: // nil if we can't process outgoing messages right now, which means it will not be selected
-			// new message outgoing
-			if !ok {
-				s.sendPacket <- &packet.ShutdownPacket{}
-				s.shutdownEvent <- shutdownMessage{sockState: sockStateClosed, permitLinger: !s.socket.isServer, reason: TerminateReasonCannotProcessOutgoing}
-				return
-			}
-
+		case msg := <-messageOut: // nil if we can't process outgoing messages right now, which means it will not be selected
 			msg.content = s.fillDataToMTU(msg.content, messageOut) // a trick to fill up the packet immediately with data (stream only)
 
 			s.processDataMsg(msg.content, msg.tim, msg.ttl, true)
@@ -171,6 +164,11 @@ func (s *udtSocketSend) goSendEvent() {
 
 		case <-s.sockClosed:
 			return
+
+		case <-s.socket.terminateSignal:
+			s.sendPacket <- &packet.ShutdownPacket{}
+			s.shutdownEvent <- shutdownMessage{sockState: sockStateClosed, permitLinger: !s.socket.isServer, reason: TerminateReasonCannotProcessOutgoing}
+			return
 		}
 	}
 }