From 32470fa04ee4b131ffca0b13d55a6b362e50b782 Mon Sep 17 00:00:00 2001
From: Mark Puha <p.mark95@gmail.com>
Date: Sat, 8 Feb 2025 19:00:00 +0100
Subject: [PATCH] add codec generation/parsing

---
 device/device.go  | 14 ++++++++++----
 device/receive.go |  8 +++++++-
 device/send.go    | 32 +++++++++++++++++++++++++++++++-
 3 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/device/device.go b/device/device.go
index 1bedce3..8178667 100644
--- a/device/device.go
+++ b/device/device.go
@@ -11,6 +11,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/amnezia-vpn/amneziawg-go/adapter"
 	"github.com/amnezia-vpn/amneziawg-go/conn"
 	"github.com/amnezia-vpn/amneziawg-go/ipc"
 	"github.com/amnezia-vpn/amneziawg-go/ratelimiter"
@@ -92,11 +93,13 @@ type Device struct {
 	closed   chan struct{}
 	log      *Logger
 
-	isASecOn abool.AtomicBool
-	aSecMux  sync.RWMutex
-	aSecCfg  aSecCfgType
-
+	isASecOn    abool.AtomicBool
+	aSecMux     sync.RWMutex
+	aSecCfg     aSecCfgType
 	junkCreator junkCreator
+
+	luaAdapter 	  *adapter.Lua
+	packetCounter atomic.Int64
 }
 
 type aSecCfgType struct {
@@ -428,6 +431,9 @@ func (device *Device) Close() {
 
 	device.resetProtocol()
 
+	if device.luaAdapter != nil {
+		device.luaAdapter.Close()
+	}
 	device.log.Verbosef("Device closed")
 	close(device.closed)
 }
diff --git a/device/receive.go b/device/receive.go
index 66c1a32..e790048 100644
--- a/device/receive.go
+++ b/device/receive.go
@@ -137,8 +137,14 @@ func (device *Device) RoutineReceiveIncoming(
 			}
 
 			// check size of packet
-
 			packet := bufsArrs[i][:size]
+			if device.luaAdapter != nil {
+				packet, err = device.luaAdapter.Parse(packet)
+				if err != nil {
+					device.log.Verbosef("Couldn't parse message; reason: %v", err)
+					continue
+				}
+			}
 			var msgType uint32
 			if device.isAdvancedSecurityOn() {
 				if assumedMsgType, ok := packetSizeToMsgType[size]; ok {
diff --git a/device/send.go b/device/send.go
index 5c54d4d..c0b9ad4 100644
--- a/device/send.go
+++ b/device/send.go
@@ -175,6 +175,10 @@ func (peer *Peer) SendHandshakeInitiation(isRetry bool) error {
 	peer.cookieGenerator.AddMacs(packet)
 	junkedHeader = append(junkedHeader, packet...)
 
+	if junkedHeader, err = peer.device.codecPacket(junkedHeader); err != nil {
+		return err
+	}
+
 	peer.timersAnyAuthenticatedPacketTraversal()
 	peer.timersAnyAuthenticatedPacketSent()
 
@@ -233,6 +237,10 @@ func (peer *Peer) SendHandshakeResponse() error {
 	peer.cookieGenerator.AddMacs(packet)
 	junkedHeader = append(junkedHeader, packet...)
 
+	if junkedHeader, err = peer.device.codecPacket(junkedHeader); err != nil {
+		return err
+	}
+
 	err = peer.BeginSymmetricSession()
 	if err != nil {
 		peer.device.log.Errorf("%v - Failed to derive keypair: %v", peer, err)
@@ -277,8 +285,13 @@ func (device *Device) SendHandshakeCookie(
 	var buf [MessageCookieReplySize]byte
 	writer := bytes.NewBuffer(buf[:0])
 	binary.Write(writer, binary.LittleEndian, reply)
+	packet := writer.Bytes()
+	if packet, err = device.codecPacket(packet); err != nil {
+		return err
+	}
+
 	// TODO: allocation could be avoided
-	device.net.bind.Send([][]byte{writer.Bytes()}, initiatingElem.endpoint)
+	device.net.bind.Send([][]byte{packet}, initiatingElem.endpoint)
 	return nil
 }
 
@@ -534,6 +547,18 @@ func calculatePaddingSize(packetSize, mtu int) int {
 	return paddedSize - lastUnit
 }
 
+func (device *Device) codecPacket(packet []byte) ([]byte, error) {
+	if device.luaAdapter != nil {
+		var err error
+		packet, err = device.luaAdapter.Generate(packet, device.packetCounter.Add(1))
+		if err != nil {
+			device.log.Errorf("%v - Failed to run codec generate: %v", device, err)
+			return nil, err
+		}
+	}
+	return packet, nil
+}
+
 /* Encrypts the elements in the queue
  * and marks them for sequential consumption (by releasing the mutex)
  *
@@ -578,6 +603,11 @@ func (device *Device) RoutineEncryption(id int) {
 				elem.packet,
 				nil,
 			)
+			// TODO: check
+			var err error
+			if elem.packet, err = device.codecPacket(elem.packet); err != nil {
+				continue
+			}
 		}
 		elemsContainer.Unlock()
 	}