From 4382acccfc743cbec670e392bd8dd3f89c23573d Mon Sep 17 00:00:00 2001
From: zjj <zhong2plus@gmail.com>
Date: Thu, 14 Nov 2024 11:00:42 +0800
Subject: [PATCH] fix: set the timeout based on the caller to avoid strange
 issues

although a ctx with deadline is set in most cases,
but passing in context.TODO() can lead to some very strange problems.
---
 prober/dns.go  |  5 +++--
 prober/icmp.go | 12 +++++++-----
 prober/tcp.go  | 19 ++++++++++---------
 3 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/prober/dns.go b/prober/dns.go
index 184fa949f..281401529 100644
--- a/prober/dns.go
+++ b/prober/dns.go
@@ -259,8 +259,9 @@ func ProbeDNS(ctx context.Context, target string, module config.Module, registry
 	msg.Question[0] = dns.Question{dns.Fqdn(module.DNS.QueryName), qt, qc}
 
 	logger.Info("Making DNS query", "target", targetIP, "dial_protocol", dialProtocol, "query", module.DNS.QueryName, "type", qt, "class", qc)
-	timeoutDeadline, _ := ctx.Deadline()
-	client.Timeout = time.Until(timeoutDeadline)
+	if timeoutDeadline, ok := ctx.Deadline(); ok {
+		client.Timeout = time.Until(timeoutDeadline)
+	}
 	requestStart := time.Now()
 	response, rtt, err := client.Exchange(msg, targetIP)
 	// The rtt value returned from client.Exchange includes only the time to
diff --git a/prober/icmp.go b/prober/icmp.go
index 883fbcf73..9f45f7249 100644
--- a/prober/icmp.go
+++ b/prober/icmp.go
@@ -293,11 +293,13 @@ func ProbeICMP(ctx context.Context, target string, module config.Module, registr
 	}
 
 	rb := make([]byte, 65536)
-	deadline, _ := ctx.Deadline()
-	if icmpConn != nil {
-		err = icmpConn.SetReadDeadline(deadline)
-	} else {
-		err = v4RawConn.SetReadDeadline(deadline)
+
+	if deadline, ok := ctx.Deadline(); ok {
+		if icmpConn != nil {
+			err = icmpConn.SetReadDeadline(deadline)
+		} else {
+			err = v4RawConn.SetReadDeadline(deadline)
+		}
 	}
 	if err != nil {
 		logger.Error("Error setting socket deadline", "err", err)
diff --git a/prober/tcp.go b/prober/tcp.go
index 98262813a..99eaf577f 100644
--- a/prober/tcp.go
+++ b/prober/tcp.go
@@ -80,8 +80,10 @@ func dialTCP(ctx context.Context, target string, module config.Module, registry
 		// via tlsConfig to enable hostname verification.
 		tlsConfig.ServerName = targetAddress
 	}
-	timeoutDeadline, _ := ctx.Deadline()
-	dialer.Deadline = timeoutDeadline
+
+	if timeoutDeadline, ok := ctx.Deadline(); ok {
+		dialer.Deadline = timeoutDeadline
+	}
 
 	logger.Info("Dialing TCP with TLS")
 	return tls.DialWithDialer(dialer, dialProtocol, dialTarget, tlsConfig)
@@ -124,7 +126,6 @@ func ProbeTCP(ctx context.Context, target string, module config.Module, registry
 		Help: "Indicates if probe failed due to regex",
 	})
 	registry.MustRegister(probeFailedDueToRegex)
-	deadline, _ := ctx.Deadline()
 
 	conn, err := dialTCP(ctx, target, module, registry, logger)
 	if err != nil {
@@ -134,13 +135,13 @@ func ProbeTCP(ctx context.Context, target string, module config.Module, registry
 	defer conn.Close()
 	logger.Info("Successfully dialed")
 
-	// Set a deadline to prevent the following code from blocking forever.
-	// If a deadline cannot be set, better fail the probe by returning an error
-	// now rather than blocking forever.
-	if err := conn.SetDeadline(deadline); err != nil {
-		logger.Error("Error setting deadline", "err", err)
-		return false
+	if deadline, ok := ctx.Deadline(); ok {
+		if err := conn.SetDeadline(deadline); err != nil {
+			logger.Error("Error setting deadline", "err", err)
+			return false
+		}
 	}
+
 	if module.TCP.TLS {
 		state := conn.(*tls.Conn).ConnectionState()
 		registry.MustRegister(probeSSLEarliestCertExpiry, probeTLSVersion, probeSSLLastChainExpiryTimestampSeconds, probeSSLLastInformation)