From 8b848b62bba80bfc12b57ca2afd63dec5539a552 Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Fri, 18 Nov 2022 18:02:46 +0800
Subject: [PATCH] fix: reset timeout in exchangeWithoutCache's singleflight

---
 dns/resolver.go | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/dns/resolver.go b/dns/resolver.go
index 4e31074c..cb186dbe 100644
--- a/dns/resolver.go
+++ b/dns/resolver.go
@@ -187,7 +187,10 @@ func (r *Resolver) ExchangeContext(ctx context.Context, m *D.Msg) (msg *D.Msg, e
 func (r *Resolver) exchangeWithoutCache(ctx context.Context, m *D.Msg) (msg *D.Msg, err error) {
 	q := m.Question[0]
 
-	ret, err, shared := r.group.Do(q.String(), func() (result any, err error) {
+	ch := r.group.DoChan(q.String(), func() (result any, err error) {
+		ctx, cancel := context.WithTimeout(context.Background(), resolver.DefaultDNSTimeout) // reset timeout in singleflight
+		defer cancel()
+
 		defer func() {
 			if err != nil {
 				return
@@ -209,6 +212,21 @@ func (r *Resolver) exchangeWithoutCache(ctx context.Context, m *D.Msg) (msg *D.M
 		return r.batchExchange(ctx, r.main, m)
 	})
 
+	var result singleflight.Result
+
+	select {
+	case result = <-ch:
+		break
+	case <-ctx.Done():
+		select {
+		case result = <-ch:
+			break
+		}
+		return nil, ctx.Err()
+	}
+
+	ret, err, shared := result.Val, result.Err, result.Shared
+
 	if err == nil {
 		msg = ret.(*D.Msg)
 		if shared {