From 1c50c2b6af88127aaf601262ce06aaacb162d942 Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Thu, 6 Feb 2025 15:16:32 +0700
Subject: [PATCH] Set deadline for custom UDP/TCP conn

Otherwise, OS resolver may hang forever if the server does not reply.

While at it, also removing unused method stopClientInfoDiscover.

Updates #344
---
 cmd/cli/prog.go | 7 -------
 resolver.go     | 2 ++
 2 files changed, 2 insertions(+), 7 deletions(-)

diff --git a/cmd/cli/prog.go b/cmd/cli/prog.go
index c7eba13..07a7592 100644
--- a/cmd/cli/prog.go
+++ b/cmd/cli/prog.go
@@ -581,13 +581,6 @@ func (p *prog) runClientInfoDiscover(ctx context.Context) {
 	}()
 }
 
-// stopClientInfoDiscover stops the current client info discover goroutine.
-// It blocks until the goroutine terminated.
-func (p *prog) stopClientInfoDiscover() {
-	p.ciTable.Stop()
-	mainLog.Load().Debug().Msg("stopped client info discover")
-}
-
 // metricsEnabled reports whether prometheus exporter is enabled/disabled.
 func (p *prog) metricsEnabled() bool {
 	return p.cfg.Service.MetricsQueryStats || p.cfg.Service.MetricsListener != ""
diff --git a/resolver.go b/resolver.go
index 01348dc..f4299e6 100644
--- a/resolver.go
+++ b/resolver.go
@@ -289,6 +289,7 @@ func customDNSExchange(ctx context.Context, msg *dns.Msg, server string, desired
 		return nil, err
 	}
 	defer udpConn.Close()
+	udpConn.SetDeadline(time.Now().Add(3 * time.Second))
 	udpDnsConn := &dns.Conn{Conn: udpConn}
 	if err = udpDnsConn.WriteMsg(msg); err != nil {
 		return nil, err
@@ -310,6 +311,7 @@ func customDNSExchange(ctx context.Context, msg *dns.Msg, server string, desired
 		return reply, nil // fallback to UDP reply if TCP dial fails.
 	}
 	defer tcpConn.Close()
+	tcpConn.SetDeadline(time.Now().Add(3 * time.Second))
 	tcpDnsConn := &dns.Conn{Conn: tcpConn}
 	if err = tcpDnsConn.WriteMsg(msg); err != nil {
 		return reply, nil // fallback if TCP write fails.