From 1896d9ba3fe8c582a54dc116af90b72c9ed7567c Mon Sep 17 00:00:00 2001
From: Yaroslav Gurov <ygurov@proton.me>
Date: Mon, 14 Jul 2025 20:39:51 +0200
Subject: [PATCH] fix: add i,j params size limitation

---
 device/awg/special_handshake_handler.go       |  6 +++---
 device/awg/tag_junk_packet_generator.go       |  4 ++++
 device/awg/tag_junk_packet_generators.go      | 10 ++++++++--
 device/awg/tag_junk_packet_generators_test.go |  2 +-
 device/device.go                              |  2 +-
 device/uapi.go                                |  8 ++++----
 6 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/device/awg/special_handshake_handler.go b/device/awg/special_handshake_handler.go
index e582d97..86b2cf4 100644
--- a/device/awg/special_handshake_handler.go
+++ b/device/awg/special_handshake_handler.go
@@ -31,12 +31,12 @@ type SpecialHandshakeHandler struct {
 	IsSet bool
 }
 
-func (handler *SpecialHandshakeHandler) Validate() error {
+func (handler *SpecialHandshakeHandler) Validate(maxSegmentSize int) error {
 	var errs []error
-	if err := handler.SpecialJunk.Validate(); err != nil {
+	if err := handler.SpecialJunk.Validate(maxSegmentSize); err != nil {
 		errs = append(errs, err)
 	}
-	if err := handler.ControlledJunk.Validate(); err != nil {
+	if err := handler.ControlledJunk.Validate(maxSegmentSize); err != nil {
 		errs = append(errs, err)
 	}
 	return errors.Join(errs...)
diff --git a/device/awg/tag_junk_packet_generator.go b/device/awg/tag_junk_packet_generator.go
index fdbebc8..ac80df8 100644
--- a/device/awg/tag_junk_packet_generator.go
+++ b/device/awg/tag_junk_packet_generator.go
@@ -57,3 +57,7 @@ func (tg *TagJunkPacketGenerator) IpcGetFields() IpcFields {
 		Value: tg.tagValue,
 	}
 }
+
+func (tg *TagJunkPacketGenerator) Size() int {
+	return tg.packetSize
+}
diff --git a/device/awg/tag_junk_packet_generators.go b/device/awg/tag_junk_packet_generators.go
index 9921eb0..0cca137 100644
--- a/device/awg/tag_junk_packet_generators.go
+++ b/device/awg/tag_junk_packet_generators.go
@@ -1,6 +1,8 @@
 package awg
 
-import "fmt"
+import (
+	"fmt"
+)
 
 type TagJunkPacketGenerators struct {
 	tagGenerators    []TagJunkPacketGenerator
@@ -20,7 +22,7 @@ func (generators *TagJunkPacketGenerators) IsDefined() bool {
 }
 
 // validate that packets were defined consecutively
-func (generators *TagJunkPacketGenerators) Validate() error {
+func (generators *TagJunkPacketGenerators) Validate(maxSegmentSize int) error {
 	seen := make([]bool, len(generators.tagGenerators))
 	for _, generator := range generators.tagGenerators {
 		index, err := generator.nameIndex()
@@ -32,6 +34,10 @@ func (generators *TagJunkPacketGenerators) Validate() error {
 		} else {
 			seen[index-1] = true
 		}
+
+		if generator.Size() > maxSegmentSize {
+			return fmt.Errorf("junk packet %s must not exceed %d bytes", generator.name, maxSegmentSize)
+		}
 	}
 
 	for _, found := range seen {
diff --git a/device/awg/tag_junk_packet_generators_test.go b/device/awg/tag_junk_packet_generators_test.go
index 6b1fd47..88aa088 100644
--- a/device/awg/tag_junk_packet_generators_test.go
+++ b/device/awg/tag_junk_packet_generators_test.go
@@ -91,7 +91,7 @@ func TestTagJunkGeneratorHandlerValidate(t *testing.T) {
 				generators.AppendGenerator(gen)
 			}
 
-			err := generators.Validate()
+			err := generators.Validate(1500)
 			if tt.wantErr {
 				require.Error(t, err)
 				require.Contains(t, err.Error(), tt.errMsg)
diff --git a/device/device.go b/device/device.go
index 1829352..51e439a 100644
--- a/device/device.go
+++ b/device/device.go
@@ -819,7 +819,7 @@ func (device *Device) handlePostConfig(tempAwg *awg.Protocol) error {
 	}
 
 	if tempAwg.HandshakeHandler.IsSet {
-		if err := tempAwg.HandshakeHandler.Validate(); err != nil {
+		if err := tempAwg.HandshakeHandler.Validate(MaxSegmentSize); err != nil {
 			errs = append(errs, ipcErrorf(
 				ipc.IpcErrorInvalid, "handshake handler validate: %w", err))
 		} else {
diff --git a/device/uapi.go b/device/uapi.go
index e9f962a..f49b186 100644
--- a/device/uapi.go
+++ b/device/uapi.go
@@ -406,12 +406,12 @@ func (device *Device) handleDeviceLine(key, value string, tempAwg *awg.Protocol)
 			return nil
 		}
 
-		generators, err := awg.Parse(key, value)
+		generator, err := awg.Parse(key, value)
 		if err != nil {
 			return ipcErrorf(ipc.IpcErrorInvalid, "invalid %s: %w", key, err)
 		}
 		device.log.Verbosef("UAPI: Updating %s", key)
-		tempAwg.HandshakeHandler.SpecialJunk.AppendGenerator(generators)
+		tempAwg.HandshakeHandler.SpecialJunk.AppendGenerator(generator)
 		tempAwg.HandshakeHandler.IsSet = true
 	case "j1", "j2", "j3":
 		if len(value) == 0 {
@@ -419,13 +419,13 @@ func (device *Device) handleDeviceLine(key, value string, tempAwg *awg.Protocol)
 			return nil
 		}
 
-		generators, err := awg.Parse(key, value)
+		generator, err := awg.Parse(key, value)
 		if err != nil {
 			return ipcErrorf(ipc.IpcErrorInvalid, "invalid %s: %w", key, err)
 		}
 		device.log.Verbosef("UAPI: Updating %s", key)
 
-		tempAwg.HandshakeHandler.ControlledJunk.AppendGenerator(generators)
+		tempAwg.HandshakeHandler.ControlledJunk.AppendGenerator(generator)
 		tempAwg.HandshakeHandler.IsSet = true
 	case "itime":
 		if len(value) == 0 {