Skip to content

Commit

Permalink
fix: set the timeout based on the caller to avoid strange issues
Browse files Browse the repository at this point in the history
although a ctx with deadline is set in most cases,
but passing in context.TODO() can lead to some very strange problems.
  • Loading branch information
zjj committed Nov 14, 2024
1 parent 3d3d4ee commit 4382acc
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 16 deletions.
5 changes: 3 additions & 2 deletions prober/dns.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 7 additions & 5 deletions prober/icmp.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
19 changes: 10 additions & 9 deletions prober/tcp.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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 {
Expand All @@ -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)
Expand Down

0 comments on commit 4382acc

Please sign in to comment.