Fix wrong context usage for reissuing expired certificate for tsh pro…

…xy kube. (#43374)

* Fix wrong context usage for reissuing expired certificate for tsh proxy kube.

* Rename context to closeContext

* Add test for request context expiration.

* Add missing context in tests.

* Remove flakiness from the test.

integration/proxy/proxy_helpers.go

@@ -567,6 +567,7 @@ func mustCreateKubeLocalProxyMiddleware(t *testing.T, teleportCluster, kubeClust
 		CertReissuer: func(ctx context.Context, teleportCluster, kubeCluster string) (tls.Certificate, error) {
 			return tls.Certificate{}, nil
 		},
+		CloseContext: context.Background(),
 	})
 }
 

lib/srv/alpnproxy/kube.go

@@ -81,7 +81,8 @@ type KubeMiddleware struct {
 	// headless controls whether proxy is working in headless login mode.
 	headless bool
 
-	logger logrus.FieldLogger
+	logger       logrus.FieldLogger
+	closeContext context.Context
 
 	// isCertReissuingRunning is used to only ever have one concurrent cert reissuing session requiring user input.
 	isCertReissuingRunning atomic.Bool
@@ -97,6 +98,7 @@ type KubeMiddlewareConfig struct {
 	Headless     bool
 	Clock        clockwork.Clock
 	Logger       logrus.FieldLogger
+	CloseContext context.Context
 }
 
 // NewKubeMiddleware creates a new KubeMiddleware.
@@ -107,6 +109,7 @@ func NewKubeMiddleware(cfg KubeMiddlewareConfig) LocalProxyHTTPMiddleware {
 		headless:     cfg.Headless,
 		clock:        cfg.Clock,
 		logger:       cfg.Logger,
+		closeContext: cfg.CloseContext,
 	}
 }
 
@@ -121,6 +124,9 @@ func (m *KubeMiddleware) CheckAndSetDefaults() error {
 	if m.logger == nil {
 		m.logger = logrus.WithField(teleport.ComponentKey, "local_proxy_kube")
 	}
+	if m.closeContext == nil {
+		return trace.BadParameter("missing close context")
+	}
 	return nil
 }
 
@@ -245,7 +251,7 @@ func (m *KubeMiddleware) reissueCertIfExpired(ctx context.Context, cert tls.Cert
 			if identity.RouteToCluster != "" {
 				cluster = identity.RouteToCluster
 			}
-			newCert, err := m.certReissuer(ctx, cluster, identity.KubernetesCluster)
+			newCert, err := m.certReissuer(m.closeContext, cluster, identity.KubernetesCluster)
 			if err == nil {
 				m.certsMu.Lock()
 				m.certs[serverName] = newCert

lib/srv/alpnproxy/local_proxy_test.go

@@ -504,9 +504,52 @@ func TestKubeMiddleware(t *testing.T) {
 	)
 
 	certReissuer = func(ctx context.Context, teleportCluster, kubeCluster string) (tls.Certificate, error) {
-		return newCert, nil
+		select {
+		case <-ctx.Done():
+			return tls.Certificate{}, ctx.Err()
+		default:
+			return newCert, nil
+		}
 	}
 
+	t.Run("expired certificate is still reissued if request context expires", func(t *testing.T) {
+		req := &http.Request{
+			TLS: &tls.ConnectionState{
+				ServerName: "kube1",
+			},
+		}
+		// we set request context to a context that will expired immediately.
+		reqCtx, cancel := context.WithDeadline(context.Background(), time.Now())
+		defer cancel()
+		req = req.WithContext(reqCtx)
+
+		km := NewKubeMiddleware(KubeMiddlewareConfig{
+			Certs:        KubeClientCerts{"kube1": kube1Cert},
+			CertReissuer: certReissuer,
+			Clock:        clockwork.NewFakeClockAt(now.Add(time.Hour * 2)),
+			CloseContext: context.Background(),
+		})
+		err := km.CheckAndSetDefaults()
+		require.NoError(t, err)
+
+		rw := responsewriters.NewMemoryResponseWriter()
+		// HandleRequest will reissue certificate if needed.
+		km.HandleRequest(rw, req)
+
+		// request timed out.
+		require.Equal(t, http.StatusInternalServerError, rw.Status())
+		require.Contains(t, rw.Buffer().String(), "context deadline exceeded")
+
+		// just let the reissuing goroutine some time to replace certs.
+		time.Sleep(10 * time.Millisecond)
+
+		// but certificate still was reissued.
+		certs, err := km.OverwriteClientCerts(req)
+		require.NoError(t, err)
+		require.Len(t, certs, 1)
+		require.Equal(t, newCert, certs[0], "certificate was not reissued")
+	})
+
 	testCases := []struct {
 		name            string
 		reqClusterName  string
@@ -559,6 +602,7 @@ func TestKubeMiddleware(t *testing.T) {
 				Certs:        tt.startCerts,
 				CertReissuer: certReissuer,
 				Clock:        tt.clock,
+				CloseContext: context.Background(),
 			})
 
 			// HandleRequest will reissue certificate if needed

lib/teleterm/gateway/kube.go

@@ -131,8 +131,9 @@ func (k *kube) makeKubeMiddleware() (alpnproxy.LocalProxyHTTPMiddleware, error)
 			cert, err := k.cfg.OnExpiredCert(ctx, k)
 			return cert, trace.Wrap(err)
 		},
-		Clock:  k.cfg.Clock,
-		Logger: k.cfg.Log,
+		Clock:        k.cfg.Clock,
+		Logger:       k.cfg.Log,
+		CloseContext: k.closeContext,
 	}), nil
 }
 

tool/tsh/common/kube_proxy.go

@@ -342,6 +342,7 @@ func makeKubeLocalProxy(cf *CLIConf, tc *client.TeleportClient, clusters kubecon
 		CertReissuer: kubeProxy.getCertReissuer(tc),
 		Headless:     cf.Headless,
 		Logger:       log,
+		CloseContext: cf.Context,
 	})
 
 	localProxy, err := alpnproxy.NewLocalProxy(