Fix datachannel data races by locking

Author: saligrama

src/datachannel/streaming.go

@@ -119,6 +119,8 @@ type DataChannel struct {
 
 	// AgentVersion received during handshake
 	agentVersion string
+
+	mutex sync.Mutex
 }
 
 type ListMessageBuffer struct {
@@ -279,6 +281,9 @@ func (dataChannel *DataChannel) SendInputDataMessage(
 	payloadType message.PayloadType,
 	inputData []byte) (err error) {
 
+	dataChannel.mutex.Lock()
+	defer dataChannel.mutex.Unlock()
+
 	var (
 		flag uint64 = 0
 		msg  []byte
@@ -343,12 +348,16 @@ func (dataChannel *DataChannel) ResendStreamDataMessageScheduler(log log.T) (err
 			streamMessageElement := dataChannel.OutgoingMessageBuffer.Messages.Front()
 			dataChannel.OutgoingMessageBuffer.Mutex.Unlock()
 
+			dataChannel.mutex.Lock()
+			localTimeout := dataChannel.RetransmissionTimeout
+			dataChannel.mutex.Unlock()
+
 			if streamMessageElement == nil {
 				continue
 			}
 
 			streamMessage := streamMessageElement.Value.(StreamingMessage)
-			if time.Since(streamMessage.LastSentTime) > dataChannel.RetransmissionTimeout {
+			if time.Since(streamMessage.LastSentTime) > localTimeout {
 				log.Debugf("Resend stream data message %d for the %d attempt.", streamMessage.SequenceNumber, *streamMessage.ResendAttempt)
 				if *streamMessage.ResendAttempt >= config.ResendMaxAttempt {
 					log.Warnf("Message %d was resent over %d times.", streamMessage.SequenceNumber, config.ResendMaxAttempt)
@@ -368,6 +377,9 @@ func (dataChannel *DataChannel) ResendStreamDataMessageScheduler(log log.T) (err
 
 // ProcessAcknowledgedMessage processes acknowledge messages by deleting them from OutgoingMessageBuffer
 func (dataChannel *DataChannel) ProcessAcknowledgedMessage(log log.T, acknowledgeMessageContent message.AcknowledgeContent) error {
+	dataChannel.mutex.Lock()
+	defer dataChannel.mutex.Unlock()
+
 	acknowledgeSequenceNumber := acknowledgeMessageContent.SequenceNumber
 	for streamMessageElement := dataChannel.OutgoingMessageBuffer.Messages.Front(); streamMessageElement != nil; streamMessageElement = streamMessageElement.Next() {
 		streamMessage := streamMessageElement.Value.(StreamingMessage)
@@ -615,6 +627,8 @@ func (dataChannel *DataChannel) HandleOutputMessage(
 	outputMessage message.ClientMessage,
 	rawMessage []byte) (err error) {
 
+	dataChannel.mutex.Lock()
+
 	// On receiving expected stream data message, send acknowledgement, process it and increment expected sequence number by 1.
 	// Further process messages from IncomingMessageBuffer
 	if outputMessage.SequenceNumber == dataChannel.ExpectedSequenceNumber {
@@ -623,40 +637,51 @@ func (dataChannel *DataChannel) HandleOutputMessage(
 		case message.HandshakeRequestPayloadType:
 			{
 				if err = SendAcknowledgeMessageCall(log, dataChannel, outputMessage); err != nil {
+					dataChannel.mutex.Unlock()
 					return err
 				}
 
 				// PayloadType is HandshakeRequest so we call our own handler instead of the provided handler
 				log.Debugf("Processing HandshakeRequest message %s", outputMessage)
+
+				// The handler will eventually request the lock in `SendInputDataMessage`, so we'll unlock here to avoid deadlock
+				dataChannel.mutex.Unlock()
 				if err = dataChannel.handleHandshakeRequest(log, outputMessage); err != nil {
 					log.Errorf("Unable to process incoming data payload, MessageType %s, "+
 						"PayloadType HandshakeRequestPayloadType, err: %s.", outputMessage.MessageType, err)
 					return err
 				}
+				dataChannel.mutex.Lock()
 			}
 		case message.HandshakeCompletePayloadType:
 			{
 				if err = SendAcknowledgeMessageCall(log, dataChannel, outputMessage); err != nil {
+					dataChannel.mutex.Unlock()
 					return err
 				}
 
+				dataChannel.mutex.Unlock()
 				if err = dataChannel.handleHandshakeComplete(log, outputMessage); err != nil {
 					log.Errorf("Unable to process incoming data payload, MessageType %s, "+
 						"PayloadType HandshakeCompletePayloadType, err: %s.", outputMessage.MessageType, err)
 					return err
 				}
+				dataChannel.mutex.Lock()
 			}
 		case message.EncChallengeRequest:
 			{
 				if err = SendAcknowledgeMessageCall(log, dataChannel, outputMessage); err != nil {
+					dataChannel.mutex.Unlock()
 					return err
 				}
 
+				dataChannel.mutex.Unlock()
 				if err = dataChannel.handleEncryptionChallengeRequest(log, outputMessage); err != nil {
 					log.Errorf("Unable to process incoming data payload, MessageType %s, "+
 						"PayloadType EncChallengeRequest, err: %s.", outputMessage.MessageType, err)
 					return err
 				}
+				dataChannel.mutex.Lock()
 			}
 		default:
 
@@ -686,11 +711,13 @@ func (dataChannel *DataChannel) HandleOutputMessage(
 			} else {
 				// Acknowledge outputMessage only if session specific handler is ready
 				if err := SendAcknowledgeMessageCall(log, dataChannel, outputMessage); err != nil {
+					dataChannel.mutex.Unlock()
 					return err
 				}
 			}
 		}
 		dataChannel.ExpectedSequenceNumber = dataChannel.ExpectedSequenceNumber + 1
+		dataChannel.mutex.Unlock()
 		return dataChannel.ProcessIncomingMessageBufferItems(log, outputMessage)
 	} else {
 		log.Debugf("Unexpected sequence message received. Received Sequence Number: %d. Expected Sequence Number: %d",
@@ -703,6 +730,7 @@ func (dataChannel *DataChannel) HandleOutputMessage(
 				outputMessage.SequenceNumber, dataChannel.ExpectedSequenceNumber)
 			if len(dataChannel.IncomingMessageBuffer.Messages) < dataChannel.IncomingMessageBuffer.Capacity {
 				if err = SendAcknowledgeMessageCall(log, dataChannel, outputMessage); err != nil {
+					dataChannel.mutex.Unlock()
 					return err
 				}
 
@@ -718,6 +746,7 @@ func (dataChannel *DataChannel) HandleOutputMessage(
 			}
 		}
 	}
+	dataChannel.mutex.Unlock()
 	return nil
 }
 
@@ -727,6 +756,9 @@ func (dataChannel *DataChannel) HandleOutputMessage(
 func (dataChannel *DataChannel) ProcessIncomingMessageBufferItems(log log.T,
 	outputMessage message.ClientMessage) (err error) {
 
+	dataChannel.mutex.Lock()
+	defer dataChannel.mutex.Unlock()
+
 	for {
 		bufferedStreamMessage := dataChannel.IncomingMessageBuffer.Messages[dataChannel.ExpectedSequenceNumber]
 		if bufferedStreamMessage.Content != nil {
@@ -900,12 +932,18 @@ func (dataChannel *DataChannel) IsSessionTypeSet() chan bool {
 
 // IsSessionEnded check if session has ended
 func (dataChannel *DataChannel) IsSessionEnded() bool {
-	return dataChannel.isSessionEnded
+	dataChannel.mutex.Lock()
+	sessionEnded := dataChannel.isSessionEnded
+	dataChannel.mutex.Unlock()
+	return sessionEnded
 }
 
 // IsSessionEnded check if session has ended
 func (dataChannel *DataChannel) EndSession() error {
+	dataChannel.mutex.Lock()
 	dataChannel.isSessionEnded = true
+	dataChannel.mutex.Unlock()
+
 	return nil
 }