chore: split dns's dialContext and listenPacket

dns/doh.go

@@ -528,14 +528,8 @@ func (doh *dnsOverHTTPS) dialQuic(ctx context.Context, addr string, tlsCfg *tls.
 		IP:   net.ParseIP(ip),
 		Port: portInt,
 	}
-	var conn net.PacketConn
+	conn, err := listenPacket(ctx, doh.proxyAdapter, "udp", addr, doh.r)
-	if wrapConn, err := dialContextExtra(ctx, doh.proxyAdapter, "udp", addr, doh.r); err == nil {
+	if err != nil {
-		if pc, ok := wrapConn.(*wrapPacketConn); ok {
-			conn = pc
-		} else {
-			return nil, fmt.Errorf("conn isn't wrapPacketConn")
-		}
-	} else {
 		return nil, err
 	}
 	return quic.DialEarlyContext(ctx, conn, &udpAddr, doh.url.Host, tlsCfg, cfg)
@@ -556,20 +550,10 @@ func (doh *dnsOverHTTPS) probeH3(
 	if err != nil {
 		return "", fmt.Errorf("failed to dial: %w", err)
 	}
+	addr = rawConn.RemoteAddr().String()
 	// It's never actually used.
 	_ = rawConn.Close()
 
-	udpConn, ok := rawConn.(*net.UDPConn)
-	if !ok {
-		if packetConn, ok := rawConn.(*wrapPacketConn); !ok {
-			return "", fmt.Errorf("not a UDP connection to %s", doh.Address())
-		} else {
-			addr = packetConn.RemoteAddr().String()
-		}
-	} else {
-		addr = udpConn.RemoteAddr().String()
-	}
-
 	// Avoid spending time on probing if this upstream only supports HTTP/3.
 	if doh.supportsH3() && !doh.supportsHTTP() {
 		return addr, nil

dns/doq.go

@@ -313,19 +313,9 @@ func (doq *dnsOverQUIC) openConnection(ctx context.Context) (conn quic.Connectio
 	if err != nil {
 		return nil, fmt.Errorf("failed to open a QUIC connection: %w", err)
 	}
+	addr := rawConn.RemoteAddr().String()
 	// It's never actually used
 	_ = rawConn.Close()
-	var addr string
-	udpConn, ok := rawConn.(*net.UDPConn)
-	if !ok {
-		if packetConn, ok := rawConn.(*wrapPacketConn); !ok {
-			return nil, fmt.Errorf("failed to open connection to %s", doq.addr)
-		} else {
-			addr = packetConn.RemoteAddr().String()
-		}
-	} else {
-		addr = udpConn.RemoteAddr().String()
-	}
 
 	ip, port, err := net.SplitHostPort(addr)
 	if err != nil {
@@ -334,14 +324,8 @@ func (doq *dnsOverQUIC) openConnection(ctx context.Context) (conn quic.Connectio
 
 	p, err := strconv.Atoi(port)
 	udpAddr := net.UDPAddr{IP: net.ParseIP(ip), Port: p}
-	var udp net.PacketConn
+	udp, err := listenPacket(ctx, doq.proxyAdapter, "udp", addr, doq.r)
-	if wrapConn, err := dialContextExtra(ctx, doq.proxyAdapter, "udp", addr, doq.r); err == nil {
+	if err != nil {
-		if pc, ok := wrapConn.(*wrapPacketConn); ok {
-			udp = pc
-		} else {
-			return nil, fmt.Errorf("quic create packet failed")
-		}
-	} else {
 		return nil, err
 	}
 

dns/util.go

@@ -168,19 +168,13 @@ func getDialHandler(r *Resolver, proxyAdapter string, opts ...dialer.Option) dia
 			opts = append(opts, dialer.WithResolver(r))
 			return dialer.DialContext(ctx, network, addr, opts...)
 		} else {
-			return dialContextExtra(ctx, proxyAdapter, network, addr, r, opts...)
-		}
-	}
-}
-
-func dialContextExtra(ctx context.Context, adapterName string, network string, addr string, r *Resolver, opts ...dialer.Option) (net.Conn, error) {
 			host, port, err := net.SplitHostPort(addr)
 			if err != nil {
 				return nil, err
 			}
-	adapter, ok := tunnel.Proxies()[adapterName]
+			adapter, ok := tunnel.Proxies()[proxyAdapter]
 			if !ok {
-		opts = append(opts, dialer.WithInterface(adapterName))
+				opts = append(opts, dialer.WithInterface(proxyAdapter))
 			}
 			if strings.Contains(network, "tcp") {
 				// tcp can resolve host by remote
@@ -207,19 +201,11 @@ func dialContextExtra(ctx context.Context, adapterName string, network string, a
 					DstPort: port,
 				}
 				if !ok {
-			packetConn, err := dialer.ListenPacket(ctx, dialer.ParseNetwork(network, dstIP), "", opts...)
+					return dialer.DialContext(ctx, network, addr, opts...)
-			if err != nil {
-				return nil, err
-			}
-
-			return &wrapPacketConn{
-				PacketConn: packetConn,
-				rAddr:      metadata.UDPAddr(),
-			}, nil
 				}
 
 				if !adapter.SupportUDP() {
-			return nil, fmt.Errorf("proxy adapter [%s] UDP is not supported", adapterName)
+					return nil, fmt.Errorf("proxy adapter [%s] UDP is not supported", proxyAdapter)
 				}
 
 				packetConn, err := adapter.ListenPacketContext(ctx, metadata, opts...)
@@ -233,6 +219,40 @@ func dialContextExtra(ctx context.Context, adapterName string, network string, a
 				}, nil
 			}
 		}
+	}
+}
+
+func listenPacket(ctx context.Context, proxyAdapter string, network string, addr string, r *Resolver, opts ...dialer.Option) (net.PacketConn, error) {
+	host, port, err := net.SplitHostPort(addr)
+	if err != nil {
+		return nil, err
+	}
+	adapter, ok := tunnel.Proxies()[proxyAdapter]
+	if !ok && len(proxyAdapter) != 0 {
+		opts = append(opts, dialer.WithInterface(proxyAdapter))
+	}
+
+	// udp must resolve host first
+	dstIP, err := resolver.ResolveIPWithResolver(ctx, host, r)
+	if err != nil {
+		return nil, err
+	}
+	metadata := &C.Metadata{
+		NetWork: C.UDP,
+		Host:    "",
+		DstIP:   dstIP,
+		DstPort: port,
+	}
+	if !ok {
+		return dialer.ListenPacket(ctx, dialer.ParseNetwork(network, dstIP), "", opts...)
+	}
+
+	if !adapter.SupportUDP() {
+		return nil, fmt.Errorf("proxy adapter [%s] UDP is not supported", proxyAdapter)
+	}
+
+	return adapter.ListenPacketContext(ctx, metadata, opts...)
+}
 
 func batchExchange(ctx context.Context, clients []dnsClient, m *D.Msg) (msg *D.Msg, err error) {
 	fast, ctx := picker.WithTimeout[*D.Msg](ctx, resolver.DefaultDNSTimeout)