fix: propagate context through MCP requests and certinfo TLS paths

Thread cancellation from MCP tool timeouts into HandleRequests,
http.NewRequestWithContext, and certinfo DialContext handshakes.
Remove parallel os.Args mutation in TestIsMCPCommand subtests.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: xenOs76

internal/certinfo/certinfo.go

@@ -1,6 +1,7 @@
 package certinfo
 
 import (
+	"context"
 	"crypto"
 	"crypto/x509"
 	"fmt"
@@ -180,7 +181,7 @@ func (c *Config) SetPrivateKeyFromFile(
 }
 
 // SetTLSEndpoint parses a host:port string and fetches the remote certificates from that endpoint.
-func (c *Config) SetTLSEndpoint(hostport string) error {
+func (c *Config) SetTLSEndpoint(ctx context.Context, hostport string) error {
 	if hostport != emptyString {
 		eHost, ePort, err := net.SplitHostPort(hostport)
 		if err != nil {
@@ -191,7 +192,7 @@ func (c *Config) SetTLSEndpoint(hostport string) error {
 		c.TLSEndpointHost = eHost
 		c.TLSEndpointPort = ePort
 
-		err = c.GetRemoteCerts()
+		err = c.GetRemoteCerts(ctx)
 		if err != nil {
 			return fmt.Errorf("unable to get endpoint certificates: %w", err)
 		}

internal/certinfo/certinfo_handlers.go

@@ -6,6 +6,7 @@ package certinfo
 
 import (
 	"cmp"
+	"context"
 	"crypto/sha256"
 	"crypto/tls"
 	"crypto/x509"
@@ -29,7 +30,7 @@ import (
 // to the provided writer in a human-readable format.
 //
 //nolint:revive
-func (c *Config) PrintData(w io.Writer) error {
+func (c *Config) PrintData(ctx context.Context, w io.Writer) error {
 	ks := style.ItemKey.PaddingBottom(0).PaddingTop(1).PaddingLeft(1)
 	sl := style.CertKeyP4.Bold(true)
 	sv := style.CertValue.Bold(false)
@@ -49,7 +50,7 @@ func (c *Config) PrintData(w io.Writer) error {
 	}
 
 	if c.TLSInfoRequested {
-		_ = c.ProbeTLSInfo()
+		_ = c.ProbeTLSInfo(ctx)
 		c.printTLSInfo(w, ks, sl, sv)
 	}
 
@@ -174,30 +175,47 @@ func (c *Config) printCACerts(w io.Writer, ks, sl, sv lipgloss.Style) error {
 	return nil
 }
 
+// dialTLS connects to serverAddr and completes a TLS handshake using ctx for cancellation.
+func dialTLS(ctx context.Context, serverAddr string, tlsConfig *tls.Config) (*tls.Conn, error) {
+	dialer := &net.Dialer{Timeout: TLSTimeout}
+
+	rawConn, err := dialer.DialContext(ctx, "tcp", serverAddr)
+	if err != nil {
+		return nil, err
+	}
+
+	conn := tls.Client(rawConn, tlsConfig)
+
+	if err = conn.HandshakeContext(ctx); err != nil {
+		_ = rawConn.Close()
+
+		return nil, err
+	}
+
+	return conn, nil
+}
+
 // GetRemoteCerts establishes a TLS connection to the configured endpoint and retrieves
 // the peer certificate chain. It also performs certificate verification unless TLSInsecure is true.
-func (c *Config) GetRemoteCerts() error {
+func (c *Config) GetRemoteCerts(ctx context.Context) error {
 	tlsConfig := &tls.Config{
 		RootCAs:            c.CACertsPool,
 		InsecureSkipVerify: c.TLSInsecure,
 	}
 
-	if c.TLSServerName != emptyString {
+	verifyName := c.TLSServerName
+	switch {
+	case c.TLSServerName != emptyString:
 		tlsConfig.ServerName = c.TLSServerName
+	case c.TLSEndpointHost != emptyString:
+		tlsConfig.ServerName = c.TLSEndpointHost
+		verifyName = c.TLSEndpointHost
+	default:
 	}
 
 	serverAddr := net.JoinHostPort(c.TLSEndpointHost, c.TLSEndpointPort)
 
-	dialer := &net.Dialer{
-		Timeout: TLSTimeout,
-	}
-
-	conn, err := tls.DialWithDialer(
-		dialer,
-		"tcp",
-		serverAddr,
-		tlsConfig,
-	)
+	conn, err := dialTLS(ctx, serverAddr, tlsConfig)
 	if err != nil {
 		return fmt.Errorf("TLS handshake failed: %w", err)
 	}
@@ -214,7 +232,7 @@ func (c *Config) GetRemoteCerts() error {
 	}
 
 	opts := x509.VerifyOptions{
-		DNSName:       c.TLSServerName,
+		DNSName:       verifyName,
 		Roots:         c.CACertsPool,
 		Intermediates: x509.NewCertPool(),
 	}
@@ -425,7 +443,7 @@ func tlsVersionToString(version uint16) string {
 }
 
 // probeProtocol tests whether the TLS endpoint supports a specific TLS protocol version.
-func (c *Config) probeProtocol(version uint16) bool {
+func (c *Config) probeProtocol(ctx context.Context, version uint16) bool {
 	tlsConfig := &tls.Config{
 		MinVersion:         version,
 		MaxVersion:         version,
@@ -434,17 +452,15 @@ func (c *Config) probeProtocol(version uint16) bool {
 
 	if c.TLSServerName != emptyString {
 		tlsConfig.ServerName = c.TLSServerName
+	} else if c.TLSEndpointHost != emptyString {
+		tlsConfig.ServerName = c.TLSEndpointHost
 	}
 
 	serverAddr := net.JoinHostPort(c.TLSEndpointHost, c.TLSEndpointPort)
 
-	dialer := &net.Dialer{
-		Timeout: TLSTimeout,
-	}
-
-	conn, err := tls.DialWithDialer(dialer, "tcp", serverAddr, tlsConfig)
+	conn, err := dialTLS(ctx, serverAddr, tlsConfig)
 	if err == nil {
-		conn.Close()
+		_ = conn.Close()
 
 		return true
 	}
@@ -453,7 +469,7 @@ func (c *Config) probeProtocol(version uint16) bool {
 }
 
 // probeCipher tests whether a specific TLS 1.0-1.2 cipher suite is supported.
-func (c *Config) probeCipher(suite *tls.CipherSuite) (bool, string) {
+func (c *Config) probeCipher(ctx context.Context, suite *tls.CipherSuite) (bool, string) {
 	tlsConfig := &tls.Config{
 		MinVersion:         tls.VersionTLS10,
 		MaxVersion:         tls.VersionTLS12,
@@ -463,19 +479,17 @@ func (c *Config) probeCipher(suite *tls.CipherSuite) (bool, string) {
 
 	if c.TLSServerName != emptyString {
 		tlsConfig.ServerName = c.TLSServerName
+	} else if c.TLSEndpointHost != emptyString {
+		tlsConfig.ServerName = c.TLSEndpointHost
 	}
 
 	serverAddr := net.JoinHostPort(c.TLSEndpointHost, c.TLSEndpointPort)
 
-	dialer := &net.Dialer{
-		Timeout: TLSTimeout,
-	}
-
-	conn, err := tls.DialWithDialer(dialer, "tcp", serverAddr, tlsConfig)
+	conn, err := dialTLS(ctx, serverAddr, tlsConfig)
 	if err == nil {
 		state := conn.ConnectionState()
 
-		conn.Close()
+		_ = conn.Close()
 
 		return true, tlsVersionToString(state.Version)
 	}
@@ -484,7 +498,7 @@ func (c *Config) probeCipher(suite *tls.CipherSuite) (bool, string) {
 }
 
 // ProbeTLSInfo concurrently scans the endpoint for supported TLS versions and cipher suites.
-func (c *Config) ProbeTLSInfo() error {
+func (c *Config) ProbeTLSInfo(ctx context.Context) error {
 	if c.TLSEndpoint == emptyString {
 		return nil
 	}
@@ -495,23 +509,27 @@ func (c *Config) ProbeTLSInfo() error {
 	versions := []uint16{tls.VersionTLS10, tls.VersionTLS11, tls.VersionTLS12, tls.VersionTLS13}
 
 	for _, v := range versions {
-		supported := c.probeProtocol(v)
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+
+		supported := c.probeProtocol(ctx, v)
 
 		c.ProbedProtocols[tlsVersionToString(v)] = supported
 	}
 
 	// 2. Probe ciphers concurrently
 	suites := append(tls.CipherSuites(), tls.InsecureCipherSuites()...)
 
-	c.ProbedCiphers = c.probeCiphersConcurrently(suites)
+	c.ProbedCiphers = c.probeCiphersConcurrently(ctx, suites)
 
 	return nil
 }
 
 // probeCiphersConcurrently manages the worker pool to concurrently scan cipher suites.
 //
 //nolint:gocognit,revive,wsl
-func (c *Config) probeCiphersConcurrently(suites []*tls.CipherSuite) []ProbedCipher {
+func (c *Config) probeCiphersConcurrently(ctx context.Context, suites []*tls.CipherSuite) []ProbedCipher {
 	type job struct {
 		suite *tls.CipherSuite
 	}
@@ -539,6 +557,10 @@ func (c *Config) probeCiphersConcurrently(suites []*tls.CipherSuite) []ProbedCip
 			defer wg.Done()
 
 			for j := range jobs {
+				if err := ctx.Err(); err != nil {
+					return
+				}
+
 				suite := j.suite
 				isTLS13 := false
 
@@ -559,7 +581,7 @@ func (c *Config) probeCiphersConcurrently(suites []*tls.CipherSuite) []ProbedCip
 					supported = c.ProbedProtocols["TLS 1.3"]
 					protoName = "TLS 1.3"
 				} else {
-					ok, name := c.probeCipher(suite)
+					ok, name := c.probeCipher(ctx, suite)
 
 					supported = ok
 					protoName = name

internal/certinfo/certinfo_handlers_test.go

@@ -2,6 +2,7 @@ package certinfo
 
 import (
 	"bytes"
+	"context"
 	"crypto/x509"
 	"crypto/x509/pkix"
 	"math/big"
@@ -145,10 +146,10 @@ func TestCertinfo_GetRemoteCerts(t *testing.T) {
 
 			cc.SetTLSServerName(tt.srvCfg.serverName)
 			cc.SetCaPoolFromFile(tt.caCertFile, inputReader)
-			cc.SetTLSEndpoint(tt.srvCfg.serverAddr)
+			cc.SetTLSEndpoint(context.Background(), tt.srvCfg.serverAddr)
 			cc.SetTLSInsecure(tt.insecure)
 
-			err = cc.GetRemoteCerts()
+			err = cc.GetRemoteCerts(context.Background())
 			if !tt.expectError {
 				require.NoError(t, err, "check error not expected")
 				require.Equal(t, tt.srvCfg.serverName, cc.TLSServerName, "check TLSServerName")
@@ -491,7 +492,7 @@ func TestCertinfo_PrintData(t *testing.T) {
 		})
 		cc.CertsBundleFilePath = "dummy"
 
-		errPrint := cc.PrintData(&buffer)
+		errPrint := cc.PrintData(context.Background(), &buffer)
 		require.Error(t, errPrint)
 		require.ErrorContains(t, errPrint, "unable to check if private key matches local certificate")
 	})
@@ -508,7 +509,7 @@ func TestCertinfo_PrintData(t *testing.T) {
 		cc.TLSEndpointHost = "localhost"
 		cc.TLSEndpointPort = "443"
 
-		errPrint := cc.PrintData(&buffer)
+		errPrint := cc.PrintData(context.Background(), &buffer)
 		require.Error(t, errPrint)
 		require.ErrorContains(t, errPrint, "unable to check if private key matches remote TLS Endpoint certificate")
 	})
@@ -520,7 +521,7 @@ func TestCertinfo_PrintData(t *testing.T) {
 
 		cc.CACertsFilePath = "non_existent_file.pem"
 
-		errPrint := cc.PrintData(&buffer)
+		errPrint := cc.PrintData(context.Background(), &buffer)
 		require.Error(t, errPrint)
 		require.ErrorContains(t, errPrint, "unable for read Root certificates")
 	})
@@ -559,15 +560,15 @@ func runPrintDataSubtest(t *testing.T, tt printDataTestCase) {
 		cc.SetTLSServerName(tt.tlsServerName)
 		cc.SetTLSInsecure(tt.tlsInsecure)
 
-		err = cc.SetTLSEndpoint(tt.tlsEndpoint)
+		err = cc.SetTLSEndpoint(context.Background(), tt.tlsEndpoint)
 		if tt.expectCertsFetchErr {
 			require.EqualError(t, err, tt.expectCertsFetcMsg)
 		} else {
 			require.NoError(t, err, "SetTLSEndpoint require NoError")
 		}
 	}
 
-	errPrint := cc.PrintData(&buffer)
+	errPrint := cc.PrintData(context.Background(), &buffer)
 	require.NoError(t, errPrint)
 
 	got := buffer.String()

internal/certinfo/certinfo_test.go

@@ -2,6 +2,7 @@ package certinfo
 
 import (
 	"bytes"
+	"context"
 	"crypto/tls"
 	"fmt"
 	"net/http"
@@ -400,7 +401,7 @@ func TestCertinfo_SetTLSEndpoint(t *testing.T) {
 			cc, errNew := New()
 			require.NoError(t, errNew)
 
-			err := cc.SetTLSEndpoint(tt.endpoint)
+			err := cc.SetTLSEndpoint(context.Background(), tt.endpoint)
 
 			if !tt.processErr {
 				// skip requiring NoError since SetTLSEndpoint will always return network errors
@@ -440,10 +441,10 @@ func TestCertinfo_ProbeTLSInfo(t *testing.T) {
 	cc.SetTLSInsecure(true)
 	cc.SetTLSServerName("example.com")
 
-	err = cc.SetTLSEndpoint(u.Host)
+	err = cc.SetTLSEndpoint(context.Background(), u.Host)
 	require.NoError(t, err)
 
-	err = cc.ProbeTLSInfo()
+	err = cc.ProbeTLSInfo(context.Background())
 	require.NoError(t, err)
 
 	// Since it's a local TLS server run by Go's httptest, it supports TLS 1.3 or TLS 1.2
@@ -470,7 +471,7 @@ func TestCertinfo_ProbeTLSInfo_NotRequested(t *testing.T) {
 	cc.SetTLSInfoRequested(false)
 	require.False(t, cc.TLSInfoRequested)
 
-	err = cc.ProbeTLSInfo()
+	err = cc.ProbeTLSInfo(context.Background())
 	require.NoError(t, err)
 	require.Empty(t, cc.NegotiatedProtocol)
 }
@@ -483,7 +484,7 @@ func TestCertinfo_ProbeTLSInfo_NoEndpoint(t *testing.T) {
 
 	cc.SetTLSInfoRequested(true)
 
-	err = cc.ProbeTLSInfo()
+	err = cc.ProbeTLSInfo(context.Background())
 	require.NoError(t, err)
 	require.Empty(t, cc.ProbedProtocols)
 }
@@ -502,7 +503,7 @@ func TestCertinfo_ProbeTLSInfo_Unreachable(t *testing.T) {
 	cc.TLSEndpointHost = "127.0.0.1"
 	cc.TLSEndpointPort = "54321"
 
-	err = cc.ProbeTLSInfo()
+	err = cc.ProbeTLSInfo(context.Background())
 	require.NoError(t, err)
 
 	// When unreachable, all scanned protocols should be unsupported
@@ -655,7 +656,7 @@ func TestCertinfo_ProbeTLSInfo_SingleCipher(t *testing.T) {
 		},
 	}
 
-	res := cc.probeCiphersConcurrently(ciphers)
+	res := cc.probeCiphersConcurrently(context.Background(), ciphers)
 	require.Len(t, res, 1)
 	require.Equal(t, "TLS_AES_128_GCM_SHA256", res[0].Name)
 	require.False(t, res[0].Supported)
@@ -673,7 +674,7 @@ func TestCertinfo_PrintData_WithTLSInfo(t *testing.T) {
 
 	var buf bytes.Buffer
 
-	err = cc.PrintData(&buf)
+	err = cc.PrintData(context.Background(), &buf)
 	require.NoError(t, err)
 	require.Contains(t, buf.String(), "Negotiated TLS Connection")
 }