From 9b6141b83f52f2ec6b6ac6e44911273b42ba2753 Mon Sep 17 00:00:00 2001
From: yuhan6665 <1588741+yuhan6665@users.noreply.github.com>
Date: Wed, 22 May 2024 08:44:02 -0400
Subject: [PATCH] Wireguard dial with context

---
 proxy/wireguard/bind.go   | 3 ++-
 proxy/wireguard/client.go | 5 +++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/proxy/wireguard/bind.go b/proxy/wireguard/bind.go
index 1fbcbc98..aaca93ff 100644
--- a/proxy/wireguard/bind.go
+++ b/proxy/wireguard/bind.go
@@ -123,12 +123,13 @@ func (bind *netBind) Close() error {
 type netBindClient struct {
 	netBind
 
+	ctx      context.Context
 	dialer   internet.Dialer
 	reserved []byte
 }
 
 func (bind *netBindClient) connectTo(endpoint *netEndpoint) error {
-	c, err := bind.dialer.Dial(context.Background(), endpoint.dst)
+	c, err := bind.dialer.Dial(bind.ctx, endpoint.dst)
 	if err != nil {
 		return err
 	}
diff --git a/proxy/wireguard/client.go b/proxy/wireguard/client.go
index 00a6fa51..04a72650 100644
--- a/proxy/wireguard/client.go
+++ b/proxy/wireguard/client.go
@@ -76,7 +76,7 @@ func New(ctx context.Context, conf *DeviceConfig) (*Handler, error) {
 	}, nil
 }
 
-func (h *Handler) processWireGuard(dialer internet.Dialer) (err error) {
+func (h *Handler) processWireGuard(ctx context.Context, dialer internet.Dialer) (err error) {
 	h.wgLock.Lock()
 	defer h.wgLock.Unlock()
 
@@ -108,6 +108,7 @@ func (h *Handler) processWireGuard(dialer internet.Dialer) (err error) {
 			},
 			workers: int(h.conf.NumWorkers),
 		},
+		ctx:      ctx,
 		dialer:   dialer,
 		reserved: h.conf.Reserved,
 	}
@@ -135,7 +136,7 @@ func (h *Handler) Process(ctx context.Context, link *transport.Link, dialer inte
 	ob.Name = "wireguard"
 	ob.CanSpliceCopy = 3
 
-	if err := h.processWireGuard(dialer); err != nil {
+	if err := h.processWireGuard(ctx, dialer); err != nil {
 		return err
 	}