0904

internal/shared/multidial.go

@@ -0,0 +1,61 @@
+package shared
+
+import (
+	"context"
+	"errors"
+	"net"
+)
+
+// DialFirstReachable 对一组已通过安全校验的候选 IP 并发拨号，返回最先成功建立的连接。
+//
+// 适用场景：SSRF 防护型出站客户端。调用方先解析域名并过滤内网地址，再以 IP 字面量拨号
+// （无二次解析，免疫 DNS rebinding），本函数负责多 IP 容错——避免固定拨单一 IP 时，
+// 恰好抽中不可达或被劣化的边缘节点（如 Cloudflare anycast 的个别边缘）导致请求稳定超时；
+// curl 等客户端默认会尝试多个 IP，因此这里保持同等的多 IP 能力。
+//
+// 成功路径会取消其余拨号并关闭任何已建立但落选的连接，避免泄漏；全部失败时返回第一个错误。
+// maxCandidates<=0 或大于候选数时按全部候选处理；ips 为空返回错误。
+func DialFirstReachable(ctx context.Context, dialer *net.Dialer, network string, ips []net.IP, port string, maxCandidates int) (net.Conn, error) {
+	if len(ips) == 0 {
+		return nil, errors.New("no candidates to dial")
+	}
+	if maxCandidates <= 0 || maxCandidates > len(ips) {
+		maxCandidates = len(ips)
+	}
+
+	dialCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	type dialResult struct {
+		conn net.Conn
+		err  error
+	}
+	results := make(chan dialResult, maxCandidates)
+	for _, ip := range ips[:maxCandidates] {
+		ip := ip
+		go func() {
+			conn, err := dialer.DialContext(dialCtx, network, net.JoinHostPort(ip.String(), port))
+			results <- dialResult{conn: conn, err: err}
+		}()
+	}
+
+	var firstErr error
+	remaining := maxCandidates
+	for remaining > 0 {
+		r := <-results
+		remaining--
+		if r.err == nil {
+			cancel() // 通知其余拨号尽快退出
+			for ; remaining > 0; remaining-- {
+				if extra := <-results; extra.err == nil && extra.conn != nil {
+					_ = extra.conn.Close()
+				}
+			}
+			return r.conn, nil
+		}
+		if firstErr == nil {
+			firstErr = r.err
+		}
+	}
+	return nil, firstErr
+}