WireGuard proxy: Give up on the endpoint lookup at the connection's deadline (#7323)

https://github.com/XTLS/Xray-core/pull/7323#issuecomment-6623261000

proxy/wireguard/bind.go

@@ -17,7 +17,7 @@ import (
 
 type bind struct {
 	ctx         context.Context
-	resolveFunc func(host string) (net.IP, error)
+	resolveFunc func(ctx context.Context, host string) (net.IP, error)
 	listenFunc  func(ctx context.Context) (net.PacketConn, error)
 	downFunc    func() error
 	reserved    []byte
@@ -143,7 +143,11 @@ func (b *bind) ParseEndpoint(s string) (conn.Endpoint, error) {
 	if port < 0 || port > 65535 {
 		return nil, errors.New("invalid port " + sport)
 	}
-	ip, err := b.resolveFunc(host)
+	// up sets the connection's context right before it hands over the endpoint
+	b.mu.Lock()
+	ctx := b.ctx
+	b.mu.Unlock()
+	ip, err := b.resolveFunc(ctx, host)
 	if err != nil {
 		return nil, err
 	}

proxy/wireguard/client.go

@@ -323,12 +323,12 @@ func (h *Handler) up(ctx context.Context) error {
 	return h.dev.Up()
 }
 
-func MakeResolveLocal(d dns.Client, streamSettings *internet.MemoryStreamConfig) func(host string) (net.IP, error) {
-	return func(host string) (net.IP, error) {
+func MakeResolveLocal(d dns.Client, streamSettings *internet.MemoryStreamConfig) func(ctx context.Context, host string) (net.IP, error) {
+	return func(ctx context.Context, host string) (net.IP, error) {
 		if ip := net.ParseIP(host); ip != nil {
 			return ip, nil
 		}
-		ips, _, err := d.LookupIP(host, dns.IPOption{IPv4Enable: true, IPv6Enable: true})
+		ips, err := lookupIPContext(ctx, d, host)
 		if err != nil {
 			return nil, err
 		}
@@ -368,6 +368,27 @@ func MakeResolveLocal(d dns.Client, streamSettings *internet.MemoryStreamConfig)
 	}
 }
 
+func lookupIPContext(ctx context.Context, d dns.Client, host string) ([]net.IP, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	type result struct {
+		ips []net.IP
+		err error
+	}
+	done := make(chan result, 1)
+	go func() {
+		ips, _, err := d.LookupIP(host, dns.IPOption{IPv4Enable: true, IPv6Enable: true})
+		done <- result{ips, err}
+	}()
+	select {
+	case r := <-done:
+		return r.ips, r.err
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+}
+
 type UDPConnClient struct {
 	net.PacketConn
 	Dest *net.UDPAddr