From c39e7650cd3fd71a7bfda8c4818528048fba220e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Mon, 6 Oct 2025 14:12:19 +0800
Subject: [PATCH] fix: wrong deadline report

---
 go.mod  |  2 +-
 go.sum  |  4 ++--
 main.go | 15 ++++++++++-----
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/go.mod b/go.mod
index d0f00d0..e020cba 100644
--- a/go.mod
+++ b/go.mod
@@ -5,7 +5,7 @@ go 1.20
 require (
 	github.com/FloatTech/ttl v0.0.0-20250224045156-012b1463287d
 	github.com/fumiama/orbyte v0.0.0-20251002065953-3bb358367eb5
-	github.com/fumiama/terasu v0.0.0-20251006044146-d81a14f4444c
+	github.com/fumiama/terasu v0.0.0-20251006061017-78a388f8731c
 	github.com/sirupsen/logrus v1.9.3
 )
 
diff --git a/go.sum b/go.sum
index b183dba..9436d77 100644
--- a/go.sum
+++ b/go.sum
@@ -7,8 +7,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/fumiama/orbyte v0.0.0-20251002065953-3bb358367eb5 h1:j9o0XVvdAeLwrBYMnh0SerrMc9CgNU6AGszbsvFzoc0=
 github.com/fumiama/orbyte v0.0.0-20251002065953-3bb358367eb5/go.mod h1:FOjdw7KdCbK2eH3gRPhwFNCoXKpu9sN5vPH4El/8e0c=
-github.com/fumiama/terasu v0.0.0-20251006044146-d81a14f4444c h1:goudQ1xbK4WNPmMCdGwyzj+pJKrSX2s5HIMYdisbRKg=
-github.com/fumiama/terasu v0.0.0-20251006044146-d81a14f4444c/go.mod h1:5wnbYtJ8Rv0GG7EIiYSqniKnGDXDvkKqCcZQehh3UCQ=
+github.com/fumiama/terasu v0.0.0-20251006061017-78a388f8731c h1:fn5vxrQ4maTkwZqwL67dYLT4aDK/q+eufGZWLDmiFNc=
+github.com/fumiama/terasu v0.0.0-20251006061017-78a388f8731c/go.mod h1:5wnbYtJ8Rv0GG7EIiYSqniKnGDXDvkKqCcZQehh3UCQ=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
diff --git a/main.go b/main.go
index 67000b5..2ed6777 100644
--- a/main.go
+++ b/main.go
@@ -32,13 +32,14 @@ var (
 		})
 	fallback *net.UDPAddr
 	forcefb  bool
+	timeout  uint
 )
 
 func main() {
 	iphost := flag.String("l", "127.0.0.1:5345", "listen DNS UDP port")
 	fbsrv := flag.String("fb", "127.0.0.1:53", "fallback to DNS UDP port")
 	debug := flag.Bool("d", false, "show debug log")
-	timeout := flag.Uint("to", 4, "dial timeout in sec")
+	flag.UintVar(&timeout, "to", 4, "dial timeout in sec")
 	flag.BoolVar(&forcefb, "ffb", false, "force using fallback")
 	flag.BoolVar(&ip.IsIPv6Available, "6", false, "use ipv6 servers")
 	frag := flag.Uint("frag", 3, "TLS first fragemt size (0 to disable)")
@@ -61,7 +62,7 @@ func main() {
 		logrus.Infoln("Set fallback server to", fallback)
 	}
 
-	dns.SetTimeout(time.Second * time.Duration(*timeout))
+	dns.SetTimeout(time.Second * time.Duration(timeout))
 
 	logrus.Infoln("Use ipv6 servers:", ip.IsIPv6Available)
 
@@ -104,6 +105,8 @@ func response(cnt uint8, conn *net.UDPConn, addr *net.UDPAddr, payload pbuf.Byte
 		err     error
 		tlsconn net.Conn
 		loopcnt = 0
+		ctx     context.Context
+		cancel  context.CancelFunc
 	)
 
 	if forcefb {
@@ -114,7 +117,9 @@ func response(cnt uint8, conn *net.UDPConn, addr *net.UDPAddr, payload pbuf.Byte
 	logrus.Debugln(addr, "Run on lock", cnt)
 
 REDAIL:
-	tlsconn, err = dialtls(cnt)
+	ctx, cancel = context.WithTimeout(context.Background(), time.Second*time.Duration(timeout))
+	defer cancel()
+	tlsconn, err = dialtls(cnt, ctx)
 	if err != nil {
 		logrus.Warnln(addr, "Dial DNS server err:", err)
 		return
@@ -238,14 +243,14 @@ func releasefree(i uint8) {
 	}
 }
 
-func dialtls(cnt uint8) (net.Conn, error) {
+func dialtls(cnt uint8, ctx context.Context) (net.Conn, error) {
 	conn := tlsconnCache.Get(cnt)
 	if conn != nil {
 		logrus.Debugln("Lock", cnt, "get cached tls conn to", conn.RemoteAddr())
 		return conn, nil
 	}
 	// dummy nw and addr
-	conn, err := dns.DefaultResolver.Dial(context.Background(), "", "")
+	conn, err := dns.DefaultResolver.Dial(ctx, "", "")
 	if err != nil {
 		return nil, err
 	}