From 675b7b7531c23520222254733f5ea56dee41a49b Mon Sep 17 00:00:00 2001
From: Mark Puha <p.mark95@gmail.com>
Date: Mon, 4 Aug 2025 18:43:37 +0200
Subject: [PATCH] fix: default wg protocol

---
 device/awg/junk_creator_test.go | 15 ++++++++++-----
 device/cookie.go                |  2 +-
 device/noise-protocol.go        | 29 +++++++++++++++++++----------
 device/send.go                  | 31 +++++++++++++++++++++++--------
 4 files changed, 53 insertions(+), 24 deletions(-)

diff --git a/device/awg/junk_creator_test.go b/device/awg/junk_creator_test.go
index 33532e4..cdf752b 100644
--- a/device/awg/junk_creator_test.go
+++ b/device/awg/junk_creator_test.go
@@ -7,6 +7,15 @@ import (
 )
 
 func setUpJunkCreator() JunkCreator {
+	mh, _ := NewMagicHeaders(
+		[]MagicHeader{
+			NewMagicHeaderSameValue(123456),
+			NewMagicHeaderSameValue(67543),
+			NewMagicHeaderSameValue(32345),
+			NewMagicHeaderSameValue(123123),
+		},
+	)
+
 	jc := NewJunkCreator(Cfg{
 		IsSet:                  true,
 		JunkPacketCount:        5,
@@ -14,11 +23,7 @@ func setUpJunkCreator() JunkCreator {
 		JunkPacketMaxSize:      1000,
 		InitHeaderJunkSize:     30,
 		ResponseHeaderJunkSize: 40,
-		// TODO
-		// InitPacketMagicHeader:      123456,
-		// ResponsePacketMagicHeader:  67543,
-		// UnderloadPacketMagicHeader: 32345,
-		// TransportPacketMagicHeader: 123123,
+		MagicHeaders:           mh,
 	})
 
 	return jc
diff --git a/device/cookie.go b/device/cookie.go
index 5e09806..6a0463c 100644
--- a/device/cookie.go
+++ b/device/cookie.go
@@ -154,7 +154,7 @@ func (st *CookieChecker) CreateReply(
 	smac1 := smac2 - blake2s.Size128
 
 	reply := new(MessageCookieReply)
-	reply.Type = msgType 
+	reply.Type = msgType
 	reply.Receiver = recv
 
 	_, err := rand.Read(reply.Nonce[:])
diff --git a/device/noise-protocol.go b/device/noise-protocol.go
index 7d15d87..6e6fe58 100644
--- a/device/noise-protocol.go
+++ b/device/noise-protocol.go
@@ -205,18 +205,22 @@ func (device *Device) CreateMessageInitiation(peer *Peer) (*MessageInitiation, e
 
 	handshake.mixHash(handshake.remoteStatic[:])
 
-	device.awg.Mux.RLock()
-	msgType, err := device.awg.GetMsgType(DefaultMessageInitiationType)
-	if err != nil {
+	msgType := DefaultMessageInitiationType
+	if device.isAWG() {
+		device.awg.Mux.RLock()
+		msgType, err = device.awg.GetMsgType(DefaultMessageInitiationType)
+		if err != nil {
+			device.awg.Mux.RUnlock()
+			return nil, fmt.Errorf("get message type: %w", err)
+		}
+
 		device.awg.Mux.RUnlock()
-		return nil, fmt.Errorf("get message type: %w", err)
 	}
 
 	msg := MessageInitiation{
 		Type:      msgType,
 		Ephemeral: handshake.localEphemeral.publicKey(),
 	}
-	device.awg.Mux.RUnlock()
 
 	handshake.mixKey(msg.Ephemeral[:])
 	handshake.mixHash(msg.Ephemeral[:])
@@ -391,14 +395,19 @@ func (device *Device) CreateMessageResponse(peer *Peer) (*MessageResponse, error
 	}
 
 	var msg MessageResponse
-	device.awg.Mux.RLock()
-	msg.Type, err = device.awg.GetMsgType(DefaultMessageResponseType)
-	if err != nil {
+	if device.isAWG() {
+		device.awg.Mux.RLock()
+		msg.Type, err = device.awg.GetMsgType(DefaultMessageResponseType)
+		if err != nil {
+			device.awg.Mux.RUnlock()
+			return nil, fmt.Errorf("get message type: %w", err)
+		}
+
 		device.awg.Mux.RUnlock()
-		return nil, fmt.Errorf("get message type: %w", err)
+	} else {
+		msg.Type = DefaultMessageResponseType
 	}
 
-	device.awg.Mux.RUnlock()
 	msg.Sender = handshake.localIndex
 	msg.Receiver = handshake.remoteIndex
 
diff --git a/device/send.go b/device/send.go
index e1aec72..89e793c 100644
--- a/device/send.go
+++ b/device/send.go
@@ -237,10 +237,17 @@ func (device *Device) SendHandshakeCookie(
 	device.log.Verbosef("Sending cookie response for denied handshake message for %v", initiatingElem.endpoint.DstToString())
 
 	sender := binary.LittleEndian.Uint32(initiatingElem.packet[4:8])
-	msgType, err := device.awg.GetMsgType(DefaultMessageCookieReplyType)
-	if err != nil {
-		device.log.Errorf("Get message type for cookie reply: %v", err)
-		return err
+	msgType := DefaultMessageCookieReplyType
+	if device.isAWG() {
+		device.awg.Mux.RLock()
+
+		var err error
+		msgType, err = device.awg.GetMsgType(DefaultMessageCookieReplyType)
+		device.awg.Mux.RUnlock()
+		if err != nil {
+			device.log.Errorf("Get message type for cookie reply: %v", err)
+			return err
+		}
 	}
 
 	reply, err := device.cookieChecker.CreateReply(
@@ -530,11 +537,19 @@ func (device *Device) RoutineEncryption(id int) {
 			fieldReceiver := header[4:8]
 			fieldNonce := header[8:16]
 
-			msgType, err := device.awg.GetMsgType(DefaultMessageTransportType)
-			if err != nil {
-				device.log.Errorf("get message type for transport: %v", err)
-				continue
+			msgType := DefaultMessageTransportType
+			if device.isAWG() {
+				device.awg.Mux.RLock()
+
+				var err error
+				msgType, err = device.awg.GetMsgType(DefaultMessageTransportType)
+				device.awg.Mux.RUnlock()
+				if err != nil {
+					device.log.Errorf("get message type for transport: %v", err)
+					continue
+				}
 			}
+
 			binary.LittleEndian.PutUint32(fieldType, msgType)
 			binary.LittleEndian.PutUint32(fieldReceiver, elem.keypair.remoteIndex)
 			binary.LittleEndian.PutUint64(fieldNonce, elem.nonce)