fix: copilot issues

Signed-off-by: Gustavo Carvalho <gustavo.carvalho@container-solutions.com>

Author: gusfcarvalho

main.go

@@ -788,6 +788,9 @@ func (e *CommandCustodianExecutor) runAWSEndpointDiagnostics(ctx context.Context
 
 	var diagnosticsErr error
 	for _, endpoint := range endpoints {
+		if err := ctx.Err(); err != nil {
+			return errors.Join(diagnosticsErr, err)
+		}
 		lookupCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
 		lookupStarted := time.Now()
 		ips, err := lookupHost(lookupCtx, endpoint.Host)
@@ -895,9 +898,13 @@ func parseNetworkDiagnosticEndpoint(value string) (networkDiagnosticEndpoint, er
 	if port == "" {
 		port = "443"
 	}
-	if _, err := strconv.Atoi(port); err != nil {
+	portNumber, err := strconv.Atoi(port)
+	if err != nil {
 		return networkDiagnosticEndpoint{}, fmt.Errorf("network diagnostic endpoint %q has invalid port %q: %w", original, port, err)
 	}
+	if portNumber < 1 || portNumber > 65535 {
+		return networkDiagnosticEndpoint{}, fmt.Errorf("network diagnostic endpoint %q has invalid port %q: must be between 1 and 65535", original, port)
+	}
 	return networkDiagnosticEndpoint{
 		Host:       strings.ToLower(host),
 		Port:       port,
@@ -934,19 +941,30 @@ type tlsProbeResult struct {
 }
 
 func defaultTLSProbeEndpoint(ctx context.Context, endpoint networkDiagnosticEndpoint) (tlsProbeResult, error) {
-	dialer := &net.Dialer{Timeout: 5 * time.Second}
+	if err := ctx.Err(); err != nil {
+		return tlsProbeResult{}, err
+	}
+	netDialer := &net.Dialer{Timeout: 5 * time.Second}
 	if deadline, ok := ctx.Deadline(); ok {
-		if remaining := time.Until(deadline); remaining > 0 && remaining < dialer.Timeout {
-			dialer.Timeout = remaining
+		if remaining := time.Until(deadline); remaining > 0 && remaining < netDialer.Timeout {
+			netDialer.Timeout = remaining
 		}
 	}
-	conn, err := tls.DialWithDialer(dialer, "tcp", net.JoinHostPort(endpoint.Host, endpoint.Port), &tls.Config{ServerName: endpoint.ServerName, MinVersion: tls.VersionTLS12})
+	tlsDialer := &tls.Dialer{
+		NetDialer: netDialer,
+		Config:    &tls.Config{ServerName: endpoint.ServerName, MinVersion: tls.VersionTLS12},
+	}
+	conn, err := tlsDialer.DialContext(ctx, "tcp", net.JoinHostPort(endpoint.Host, endpoint.Port))
 	if err != nil {
 		return tlsProbeResult{}, err
 	}
 	defer conn.Close()
 
-	state := conn.ConnectionState()
+	tlsConn, ok := conn.(*tls.Conn)
+	if !ok {
+		return tlsProbeResult{}, errors.New("TLS dial returned non-TLS connection")
+	}
+	state := tlsConn.ConnectionState()
 	return tlsProbeResult{
 		RemoteAddr: conn.RemoteAddr().String(),
 		TLSVersion: tlsVersionString(state.Version),

main_test.go

@@ -616,6 +616,35 @@ printf '[]' > "$out/test-policy/resources.json"
 		}
 	})
 
+	t.Run("network diagnostics stop when context is canceled", func(t *testing.T) {
+		stubNetworkDiagnostics(
+			t,
+			func(ctx context.Context, host string) ([]string, error) {
+				t.Fatalf("did not expect DNS probe after context cancellation, got host %s", host)
+				return nil, nil
+			},
+			func(ctx context.Context, endpoint networkDiagnosticEndpoint) (tlsProbeResult, error) {
+				t.Fatalf("did not expect TLS probe after context cancellation, got endpoint %#v", endpoint)
+				return tlsProbeResult{}, nil
+			},
+		)
+		executor := &CommandCustodianExecutor{Logger: hclog.NewNullLogger()}
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		err := executor.runAWSEndpointDiagnostics(ctx, CustodianExecutionRequest{
+			Check: CustodianCheck{
+				Name:     "test-policy",
+				Resource: "aws.backup-vault",
+				Provider: "aws",
+			},
+			NetworkDiagnosticEndpoints: []string{"vpce-123.backup.eu-west-1.vpce.amazonaws.com"},
+		})
+		if !errors.Is(err, context.Canceled) {
+			t.Fatalf("expected context canceled error, got %v", err)
+		}
+	})
+
 	t.Run("network diagnostics skip service probes when regions are not concrete", func(t *testing.T) {
 		stubNetworkDiagnostics(
 			t,
@@ -1094,6 +1123,27 @@ func TestDiagnosticHelpers(t *testing.T) {
 		if err == nil {
 			t.Fatalf("expected invalid endpoint port error")
 		}
+		for _, endpoint := range []string{
+			"vpce-123.backup.eu-west-1.vpce.amazonaws.com:0",
+			"vpce-123.backup.eu-west-1.vpce.amazonaws.com:65536",
+		} {
+			if _, _, err := awsDiagnosticEndpointsForCheck("aws.backup-vault", nil, []string{endpoint}); err == nil {
+				t.Fatalf("expected invalid endpoint port error for %s", endpoint)
+			}
+		}
+	})
+
+	t.Run("tls probe returns immediately when context is canceled", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		_, err := defaultTLSProbeEndpoint(ctx, networkDiagnosticEndpoint{
+			Host:       "example.invalid",
+			Port:       "443",
+			ServerName: "example.invalid",
+		})
+		if !errors.Is(err, context.Canceled) {
+			t.Fatalf("expected context canceled error, got %v", err)
+		}
 	})
 
 	t.Run("reports unknown resource types", func(t *testing.T) {
@@ -1801,7 +1851,7 @@ func TestAWSResourceExplorerURL(t *testing.T) {
 		}
 	})
 
-	t.Run("logs common non arn skipped link generation at debug", func(t *testing.T) {
+	t.Run("does not warn for common non arn skipped link generation", func(t *testing.T) {
 		var logs bytes.Buffer
 		logger := hclog.New(&hclog.LoggerOptions{
 			Name:   "test",