From b9d4759cef9c52693c67d879f28b81f6bcc29403 Mon Sep 17 00:00:00 2001
From: Mark Puha <marko10@inf.elte.hu>
Date: Tue, 12 Sep 2023 06:14:19 +0200
Subject: [PATCH] add remaning guarding

Signed-off-by: Mark Puha <marko10@inf.elte.hu>
---
 device/noise-protocol.go | 10 ++++++++++
 device/send.go           | 29 ++++++++++++++++++-----------
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/device/noise-protocol.go b/device/noise-protocol.go
index e7ad927..0a01359 100644
--- a/device/noise-protocol.go
+++ b/device/noise-protocol.go
@@ -199,10 +199,12 @@ func (device *Device) CreateMessageInitiation(
 
 	handshake.mixHash(handshake.remoteStatic[:])
 
+	device.aSecMux.RLock()
 	msg := MessageInitiation{
 		Type:      MessageInitiationType,
 		Ephemeral: handshake.localEphemeral.publicKey(),
 	}
+	device.aSecMux.RUnlock()
 
 	handshake.mixKey(msg.Ephemeral[:])
 	handshake.mixHash(msg.Ephemeral[:])
@@ -261,9 +263,12 @@ func (device *Device) ConsumeMessageInitiation(msg *MessageInitiation) *Peer {
 		chainKey [blake2s.Size]byte
 	)
 
+	device.aSecMux.RLock()
 	if msg.Type != MessageInitiationType {
+		device.aSecMux.RUnlock()
 		return nil
 	}
+	device.aSecMux.RUnlock()
 
 	device.staticIdentity.RLock()
 	defer device.staticIdentity.RUnlock()
@@ -392,7 +397,9 @@ func (device *Device) CreateMessageResponse(
 	}
 
 	var msg MessageResponse
+	device.aSecMux.RLock()
 	msg.Type = MessageResponseType
+	device.aSecMux.RUnlock()
 	msg.Sender = handshake.localIndex
 	msg.Receiver = handshake.remoteIndex
 
@@ -442,9 +449,12 @@ func (device *Device) CreateMessageResponse(
 }
 
 func (device *Device) ConsumeMessageResponse(msg *MessageResponse) *Peer {
+	device.aSecMux.RLock()
 	if msg.Type != MessageResponseType {
+		device.aSecMux.RUnlock()
 		return nil
 	}
+	device.aSecMux.RUnlock()
 
 	// lookup handshake by receiver
 
diff --git a/device/send.go b/device/send.go
index 15d0a14..224cc48 100644
--- a/device/send.go
+++ b/device/send.go
@@ -128,8 +128,10 @@ func (peer *Peer) SendHandshakeInitiation(isRetry bool) error {
 	// so only packet processed for cookie generation
 	var junkedHeader []byte
 	if peer.device.isAdvancedSecurityOn() {
+		peer.device.aSecMux.RLock()
 		err = peer.sendJunkPackets()
 		if err != nil {
+			peer.device.aSecMux.RUnlock()
 			peer.device.log.Errorf("%v - %v", peer, err)
 			return err
 		}
@@ -138,11 +140,13 @@ func (peer *Peer) SendHandshakeInitiation(isRetry bool) error {
 			writer := bytes.NewBuffer(buf[:0])
 			err = appendJunk(writer, peer.device.aSecCfg.initPacketJunkSize)
 			if err != nil {
+				peer.device.aSecMux.RUnlock()
 				peer.device.log.Errorf("%v - %v", peer, err)
 				return err
 			}
 			junkedHeader = writer.Bytes()
 		}
+		peer.device.aSecMux.RUnlock()
 	}
 	var buf [MessageInitiationSize]byte
 	writer := bytes.NewBuffer(buf[:0])
@@ -184,17 +188,20 @@ func (peer *Peer) SendHandshakeResponse() error {
 		return err
 	}
 	var junkedHeader []byte
-	if peer.device.isAdvancedSecurityOn() &&
-		peer.device.aSecCfg.responsePacketJunkSize != 0 {
-
-		buf := make([]byte, 0, peer.device.aSecCfg.responsePacketJunkSize)
-		writer := bytes.NewBuffer(buf[:0])
-		err = appendJunk(writer, peer.device.aSecCfg.responsePacketJunkSize)
-		if err != nil {
-			peer.device.log.Errorf("%v - %v", peer, err)
-			return err
-		}
-		junkedHeader = writer.Bytes()
+	if peer.device.isAdvancedSecurityOn() {
+		peer.device.aSecMux.RLock()
+		if peer.device.aSecCfg.responsePacketJunkSize != 0 {
+			buf := make([]byte, 0, peer.device.aSecCfg.responsePacketJunkSize)
+			writer := bytes.NewBuffer(buf[:0])
+			err = appendJunk(writer, peer.device.aSecCfg.responsePacketJunkSize)
+			if err != nil {
+				peer.device.aSecMux.RUnlock()
+				peer.device.log.Errorf("%v - %v", peer, err)
+				return err
+			}
+			junkedHeader = writer.Bytes()
+		} 
+		peer.device.aSecMux.RUnlock()
 	}
 	var buf [MessageResponseSize]byte
 	writer := bytes.NewBuffer(buf[:0])