feat: awg 2.0 (#91)

* feat: ranged H1-H4
* feat: S3, S4 support
* chore: updated awg-tools version

---------

Co-authored-by: Yaroslav Gurov <[email protected]>

Author: marko1777

Dockerfile

@@ -6,7 +6,7 @@ RUN go mod download && \
     go build -ldflags '-linkmode external -extldflags "-fno-PIC -static"' -v -o /usr/bin
 
 FROM alpine:3.19
-ARG AWGTOOLS_RELEASE="1.0.20241018"
+ARG AWGTOOLS_RELEASE="1.0.20250901"
 
 RUN apk --no-cache add iproute2 iptables bash && \
     cd /usr/bin/ && \

device/awg/awg.go

@@ -3,142 +3,88 @@ package awg
 import (
 	"bytes"
 	"fmt"
-	"slices"
-	"strconv"
-	"strings"
 	"sync"
 
 	"github.com/tevino/abool"
 )
 
-type aSecCfgType struct {
-	IsSet                      bool
-	JunkPacketCount            int
-	JunkPacketMinSize          int
-	JunkPacketMaxSize          int
-	InitHeaderJunkSize         int
-	ResponseHeaderJunkSize     int
-	CookieReplyHeaderJunkSize  int
-	TransportHeaderJunkSize    int
-	InitPacketMagicHeader      uint32
-	ResponsePacketMagicHeader  uint32
-	UnderloadPacketMagicHeader uint32
-	TransportPacketMagicHeader uint32
-	// InitPacketMagicHeader      Limit
-	// ResponsePacketMagicHeader  Limit
-	// UnderloadPacketMagicHeader Limit
-	// TransportPacketMagicHeader Limit
-}
-
-type Limit struct {
-	Min        uint32
-	Max        uint32
-	HeaderType uint32
-}
-
-func NewLimit(min, max, headerType uint32) (Limit, error) {
-	if min > max {
-		return Limit{}, fmt.Errorf("min (%d) cannot be greater than max (%d)", min, max)
-	}
-
-	return Limit{
-		Min:        min,
-		Max:        max,
-		HeaderType: headerType,
-	}, nil
-}
-
-func ParseMagicHeader(key, value string, defaultHeaderType uint32) (Limit, error) {
-	// tempAwg.ASecCfg.InitPacketMagicHeader, err = awg.NewLimit(uint32(initPacketMagicHeaderMin), uint32(initPacketMagicHeaderMax), DNewLimit(min, max, headerType)efaultMessageInitiationType)
-	// var min, max, headerType uint32
-	// _, err := fmt.Sscanf(value, "%d-%d:%d", &min, &max, &headerType)
-	// if err != nil {
-	// 	return Limit{}, fmt.Errorf("invalid magic header format: %s", value)
-	// }
-
-	limits := strings.Split(value, "-")
-	if len(limits) != 2 {
-		return Limit{}, fmt.Errorf("invalid format for key: %s; %s", key, value)
-	}
-
-	min, err := strconv.ParseUint(limits[0], 10, 32)
-	if err != nil {
-		return Limit{}, fmt.Errorf("parse min key: %s; value: ; %w", key, limits[0], err)
-	}
-
-	max, err := strconv.ParseUint(limits[1], 10, 32)
-	if err != nil {
-		return Limit{}, fmt.Errorf("parse max key: %s; value: ; %w", key, limits[0], err)
-	}
-
-	limit, err := NewLimit(uint32(min), uint32(max), defaultHeaderType)
-	if err != nil {
-		return Limit{}, fmt.Errorf("new lmit key: %s; value: ; %w", key, limits[0], err)
-	}
-
-	return limit, nil
-}
-
-type Limits []Limit
-
-func NewLimits(limits []Limit) Limits {
-	slices.SortFunc(limits, func(a, b Limit) int {
-		if a.Min < b.Min {
-			return -1
-		} else if a.Min > b.Min {
-			return 1
-		}
-		return 0
-	})
-
-	return Limits(limits)
+type Cfg struct {
+	IsSet                     bool
+	JunkPacketCount           int
+	JunkPacketMinSize         int
+	JunkPacketMaxSize         int
+	InitHeaderJunkSize        int
+	ResponseHeaderJunkSize    int
+	CookieReplyHeaderJunkSize int
+	TransportHeaderJunkSize   int
+
+	MagicHeaders MagicHeaders
 }
 
 type Protocol struct {
-	IsASecOn abool.AtomicBool
+	IsOn abool.AtomicBool
 	// TODO: revision the need of the mutex
-	ASecMux     sync.RWMutex
-	ASecCfg     aSecCfgType
-	JunkCreator junkCreator
+	Mux         sync.RWMutex
+	Cfg         Cfg
+	JunkCreator JunkCreator
 
 	HandshakeHandler SpecialHandshakeHandler
 }
 
 func (protocol *Protocol) CreateInitHeaderJunk() ([]byte, error) {
-	return protocol.createHeaderJunk(protocol.ASecCfg.InitHeaderJunkSize)
+	protocol.Mux.RLock()
+	defer protocol.Mux.RUnlock()
+
+	return protocol.createHeaderJunk(protocol.Cfg.InitHeaderJunkSize, 0)
 }
 
 func (protocol *Protocol) CreateResponseHeaderJunk() ([]byte, error) {
-	return protocol.createHeaderJunk(protocol.ASecCfg.ResponseHeaderJunkSize)
+	protocol.Mux.RLock()
+	defer protocol.Mux.RUnlock()
+
+	return protocol.createHeaderJunk(protocol.Cfg.ResponseHeaderJunkSize, 0)
 }
 
 func (protocol *Protocol) CreateCookieReplyHeaderJunk() ([]byte, error) {
-	return protocol.createHeaderJunk(protocol.ASecCfg.CookieReplyHeaderJunkSize)
+	protocol.Mux.RLock()
+	defer protocol.Mux.RUnlock()
+
+	return protocol.createHeaderJunk(protocol.Cfg.CookieReplyHeaderJunkSize, 0)
 }
 
 func (protocol *Protocol) CreateTransportHeaderJunk(packetSize int) ([]byte, error) {
-	return protocol.createHeaderJunk(protocol.ASecCfg.TransportHeaderJunkSize, packetSize)
+	protocol.Mux.RLock()
+	defer protocol.Mux.RUnlock()
+
+	return protocol.createHeaderJunk(protocol.Cfg.TransportHeaderJunkSize, packetSize)
 }
 
-func (protocol *Protocol) createHeaderJunk(junkSize int, optExtraSize ...int) ([]byte, error) {
-	extraSize := 0
-	if len(optExtraSize) == 1 {
-		extraSize = optExtraSize[0]
+func (protocol *Protocol) createHeaderJunk(junkSize int, extraSize int) ([]byte, error) {
+	if junkSize == 0 {
+		return nil, nil
+	}
+
+	buf := make([]byte, 0, junkSize+extraSize)
+	writer := bytes.NewBuffer(buf[:0])
+
+	err := protocol.JunkCreator.AppendJunk(writer, junkSize)
+	if err != nil {
+		return nil, fmt.Errorf("append junk: %w", err)
 	}
 
-	var junk []byte
-	protocol.ASecMux.RLock()
-	if junkSize != 0 {
-		buf := make([]byte, 0, junkSize+extraSize)
-		writer := bytes.NewBuffer(buf[:0])
-		err := protocol.JunkCreator.AppendJunk(writer, junkSize)
-		if err != nil {
-			protocol.ASecMux.RUnlock()
-			return nil, err
+	return writer.Bytes(), nil
+}
+
+func (protocol *Protocol) GetMagicHeaderMinFor(msgType uint32) (uint32, error) {
+	for _, magicHeader := range protocol.Cfg.MagicHeaders.Values {
+		if magicHeader.Min <= msgType && msgType <= magicHeader.Max {
+			return magicHeader.Min, nil
 		}
-		junk = writer.Bytes()
 	}
-	protocol.ASecMux.RUnlock()
 
-	return junk, nil
+	return 0, fmt.Errorf("no header for value: %d", msgType)
+}
+
+func (protocol *Protocol) GetMsgType(defaultMsgType uint32) (uint32, error) {
+	return protocol.Cfg.MagicHeaders.Get(defaultMsgType)
 }

device/awg/junk_creator.go

@@ -2,69 +2,49 @@ package awg
 
 import (
 	"bytes"
-	crand "crypto/rand"
 	"fmt"
-	v2 "math/rand/v2"
 )
 
-type junkCreator struct {
-	aSecCfg  aSecCfgType
-	cha8Rand *v2.ChaCha8
+type JunkCreator struct {
+	cfg             Cfg
+	randomGenerator PRNG[int]
 }
 
 // TODO: refactor param to only pass the junk related params
-func NewJunkCreator(aSecCfg aSecCfgType) (junkCreator, error) {
-	buf := make([]byte, 32)
-	_, err := crand.Read(buf)
-	if err != nil {
-		return junkCreator{}, err
-	}
-	return junkCreator{aSecCfg: aSecCfg, cha8Rand: v2.NewChaCha8([32]byte(buf))}, nil
+func NewJunkCreator(cfg Cfg) JunkCreator {
+	return JunkCreator{cfg: cfg, randomGenerator: NewPRNG[int]()}
 }
 
-// Should be called with aSecMux RLocked
-func (jc *junkCreator) CreateJunkPackets(junks *[][]byte) error {
-	if jc.aSecCfg.JunkPacketCount == 0 {
-		return nil
+// Should be called with awg mux RLocked
+func (jc *JunkCreator) CreateJunkPackets(junks *[][]byte) {
+	if jc.cfg.JunkPacketCount == 0 {
+		return
 	}
 
-	for range jc.aSecCfg.JunkPacketCount {
+	for range jc.cfg.JunkPacketCount {
 		packetSize := jc.randomPacketSize()
-		junk, err := jc.randomJunkWithSize(packetSize)
-		if err != nil {
-			return fmt.Errorf("create junk packet: %v", err)
-		}
+		junk := jc.randomJunkWithSize(packetSize)
 		*junks = append(*junks, junk)
 	}
-	return nil
+	return
 }
 
-// Should be called with aSecMux RLocked
-func (jc *junkCreator) randomPacketSize() int {
-	return int(
-		jc.cha8Rand.Uint64()%uint64(
-			jc.aSecCfg.JunkPacketMaxSize-jc.aSecCfg.JunkPacketMinSize,
-		),
-	) + jc.aSecCfg.JunkPacketMinSize
+// Should be called with awg mux RLocked
+func (jc *JunkCreator) randomPacketSize() int {
+	return jc.randomGenerator.RandomSizeInRange(jc.cfg.JunkPacketMinSize, jc.cfg.JunkPacketMaxSize)
 }
 
-// Should be called with aSecMux RLocked
-func (jc *junkCreator) AppendJunk(writer *bytes.Buffer, size int) error {
-	headerJunk, err := jc.randomJunkWithSize(size)
-	if err != nil {
-		return fmt.Errorf("create header junk: %v", err)
-	}
-	_, err = writer.Write(headerJunk)
+// Should be called with awg mux RLocked
+func (jc *JunkCreator) AppendJunk(writer *bytes.Buffer, size int) error {
+	headerJunk := jc.randomJunkWithSize(size)
+	_, err := writer.Write(headerJunk)
 	if err != nil {
 		return fmt.Errorf("write header junk: %v", err)
 	}
 	return nil
 }
 
-// Should be called with aSecMux RLocked
-func (jc *junkCreator) randomJunkWithSize(size int) ([]byte, error) {
-	// TODO: use a memory pool to allocate
-	junk := make([]byte, size)
-	_, err := jc.cha8Rand.Read(junk)
-	return junk, err
+// Should be called with awg mux RLocked
+func (jc *JunkCreator) randomJunkWithSize(size int) []byte {
+	return jc.randomGenerator.ReadSize(size)
 }