From 35886b88d7ce8689a9dbd648d155de4cffc55c6a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Wed, 26 Oct 2022 20:21:07 +0800
Subject: [PATCH] Add option for custom wireguard reserved bytes

---
 option/wireguard.go                |  1 +
 outbound/wireguard.go              |  9 ++++++++-
 transport/wireguard/client_bind.go | 14 +++++++++++++-
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/option/wireguard.go b/option/wireguard.go
index 3b1ffb3f..978c130d 100644
--- a/option/wireguard.go
+++ b/option/wireguard.go
@@ -9,6 +9,7 @@ type WireGuardOutboundOptions struct {
 	PrivateKey      string                 `json:"private_key"`
 	PeerPublicKey   string                 `json:"peer_public_key"`
 	PreSharedKey    string                 `json:"pre_shared_key,omitempty"`
+	Reserved        []uint8                `json:"reserved,omitempty"`
 	MTU             uint32                 `json:"mtu,omitempty"`
 	Network         NetworkList            `json:"network,omitempty"`
 }
diff --git a/outbound/wireguard.go b/outbound/wireguard.go
index 0a8d245d..a0c1e933 100644
--- a/outbound/wireguard.go
+++ b/outbound/wireguard.go
@@ -45,8 +45,15 @@ func NewWireGuard(ctx context.Context, router adapter.Router, logger log.Context
 			tag:      tag,
 		},
 	}
+	var reserved [3]uint8
+	if len(options.Reserved) > 0 {
+		if len(options.Reserved) != 3 {
+			return nil, E.New("invalid reserved value, required 3 bytes, got ", len(options.Reserved))
+		}
+		copy(reserved[:], options.Reserved)
+	}
 	peerAddr := options.ServerOptions.Build()
-	outbound.bind = wireguard.NewClientBind(ctx, dialer.New(router, options.DialerOptions), peerAddr)
+	outbound.bind = wireguard.NewClientBind(ctx, dialer.New(router, options.DialerOptions), peerAddr, reserved)
 	localPrefixes := common.Map(options.LocalAddress, option.ListenPrefix.Build)
 	if len(localPrefixes) == 0 {
 		return nil, E.New("missing local address")
diff --git a/transport/wireguard/client_bind.go b/transport/wireguard/client_bind.go
index 270e09fd..1ecbda23 100644
--- a/transport/wireguard/client_bind.go
+++ b/transport/wireguard/client_bind.go
@@ -18,16 +18,18 @@ type ClientBind struct {
 	ctx        context.Context
 	dialer     N.Dialer
 	peerAddr   M.Socksaddr
+	reserved   [3]uint8
 	connAccess sync.Mutex
 	conn       *wireConn
 	done       chan struct{}
 }
 
-func NewClientBind(ctx context.Context, dialer N.Dialer, peerAddr M.Socksaddr) *ClientBind {
+func NewClientBind(ctx context.Context, dialer N.Dialer, peerAddr M.Socksaddr, reserved [3]uint8) *ClientBind {
 	return &ClientBind{
 		ctx:      ctx,
 		dialer:   dialer,
 		peerAddr: peerAddr,
+		reserved: reserved,
 	}
 }
 
@@ -89,6 +91,11 @@ func (c *ClientBind) receive(b []byte) (n int, ep conn.Endpoint, err error) {
 		}
 		return
 	}
+	if n > 3 {
+		b[1] = 0
+		b[2] = 0
+		b[3] = 0
+	}
 	ep = Endpoint(c.peerAddr)
 	return
 }
@@ -119,6 +126,11 @@ func (c *ClientBind) Send(b []byte, ep conn.Endpoint) error {
 	if err != nil {
 		return err
 	}
+	if len(b) > 3 {
+		b[1] = c.reserved[0]
+		b[2] = c.reserved[1]
+		b[3] = c.reserved[2]
+	}
 	_, err = udpConn.Write(b)
 	if err != nil {
 		udpConn.Close()