feat: awg 2.0 (#91)

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

---------

Co-authored-by: Yaroslav Gurov <ygurov@proton.me>

Dockerfile

@@ -6,7 +6,8 @@ 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.20250704"
+ARG AWGTOOLS_RELEASE="1.0.20250901"
+
 RUN apk --no-cache add iproute2 iptables bash && \
     cd /usr/bin/ && \
     wget https://github.com/amnezia-vpn/amneziawg-tools/releases/download/v${AWGTOOLS_RELEASE}/alpine-3.19-amneziawg-tools.zip && \

README.md

@@ -50,3 +50,4 @@ $ git clone https://github.com/amnezia-vpn/amneziawg-go
 $ cd amneziawg-go
 $ make
 ```
+ 

conn/bind_std.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/bind_windows.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/bindtest/bindtest.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package bindtest

conn/boundif_android.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/conn.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 // Package conn implements WireGuard's network connections.

conn/conn_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/controlfns.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/controlfns_linux.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn
@@ -13,6 +13,35 @@ import (
 	"golang.org/x/sys/unix"
 )
 
+// Taken from go/src/internal/syscall/unix/kernel_version_linux.go
+func kernelVersion() (major, minor int) {
+	var uname unix.Utsname
+	if err := unix.Uname(&uname); err != nil {
+		return
+	}
+
+	var (
+		values    [2]int
+		value, vi int
+	)
+	for _, c := range uname.Release {
+		if '0' <= c && c <= '9' {
+			value = (value * 10) + int(c-'0')
+		} else {
+			// Note that we're assuming N.N.N here.
+			// If we see anything else, we are likely to mis-parse it.
+			values[vi] = value
+			vi++
+			if vi >= len(values) {
+				break
+			}
+			value = 0
+		}
+	}
+
+	return values[0], values[1]
+}
+
 func init() {
 	controlFns = append(controlFns,
 
@@ -57,5 +86,24 @@ func init() {
 			}
 			return err
 		},
+
+		// Attempt to enable UDP_GRO
+		func(network, address string, c syscall.RawConn) error {
+			// Kernels below 5.12 are missing 98184612aca0 ("net:
+			// udp: Add support for getsockopt(..., ..., UDP_GRO,
+			// ..., ...);"), which means we can't read this back
+			// later. We could pipe the return value through to
+			// the rest of the code, but UDP_GRO is kind of buggy
+			// anyway, so just gate this here.
+			major, minor := kernelVersion()
+			if major < 5 || (major == 5 && minor < 12) {
+				return nil
+			}
+
+			c.Control(func(fd uintptr) {
+				_ = unix.SetsockoptInt(int(fd), unix.IPPROTO_UDP, unix.UDP_GRO, 1)
+			})
+			return nil
+		},
 	)
 }

conn/controlfns_unix.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/controlfns_windows.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/default.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/errors_default.go

@@ -2,11 +2,11 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn
 
-func errShouldDisableUDPGSO(err error) bool {
+func errShouldDisableUDPGSO(_ error) bool {
 	return false
 }

conn/errors_linux.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/features_default.go

@@ -3,13 +3,13 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn
 
 import "net"
 
-func supportsUDPOffload(conn *net.UDPConn) (txOffload, rxOffload bool) {
+func supportsUDPOffload(_ *net.UDPConn) (txOffload, rxOffload bool) {
 	return
 }

conn/features_linux.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/gso_default.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/gso_linux.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/mark_default.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/mark_unix.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/sticky_default.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/sticky_linux.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/sticky_linux_test.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package conn

conn/winrio/rio_windows.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package winrio

device/allowedips.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
@@ -223,6 +223,60 @@ func (table *AllowedIPs) EntriesForPeer(peer *Peer, cb func(prefix netip.Prefix)
 	}
 }
 
+func (node *trieEntry) remove() {
+	node.removeFromPeerEntries()
+	node.peer = nil
+	if node.child[0] != nil && node.child[1] != nil {
+		return
+	}
+	bit := 0
+	if node.child[0] == nil {
+		bit = 1
+	}
+	child := node.child[bit]
+	if child != nil {
+		child.parent = node.parent
+	}
+	*node.parent.parentBit = child
+	if node.child[0] != nil || node.child[1] != nil || node.parent.parentBitType > 1 {
+		node.zeroizePointers()
+		return
+	}
+	parent := (*trieEntry)(unsafe.Pointer(uintptr(unsafe.Pointer(node.parent.parentBit)) - unsafe.Offsetof(node.child) - unsafe.Sizeof(node.child[0])*uintptr(node.parent.parentBitType)))
+	if parent.peer != nil {
+		node.zeroizePointers()
+		return
+	}
+	child = parent.child[node.parent.parentBitType^1]
+	if child != nil {
+		child.parent = parent.parent
+	}
+	*parent.parent.parentBit = child
+	node.zeroizePointers()
+	parent.zeroizePointers()
+}
+
+func (table *AllowedIPs) Remove(prefix netip.Prefix, peer *Peer) {
+	table.mutex.Lock()
+	defer table.mutex.Unlock()
+	var node *trieEntry
+	var exact bool
+
+	if prefix.Addr().Is6() {
+		ip := prefix.Addr().As16()
+		node, exact = table.IPv6.nodePlacement(ip[:], uint8(prefix.Bits()))
+	} else if prefix.Addr().Is4() {
+		ip := prefix.Addr().As4()
+		node, exact = table.IPv4.nodePlacement(ip[:], uint8(prefix.Bits()))
+	} else {
+		panic(errors.New("removing unknown address type"))
+	}
+	if !exact || node == nil || peer != node.peer {
+		return
+	}
+	node.remove()
+}
+
 func (table *AllowedIPs) RemoveByPeer(peer *Peer) {
 	table.mutex.Lock()
 	defer table.mutex.Unlock()
@@ -230,38 +284,7 @@ func (table *AllowedIPs) RemoveByPeer(peer *Peer) {
 	var next *list.Element
 	for elem := peer.trieEntries.Front(); elem != nil; elem = next {
 		next = elem.Next()
-		node := elem.Value.(*trieEntry)
-
-		node.removeFromPeerEntries()
-		node.peer = nil
-		if node.child[0] != nil && node.child[1] != nil {
-			continue
-		}
-		bit := 0
-		if node.child[0] == nil {
-			bit = 1
-		}
-		child := node.child[bit]
-		if child != nil {
-			child.parent = node.parent
-		}
-		*node.parent.parentBit = child
-		if node.child[0] != nil || node.child[1] != nil || node.parent.parentBitType > 1 {
-			node.zeroizePointers()
-			continue
-		}
-		parent := (*trieEntry)(unsafe.Pointer(uintptr(unsafe.Pointer(node.parent.parentBit)) - unsafe.Offsetof(node.child) - unsafe.Sizeof(node.child[0])*uintptr(node.parent.parentBitType)))
-		if parent.peer != nil {
-			node.zeroizePointers()
-			continue
-		}
-		child = parent.child[node.parent.parentBitType^1]
-		if child != nil {
-			child.parent = parent.parent
-		}
-		*parent.parent.parentBit = child
-		node.zeroizePointers()
-		parent.zeroizePointers()
+		elem.Value.(*trieEntry).remove()
 	}
 }
 

device/allowedips_rand_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
@@ -83,7 +83,7 @@ func TestTrieRandom(t *testing.T) {
 	var peers []*Peer
 	var allowedIPs AllowedIPs
 
-	rand.Seed(1)
+	rng := rand.New(rand.NewSource(1))
 
 	for n := 0; n < NumberOfPeers; n++ {
 		peers = append(peers, &Peer{})
@@ -91,14 +91,14 @@ func TestTrieRandom(t *testing.T) {
 
 	for n := 0; n < NumberOfAddresses; n++ {
 		var addr4 [4]byte
-		rand.Read(addr4[:])
+		rng.Read(addr4[:])
 		cidr := uint8(rand.Intn(32) + 1)
 		index := rand.Intn(NumberOfPeers)
 		allowedIPs.Insert(netip.PrefixFrom(netip.AddrFrom4(addr4), int(cidr)), peers[index])
 		slow4 = slow4.Insert(addr4[:], cidr, peers[index])
 
 		var addr6 [16]byte
-		rand.Read(addr6[:])
+		rng.Read(addr6[:])
 		cidr = uint8(rand.Intn(128) + 1)
 		index = rand.Intn(NumberOfPeers)
 		allowedIPs.Insert(netip.PrefixFrom(netip.AddrFrom16(addr6), int(cidr)), peers[index])
@@ -109,7 +109,7 @@ func TestTrieRandom(t *testing.T) {
 	for p = 0; ; p++ {
 		for n := 0; n < NumberOfTests; n++ {
 			var addr4 [4]byte
-			rand.Read(addr4[:])
+			rng.Read(addr4[:])
 			peer1 := slow4.Lookup(addr4[:])
 			peer2 := allowedIPs.Lookup(addr4[:])
 			if peer1 != peer2 {
@@ -117,7 +117,7 @@ func TestTrieRandom(t *testing.T) {
 			}
 
 			var addr6 [16]byte
-			rand.Read(addr6[:])
+			rng.Read(addr6[:])
 			peer1 = slow6.Lookup(addr6[:])
 			peer2 = allowedIPs.Lookup(addr6[:])
 			if peer1 != peer2 {

device/allowedips_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
@@ -39,12 +39,12 @@ func TestCommonBits(t *testing.T) {
 	}
 }
 
-func benchmarkTrie(peerNumber, addressNumber, addressLength int, b *testing.B) {
+func benchmarkTrie(peerNumber, addressNumber, _ int, b *testing.B) {
 	var trie *trieEntry
 	var peers []*Peer
 	root := parentIndirection{&trie, 2}
 
-	rand.Seed(1)
+	rng := rand.New(rand.NewSource(1))
 
 	const AddressLength = 4
 
@@ -54,15 +54,15 @@ func benchmarkTrie(peerNumber, addressNumber, addressLength int, b *testing.B) {
 
 	for n := 0; n < addressNumber; n++ {
 		var addr [AddressLength]byte
-		rand.Read(addr[:])
-		cidr := uint8(rand.Uint32() % (AddressLength * 8))
-		index := rand.Int() % peerNumber
+		rng.Read(addr[:])
+		cidr := uint8(rng.Uint32() % (AddressLength * 8))
+		index := rng.Int() % peerNumber
 		root.insert(addr[:], cidr, peers[index])
 	}
 
 	for n := 0; n < b.N; n++ {
 		var addr [AddressLength]byte
-		rand.Read(addr[:])
+		rng.Read(addr[:])
 		trie.lookup(addr[:])
 	}
 }
@@ -101,6 +101,10 @@ func TestTrieIPv4(t *testing.T) {
 		allowedIPs.Insert(netip.PrefixFrom(netip.AddrFrom4([4]byte{a, b, c, d}), int(cidr)), peer)
 	}
 
+	remove := func(peer *Peer, a, b, c, d byte, cidr uint8) {
+		allowedIPs.Remove(netip.PrefixFrom(netip.AddrFrom4([4]byte{a, b, c, d}), int(cidr)), peer)
+	}
+
 	assertEQ := func(peer *Peer, a, b, c, d byte) {
 		p := allowedIPs.Lookup([]byte{a, b, c, d})
 		if p != peer {
@@ -176,6 +180,21 @@ func TestTrieIPv4(t *testing.T) {
 	allowedIPs.RemoveByPeer(a)
 
 	assertNEQ(a, 192, 168, 0, 1)
+
+	insert(a, 1, 0, 0, 0, 32)
+	insert(a, 192, 0, 0, 0, 24)
+	assertEQ(a, 1, 0, 0, 0)
+	assertEQ(a, 192, 0, 0, 1)
+	remove(a, 192, 0, 0, 0, 32)
+	assertEQ(a, 192, 0, 0, 1)
+	remove(nil, 192, 0, 0, 0, 24)
+	assertEQ(a, 192, 0, 0, 1)
+	remove(b, 192, 0, 0, 0, 24)
+	assertEQ(a, 192, 0, 0, 1)
+	remove(a, 192, 0, 0, 0, 24)
+	assertNEQ(a, 192, 0, 0, 1)
+	remove(a, 1, 0, 0, 0, 32)
+	assertNEQ(a, 1, 0, 0, 0)
 }
 
 /* Test ported from kernel implementation:
@@ -211,6 +230,15 @@ func TestTrieIPv6(t *testing.T) {
 		allowedIPs.Insert(netip.PrefixFrom(netip.AddrFrom16(*(*[16]byte)(addr)), int(cidr)), peer)
 	}
 
+	remove := func(peer *Peer, a, b, c, d uint32, cidr uint8) {
+		var addr []byte
+		addr = append(addr, expand(a)...)
+		addr = append(addr, expand(b)...)
+		addr = append(addr, expand(c)...)
+		addr = append(addr, expand(d)...)
+		allowedIPs.Remove(netip.PrefixFrom(netip.AddrFrom16(*(*[16]byte)(addr)), int(cidr)), peer)
+	}
+
 	assertEQ := func(peer *Peer, a, b, c, d uint32) {
 		var addr []byte
 		addr = append(addr, expand(a)...)
@@ -223,6 +251,18 @@ func TestTrieIPv6(t *testing.T) {
 		}
 	}
 
+	assertNEQ := func(peer *Peer, a, b, c, d uint32) {
+		var addr []byte
+		addr = append(addr, expand(a)...)
+		addr = append(addr, expand(b)...)
+		addr = append(addr, expand(c)...)
+		addr = append(addr, expand(d)...)
+		p := allowedIPs.Lookup(addr)
+		if p == peer {
+			t.Error("Assert NEQ failed")
+		}
+	}
+
 	insert(d, 0x26075300, 0x60006b00, 0, 0xc05f0543, 128)
 	insert(c, 0x26075300, 0x60006b00, 0, 0, 64)
 	insert(e, 0, 0, 0, 0, 0)
@@ -244,4 +284,21 @@ func TestTrieIPv6(t *testing.T) {
 	assertEQ(h, 0x24046800, 0x40040800, 0, 0)
 	assertEQ(h, 0x24046800, 0x40040800, 0x10101010, 0x10101010)
 	assertEQ(a, 0x24046800, 0x40040800, 0xdeadbeef, 0xdeadbeef)
+
+	insert(a, 0x24446801, 0x40e40800, 0xdeaebeef, 0xdefbeef, 128)
+	insert(a, 0x24446800, 0xf0e40800, 0xeeaebeef, 0, 98)
+	assertEQ(a, 0x24446801, 0x40e40800, 0xdeaebeef, 0xdefbeef)
+	assertEQ(a, 0x24446800, 0xf0e40800, 0xeeaebeef, 0x10101010)
+	remove(a, 0x24446801, 0x40e40800, 0xdeaebeef, 0xdefbeef, 96)
+	assertEQ(a, 0x24446801, 0x40e40800, 0xdeaebeef, 0xdefbeef)
+	remove(nil, 0x24446801, 0x40e40800, 0xdeaebeef, 0xdefbeef, 128)
+	assertEQ(a, 0x24446801, 0x40e40800, 0xdeaebeef, 0xdefbeef)
+	remove(b, 0x24446801, 0x40e40800, 0xdeaebeef, 0xdefbeef, 128)
+	assertEQ(a, 0x24446801, 0x40e40800, 0xdeaebeef, 0xdefbeef)
+	remove(a, 0x24446801, 0x40e40800, 0xdeaebeef, 0xdefbeef, 128)
+	assertNEQ(a, 0x24446801, 0x40e40800, 0xdeaebeef, 0xdefbeef)
+	remove(b, 0x24446800, 0xf0e40800, 0xeeaebeef, 0, 98)
+	assertEQ(a, 0x24446800, 0xf0e40800, 0xeeaebeef, 0x10101010)
+	remove(a, 0x24446800, 0xf0e40800, 0xeeaebeef, 0, 98)
+	assertNEQ(a, 0x24446800, 0xf0e40800, 0xeeaebeef, 0x10101010)
 }

device/awg/awg.go

@@ -1,30 +1,90 @@
 package awg
 
 import (
+	"bytes"
+	"fmt"
 	"sync"
 
 	"github.com/tevino/abool"
 )
 
+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
 }
 
-type aSecCfgType struct {
-	IsSet                      bool
-	JunkPacketCount            int
-	JunkPacketMinSize          int
-	JunkPacketMaxSize          int
-	InitPacketJunkSize         int
-	ResponsePacketJunkSize     int
-	InitPacketMagicHeader      uint32
-	ResponsePacketMagicHeader  uint32
-	UnderloadPacketMagicHeader uint32
-	TransportPacketMagicHeader uint32
+func (protocol *Protocol) CreateInitHeaderJunk() ([]byte, error) {
+	protocol.Mux.RLock()
+	defer protocol.Mux.RUnlock()
+
+	return protocol.createHeaderJunk(protocol.Cfg.InitHeaderJunkSize, 0)
+}
+
+func (protocol *Protocol) CreateResponseHeaderJunk() ([]byte, error) {
+	protocol.Mux.RLock()
+	defer protocol.Mux.RUnlock()
+
+	return protocol.createHeaderJunk(protocol.Cfg.ResponseHeaderJunkSize, 0)
+}
+
+func (protocol *Protocol) CreateCookieReplyHeaderJunk() ([]byte, error) {
+	protocol.Mux.RLock()
+	defer protocol.Mux.RUnlock()
+
+	return protocol.createHeaderJunk(protocol.Cfg.CookieReplyHeaderJunkSize, 0)
+}
+
+func (protocol *Protocol) CreateTransportHeaderJunk(packetSize int) ([]byte, error) {
+	protocol.Mux.RLock()
+	defer protocol.Mux.RUnlock()
+
+	return protocol.createHeaderJunk(protocol.Cfg.TransportHeaderJunkSize, packetSize)
+}
+
+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)
+	}
+
+	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
+		}
+	}
+
+	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,68 +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]
 }
 
-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
+// TODO: refactor param to only pass the junk related params
+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)
 }

device/awg/junk_creator_test.go

@@ -6,43 +6,34 @@ import (
 	"testing"
 )
 
-func setUpJunkCreator(t *testing.T) (junkCreator, error) {
-	jc, err := NewJunkCreator(aSecCfgType{
-		IsSet:                      true,
-		JunkPacketCount:            5,
-		JunkPacketMinSize:          500,
-		JunkPacketMaxSize:          1000,
-		InitPacketJunkSize:         30,
-		ResponsePacketJunkSize:     40,
-		InitPacketMagicHeader:      123456,
-		ResponsePacketMagicHeader:  67543,
-		UnderloadPacketMagicHeader: 32345,
-		TransportPacketMagicHeader: 123123,
+func setUpJunkCreator() JunkCreator {
+	mh, _ := NewMagicHeaders(
+		[]MagicHeader{
+			NewMagicHeaderSameValue(123456),
+			NewMagicHeaderSameValue(67543),
+			NewMagicHeaderSameValue(32345),
+			NewMagicHeaderSameValue(123123),
+		},
+	)
+
+	jc := NewJunkCreator(Cfg{
+		IsSet:                  true,
+		JunkPacketCount:        5,
+		JunkPacketMinSize:      500,
+		JunkPacketMaxSize:      1000,
+		InitHeaderJunkSize:     30,
+		ResponseHeaderJunkSize: 40,
+		MagicHeaders:           mh,
 	})
 
-	if err != nil {
-		t.Errorf("failed to create junk creator %v", err)
-		return junkCreator{}, err
-	}
-
-	return jc, nil
+	return jc
 }
 
 func Test_junkCreator_createJunkPackets(t *testing.T) {
-	jc, err := setUpJunkCreator(t)
-	if err != nil {
-		return
-	}
+	jc := setUpJunkCreator()
 	t.Run("valid", func(t *testing.T) {
-		got := make([][]byte, 0, jc.aSecCfg.JunkPacketCount)
-		err := jc.CreateJunkPackets(&got)
-		if err != nil {
-			t.Errorf(
-				"junkCreator.createJunkPackets() = %v; failed",
-				err,
-			)
-			return
-		}
+		got := make([][]byte, 0, jc.cfg.JunkPacketCount)
+		jc.CreateJunkPackets(&got)
 		seen := make(map[string]bool)
 		for _, junk := range got {
 			key := string(junk)
@@ -61,34 +52,28 @@ func Test_junkCreator_createJunkPackets(t *testing.T) {
 
 func Test_junkCreator_randomJunkWithSize(t *testing.T) {
 	t.Run("valid", func(t *testing.T) {
-		jc, err := setUpJunkCreator(t)
-		if err != nil {
-			return
-		}
-		r1, _ := jc.randomJunkWithSize(10)
-		r2, _ := jc.randomJunkWithSize(10)
+		jc := setUpJunkCreator()
+		r1 := jc.randomJunkWithSize(10)
+		r2 := jc.randomJunkWithSize(10)
 		fmt.Printf("%v\n%v\n", r1, r2)
 		if bytes.Equal(r1, r2) {
-			t.Errorf("same junks %v", err)
+			t.Errorf("same junks")
 			return
 		}
 	})
 }
 
 func Test_junkCreator_randomPacketSize(t *testing.T) {
-	jc, err := setUpJunkCreator(t)
-	if err != nil {
-		return
-	}
+	jc := setUpJunkCreator()
 	for range [30]struct{}{} {
 		t.Run("valid", func(t *testing.T) {
-			if got := jc.randomPacketSize(); jc.aSecCfg.JunkPacketMinSize > got ||
-				got > jc.aSecCfg.JunkPacketMaxSize {
+			if got := jc.randomPacketSize(); jc.cfg.JunkPacketMinSize > got ||
+				got > jc.cfg.JunkPacketMaxSize {
 				t.Errorf(
 					"junkCreator.randomPacketSize() = %v, not between range [%v,%v]",
 					got,
-					jc.aSecCfg.JunkPacketMinSize,
-					jc.aSecCfg.JunkPacketMaxSize,
+					jc.cfg.JunkPacketMinSize,
+					jc.cfg.JunkPacketMaxSize,
 				)
 			}
 		})
@@ -96,10 +81,7 @@ func Test_junkCreator_randomPacketSize(t *testing.T) {
 }
 
 func Test_junkCreator_appendJunk(t *testing.T) {
-	jc, err := setUpJunkCreator(t)
-	if err != nil {
-		return
-	}
+	jc := setUpJunkCreator()
 	t.Run("valid", func(t *testing.T) {
 		s := "apple"
 		buffer := bytes.NewBuffer([]byte(s))

device/awg/magic_header.go

@@ -0,0 +1,97 @@
+package awg
+
+import (
+	"cmp"
+	"fmt"
+	"slices"
+	"strconv"
+	"strings"
+)
+
+type MagicHeader struct {
+	Min uint32
+	Max uint32
+}
+
+func NewMagicHeaderSameValue(value uint32) MagicHeader {
+	return MagicHeader{Min: value, Max: value}
+}
+
+func NewMagicHeader(min, max uint32) (MagicHeader, error) {
+	if min > max {
+		return MagicHeader{}, fmt.Errorf("min (%d) cannot be greater than max (%d)", min, max)
+	}
+
+	return MagicHeader{Min: min, Max: max}, nil
+}
+
+func ParseMagicHeader(key, value string) (MagicHeader, error) {
+	hyphenIdx := strings.Index(value, "-")
+	if hyphenIdx == -1 {
+		// if there is no hyphen, we treat it as single magic header value
+		magicHeader, err := strconv.ParseUint(value, 10, 32)
+		if err != nil {
+			return MagicHeader{}, fmt.Errorf("parse key: %s; value: %s; %w", key, value, err)
+		}
+
+		return NewMagicHeader(uint32(magicHeader), uint32(magicHeader))
+	}
+
+	minStr := value[:hyphenIdx]
+	maxStr := value[hyphenIdx+1:]
+	if len(minStr) == 0 || len(maxStr) == 0 {
+		return MagicHeader{}, fmt.Errorf("invalid value for key: %s; value: %s; expected format: min-max", key, value)
+	}
+
+	min, err := strconv.ParseUint(minStr, 10, 32)
+	if err != nil {
+		return MagicHeader{}, fmt.Errorf("parse min key: %s; value: %s; %w", key, minStr, err)
+	}
+
+	max, err := strconv.ParseUint(maxStr, 10, 32)
+	if err != nil {
+		return MagicHeader{}, fmt.Errorf("parse max key: %s; value: %s; %w", key, maxStr, err)
+	}
+
+	magicHeader, err := NewMagicHeader(uint32(min), uint32(max))
+	if err != nil {
+		return MagicHeader{}, fmt.Errorf("new magicHeader key: %s; value: %s-%s; %w", key, minStr, maxStr, err)
+	}
+
+	return magicHeader, nil
+}
+
+type MagicHeaders struct {
+	Values          []MagicHeader
+	randomGenerator RandomNumberGenerator[uint32]
+}
+
+func NewMagicHeaders(headerValues []MagicHeader) (MagicHeaders, error) {
+	if len(headerValues) != 4 {
+		return MagicHeaders{}, fmt.Errorf("all header types should be included: %v", headerValues)
+	}
+
+	sortedMagicHeaders := slices.SortedFunc(slices.Values(headerValues), func(lhs MagicHeader, rhs MagicHeader) int {
+		return cmp.Compare(lhs.Min, rhs.Min)
+	})
+
+	for i := range 3 {
+		if sortedMagicHeaders[i].Max >= sortedMagicHeaders[i+1].Min {
+			return MagicHeaders{}, fmt.Errorf(
+				"magic headers shouldn't overlap; %v > %v",
+				sortedMagicHeaders[i].Max,
+				sortedMagicHeaders[i+1].Min,
+			)
+		}
+	}
+
+	return MagicHeaders{Values: headerValues, randomGenerator: NewPRNG[uint32]()}, nil
+}
+
+func (mh *MagicHeaders) Get(defaultMsgType uint32) (uint32, error) {
+	if defaultMsgType == 0 || defaultMsgType > 4 {
+		return 0, fmt.Errorf("invalid msg type: %d", defaultMsgType)
+	}
+
+	return mh.randomGenerator.RandomSizeInRange(mh.Values[defaultMsgType-1].Min, mh.Values[defaultMsgType-1].Max), nil
+}

device/awg/magic_header_test.go

@@ -0,0 +1,488 @@
+package awg
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewMagicHeaderSameValue(t *testing.T) {
+	tests := []struct {
+		name     string
+		value    uint32
+		expected MagicHeader
+	}{
+		{
+			name:     "zero value",
+			value:    0,
+			expected: MagicHeader{Min: 0, Max: 0},
+		},
+		{
+			name:     "small value",
+			value:    1,
+			expected: MagicHeader{Min: 1, Max: 1},
+		},
+		{
+			name:     "large value",
+			value:    4294967295, // max uint32
+			expected: MagicHeader{Min: 4294967295, Max: 4294967295},
+		},
+		{
+			name:     "medium value",
+			value:    1000,
+			expected: MagicHeader{Min: 1000, Max: 1000},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			result := NewMagicHeaderSameValue(tt.value)
+			require.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+func TestNewMagicHeader(t *testing.T) {
+	tests := []struct {
+		name     string
+		min      uint32
+		max      uint32
+		expected MagicHeader
+		errorMsg string
+	}{
+		{
+			name:     "valid range",
+			min:      1,
+			max:      10,
+			expected: MagicHeader{Min: 1, Max: 10},
+		},
+		{
+			name:     "equal values",
+			min:      5,
+			max:      5,
+			expected: MagicHeader{Min: 5, Max: 5},
+		},
+		{
+			name:     "zero range",
+			min:      0,
+			max:      0,
+			expected: MagicHeader{Min: 0, Max: 0},
+		},
+		{
+			name:     "max uint32 range",
+			min:      4294967294,
+			max:      4294967295,
+			expected: MagicHeader{Min: 4294967294, Max: 4294967295},
+		},
+		{
+			name:     "min greater than max",
+			min:      10,
+			max:      5,
+			expected: MagicHeader{},
+			errorMsg: "min (10) cannot be greater than max (5)",
+		},
+		{
+			name:     "large min greater than max",
+			min:      4294967295,
+			max:      1,
+			expected: MagicHeader{},
+			errorMsg: "min (4294967295) cannot be greater than max (1)",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			result, err := NewMagicHeader(tt.min, tt.max)
+
+			if tt.errorMsg != "" {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tt.errorMsg)
+				require.Equal(t, MagicHeader{}, result)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tt.expected, result)
+			}
+		})
+	}
+}
+
+func TestParseMagicHeader(t *testing.T) {
+	tests := []struct {
+		name     string
+		key      string
+		value    string
+		expected MagicHeader
+		errorMsg string
+	}{
+		{
+			name:     "single value",
+			key:      "header1",
+			value:    "100",
+			expected: MagicHeader{Min: 100, Max: 100},
+		},
+		{
+			name:     "valid range",
+			key:      "header2",
+			value:    "10-20",
+			expected: MagicHeader{Min: 10, Max: 20},
+		},
+		{
+			name:     "zero single value",
+			key:      "header3",
+			value:    "0",
+			expected: MagicHeader{Min: 0, Max: 0},
+		},
+		{
+			name:     "zero range",
+			key:      "header4",
+			value:    "0-0",
+			expected: MagicHeader{Min: 0, Max: 0},
+		},
+		{
+			name:     "max uint32 single",
+			key:      "header5",
+			value:    "4294967295",
+			expected: MagicHeader{Min: 4294967295, Max: 4294967295},
+		},
+		{
+			name:     "max uint32 range",
+			key:      "header6",
+			value:    "4294967294-4294967295",
+			expected: MagicHeader{Min: 4294967294, Max: 4294967295},
+		},
+		{
+			name:     "invalid single value - not number",
+			key:      "header7",
+			value:    "abc",
+			expected: MagicHeader{},
+			errorMsg: "parse key: header7; value: abc;",
+		},
+		{
+			name:     "invalid single value - negative",
+			key:      "header8",
+			value:    "-5",
+			expected: MagicHeader{},
+			errorMsg: "invalid value for key: header8; value: -5;",
+		},
+		{
+			name:     "invalid single value - too large",
+			key:      "header9",
+			value:    "4294967296",
+			expected: MagicHeader{},
+			errorMsg: "parse key: header9; value: 4294967296;",
+		},
+		{
+			name:     "invalid range - min not number",
+			key:      "header10",
+			value:    "abc-10",
+			expected: MagicHeader{},
+			errorMsg: "parse min key: header10; value: abc;",
+		},
+		{
+			name:     "invalid range - max not number",
+			key:      "header11",
+			value:    "10-abc",
+			expected: MagicHeader{},
+			errorMsg: "parse max key: header11; value: abc;",
+		},
+		{
+			name:     "invalid range - min greater than max",
+			key:      "header12",
+			value:    "20-10",
+			expected: MagicHeader{},
+			errorMsg: "new magicHeader key: header12; value: 20-10;",
+		},
+		{
+			name:     "invalid range - too many parts",
+			key:      "header13",
+			value:    "10-20-30",
+			expected: MagicHeader{},
+			errorMsg: "parse key: header13; value: 10-20-30;",
+		},
+		{
+			name:     "empty value",
+			key:      "header14",
+			value:    "",
+			expected: MagicHeader{},
+			errorMsg: "parse key: header14; value: ;",
+		},
+		{
+			name:     "hyphen only",
+			key:      "header15",
+			value:    "-",
+			expected: MagicHeader{},
+			errorMsg: "invalid value for key: header15; value: -;",
+		},
+		{
+			name:     "empty min",
+			key:      "header16",
+			value:    "-10",
+			expected: MagicHeader{},
+			errorMsg: "invalid value for key: header16; value: -10;",
+		},
+		{
+			name:     "empty max",
+			key:      "header17",
+			value:    "10-",
+			expected: MagicHeader{},
+			errorMsg: "invalid value for key: header17; value: 10-;",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			result, err := ParseMagicHeader(tt.key, tt.value)
+
+			if tt.errorMsg != "" {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tt.errorMsg)
+				require.Equal(t, MagicHeader{}, result)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tt.expected, result)
+			}
+		})
+	}
+}
+
+func TestNewMagicHeaders(t *testing.T) {
+	tests := []struct {
+		name         string
+		magicHeaders []MagicHeader
+		errorMsg     string
+	}{
+		{
+			name: "valid non-overlapping headers",
+			magicHeaders: []MagicHeader{
+				{Min: 1, Max: 10},
+				{Min: 11, Max: 20},
+				{Min: 21, Max: 30},
+				{Min: 31, Max: 40},
+			},
+		},
+		{
+			name: "valid adjacent headers",
+			magicHeaders: []MagicHeader{
+				{Min: 1, Max: 1},
+				{Min: 2, Max: 2},
+				{Min: 3, Max: 3},
+				{Min: 4, Max: 4},
+			},
+		},
+		{
+			name: "valid zero-based headers",
+			magicHeaders: []MagicHeader{
+				{Min: 0, Max: 0},
+				{Min: 1, Max: 1},
+				{Min: 2, Max: 2},
+				{Min: 3, Max: 3},
+			},
+		},
+		{
+			name: "valid large value headers",
+			magicHeaders: []MagicHeader{
+				{Min: 4294967290, Max: 4294967291},
+				{Min: 4294967292, Max: 4294967293},
+				{Min: 4294967294, Max: 4294967294},
+				{Min: 4294967295, Max: 4294967295},
+			},
+		},
+		{
+			name: "too few headers",
+			magicHeaders: []MagicHeader{
+				{Min: 1, Max: 10},
+				{Min: 11, Max: 20},
+				{Min: 21, Max: 30},
+			},
+			errorMsg: "all header types should be included:",
+		},
+		{
+			name: "too many headers",
+			magicHeaders: []MagicHeader{
+				{Min: 1, Max: 10},
+				{Min: 11, Max: 20},
+				{Min: 21, Max: 30},
+				{Min: 31, Max: 40},
+				{Min: 41, Max: 50},
+			},
+			errorMsg: "all header types should be included:",
+		},
+		{
+			name:         "empty headers",
+			magicHeaders: []MagicHeader{},
+			errorMsg:     "all header types should be included:",
+		},
+		{
+			name: "overlapping headers",
+			magicHeaders: []MagicHeader{
+				{Min: 1, Max: 15},
+				{Min: 10, Max: 20},
+				{Min: 25, Max: 30},
+				{Min: 35, Max: 40},
+			},
+			errorMsg: "magic headers shouldn't overlap;",
+		},
+		{
+			name: "overlapping headers at limit-first",
+			magicHeaders: []MagicHeader{
+				{Min: 1, Max: 10},
+				{Min: 10, Max: 20},
+				{Min: 25, Max: 30},
+				{Min: 35, Max: 40},
+			},
+			errorMsg: "magic headers shouldn't overlap;",
+		},
+		{
+			name: "overlapping headers at limit-second",
+			magicHeaders: []MagicHeader{
+				{Min: 1, Max: 10},
+				{Min: 15, Max: 25},
+				{Min: 25, Max: 30},
+				{Min: 35, Max: 40},
+			},
+			errorMsg: "magic headers shouldn't overlap;",
+		},
+		{
+			name: "overlapping headers at limit-third",
+			magicHeaders: []MagicHeader{
+				{Min: 1, Max: 10},
+				{Min: 15, Max: 25},
+				{Min: 30, Max: 35},
+				{Min: 35, Max: 40},
+			},
+			errorMsg: "magic headers shouldn't overlap;",
+		},
+		{
+			name: "identical ranges",
+			magicHeaders: []MagicHeader{
+				{Min: 10, Max: 20},
+				{Min: 10, Max: 20},
+				{Min: 25, Max: 30},
+				{Min: 35, Max: 40},
+			},
+			errorMsg: "magic headers shouldn't overlap;",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			result, err := NewMagicHeaders(tt.magicHeaders)
+
+			if tt.errorMsg != "" {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tt.errorMsg)
+				require.Equal(t, MagicHeaders{}, result)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tt.magicHeaders, result.Values)
+				require.NotNil(t, result.randomGenerator)
+			}
+		})
+	}
+}
+
+// Mock PRNG for testing
+type mockPRNG struct {
+	returnValue uint32
+}
+
+func (m *mockPRNG) RandomSizeInRange(min, max uint32) uint32 {
+	return m.returnValue
+}
+
+func (m *mockPRNG) Get() uint64 {
+	return 0
+}
+func (m *mockPRNG) ReadSize(size int) []byte {
+	return make([]byte, size)
+}
+
+func TestMagicHeaders_Get(t *testing.T) {
+	// Create test headers
+	headers := []MagicHeader{
+		{Min: 1, Max: 10},
+		{Min: 11, Max: 20},
+		{Min: 21, Max: 30},
+		{Min: 31, Max: 40},
+	}
+
+	tests := []struct {
+		name           string
+		defaultMsgType uint32
+		mockValue      uint32
+		expectedValue  uint32
+		errorMsg       string
+	}{
+		{
+			name:           "valid type 1",
+			defaultMsgType: 1,
+			mockValue:      5,
+			expectedValue:  5,
+		},
+		{
+			name:           "valid type 2",
+			defaultMsgType: 2,
+			mockValue:      15,
+			expectedValue:  15,
+		},
+		{
+			name:           "valid type 3",
+			defaultMsgType: 3,
+			mockValue:      25,
+			expectedValue:  25,
+		},
+		{
+			name:           "valid type 4",
+			defaultMsgType: 4,
+			mockValue:      35,
+			expectedValue:  35,
+		},
+		{
+			name:           "invalid type 0",
+			defaultMsgType: 0,
+			mockValue:      0,
+			expectedValue:  0,
+			errorMsg:       "invalid msg type: 0",
+		},
+		{
+			name:           "invalid type 5",
+			defaultMsgType: 5,
+			mockValue:      0,
+			expectedValue:  0,
+			errorMsg:       "invalid msg type: 5",
+		},
+		{
+			name:           "invalid type max uint32",
+			defaultMsgType: 4294967295,
+			mockValue:      0,
+			expectedValue:  0,
+			errorMsg:       "invalid msg type: 4294967295",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			// Create a new instance with mock PRNG for each test
+			testMagicHeaders := MagicHeaders{
+				Values:          headers,
+				randomGenerator: &mockPRNG{returnValue: tt.mockValue},
+			}
+
+			result, err := testMagicHeaders.Get(tt.defaultMsgType)
+
+			if tt.errorMsg != "" {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tt.errorMsg)
+				require.Equal(t, uint32(0), result)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tt.expectedValue, result)
+			}
+		})
+	}
+}

device/awg/prng.go

@@ -0,0 +1,50 @@
+package awg
+
+import (
+	crand "crypto/rand"
+	v2 "math/rand/v2"
+
+	"golang.org/x/exp/constraints"
+)
+
+type RandomNumberGenerator[T constraints.Integer] interface {
+	RandomSizeInRange(min, max T) T
+	Get() uint64
+	ReadSize(size int) []byte
+}
+
+type PRNG[T constraints.Integer] struct {
+	cha8Rand *v2.ChaCha8
+}
+
+func NewPRNG[T constraints.Integer]() PRNG[T] {
+	buf := make([]byte, 32)
+	_, _ = crand.Read(buf)
+
+	return PRNG[T]{
+		cha8Rand: v2.NewChaCha8([32]byte(buf)),
+	}
+}
+
+func (p PRNG[T]) RandomSizeInRange(min, max T) T {
+	if min > max {
+		panic("min must be less than max")
+	}
+
+	if min == max {
+		return min
+	}
+
+	return T(p.Get()%uint64(max-min)) + min
+}
+
+func (p PRNG[T]) Get() uint64 {
+	return p.cha8Rand.Uint64()
+}
+
+func (p PRNG[T]) ReadSize(size int) []byte {
+	// TODO: use a memory pool to allocate
+	buf := make([]byte, size)
+	_, _ = p.cha8Rand.Read(buf)
+	return buf
+}

device/awg/special_handshake_handler.go

@@ -1,9 +1,6 @@
 package awg
 
 import (
-	"errors"
-	"time"
-
 	"github.com/tevino/abool"
 	"go.uber.org/atomic"
 )
@@ -21,25 +18,13 @@ var WaitResponse = struct {
 }
 
 type SpecialHandshakeHandler struct {
-	isFirstDone    bool
-	SpecialJunk    TagJunkPacketGenerators
-	ControlledJunk TagJunkPacketGenerators
-
-	nextItime time.Time
-	ITimeout  time.Duration // seconds
+	SpecialJunk TagJunkPacketGenerators
 
 	IsSet bool
 }
 
 func (handler *SpecialHandshakeHandler) Validate() error {
-	var errs []error
-	if err := handler.SpecialJunk.Validate(); err != nil {
-		errs = append(errs, err)
-	}
-	if err := handler.ControlledJunk.Validate(); err != nil {
-		errs = append(errs, err)
-	}
-	return errors.Join(errs...)
+	return handler.SpecialJunk.Validate()
 }
 
 func (handler *SpecialHandshakeHandler) GenerateSpecialJunk() [][]byte {
@@ -47,27 +32,5 @@ func (handler *SpecialHandshakeHandler) GenerateSpecialJunk() [][]byte {
 		return nil
 	}
 
-	// TODO: create tests
-	if !handler.isFirstDone {
-		handler.isFirstDone = true
-	} else if !handler.isTimeToSendSpecial() {
-		return nil
-	}
-
-	rv := handler.SpecialJunk.GeneratePackets()
-	handler.nextItime = time.Now().Add(handler.ITimeout)
-
-	return rv
-}
-
-func (handler *SpecialHandshakeHandler) isTimeToSendSpecial() bool {
-	return time.Now().After(handler.nextItime)
-}
-
-func (handler *SpecialHandshakeHandler) GenerateControlledJunk() [][]byte {
-	if !handler.ControlledJunk.IsDefined() {
-		return nil
-	}
-
-	return handler.ControlledJunk.GeneratePackets()
+	return handler.SpecialJunk.GeneratePackets()
 }

device/awg/tag_generator.go

@@ -59,43 +59,110 @@ func hexToBytes(hexStr string) ([]byte, error) {
 	return hex.DecodeString(hexStr)
 }
 
-type RandomPacketGenerator struct {
+type randomGeneratorBase struct {
 	cha8Rand *v2.ChaCha8
 	size     int
 }
 
-func (rpg *RandomPacketGenerator) Generate() []byte {
-	junk := make([]byte, rpg.size)
-	rpg.cha8Rand.Read(junk)
-	return junk
-}
-
-func (rpg *RandomPacketGenerator) Size() int {
-	return rpg.size
-}
-
-func newRandomPacketGenerator(param string) (Generator, error) {
+func newRandomGeneratorBase(param string) (*randomGeneratorBase, error) {
 	size, err := strconv.Atoi(param)
 	if err != nil {
-		return nil, fmt.Errorf("random packet parse int: %w", err)
+		return nil, fmt.Errorf("parse int: %w", err)
 	}
 
 	if size > 1000 {
-		return nil, fmt.Errorf("random packet size must be less than 1000")
+		return nil, fmt.Errorf("size must be less than 1000")
 	}
 
 	buf := make([]byte, 32)
 	_, err = crand.Read(buf)
 	if err != nil {
-		return nil, fmt.Errorf("random packet crand read: %w", err)
+		return nil, fmt.Errorf("crand read: %w", err)
 	}
 
-	return &RandomPacketGenerator{
+	return &randomGeneratorBase{
 		cha8Rand: v2.NewChaCha8([32]byte(buf)),
 		size:     size,
 	}, nil
 }
 
+func (rpg *randomGeneratorBase) generate() []byte {
+	junk := make([]byte, rpg.size)
+	rpg.cha8Rand.Read(junk)
+	return junk
+}
+
+func (rpg *randomGeneratorBase) Size() int {
+	return rpg.size
+}
+
+type RandomBytesGenerator struct {
+	*randomGeneratorBase
+}
+
+func newRandomBytesGenerator(param string) (Generator, error) {
+	rpgBase, err := newRandomGeneratorBase(param)
+	if err != nil {
+		return nil, fmt.Errorf("new random bytes generator: %w", err)
+	}
+
+	return &RandomBytesGenerator{randomGeneratorBase: rpgBase}, nil
+}
+
+func (rpg *RandomBytesGenerator) Generate() []byte {
+	return rpg.generate()
+}
+
+const alphanumericChars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
+
+type RandomASCIIGenerator struct {
+	*randomGeneratorBase
+}
+
+func newRandomASCIIGenerator(param string) (Generator, error) {
+	rpgBase, err := newRandomGeneratorBase(param)
+	if err != nil {
+		return nil, fmt.Errorf("new random ascii generator: %w", err)
+	}
+
+	return &RandomASCIIGenerator{randomGeneratorBase: rpgBase}, nil
+}
+
+func (rpg *RandomASCIIGenerator) Generate() []byte {
+	junk := rpg.generate()
+
+	result := make([]byte, rpg.size)
+	for i, b := range junk {
+		result[i] = alphanumericChars[b%byte(len(alphanumericChars))]
+	}
+
+	return result
+}
+
+type RandomDigitGenerator struct {
+	*randomGeneratorBase
+}
+
+func newRandomDigitGenerator(param string) (Generator, error) {
+	rpgBase, err := newRandomGeneratorBase(param)
+	if err != nil {
+		return nil, fmt.Errorf("new random digit generator: %w", err)
+	}
+
+	return &RandomDigitGenerator{randomGeneratorBase: rpgBase}, nil
+}
+
+func (rpg *RandomDigitGenerator) Generate() []byte {
+	junk := rpg.generate()
+
+	result := make([]byte, rpg.size)
+	for i, b := range junk {
+		result[i] = '0' + (b % 10) // Convert to digit character
+	}
+
+	return result
+}
+
 type TimestampGenerator struct {
 }
 
@@ -117,34 +184,6 @@ func newTimestampGenerator(param string) (Generator, error) {
 	return &TimestampGenerator{}, nil
 }
 
-type WaitTimeoutGenerator struct {
-	waitTimeout time.Duration
-}
-
-func (wtg *WaitTimeoutGenerator) Generate() []byte {
-	time.Sleep(wtg.waitTimeout)
-	return []byte{}
-}
-
-func (wtg *WaitTimeoutGenerator) Size() int {
-	return 0
-}
-
-func newWaitTimeoutGenerator(param string) (Generator, error) {
-	timeout, err := strconv.Atoi(param)
-	if err != nil {
-		return nil, fmt.Errorf("timeout parse int: %w", err)
-	}
-
-	if timeout > 5000 {
-		return nil, fmt.Errorf("timeout must be less than 5000ms")
-	}
-
-	return &WaitTimeoutGenerator{
-		waitTimeout: time.Duration(timeout) * time.Millisecond,
-	}, nil
-}
-
 type PacketCounterGenerator struct {
 }
 

device/awg/tag_generator_test.go

@@ -8,7 +8,9 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func Test_newBytesGenerator(t *testing.T) {
+func TestNewBytesGenerator(t *testing.T) {
+	t.Parallel()
+
 	type args struct {
 		param string
 	}
@@ -63,6 +65,8 @@ func Test_newBytesGenerator(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
 			got, err := newBytesGenerator(tt.args.param)
 
 			if tt.wantErr != nil {
@@ -80,7 +84,9 @@ func Test_newBytesGenerator(t *testing.T) {
 	}
 }
 
-func Test_newRandomPacketGenerator(t *testing.T) {
+func TestNewRandomBytesGenerator(t *testing.T) {
+	t.Parallel()
+
 	type args struct {
 		param string
 	}
@@ -117,9 +123,134 @@ func Test_newRandomPacketGenerator(t *testing.T) {
 			},
 		},
 	}
+
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := newRandomPacketGenerator(tt.args.param)
+			t.Parallel()
+
+			got, err := newRandomBytesGenerator(tt.args.param)
+			if tt.wantErr != nil {
+				require.ErrorAs(t, err, &tt.wantErr)
+				require.Nil(t, got)
+				return
+			}
+
+			require.Nil(t, err)
+			require.NotNil(t, got)
+			first := got.Generate()
+
+			second := got.Generate()
+			require.NotEqual(t, first, second)
+		})
+	}
+}
+
+func TestNewRandomASCIIGenerator(t *testing.T) {
+	t.Parallel()
+
+	type args struct {
+		param string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		wantErr error
+	}{
+		{
+			name: "empty",
+			args: args{
+				param: "",
+			},
+			wantErr: fmt.Errorf("parse int"),
+		},
+		{
+			name: "not an int",
+			args: args{
+				param: "x",
+			},
+			wantErr: fmt.Errorf("parse int"),
+		},
+		{
+			name: "too large",
+			args: args{
+				param: "1001",
+			},
+			wantErr: fmt.Errorf("random packet size must be less than 1000"),
+		},
+		{
+			name: "valid",
+			args: args{
+				param: "12",
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			got, err := newRandomASCIIGenerator(tt.args.param)
+			if tt.wantErr != nil {
+				require.ErrorAs(t, err, &tt.wantErr)
+				require.Nil(t, got)
+				return
+			}
+
+			require.Nil(t, err)
+			require.NotNil(t, got)
+			first := got.Generate()
+
+			second := got.Generate()
+			require.NotEqual(t, first, second)
+		})
+	}
+}
+
+func TestNewRandomDigitGenerator(t *testing.T) {
+	t.Parallel()
+
+	type args struct {
+		param string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		wantErr error
+	}{
+		{
+			name: "empty",
+			args: args{
+				param: "",
+			},
+			wantErr: fmt.Errorf("parse int"),
+		},
+		{
+			name: "not an int",
+			args: args{
+				param: "x",
+			},
+			wantErr: fmt.Errorf("parse int"),
+		},
+		{
+			name: "too large",
+			args: args{
+				param: "1001",
+			},
+			wantErr: fmt.Errorf("random packet size must be less than 1000"),
+		},
+		{
+			name: "valid",
+			args: args{
+				param: "12",
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			got, err := newRandomDigitGenerator(tt.args.param)
 			if tt.wantErr != nil {
 				require.ErrorAs(t, err, &tt.wantErr)
 				require.Nil(t, got)
@@ -137,6 +268,8 @@ func Test_newRandomPacketGenerator(t *testing.T) {
 }
 
 func TestPacketCounterGenerator(t *testing.T) {
+	t.Parallel()
+
 	tests := []struct {
 		name    string
 		param   string
@@ -155,7 +288,6 @@ func TestPacketCounterGenerator(t *testing.T) {
 	}
 
 	for _, tc := range tests {
-		tc := tc // capture range variable
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 

device/awg/tag_parser.go

@@ -12,21 +12,21 @@ type IpcFields struct{ Key, Value string }
 type EnumTag string
 
 const (
-	BytesEnumTag        EnumTag = "b"
-	CounterEnumTag      EnumTag = "c"
-	TimestampEnumTag    EnumTag = "t"
-	RandomBytesEnumTag  EnumTag = "r"
-	WaitTimeoutEnumTag  EnumTag = "wt"
-	WaitResponseEnumTag EnumTag = "wr"
+	BytesEnumTag       EnumTag = "b"
+	CounterEnumTag     EnumTag = "c"
+	TimestampEnumTag   EnumTag = "t"
+	RandomBytesEnumTag EnumTag = "r"
+	RandomASCIIEnumTag EnumTag = "rc"
+	RandomDigitEnumTag EnumTag = "rd"
 )
 
 var generatorCreator = map[EnumTag]newGenerator{
 	BytesEnumTag:       newBytesGenerator,
 	CounterEnumTag:     newPacketCounterGenerator,
 	TimestampEnumTag:   newTimestampGenerator,
-	RandomBytesEnumTag: newRandomPacketGenerator,
-	WaitTimeoutEnumTag: newWaitTimeoutGenerator,
-	// WaitResponseEnumTag: newWaitResponseGenerator,
+	RandomBytesEnumTag: newRandomBytesGenerator,
+	RandomASCIIEnumTag: newRandomASCIIGenerator,
+	RandomDigitEnumTag: newRandomDigitGenerator,
 }
 
 // helper map to determine enumTags are unique
@@ -55,7 +55,7 @@ func parseTag(input string) (Tag, error) {
 	return tag, nil
 }
 
-func Parse(name, input string) (TagJunkPacketGenerator, error) {
+func ParseTagJunkGenerator(name, input string) (TagJunkPacketGenerator, error) {
 	inputSlice := strings.Split(input, "<")
 	if len(inputSlice) <= 1 {
 		return TagJunkPacketGenerator{}, fmt.Errorf("empty input: %s", input)

device/awg/tag_parser_test.go

@@ -64,7 +64,7 @@ func TestParse(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			_, err := Parse(tt.args.name, tt.args.input)
+			_, err := ParseTagJunkGenerator(tt.args.name, tt.args.input)
 
 			// TODO:  ErrorAs doesn't work as you think
 			if tt.wantErr != nil {

device/bind_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/channels.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/constants.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/cookie.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
@@ -118,6 +118,7 @@ func (st *CookieChecker) CreateReply(
 	msg []byte,
 	recv uint32,
 	src []byte,
+	msgType uint32,
 ) (*MessageCookieReply, error) {
 	st.RLock()
 
@@ -153,7 +154,7 @@ func (st *CookieChecker) CreateReply(
 	smac1 := smac2 - blake2s.Size128
 
 	reply := new(MessageCookieReply)
-	reply.Type = MessageCookieReplyType
+	reply.Type = msgType
 	reply.Receiver = recv
 
 	_, err := rand.Read(reply.Nonce[:])

device/cookie_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
@@ -99,7 +99,7 @@ func TestCookieMAC1(t *testing.T) {
 			0x8c, 0xe1, 0xe8, 0xfa, 0x67, 0x20, 0x80, 0x6d,
 		}
 		generator.AddMacs(msg)
-		reply, err := checker.CreateReply(msg, 1377, src)
+		reply, err := checker.CreateReply(msg, 1377, src, DefaultMessageCookieReplyType)
 		if err != nil {
 			t.Fatal("Failed to create cookie reply:", err)
 		}

device/device.go

@@ -1,12 +1,14 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
 
 import (
+	"encoding/binary"
 	"errors"
+	"fmt"
 	"runtime"
 	"sync"
 	"sync/atomic"
@@ -578,6 +580,7 @@ func (device *Device) BindClose() error {
 	device.net.Unlock()
 	return err
 }
+
 func (device *Device) isAWG() bool {
 	return device.version >= VersionAwg
 }
@@ -591,134 +594,126 @@ func (device *Device) resetProtocol() {
 }
 
 func (device *Device) handlePostConfig(tempAwg *awg.Protocol) error {
-	if !tempAwg.ASecCfg.IsSet && !tempAwg.HandshakeHandler.IsSet {
+	if !tempAwg.Cfg.IsSet && !tempAwg.HandshakeHandler.IsSet {
 		return nil
 	}
 
 	var errs []error
 
-	isASecOn := false
-	device.awg.ASecMux.Lock()
-	if tempAwg.ASecCfg.JunkPacketCount < 0 {
+	isAwgOn := false
+	device.awg.Mux.Lock()
+	if tempAwg.Cfg.JunkPacketCount < 0 {
 		errs = append(errs, ipcErrorf(
 			ipc.IpcErrorInvalid,
 			"JunkPacketCount should be non negative",
 		),
 		)
 	}
-	device.awg.ASecCfg.JunkPacketCount = tempAwg.ASecCfg.JunkPacketCount
-	if tempAwg.ASecCfg.JunkPacketCount != 0 {
-		isASecOn = true
+	device.awg.Cfg.JunkPacketCount = tempAwg.Cfg.JunkPacketCount
+	if tempAwg.Cfg.JunkPacketCount != 0 {
+		isAwgOn = true
 	}
 
-	device.awg.ASecCfg.JunkPacketMinSize = tempAwg.ASecCfg.JunkPacketMinSize
-	if tempAwg.ASecCfg.JunkPacketMinSize != 0 {
-		isASecOn = true
+	device.awg.Cfg.JunkPacketMinSize = tempAwg.Cfg.JunkPacketMinSize
+	if tempAwg.Cfg.JunkPacketMinSize != 0 {
+		isAwgOn = true
 	}
 
-	if device.awg.ASecCfg.JunkPacketCount > 0 &&
-		tempAwg.ASecCfg.JunkPacketMaxSize == tempAwg.ASecCfg.JunkPacketMinSize {
+	if device.awg.Cfg.JunkPacketCount > 0 &&
+		tempAwg.Cfg.JunkPacketMaxSize == tempAwg.Cfg.JunkPacketMinSize {
 
-		tempAwg.ASecCfg.JunkPacketMaxSize++ // to make rand gen work
+		tempAwg.Cfg.JunkPacketMaxSize++ // to make rand gen work
 	}
 
-	if tempAwg.ASecCfg.JunkPacketMaxSize >= MaxSegmentSize {
-		device.awg.ASecCfg.JunkPacketMinSize = 0
-		device.awg.ASecCfg.JunkPacketMaxSize = 1
+	if tempAwg.Cfg.JunkPacketMaxSize >= MaxSegmentSize {
+		device.awg.Cfg.JunkPacketMinSize = 0
+		device.awg.Cfg.JunkPacketMaxSize = 1
 		errs = append(errs, ipcErrorf(
 			ipc.IpcErrorInvalid,
 			"JunkPacketMaxSize: %d; should be smaller than maxSegmentSize: %d",
-			tempAwg.ASecCfg.JunkPacketMaxSize,
+			tempAwg.Cfg.JunkPacketMaxSize,
 			MaxSegmentSize,
 		))
-	} else if tempAwg.ASecCfg.JunkPacketMaxSize < tempAwg.ASecCfg.JunkPacketMinSize {
+	} else if tempAwg.Cfg.JunkPacketMaxSize < tempAwg.Cfg.JunkPacketMinSize {
 		errs = append(errs, ipcErrorf(
 			ipc.IpcErrorInvalid,
 			"maxSize: %d; should be greater than minSize: %d",
-			tempAwg.ASecCfg.JunkPacketMaxSize,
-			tempAwg.ASecCfg.JunkPacketMinSize,
+			tempAwg.Cfg.JunkPacketMaxSize,
+			tempAwg.Cfg.JunkPacketMinSize,
 		))
 	} else {
-		device.awg.ASecCfg.JunkPacketMaxSize = tempAwg.ASecCfg.JunkPacketMaxSize
+		device.awg.Cfg.JunkPacketMaxSize = tempAwg.Cfg.JunkPacketMaxSize
 	}
 
-	if tempAwg.ASecCfg.JunkPacketMaxSize != 0 {
-		isASecOn = true
+	if tempAwg.Cfg.JunkPacketMaxSize != 0 {
+		isAwgOn = true
 	}
 
-	if MessageInitiationSize+tempAwg.ASecCfg.InitPacketJunkSize >= MaxSegmentSize {
-		errs = append(errs, ipcErrorf(
+	magicHeaders := make([]awg.MagicHeader, 4)
+
+	if len(tempAwg.Cfg.MagicHeaders.Values) != 4 {
+		return ipcErrorf(
 			ipc.IpcErrorInvalid,
-			`init header size(148) + junkSize:%d; should be smaller than maxSegmentSize: %d`,
-			tempAwg.ASecCfg.InitPacketJunkSize,
-			MaxSegmentSize,
-		),
+			"magic headers should have 4 values; got: %d",
+			len(tempAwg.Cfg.MagicHeaders.Values),
 		)
-	} else {
-		device.awg.ASecCfg.InitPacketJunkSize = tempAwg.ASecCfg.InitPacketJunkSize
 	}
 
-	if tempAwg.ASecCfg.InitPacketJunkSize != 0 {
-		isASecOn = true
-	}
-
-	if MessageResponseSize+tempAwg.ASecCfg.ResponsePacketJunkSize >= MaxSegmentSize {
-		errs = append(errs, ipcErrorf(
-			ipc.IpcErrorInvalid,
-			`response header size(92) + junkSize:%d; should be smaller than maxSegmentSize: %d`,
-			tempAwg.ASecCfg.ResponsePacketJunkSize,
-			MaxSegmentSize,
-		),
-		)
-	} else {
-		device.awg.ASecCfg.ResponsePacketJunkSize = tempAwg.ASecCfg.ResponsePacketJunkSize
-	}
-
-	if tempAwg.ASecCfg.ResponsePacketJunkSize != 0 {
-		isASecOn = true
-	}
-
-	if tempAwg.ASecCfg.InitPacketMagicHeader > 4 {
-		isASecOn = true
+	if tempAwg.Cfg.MagicHeaders.Values[0].Min > 4 {
+		isAwgOn = true
 		device.log.Verbosef("UAPI: Updating init_packet_magic_header")
-		device.awg.ASecCfg.InitPacketMagicHeader = tempAwg.ASecCfg.InitPacketMagicHeader
-		MessageInitiationType = device.awg.ASecCfg.InitPacketMagicHeader
+		magicHeaders[0] = tempAwg.Cfg.MagicHeaders.Values[0]
+
+		MessageInitiationType = magicHeaders[0].Min
 	} else {
 		device.log.Verbosef("UAPI: Using default init type")
 		MessageInitiationType = DefaultMessageInitiationType
+		magicHeaders[0] = awg.NewMagicHeaderSameValue(DefaultMessageInitiationType)
 	}
 
-	if tempAwg.ASecCfg.ResponsePacketMagicHeader > 4 {
-		isASecOn = true
+	if tempAwg.Cfg.MagicHeaders.Values[1].Min > 4 {
+		isAwgOn = true
+
 		device.log.Verbosef("UAPI: Updating response_packet_magic_header")
-		device.awg.ASecCfg.ResponsePacketMagicHeader = tempAwg.ASecCfg.ResponsePacketMagicHeader
-		MessageResponseType = device.awg.ASecCfg.ResponsePacketMagicHeader
+		magicHeaders[1] = tempAwg.Cfg.MagicHeaders.Values[1]
+		MessageResponseType = magicHeaders[1].Min
 	} else {
 		device.log.Verbosef("UAPI: Using default response type")
 		MessageResponseType = DefaultMessageResponseType
+		magicHeaders[1] = awg.NewMagicHeaderSameValue(DefaultMessageResponseType)
 	}
 
-	if tempAwg.ASecCfg.UnderloadPacketMagicHeader > 4 {
-		isASecOn = true
+	if tempAwg.Cfg.MagicHeaders.Values[2].Min > 4 {
+		isAwgOn = true
+
 		device.log.Verbosef("UAPI: Updating underload_packet_magic_header")
-		device.awg.ASecCfg.UnderloadPacketMagicHeader = tempAwg.ASecCfg.UnderloadPacketMagicHeader
-		MessageCookieReplyType = device.awg.ASecCfg.UnderloadPacketMagicHeader
+		magicHeaders[2] = tempAwg.Cfg.MagicHeaders.Values[2]
+		MessageCookieReplyType = magicHeaders[2].Min
 	} else {
 		device.log.Verbosef("UAPI: Using default underload type")
 		MessageCookieReplyType = DefaultMessageCookieReplyType
+		magicHeaders[2] = awg.NewMagicHeaderSameValue(DefaultMessageCookieReplyType)
 	}
 
-	if tempAwg.ASecCfg.TransportPacketMagicHeader > 4 {
-		isASecOn = true
+	if tempAwg.Cfg.MagicHeaders.Values[3].Min > 4 {
+		isAwgOn = true
+
 		device.log.Verbosef("UAPI: Updating transport_packet_magic_header")
-		device.awg.ASecCfg.TransportPacketMagicHeader = tempAwg.ASecCfg.TransportPacketMagicHeader
-		MessageTransportType = device.awg.ASecCfg.TransportPacketMagicHeader
+		magicHeaders[3] = tempAwg.Cfg.MagicHeaders.Values[3]
+		MessageTransportType = magicHeaders[3].Min
 	} else {
 		device.log.Verbosef("UAPI: Using default transport type")
 		MessageTransportType = DefaultMessageTransportType
+		magicHeaders[3] = awg.NewMagicHeaderSameValue(DefaultMessageTransportType)
 	}
 
-	isSameMap := map[uint32]struct{}{
+	var err error
+	device.awg.Cfg.MagicHeaders, err = awg.NewMagicHeaders(magicHeaders)
+	if err != nil {
+		errs = append(errs, ipcErrorf(ipc.IpcErrorInvalid, "new magic headers: %w", err))
+	}
+
+	isSameHeaderMap := map[uint32]struct{}{
 		MessageInitiationType:  {},
 		MessageResponseType:    {},
 		MessageCookieReplyType: {},
@@ -726,7 +721,7 @@ func (device *Device) handlePostConfig(tempAwg *awg.Protocol) error {
 	}
 
 	// size will be different if same values
-	if len(isSameMap) != 4 {
+	if len(isSameHeaderMap) != 4 {
 		errs = append(errs, ipcErrorf(
 			ipc.IpcErrorInvalid,
 			`magic headers should differ; got: init:%d; recv:%d; unde:%d; tran:%d`,
@@ -738,39 +733,113 @@ func (device *Device) handlePostConfig(tempAwg *awg.Protocol) error {
 		)
 	}
 
-	newInitSize := MessageInitiationSize + device.awg.ASecCfg.InitPacketJunkSize
-	newResponseSize := MessageResponseSize + device.awg.ASecCfg.ResponsePacketJunkSize
+	newInitSize := MessageInitiationSize + tempAwg.Cfg.InitHeaderJunkSize
 
-	if newInitSize == newResponseSize {
+	if newInitSize >= MaxSegmentSize {
 		errs = append(errs, ipcErrorf(
 			ipc.IpcErrorInvalid,
-			`new init size:%d; and new response size:%d; should differ`,
-			newInitSize,
-			newResponseSize,
+			`init header size(148) + junkSize:%d; should be smaller than maxSegmentSize: %d`,
+			tempAwg.Cfg.InitHeaderJunkSize,
+			MaxSegmentSize,
 		),
 		)
 	} else {
-		packetSizeToMsgType = map[int]uint32{
-			newInitSize:            MessageInitiationType,
-			newResponseSize:        MessageResponseType,
-			MessageCookieReplySize: MessageCookieReplyType,
-			MessageTransportSize:   MessageTransportType,
-		}
+		device.awg.Cfg.InitHeaderJunkSize = tempAwg.Cfg.InitHeaderJunkSize
+	}
 
+	if tempAwg.Cfg.InitHeaderJunkSize != 0 {
+		isAwgOn = true
+	}
+
+	newResponseSize := MessageResponseSize + tempAwg.Cfg.ResponseHeaderJunkSize
+
+	if newResponseSize >= MaxSegmentSize {
+		errs = append(errs, ipcErrorf(
+			ipc.IpcErrorInvalid,
+			`response header size(92) + junkSize:%d; should be smaller than maxSegmentSize: %d`,
+			tempAwg.Cfg.ResponseHeaderJunkSize,
+			MaxSegmentSize,
+		),
+		)
+	} else {
+		device.awg.Cfg.ResponseHeaderJunkSize = tempAwg.Cfg.ResponseHeaderJunkSize
+	}
+
+	if tempAwg.Cfg.ResponseHeaderJunkSize != 0 {
+		isAwgOn = true
+	}
+
+	newCookieSize := MessageCookieReplySize + tempAwg.Cfg.CookieReplyHeaderJunkSize
+
+	if newCookieSize >= MaxSegmentSize {
+		errs = append(errs, ipcErrorf(
+			ipc.IpcErrorInvalid,
+			`cookie reply size(92) + junkSize:%d; should be smaller than maxSegmentSize: %d`,
+			tempAwg.Cfg.CookieReplyHeaderJunkSize,
+			MaxSegmentSize,
+		),
+		)
+	} else {
+		device.awg.Cfg.CookieReplyHeaderJunkSize = tempAwg.Cfg.CookieReplyHeaderJunkSize
+	}
+
+	if tempAwg.Cfg.CookieReplyHeaderJunkSize != 0 {
+		isAwgOn = true
+	}
+
+	newTransportSize := MessageTransportSize + tempAwg.Cfg.TransportHeaderJunkSize
+
+	if newTransportSize >= MaxSegmentSize {
+		errs = append(errs, ipcErrorf(
+			ipc.IpcErrorInvalid,
+			`transport size(92) + junkSize:%d; should be smaller than maxSegmentSize: %d`,
+			tempAwg.Cfg.TransportHeaderJunkSize,
+			MaxSegmentSize,
+		),
+		)
+	} else {
+		device.awg.Cfg.TransportHeaderJunkSize = tempAwg.Cfg.TransportHeaderJunkSize
+	}
+
+	if tempAwg.Cfg.TransportHeaderJunkSize != 0 {
+		isAwgOn = true
+	}
+
+	isSameSizeMap := map[int]struct{}{
+		newInitSize:      {},
+		newResponseSize:  {},
+		newCookieSize:    {},
+		newTransportSize: {},
+	}
+
+	if len(isSameSizeMap) != 4 {
+		errs = append(errs, ipcErrorf(
+			ipc.IpcErrorInvalid,
+			`new sizes should differ; init: %d; response: %d; cookie: %d; trans: %d`,
+			newInitSize,
+			newResponseSize,
+			newCookieSize,
+			newTransportSize,
+		),
+		)
+	} else {
 		msgTypeToJunkSize = map[uint32]int{
-			MessageInitiationType:  device.awg.ASecCfg.InitPacketJunkSize,
-			MessageResponseType:    device.awg.ASecCfg.ResponsePacketJunkSize,
-			MessageCookieReplyType: 0,
-			MessageTransportType:   0,
+			MessageInitiationType:  device.awg.Cfg.InitHeaderJunkSize,
+			MessageResponseType:    device.awg.Cfg.ResponseHeaderJunkSize,
+			MessageCookieReplyType: device.awg.Cfg.CookieReplyHeaderJunkSize,
+			MessageTransportType:   device.awg.Cfg.TransportHeaderJunkSize,
+		}
+
+		packetSizeToMsgType = map[int]uint32{
+			newInitSize:      MessageInitiationType,
+			newResponseSize:  MessageResponseType,
+			newCookieSize:    MessageCookieReplyType,
+			newTransportSize: MessageTransportType,
 		}
 	}
 
-	device.awg.IsASecOn.SetTo(isASecOn)
-	var err error
-	device.awg.JunkCreator, err = awg.NewJunkCreator(device.awg.ASecCfg)
-	if err != nil {
-		errs = append(errs, err)
-	}
+	device.awg.IsOn.SetTo(isAwgOn)
+	device.awg.JunkCreator = awg.NewJunkCreator(device.awg.Cfg)
 
 	if tempAwg.HandshakeHandler.IsSet {
 		if err := tempAwg.HandshakeHandler.Validate(); err != nil {
@@ -778,15 +847,91 @@ func (device *Device) handlePostConfig(tempAwg *awg.Protocol) error {
 				ipc.IpcErrorInvalid, "handshake handler validate: %w", err))
 		} else {
 			device.awg.HandshakeHandler = tempAwg.HandshakeHandler
-			device.awg.HandshakeHandler.ControlledJunk.DefaultJunkCount = tempAwg.ASecCfg.JunkPacketCount
-			device.awg.HandshakeHandler.SpecialJunk.DefaultJunkCount = tempAwg.ASecCfg.JunkPacketCount
+			device.awg.HandshakeHandler.SpecialJunk.DefaultJunkCount = tempAwg.Cfg.JunkPacketCount
 			device.version = VersionAwgSpecialHandshake
 		}
 	} else {
 		device.version = VersionAwg
 	}
 
-	device.awg.ASecMux.Unlock()
+	device.awg.Mux.Unlock()
 
 	return errors.Join(errs...)
 }
+
+func (device *Device) ProcessAWGPacket(size int, packet *[]byte, buffer *[MaxMessageSize]byte) (uint32, error) {
+	// TODO:
+	// if awg.WaitResponse.ShouldWait.IsSet() {
+	// 	awg.WaitResponse.Channel <- struct{}{}
+	// }
+
+	expectedMsgType, isKnownSize := packetSizeToMsgType[size]
+	if !isKnownSize {
+		msgType, err := device.handleTransport(size, packet, buffer)
+
+		if err != nil {
+			return 0, fmt.Errorf("handle transport: %w", err)
+		}
+
+		return msgType, nil
+	}
+
+	junkSize := msgTypeToJunkSize[expectedMsgType]
+
+	// transport size can align with other header types;
+	// making sure we have the right actualMsgType
+	actualMsgType, err := device.getMsgType(packet, junkSize)
+	if err != nil {
+		return 0, fmt.Errorf("get msg type: %w", err)
+	}
+
+	if actualMsgType == expectedMsgType {
+		*packet = (*packet)[junkSize:]
+		return actualMsgType, nil
+	}
+
+	device.log.Verbosef("awg: transport packet lined up with another msg type")
+
+	msgType, err := device.handleTransport(size, packet, buffer)
+	if err != nil {
+		return 0, fmt.Errorf("handle transport: %w", err)
+	}
+
+	return msgType, nil
+}
+
+func (device *Device) getMsgType(packet *[]byte, junkSize int) (uint32, error) {
+	msgTypeValue := binary.LittleEndian.Uint32((*packet)[junkSize : junkSize+4])
+	msgType, err := device.awg.GetMagicHeaderMinFor(msgTypeValue)
+
+	if err != nil {
+		return 0, fmt.Errorf("get magic header min: %w", err)
+	}
+
+	return msgType, nil
+}
+
+func (device *Device) handleTransport(size int, packet *[]byte, buffer *[MaxMessageSize]byte) (uint32, error) {
+	junkSize := device.awg.Cfg.TransportHeaderJunkSize
+
+	msgType, err := device.getMsgType(packet, junkSize)
+	if err != nil {
+		return 0, fmt.Errorf("get msg type: %w", err)
+	}
+
+	if msgType != MessageTransportType {
+		// probably a junk packet
+		return 0, fmt.Errorf("Received message with unknown type: %d", msgType)
+	}
+
+	if junkSize > 0 {
+		// remove junk from buffer by shifting the packet
+		// this buffer is also used for decryption, so it needs to be corrected
+		copy((*buffer)[:size], (*packet)[junkSize:])
+		size -= junkSize
+		// need to reinitialize packet as well
+		(*packet) = (*packet)[:size]
+	}
+
+	return msgType, nil
+}

device/device_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
@@ -136,6 +136,7 @@ func (pair *testPair) Send(
 		// pong is the new ping
 		p0, p1 = p1, p0
 	}
+
 	msg := tuntest.Ping(p0.ip, p1.ip)
 	p1.tun.Outbound <- msg
 	timer := time.NewTimer(6 * time.Second)
@@ -226,16 +227,19 @@ func TestTwoDevicePing(t *testing.T) {
 // Run test with -race=false to avoid the race for setting the default msgTypes 2 times
 func TestAWGDevicePing(t *testing.T) {
 	goroutineLeakCheck(t)
+
 	pair := genTestPair(t, true,
 		"jc", "5",
 		"jmin", "500",
 		"jmax", "1000",
-		"s1", "30",
-		"s2", "40",
-		"h1", "123456",
-		"h2", "67543",
-		"h4", "32345",
-		"h3", "123123",
+		"s1", "15",
+		"s2", "18",
+		"s3", "20",
+		"s4", "25",
+		"h1", "123456-123500",
+		"h2", "67543-67550",
+		"h3", "123123-123200",
+		"h4", "32345-32350",
 	)
 	t.Run("ping 1.0.0.1", func(t *testing.T) {
 		pair.Send(t, Ping, nil)
@@ -260,12 +264,10 @@ func TestAWGHandshakeDevicePing(t *testing.T) {
 
 	goroutineLeakCheck(t)
 	pair := genTestPair(t, true,
-		"i1", "<b 0xf6ab3267fa><c><b 0xf6ab><t><r 10><wt 10>",
-		"i2", "<b 0xf6ab3267fa><r 100>",
-		"j1", "<b 0xffffffff><c><b 0xf6ab><t><r 10>",
-		"j2", "<c><b 0xf6ab><t><wt 1000>",
-		"j3", "<t><b 0xf6ab><c><r 10>",
-		"itime", "60",
+		"i1", "<b 0xf6ab3267fa><c><b 0xf6ab><t><r 10>",
+		"i2", "<b 0xf6ab3267fa><c><b 0xf6ab><t><rc 10>",
+		"i3", "<b 0xf6ab3267fa><c><b 0xf6ab><t><rd 10>",
+		"i4", "<b 0xf6ab3267fa><r 100>",
 		// "jc", "1",
 		// "jmin", "500",
 		// "jmax", "1000",

device/endpoint_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/indextable.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/ip.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/kdf_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/keypair.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/logger.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/mobilequirks.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/noise-helpers.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/noise-protocol.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
@@ -82,9 +82,10 @@ const (
 	MessageTransportOffsetContent  = 16
 )
 
-var packetSizeToMsgType map[int]uint32
-
-var msgTypeToJunkSize map[uint32]int
+var (
+	packetSizeToMsgType map[int]uint32
+	msgTypeToJunkSize   map[uint32]int
+)
 
 /* Type is an 8-bit field, followed by 3 nul bytes,
  * by marshalling the messages in little-endian byteorder
@@ -204,12 +205,22 @@ func (device *Device) CreateMessageInitiation(peer *Peer) (*MessageInitiation, e
 
 	handshake.mixHash(handshake.remoteStatic[:])
 
-	device.awg.ASecMux.RLock()
+	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()
+	}
+
 	msg := MessageInitiation{
-		Type:      MessageInitiationType,
+		Type:      msgType,
 		Ephemeral: handshake.localEphemeral.publicKey(),
 	}
-	device.awg.ASecMux.RUnlock()
 
 	handshake.mixKey(msg.Ephemeral[:])
 	handshake.mixHash(msg.Ephemeral[:])
@@ -263,12 +274,13 @@ func (device *Device) ConsumeMessageInitiation(msg *MessageInitiation) *Peer {
 		chainKey [blake2s.Size]byte
 	)
 
-	device.awg.ASecMux.RLock()
+	device.awg.Mux.RLock()
+
 	if msg.Type != MessageInitiationType {
-		device.awg.ASecMux.RUnlock()
+		device.awg.Mux.RUnlock()
 		return nil
 	}
-	device.awg.ASecMux.RUnlock()
+	device.awg.Mux.RUnlock()
 
 	device.staticIdentity.RLock()
 	defer device.staticIdentity.RUnlock()
@@ -383,9 +395,19 @@ func (device *Device) CreateMessageResponse(peer *Peer) (*MessageResponse, error
 	}
 
 	var msg MessageResponse
-	device.awg.ASecMux.RLock()
-	msg.Type = MessageResponseType
-	device.awg.ASecMux.RUnlock()
+	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()
+	} else {
+		msg.Type = DefaultMessageResponseType
+	}
+
 	msg.Sender = handshake.localIndex
 	msg.Receiver = handshake.remoteIndex
 
@@ -435,12 +457,13 @@ func (device *Device) CreateMessageResponse(peer *Peer) (*MessageResponse, error
 }
 
 func (device *Device) ConsumeMessageResponse(msg *MessageResponse) *Peer {
-	device.awg.ASecMux.RLock()
+	device.awg.Mux.RLock()
+
 	if msg.Type != MessageResponseType {
-		device.awg.ASecMux.RUnlock()
+		device.awg.Mux.RUnlock()
 		return nil
 	}
-	device.awg.ASecMux.RUnlock()
+	device.awg.Mux.RUnlock()
 
 	// lookup handshake by receiver
 

device/noise-types.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/noise_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/peer.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
@@ -114,6 +114,16 @@ func (device *Device) NewPeer(pk NoisePublicKey) (*Peer, error) {
 	return peer, nil
 }
 
+func (peer *Peer) SendAndCountBuffers(buffers [][]byte) error {
+	err := peer.SendBuffers(buffers)
+	if err == nil {
+		awg.PacketCounter.Add(uint64(len(buffers)))
+		return nil
+	}
+
+	return err
+}
+
 func (peer *Peer) SendBuffers(buffers [][]byte) error {
 	peer.device.net.RLock()
 	defer peer.device.net.RUnlock()
@@ -145,16 +155,6 @@ func (peer *Peer) SendBuffers(buffers [][]byte) error {
 	return err
 }
 
-func (peer *Peer) SendAndCountBuffers(buffers [][]byte) error {
-	err := peer.SendBuffers(buffers)
-	if err == nil {
-		awg.PacketCounter.Add(uint64(len(buffers)))
-		return nil
-	}
-
-	return err
-}
-
 func (peer *Peer) String() string {
 	// The awful goo that follows is identical to:
 	//

device/pools.go

@@ -1,20 +1,19 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
 
 import (
 	"sync"
-	"sync/atomic"
 )
 
 type WaitPool struct {
 	pool  sync.Pool
 	cond  sync.Cond
 	lock  sync.Mutex
-	count atomic.Uint32
+	count uint32 // Get calls not yet Put back
 	max   uint32
 }
 
@@ -27,10 +26,10 @@ func NewWaitPool(max uint32, new func() any) *WaitPool {
 func (p *WaitPool) Get() any {
 	if p.max != 0 {
 		p.lock.Lock()
-		for p.count.Load() >= p.max {
+		for p.count >= p.max {
 			p.cond.Wait()
 		}
-		p.count.Add(1)
+		p.count++
 		p.lock.Unlock()
 	}
 	return p.pool.Get()
@@ -41,7 +40,9 @@ func (p *WaitPool) Put(x any) {
 	if p.max == 0 {
 		return
 	}
-	p.count.Add(^uint32(0))
+	p.lock.Lock()
+	defer p.lock.Unlock()
+	p.count--
 	p.cond.Signal()
 }
 

device/pools_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
@@ -32,7 +32,9 @@ func TestWaitPool(t *testing.T) {
 	wg.Add(workers)
 	var max atomic.Uint32
 	updateMax := func() {
-		count := p.count.Load()
+		p.lock.Lock()
+		count := p.count
+		p.lock.Unlock()
 		if count > p.max {
 			t.Errorf("count (%d) > max (%d)", count, p.max)
 		}

device/queueconstants_android.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/queueconstants_default.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/queueconstants_ios.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/queueconstants_windows.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/race_disabled_test.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/race_enabled_test.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/receive.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
@@ -129,7 +129,7 @@ func (device *Device) RoutineReceiveIncoming(
 		}
 		deathSpiral = 0
 
-		device.awg.ASecMux.RLock()
+		device.awg.Mux.RLock()
 		// handle each packet in the batch
 		for i, size := range sizes[:count] {
 			if size < MinMessageSize {
@@ -140,33 +140,16 @@ func (device *Device) RoutineReceiveIncoming(
 			packet := bufsArrs[i][:size]
 			var msgType uint32
 			if device.isAWG() {
-				// TODO:
-				// if awg.WaitResponse.ShouldWait.IsSet() {
-				// 	awg.WaitResponse.Channel <- struct{}{}
-				// }
+				msgType, err = device.ProcessAWGPacket(size, &packet, bufsArrs[i])
 
-				if assumedMsgType, ok := packetSizeToMsgType[size]; ok {
-					junkSize := msgTypeToJunkSize[assumedMsgType]
-					// transport size can align with other header types;
-					// making sure we have the right msgType
-					msgType = binary.LittleEndian.Uint32(packet[junkSize : junkSize+4])
-					if msgType == assumedMsgType {
-						packet = packet[junkSize:]
-					} else {
-						device.log.Verbosef("transport packet lined up with another msg type")
-						msgType = binary.LittleEndian.Uint32(packet[:4])
-					}
-				} else {
-					msgType = binary.LittleEndian.Uint32(packet[:4])
-					if msgType != MessageTransportType {
-						// probably a junk packet
-						device.log.Verbosef("aSec: Received message with unknown type: %d", msgType)
-						continue
-					}
+				if err != nil {
+					device.log.Verbosef("awg: process packet: %v", err)
+					continue
 				}
 			} else {
 				msgType = binary.LittleEndian.Uint32(packet[:4])
 			}
+
 			switch msgType {
 
 			// check if transport
@@ -250,7 +233,7 @@ func (device *Device) RoutineReceiveIncoming(
 			default:
 			}
 		}
-		device.awg.ASecMux.RUnlock()
+		device.awg.Mux.RUnlock()
 		for peer, elemsContainer := range elemsByPeer {
 			if peer.isRunning.Load() {
 				peer.queue.inbound.c <- elemsContainer
@@ -309,7 +292,7 @@ func (device *Device) RoutineHandshake(id int) {
 
 	for elem := range device.queue.handshake.c {
 
-		device.awg.ASecMux.RLock()
+		device.awg.Mux.RLock()
 
 		// handle cookie fields and ratelimiting
 
@@ -396,6 +379,9 @@ func (device *Device) RoutineHandshake(id int) {
 				goto skip
 			}
 
+			// have to reassign msgType for ranged msgType to work
+			msg.Type = elem.msgType
+
 			// consume initiation
 			peer := device.ConsumeMessageInitiation(&msg)
 			if peer == nil {
@@ -428,6 +414,9 @@ func (device *Device) RoutineHandshake(id int) {
 				goto skip
 			}
 
+			// have to reassign msgType for ranged msgType to work
+			msg.Type = elem.msgType
+
 			// consume response
 
 			peer := device.ConsumeMessageResponse(&msg)
@@ -461,7 +450,7 @@ func (device *Device) RoutineHandshake(id int) {
 			peer.SendKeepalive()
 		}
 	skip:
-		device.awg.ASecMux.RUnlock()
+		device.awg.Mux.RUnlock()
 		device.PutMessageBuffer(elem.buffer)
 	}
 }

device/send.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
@@ -124,35 +124,25 @@ func (peer *Peer) SendHandshakeInitiation(isRetry bool) error {
 		return err
 	}
 	var sendBuffer [][]byte
+
 	// so only packet processed for cookie generation
 	var junkedHeader []byte
-
 	if peer.device.version >= VersionAwg {
 		var junks [][]byte
 		if peer.device.version == VersionAwgSpecialHandshake {
-			peer.device.awg.ASecMux.RLock()
+			peer.device.awg.Mux.RLock()
 			// set junks depending on packet type
 			junks = peer.device.awg.HandshakeHandler.GenerateSpecialJunk()
-			if junks == nil {
-				junks = peer.device.awg.HandshakeHandler.GenerateControlledJunk()
-				if junks != nil {
-					peer.device.log.Verbosef("%v - Controlled junks sent", peer)
-				}
-			} else {
+			if junks != nil {
 				peer.device.log.Verbosef("%v - Special junks sent", peer)
 			}
-			peer.device.awg.ASecMux.RUnlock()
+			peer.device.awg.Mux.RUnlock()
 		} else {
-			junks = make([][]byte, 0, peer.device.awg.ASecCfg.JunkPacketCount)
-		}
-		peer.device.awg.ASecMux.RLock()
-		err := peer.device.awg.JunkCreator.CreateJunkPackets(&junks)
-		peer.device.awg.ASecMux.RUnlock()
-
-		if err != nil {
-			peer.device.log.Errorf("%v - %v", peer, err)
-			return err
+			junks = make([][]byte, 0, peer.device.awg.Cfg.JunkPacketCount)
 		}
+		peer.device.awg.Mux.RLock()
+		peer.device.awg.JunkCreator.CreateJunkPackets(&junks)
+		peer.device.awg.Mux.RUnlock()
 
 		if len(junks) > 0 {
 			err = peer.SendBuffers(junks)
@@ -163,19 +153,11 @@ func (peer *Peer) SendHandshakeInitiation(isRetry bool) error {
 			}
 		}
 
-		peer.device.awg.ASecMux.RLock()
-		if peer.device.awg.ASecCfg.InitPacketJunkSize != 0 {
-			buf := make([]byte, 0, peer.device.awg.ASecCfg.InitPacketJunkSize)
-			writer := bytes.NewBuffer(buf[:0])
-			err = peer.device.awg.JunkCreator.AppendJunk(writer, peer.device.awg.ASecCfg.InitPacketJunkSize)
-			if err != nil {
-				peer.device.log.Errorf("%v - %v", peer, err)
-				peer.device.awg.ASecMux.RUnlock()
-				return err
-			}
-			junkedHeader = writer.Bytes()
+		junkedHeader, err = peer.device.awg.CreateInitHeaderJunk()
+		if err != nil {
+			peer.device.log.Errorf("%v - %v", peer, err)
+			return err
 		}
-		peer.device.awg.ASecMux.RUnlock()
 	}
 
 	var buf [MessageInitiationSize]byte
@@ -211,22 +193,13 @@ func (peer *Peer) SendHandshakeResponse() error {
 		peer.device.log.Errorf("%v - Failed to create response message: %v", peer, err)
 		return err
 	}
-	var junkedHeader []byte
-	if peer.device.isAWG() {
-		peer.device.awg.ASecMux.RLock()
-		if peer.device.awg.ASecCfg.ResponsePacketJunkSize != 0 {
-			buf := make([]byte, 0, peer.device.awg.ASecCfg.ResponsePacketJunkSize)
-			writer := bytes.NewBuffer(buf[:0])
-			err = peer.device.awg.JunkCreator.AppendJunk(writer, peer.device.awg.ASecCfg.ResponsePacketJunkSize)
-			if err != nil {
-				peer.device.awg.ASecMux.RUnlock()
-				peer.device.log.Errorf("%v - %v", peer, err)
-				return err
-			}
-			junkedHeader = writer.Bytes()
-		}
-		peer.device.awg.ASecMux.RUnlock()
+
+	junkedHeader, err := peer.device.awg.CreateResponseHeaderJunk()
+	if err != nil {
+		peer.device.log.Errorf("%v - %v", peer, err)
+		return err
 	}
+
 	var buf [MessageResponseSize]byte
 	writer := bytes.NewBuffer(buf[:0])
 
@@ -259,21 +232,43 @@ 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 := 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(
 		initiatingElem.packet,
 		sender,
 		initiatingElem.endpoint.DstToBytes(),
+		msgType,
 	)
 	if err != nil {
 		device.log.Errorf("Failed to create cookie reply: %v", err)
 		return err
 	}
 
+	junkedHeader, err := device.awg.CreateCookieReplyHeaderJunk()
+	if err != nil {
+		device.log.Errorf("%v - %v", device, err)
+		return err
+	}
+
 	var buf [MessageCookieReplySize]byte
 	writer := bytes.NewBuffer(buf[:0])
 	binary.Write(writer, binary.LittleEndian, reply)
+
+	junkedHeader = append(junkedHeader, writer.Bytes()...)
 	// TODO: allocation could be avoided
-	device.net.bind.Send([][]byte{writer.Bytes()}, initiatingElem.endpoint)
+	device.net.bind.Send([][]byte{junkedHeader}, initiatingElem.endpoint)
 	return nil
 }
 
@@ -537,7 +532,20 @@ func (device *Device) RoutineEncryption(id int) {
 			fieldReceiver := header[4:8]
 			fieldNonce := header[8:16]
 
-			binary.LittleEndian.PutUint32(fieldType, MessageTransportType)
+			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)
 
@@ -586,6 +594,7 @@ func (peer *Peer) RoutineSequentialSender(maxBatchSize int) {
 				device.PutMessageBuffer(elem.buffer)
 				device.PutOutboundElement(elem)
 			}
+			device.PutOutboundElementsContainer(elemsContainer)
 			continue
 		}
 		dataSent := false
@@ -593,6 +602,14 @@ func (peer *Peer) RoutineSequentialSender(maxBatchSize int) {
 		for _, elem := range elemsContainer.elems {
 			if len(elem.packet) != MessageKeepaliveSize {
 				dataSent = true
+
+				junkedHeader, err := device.awg.CreateTransportHeaderJunk(len(elem.packet))
+				if err != nil {
+					device.log.Errorf("%v - %v", device, err)
+					continue
+				}
+
+				elem.packet = append(junkedHeader, elem.packet...)
 			}
 			bufs = append(bufs, elem.packet)
 		}
@@ -604,6 +621,7 @@ func (peer *Peer) RoutineSequentialSender(maxBatchSize int) {
 		if dataSent {
 			peer.timersDataSent()
 		}
+
 		for _, elem := range elemsContainer.elems {
 			device.PutMessageBuffer(elem.buffer)
 			device.PutOutboundElement(elem)

device/sticky_default.go

@@ -7,6 +7,6 @@ import (
 	"github.com/amnezia-vpn/amneziawg-go/rwcancel"
 )
 
-func (device *Device) startRouteListener(bind conn.Bind) (*rwcancel.RWCancel, error) {
+func (device *Device) startRouteListener(_ conn.Bind) (*rwcancel.RWCancel, error) {
 	return nil, nil
 }

device/sticky_linux.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  *
  * This implements userspace semantics of "sticky sockets", modeled after
  * WireGuard's kernelspace implementation. This is more or less a straight port
@@ -9,7 +9,7 @@
  *
  * Currently there is no way to achieve this within the net package:
  * See e.g. https://github.com/golang/go/issues/17930
- * So this code is remains platform dependent.
+ * So this code remains platform dependent.
  */
 
 package device
@@ -47,7 +47,7 @@ func (device *Device) startRouteListener(bind conn.Bind) (*rwcancel.RWCancel, er
 	return netlinkCancel, nil
 }
 
-func (device *Device) routineRouteListener(bind conn.Bind, netlinkSock int, netlinkCancel *rwcancel.RWCancel) {
+func (device *Device) routineRouteListener(_ conn.Bind, netlinkSock int, netlinkCancel *rwcancel.RWCancel) {
 	type peerEndpointPtr struct {
 		peer     *Peer
 		endpoint *conn.Endpoint

device/timers.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  *
  * This is based heavily on timers.c from the kernel implementation.
  */

device/tun.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device

device/uapi.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package device
@@ -99,45 +99,42 @@ func (device *Device) IpcGetOperation(w io.Writer) error {
 		}
 
 		if device.isAWG() {
-			if device.awg.ASecCfg.JunkPacketCount != 0 {
-				sendf("jc=%d", device.awg.ASecCfg.JunkPacketCount)
+			if device.awg.Cfg.JunkPacketCount != 0 {
+				sendf("jc=%d", device.awg.Cfg.JunkPacketCount)
 			}
-			if device.awg.ASecCfg.JunkPacketMinSize != 0 {
-				sendf("jmin=%d", device.awg.ASecCfg.JunkPacketMinSize)
+			if device.awg.Cfg.JunkPacketMinSize != 0 {
+				sendf("jmin=%d", device.awg.Cfg.JunkPacketMinSize)
 			}
-			if device.awg.ASecCfg.JunkPacketMaxSize != 0 {
-				sendf("jmax=%d", device.awg.ASecCfg.JunkPacketMaxSize)
+			if device.awg.Cfg.JunkPacketMaxSize != 0 {
+				sendf("jmax=%d", device.awg.Cfg.JunkPacketMaxSize)
 			}
-			if device.awg.ASecCfg.InitPacketJunkSize != 0 {
-				sendf("s1=%d", device.awg.ASecCfg.InitPacketJunkSize)
+			if device.awg.Cfg.InitHeaderJunkSize != 0 {
+				sendf("s1=%d", device.awg.Cfg.InitHeaderJunkSize)
 			}
-			if device.awg.ASecCfg.ResponsePacketJunkSize != 0 {
-				sendf("s2=%d", device.awg.ASecCfg.ResponsePacketJunkSize)
+			if device.awg.Cfg.ResponseHeaderJunkSize != 0 {
+				sendf("s2=%d", device.awg.Cfg.ResponseHeaderJunkSize)
 			}
-			if device.awg.ASecCfg.InitPacketMagicHeader != 0 {
-				sendf("h1=%d", device.awg.ASecCfg.InitPacketMagicHeader)
+			if device.awg.Cfg.CookieReplyHeaderJunkSize != 0 {
+				sendf("s3=%d", device.awg.Cfg.CookieReplyHeaderJunkSize)
 			}
-			if device.awg.ASecCfg.ResponsePacketMagicHeader != 0 {
-				sendf("h2=%d", device.awg.ASecCfg.ResponsePacketMagicHeader)
+			if device.awg.Cfg.TransportHeaderJunkSize != 0 {
+				sendf("s4=%d", device.awg.Cfg.TransportHeaderJunkSize)
 			}
-			if device.awg.ASecCfg.UnderloadPacketMagicHeader != 0 {
-				sendf("h3=%d", device.awg.ASecCfg.UnderloadPacketMagicHeader)
-			}
-			if device.awg.ASecCfg.TransportPacketMagicHeader != 0 {
-				sendf("h4=%d", device.awg.ASecCfg.TransportPacketMagicHeader)
+			for i, magicHeader := range device.awg.Cfg.MagicHeaders.Values {
+				if magicHeader.Min > 4 {
+					if magicHeader.Min == magicHeader.Max {
+						sendf("h%d=%d", i+1, magicHeader.Min)
+						continue
+					}
+
+					sendf("h%d=%d-%d", i+1, magicHeader.Min, magicHeader.Max)
+				}
 			}
 
 			specialJunkIpcFields := device.awg.HandshakeHandler.SpecialJunk.IpcGetFields()
 			for _, field := range specialJunkIpcFields {
 				sendf("%s=%s", field.Key, field.Value)
 			}
-			controlledJunkIpcFields := device.awg.HandshakeHandler.ControlledJunk.IpcGetFields()
-			for _, field := range controlledJunkIpcFields {
-				sendf("%s=%s", field.Key, field.Value)
-			}
-			if device.awg.HandshakeHandler.ITimeout != 0 {
-				sendf("itime=%d", device.awg.HandshakeHandler.ITimeout/time.Second)
-			}
 		}
 
 		for _, peer := range device.peers.keyMap {
@@ -194,6 +191,8 @@ func (device *Device) IpcSetOperation(r io.Reader) (err error) {
 	deviceConfig := true
 
 	tempAwg := awg.Protocol{}
+	tempAwg.Cfg.MagicHeaders.Values = make([]awg.MagicHeader, 4)
+
 	scanner := bufio.NewScanner(r)
 	for scanner.Scan() {
 		line := scanner.Text()
@@ -306,8 +305,8 @@ func (device *Device) handleDeviceLine(key, value string, tempAwg *awg.Protocol)
 			return ipcErrorf(ipc.IpcErrorInvalid, "parse junk_packet_count %w", err)
 		}
 		device.log.Verbosef("UAPI: Updating junk_packet_count")
-		tempAwg.ASecCfg.JunkPacketCount = junkPacketCount
-		tempAwg.ASecCfg.IsSet = true
+		tempAwg.Cfg.JunkPacketCount = junkPacketCount
+		tempAwg.Cfg.IsSet = true
 
 	case "jmin":
 		junkPacketMinSize, err := strconv.Atoi(value)
@@ -315,8 +314,8 @@ func (device *Device) handleDeviceLine(key, value string, tempAwg *awg.Protocol)
 			return ipcErrorf(ipc.IpcErrorInvalid, "parse junk_packet_min_size %w", err)
 		}
 		device.log.Verbosef("UAPI: Updating junk_packet_min_size")
-		tempAwg.ASecCfg.JunkPacketMinSize = junkPacketMinSize
-		tempAwg.ASecCfg.IsSet = true
+		tempAwg.Cfg.JunkPacketMinSize = junkPacketMinSize
+		tempAwg.Cfg.IsSet = true
 
 	case "jmax":
 		junkPacketMaxSize, err := strconv.Atoi(value)
@@ -324,8 +323,8 @@ func (device *Device) handleDeviceLine(key, value string, tempAwg *awg.Protocol)
 			return ipcErrorf(ipc.IpcErrorInvalid, "parse junk_packet_max_size %w", err)
 		}
 		device.log.Verbosef("UAPI: Updating junk_packet_max_size")
-		tempAwg.ASecCfg.JunkPacketMaxSize = junkPacketMaxSize
-		tempAwg.ASecCfg.IsSet = true
+		tempAwg.Cfg.JunkPacketMaxSize = junkPacketMaxSize
+		tempAwg.Cfg.IsSet = true
 
 	case "s1":
 		initPacketJunkSize, err := strconv.Atoi(value)
@@ -333,8 +332,8 @@ func (device *Device) handleDeviceLine(key, value string, tempAwg *awg.Protocol)
 			return ipcErrorf(ipc.IpcErrorInvalid, "parse init_packet_junk_size %w", err)
 		}
 		device.log.Verbosef("UAPI: Updating init_packet_junk_size")
-		tempAwg.ASecCfg.InitPacketJunkSize = initPacketJunkSize
-		tempAwg.ASecCfg.IsSet = true
+		tempAwg.Cfg.InitHeaderJunkSize = initPacketJunkSize
+		tempAwg.Cfg.IsSet = true
 
 	case "s2":
 		responsePacketJunkSize, err := strconv.Atoi(value)
@@ -342,76 +341,71 @@ func (device *Device) handleDeviceLine(key, value string, tempAwg *awg.Protocol)
 			return ipcErrorf(ipc.IpcErrorInvalid, "parse response_packet_junk_size %w", err)
 		}
 		device.log.Verbosef("UAPI: Updating response_packet_junk_size")
-		tempAwg.ASecCfg.ResponsePacketJunkSize = responsePacketJunkSize
-		tempAwg.ASecCfg.IsSet = true
+		tempAwg.Cfg.ResponseHeaderJunkSize = responsePacketJunkSize
+		tempAwg.Cfg.IsSet = true
 
+	case "s3":
+		cookieReplyPacketJunkSize, err := strconv.Atoi(value)
+		if err != nil {
+			return ipcErrorf(ipc.IpcErrorInvalid, "parse cookie_reply_packet_junk_size %w", err)
+		}
+		device.log.Verbosef("UAPI: Updating cookie_reply_packet_junk_size")
+		tempAwg.Cfg.CookieReplyHeaderJunkSize = cookieReplyPacketJunkSize
+		tempAwg.Cfg.IsSet = true
+
+	case "s4":
+		transportPacketJunkSize, err := strconv.Atoi(value)
+		if err != nil {
+			return ipcErrorf(ipc.IpcErrorInvalid, "parse transport_packet_junk_size %w", err)
+		}
+		device.log.Verbosef("UAPI: Updating transport_packet_junk_size")
+		tempAwg.Cfg.TransportHeaderJunkSize = transportPacketJunkSize
+		tempAwg.Cfg.IsSet = true
 	case "h1":
-		initPacketMagicHeader, err := strconv.ParseUint(value, 10, 32)
+		initMagicHeader, err := awg.ParseMagicHeader(key, value)
 		if err != nil {
-			return ipcErrorf(ipc.IpcErrorInvalid, "parse init_packet_magic_header %w", err)
+			return ipcErrorf(ipc.IpcErrorInvalid, "uapi: %w", err)
 		}
-		tempAwg.ASecCfg.InitPacketMagicHeader = uint32(initPacketMagicHeader)
-		tempAwg.ASecCfg.IsSet = true
 
+		tempAwg.Cfg.MagicHeaders.Values[0] = initMagicHeader
+		tempAwg.Cfg.IsSet = true
 	case "h2":
-		responsePacketMagicHeader, err := strconv.ParseUint(value, 10, 32)
+		responseMagicHeader, err := awg.ParseMagicHeader(key, value)
 		if err != nil {
-			return ipcErrorf(ipc.IpcErrorInvalid, "parse response_packet_magic_header %w", err)
+			return ipcErrorf(ipc.IpcErrorInvalid, "uapi: %w", err)
 		}
-		tempAwg.ASecCfg.ResponsePacketMagicHeader = uint32(responsePacketMagicHeader)
-		tempAwg.ASecCfg.IsSet = true
 
+		tempAwg.Cfg.MagicHeaders.Values[1] = responseMagicHeader
+		tempAwg.Cfg.IsSet = true
 	case "h3":
-		underloadPacketMagicHeader, err := strconv.ParseUint(value, 10, 32)
+		cookieReplyMagicHeader, err := awg.ParseMagicHeader(key, value)
 		if err != nil {
-			return ipcErrorf(ipc.IpcErrorInvalid, "parse underload_packet_magic_header %w", err)
+			return ipcErrorf(ipc.IpcErrorInvalid, "uapi: %w", err)
 		}
-		tempAwg.ASecCfg.UnderloadPacketMagicHeader = uint32(underloadPacketMagicHeader)
-		tempAwg.ASecCfg.IsSet = true
 
+		tempAwg.Cfg.MagicHeaders.Values[2] = cookieReplyMagicHeader
+		tempAwg.Cfg.IsSet = true
 	case "h4":
-		transportPacketMagicHeader, err := strconv.ParseUint(value, 10, 32)
+		transportMagicHeader, err := awg.ParseMagicHeader(key, value)
 		if err != nil {
-			return ipcErrorf(ipc.IpcErrorInvalid, "parse transport_packet_magic_header %w", err)
+			return ipcErrorf(ipc.IpcErrorInvalid, "uapi: %w", err)
 		}
-		tempAwg.ASecCfg.TransportPacketMagicHeader = uint32(transportPacketMagicHeader)
-		tempAwg.ASecCfg.IsSet = true
+
+		tempAwg.Cfg.MagicHeaders.Values[3] = transportMagicHeader
+		tempAwg.Cfg.IsSet = true
 	case "i1", "i2", "i3", "i4", "i5":
 		if len(value) == 0 {
 			device.log.Verbosef("UAPI: received empty %s", key)
 			return nil
 		}
 
-		generators, err := awg.Parse(key, value)
+		generators, err := awg.ParseTagJunkGenerator(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.IsSet = true
-	case "j1", "j2", "j3":
-		if len(value) == 0 {
-			device.log.Verbosef("UAPI: received empty %s", key)
-			return nil
-		}
-
-		generators, 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.IsSet = true
-	case "itime":
-		itime, err := strconv.ParseInt(value, 10, 64)
-		if err != nil {
-			return ipcErrorf(ipc.IpcErrorInvalid, "parse itime %w", err)
-		}
-		device.log.Verbosef("UAPI: Updating itime")
-
-		tempAwg.HandshakeHandler.ITimeout = time.Duration(itime) * time.Second
-		tempAwg.HandshakeHandler.IsSet = true
 	default:
 		return ipcErrorf(ipc.IpcErrorInvalid, "invalid UAPI device key: %v", key)
 	}
@@ -561,7 +555,14 @@ func (device *Device) handlePeerLine(
 		device.allowedips.RemoveByPeer(peer.Peer)
 
 	case "allowed_ip":
-		device.log.Verbosef("%v - UAPI: Adding allowedip", peer.Peer)
+		add := true
+		verb := "Adding"
+		if len(value) > 0 && value[0] == '-' {
+			add = false
+			verb = "Removing"
+			value = value[1:]
+		}
+		device.log.Verbosef("%v - UAPI: %s allowedip", peer.Peer, verb)
 		prefix, err := netip.ParsePrefix(value)
 		if err != nil {
 			return ipcErrorf(ipc.IpcErrorInvalid, "failed to set allowed ip: %w", err)
@@ -569,7 +570,11 @@ func (device *Device) handlePeerLine(
 		if peer.dummy {
 			return nil
 		}
-		device.allowedips.Insert(prefix, peer.Peer)
+		if add {
+			device.allowedips.Insert(prefix, peer.Peer)
+		} else {
+			device.allowedips.Remove(prefix, peer.Peer)
+		}
 
 	case "protocol_version":
 		if value != "1" {

format_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 package main
 

go.mod

@@ -7,10 +7,11 @@ require (
 	github.com/tevino/abool v1.2.0
 	go.uber.org/atomic v1.11.0
 	golang.org/x/crypto v0.39.0
+	golang.org/x/exp v0.0.0-20230725093048-515e97ebf090
 	golang.org/x/net v0.41.0
 	golang.org/x/sys v0.33.0
 	golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2
-	gvisor.dev/gvisor v0.0.0-20250606233247-e3c4c4cad86f
+	gvisor.dev/gvisor v0.0.0-20231202080848-1f7806d17489
 )
 
 require (

go.sum

@@ -2,8 +2,6 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/google/btree v1.1.3 h1:CVpQJjYgC4VbzxeGVHfvZrv1ctoYCAI8vbl07Fcxlyg=
 github.com/google/btree v1.1.3/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4=
-github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
-github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
@@ -14,8 +12,8 @@ go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE=
 go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
 golang.org/x/crypto v0.39.0 h1:SHs+kF4LP+f+p14esP5jAoDpHU8Gu/v9lFRK6IT5imM=
 golang.org/x/crypto v0.39.0/go.mod h1:L+Xg3Wf6HoL4Bn4238Z6ft6KfEpN0tJGo53AAPC632U=
-golang.org/x/mod v0.21.0 h1:vvrHzRwRfVKSiLrG+d4FMl/Qi4ukBCE6kZlTUkDYRT0=
-golang.org/x/mod v0.21.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY=
+golang.org/x/exp v0.0.0-20230725093048-515e97ebf090 h1:Di6/M8l0O2lCLc6VVRWhgCiApHV8MnQurBnFSHsQtNY=
+golang.org/x/exp v0.0.0-20230725093048-515e97ebf090/go.mod h1:FXUEEKJgO7OQYeo8N01OfiKP8RXMtf6e8aTskBGqWdc=
 golang.org/x/net v0.41.0 h1:vBTly1HeNPEn3wtREYfy4GZ/NECgw2Cnl+nK6Nz3uvw=
 golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA=
 golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw=
@@ -28,5 +26,5 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
-gvisor.dev/gvisor v0.0.0-20250606233247-e3c4c4cad86f h1:zmc4cHEcCudRt2O8VsCW7nYLfAsbVY2i910/DAop1TM=
-gvisor.dev/gvisor v0.0.0-20250606233247-e3c4c4cad86f/go.mod h1:3r5CMtNQMKIvBlrmM9xWUNamjKBYPOWyXOjmg5Kts3g=
+gvisor.dev/gvisor v0.0.0-20231202080848-1f7806d17489 h1:ze1vwAdliUAr68RQ5NtufWaXaOg8WUO2OACzEV+TNdE=
+gvisor.dev/gvisor v0.0.0-20231202080848-1f7806d17489/go.mod h1:10sU+Uh5KKNv1+2x2A0Gvzt8FjD3ASIhorV3YsauXhk=

ipc/uapi_bsd.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package ipc

ipc/uapi_linux.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package ipc

ipc/uapi_unix.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package ipc

ipc/uapi_wasm.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package ipc

ipc/uapi_windows.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package ipc

main.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package main

main_windows.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package main

ratelimiter/ratelimiter.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package ratelimiter

ratelimiter/ratelimiter_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package ratelimiter

replay/replay.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 // Package replay implements an efficient anti-replay algorithm as specified in RFC 6479.

replay/replay_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package replay

rwcancel/rwcancel.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 // Package rwcancel implements cancelable read/write operations on
@@ -64,7 +64,7 @@ func (rw *RWCancel) ReadyRead() bool {
 
 func (rw *RWCancel) ReadyWrite() bool {
 	closeFd := int32(rw.closingReader.Fd())
-	pollFds := []unix.PollFd{{Fd: int32(rw.fd), Events: unix.POLLOUT}, {Fd: closeFd, Events: unix.POLLOUT}}
+	pollFds := []unix.PollFd{{Fd: int32(rw.fd), Events: unix.POLLOUT}, {Fd: closeFd, Events: unix.POLLIN}}
 	var err error
 	for {
 		_, err = unix.Poll(pollFds, -1)

tai64n/tai64n.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package tai64n

tai64n/tai64n_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package tai64n

tun/alignment_windows_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package tun

tun/checksum.go

@@ -1,102 +1,86 @@
 package tun
 
-import "encoding/binary"
+import (
+	"encoding/binary"
+	"math/bits"
+)
 
 // TODO: Explore SIMD and/or other assembly optimizations.
-// TODO: Test native endian loads. See RFC 1071 section 2 part B.
 func checksumNoFold(b []byte, initial uint64) uint64 {
-	ac := initial
+	tmp := make([]byte, 8)
+	binary.NativeEndian.PutUint64(tmp, initial)
+	ac := binary.BigEndian.Uint64(tmp)
+	var carry uint64
 
 	for len(b) >= 128 {
-		ac += uint64(binary.BigEndian.Uint32(b[:4]))
-		ac += uint64(binary.BigEndian.Uint32(b[4:8]))
-		ac += uint64(binary.BigEndian.Uint32(b[8:12]))
-		ac += uint64(binary.BigEndian.Uint32(b[12:16]))
-		ac += uint64(binary.BigEndian.Uint32(b[16:20]))
-		ac += uint64(binary.BigEndian.Uint32(b[20:24]))
-		ac += uint64(binary.BigEndian.Uint32(b[24:28]))
-		ac += uint64(binary.BigEndian.Uint32(b[28:32]))
-		ac += uint64(binary.BigEndian.Uint32(b[32:36]))
-		ac += uint64(binary.BigEndian.Uint32(b[36:40]))
-		ac += uint64(binary.BigEndian.Uint32(b[40:44]))
-		ac += uint64(binary.BigEndian.Uint32(b[44:48]))
-		ac += uint64(binary.BigEndian.Uint32(b[48:52]))
-		ac += uint64(binary.BigEndian.Uint32(b[52:56]))
-		ac += uint64(binary.BigEndian.Uint32(b[56:60]))
-		ac += uint64(binary.BigEndian.Uint32(b[60:64]))
-		ac += uint64(binary.BigEndian.Uint32(b[64:68]))
-		ac += uint64(binary.BigEndian.Uint32(b[68:72]))
-		ac += uint64(binary.BigEndian.Uint32(b[72:76]))
-		ac += uint64(binary.BigEndian.Uint32(b[76:80]))
-		ac += uint64(binary.BigEndian.Uint32(b[80:84]))
-		ac += uint64(binary.BigEndian.Uint32(b[84:88]))
-		ac += uint64(binary.BigEndian.Uint32(b[88:92]))
-		ac += uint64(binary.BigEndian.Uint32(b[92:96]))
-		ac += uint64(binary.BigEndian.Uint32(b[96:100]))
-		ac += uint64(binary.BigEndian.Uint32(b[100:104]))
-		ac += uint64(binary.BigEndian.Uint32(b[104:108]))
-		ac += uint64(binary.BigEndian.Uint32(b[108:112]))
-		ac += uint64(binary.BigEndian.Uint32(b[112:116]))
-		ac += uint64(binary.BigEndian.Uint32(b[116:120]))
-		ac += uint64(binary.BigEndian.Uint32(b[120:124]))
-		ac += uint64(binary.BigEndian.Uint32(b[124:128]))
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[:8]), 0)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[8:16]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[16:24]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[24:32]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[32:40]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[40:48]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[48:56]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[56:64]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[64:72]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[72:80]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[80:88]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[88:96]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[96:104]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[104:112]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[112:120]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[120:128]), carry)
+		ac += carry
 		b = b[128:]
 	}
 	if len(b) >= 64 {
-		ac += uint64(binary.BigEndian.Uint32(b[:4]))
-		ac += uint64(binary.BigEndian.Uint32(b[4:8]))
-		ac += uint64(binary.BigEndian.Uint32(b[8:12]))
-		ac += uint64(binary.BigEndian.Uint32(b[12:16]))
-		ac += uint64(binary.BigEndian.Uint32(b[16:20]))
-		ac += uint64(binary.BigEndian.Uint32(b[20:24]))
-		ac += uint64(binary.BigEndian.Uint32(b[24:28]))
-		ac += uint64(binary.BigEndian.Uint32(b[28:32]))
-		ac += uint64(binary.BigEndian.Uint32(b[32:36]))
-		ac += uint64(binary.BigEndian.Uint32(b[36:40]))
-		ac += uint64(binary.BigEndian.Uint32(b[40:44]))
-		ac += uint64(binary.BigEndian.Uint32(b[44:48]))
-		ac += uint64(binary.BigEndian.Uint32(b[48:52]))
-		ac += uint64(binary.BigEndian.Uint32(b[52:56]))
-		ac += uint64(binary.BigEndian.Uint32(b[56:60]))
-		ac += uint64(binary.BigEndian.Uint32(b[60:64]))
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[:8]), 0)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[8:16]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[16:24]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[24:32]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[32:40]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[40:48]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[48:56]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[56:64]), carry)
+		ac += carry
 		b = b[64:]
 	}
 	if len(b) >= 32 {
-		ac += uint64(binary.BigEndian.Uint32(b[:4]))
-		ac += uint64(binary.BigEndian.Uint32(b[4:8]))
-		ac += uint64(binary.BigEndian.Uint32(b[8:12]))
-		ac += uint64(binary.BigEndian.Uint32(b[12:16]))
-		ac += uint64(binary.BigEndian.Uint32(b[16:20]))
-		ac += uint64(binary.BigEndian.Uint32(b[20:24]))
-		ac += uint64(binary.BigEndian.Uint32(b[24:28]))
-		ac += uint64(binary.BigEndian.Uint32(b[28:32]))
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[:8]), 0)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[8:16]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[16:24]), carry)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[24:32]), carry)
+		ac += carry
 		b = b[32:]
 	}
 	if len(b) >= 16 {
-		ac += uint64(binary.BigEndian.Uint32(b[:4]))
-		ac += uint64(binary.BigEndian.Uint32(b[4:8]))
-		ac += uint64(binary.BigEndian.Uint32(b[8:12]))
-		ac += uint64(binary.BigEndian.Uint32(b[12:16]))
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[:8]), 0)
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[8:16]), carry)
+		ac += carry
 		b = b[16:]
 	}
 	if len(b) >= 8 {
-		ac += uint64(binary.BigEndian.Uint32(b[:4]))
-		ac += uint64(binary.BigEndian.Uint32(b[4:8]))
+		ac, carry = bits.Add64(ac, binary.NativeEndian.Uint64(b[:8]), 0)
+		ac += carry
 		b = b[8:]
 	}
 	if len(b) >= 4 {
-		ac += uint64(binary.BigEndian.Uint32(b))
+		ac, carry = bits.Add64(ac, uint64(binary.NativeEndian.Uint32(b[:4])), 0)
+		ac += carry
 		b = b[4:]
 	}
 	if len(b) >= 2 {
-		ac += uint64(binary.BigEndian.Uint16(b))
+		ac, carry = bits.Add64(ac, uint64(binary.NativeEndian.Uint16(b[:2])), 0)
+		ac += carry
 		b = b[2:]
 	}
 	if len(b) == 1 {
-		ac += uint64(b[0]) << 8
+		tmp := binary.NativeEndian.Uint16([]byte{b[0], 0})
+		ac, carry = bits.Add64(ac, uint64(tmp), 0)
+		ac += carry
 	}
 
-	return ac
+	binary.NativeEndian.PutUint64(tmp, ac)
+	return binary.BigEndian.Uint64(tmp)
 }
 
 func checksum(b []byte, initial uint64) uint16 {

tun/checksum_test.go

@@ -1,11 +1,74 @@
 package tun
 
 import (
+	"encoding/binary"
 	"fmt"
 	"math/rand"
 	"testing"
+
+	"golang.org/x/sys/unix"
 )
 
+func checksumRef(b []byte, initial uint16) uint16 {
+	ac := uint64(initial)
+
+	for len(b) >= 2 {
+		ac += uint64(binary.BigEndian.Uint16(b))
+		b = b[2:]
+	}
+	if len(b) == 1 {
+		ac += uint64(b[0]) << 8
+	}
+
+	for (ac >> 16) > 0 {
+		ac = (ac >> 16) + (ac & 0xffff)
+	}
+	return uint16(ac)
+}
+
+func pseudoHeaderChecksumRefNoFold(protocol uint8, srcAddr, dstAddr []byte, totalLen uint16) uint16 {
+	sum := checksumRef(srcAddr, 0)
+	sum = checksumRef(dstAddr, sum)
+	sum = checksumRef([]byte{0, protocol}, sum)
+	tmp := make([]byte, 2)
+	binary.BigEndian.PutUint16(tmp, totalLen)
+	return checksumRef(tmp, sum)
+}
+
+func TestChecksum(t *testing.T) {
+	for length := 0; length <= 9001; length++ {
+		buf := make([]byte, length)
+		rng := rand.New(rand.NewSource(1))
+		rng.Read(buf)
+		csum := checksum(buf, 0x1234)
+		csumRef := checksumRef(buf, 0x1234)
+		if csum != csumRef {
+			t.Error("Expected checksum", csumRef, "got", csum)
+		}
+	}
+}
+
+func TestPseudoHeaderChecksum(t *testing.T) {
+	for _, addrLen := range []int{4, 16} {
+		for length := 0; length <= 9001; length++ {
+			srcAddr := make([]byte, addrLen)
+			dstAddr := make([]byte, addrLen)
+			buf := make([]byte, length)
+			rng := rand.New(rand.NewSource(1))
+			rng.Read(srcAddr)
+			rng.Read(dstAddr)
+			rng.Read(buf)
+			phSum := pseudoHeaderChecksumNoFold(unix.IPPROTO_TCP, srcAddr, dstAddr, uint16(length))
+			csum := checksum(buf, phSum)
+			phSumRef := pseudoHeaderChecksumRefNoFold(unix.IPPROTO_TCP, srcAddr, dstAddr, uint16(length))
+			csumRef := checksumRef(buf, phSumRef)
+			if csum != csumRef {
+				t.Error("Expected checksumRef", csumRef, "got", csum)
+			}
+		}
+	}
+}
+
 func BenchmarkChecksum(b *testing.B) {
 	lengths := []int{
 		64,

tun/netstack/examples/http_client.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package main

tun/netstack/examples/http_server.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package main

tun/netstack/examples/ping_client.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package main

tun/netstack/tun.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package netstack
@@ -43,6 +43,7 @@ type netTun struct {
 	ep             *channel.Endpoint
 	stack          *stack.Stack
 	events         chan tun.Event
+	notifyHandle   *channel.NotificationHandle
 	incomingPacket chan *buffer.View
 	mtu            int
 	dnsServers     []netip.Addr
@@ -70,7 +71,7 @@ func CreateNetTUN(localAddresses, dnsServers []netip.Addr, mtu int) (tun.Device,
 	if tcpipErr != nil {
 		return nil, nil, fmt.Errorf("could not enable TCP SACK: %v", tcpipErr)
 	}
-	dev.ep.AddNotify(dev)
+	dev.notifyHandle = dev.ep.AddNotify(dev)
 	tcpipErr = dev.stack.CreateNIC(1, dev.ep)
 	if tcpipErr != nil {
 		return nil, nil, fmt.Errorf("CreateNIC: %v", tcpipErr)
@@ -155,7 +156,7 @@ func (tun *netTun) Write(buf [][]byte, offset int) (int, error) {
 
 func (tun *netTun) WriteNotify() {
 	pkt := tun.ep.Read()
-	if pkt.IsNil() {
+	if pkt == nil {
 		return
 	}
 
@@ -167,13 +168,14 @@ func (tun *netTun) WriteNotify() {
 
 func (tun *netTun) Close() error {
 	tun.stack.RemoveNIC(1)
+	tun.stack.Close()
+	tun.ep.RemoveNotify(tun.notifyHandle)
+	tun.ep.Close()
 
 	if tun.events != nil {
 		close(tun.events)
 	}
 
-	tun.ep.Close()
-
 	if tun.incomingPacket != nil {
 		close(tun.incomingPacket)
 	}

tun/offload_linux.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package tun

tun/offload_linux_test.go

@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package tun

tun/operateonfd.go

@@ -2,7 +2,7 @@
 
 /* SPDX-License-Identifier: MIT
  *
- * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
  */
 
 package tun