locate: preserve connector TLS for store liveness probes

README.md

@@ -26,6 +26,17 @@ There are 2 major versions of the `client-go` package.
 
 Explanation of related concepts and sample code can be found in the [Client-Go Wiki](https://github.com/tikv/client-go/wiki). There are also executable examples in [examples directory](https://github.com/tikv/client-go/tree/master/examples).
 
+### Connector-specific recovery TLS
+
+When constructing a transactional store with an RPC client configured using
+`tikv.WithSecurity(security)`, also pass
+`tikv.WithStoreLivenessSecurity(security)` to `tikv.NewKVStore`. Store recovery
+health probes then use that connector's CA and client certificate/key, with
+normal TLS hostname verification, without changing SDK global configuration. An explicitly empty
+security policy selects plaintext for that connector; omitting the store option
+retains the legacy global security default. In-flight health probes are isolated
+between connectors, including connectors that address the same endpoint.
+
 ## Contributing to client-go
 
 Pull Requests and issues are welcomed. Please check [CONTRIBUTING.md](./CONTRIBUTING.md).

internal/locate/region_cache.go

@@ -718,6 +718,13 @@ func WithRequestHealthFeedbackCallback(callback func(ctx context.Context, addr s
 	}
 }
 
+// SetStoreLivenessSecurity selects connector-local TLS credentials for store
+// recovery probes. An empty policy explicitly selects plaintext. Changing the
+// policy starts independent in-flight probes and cannot reuse an old result.
+func (c *RegionCache) SetStoreLivenessSecurity(security config.Security) {
+	c.stores.setStoreLivenessSecurity(security)
+}
+
 // NewRegionCache creates a RegionCache.
 func NewRegionCache(pdClient pd.Client, opt ...RegionCacheOpt) *RegionCache {
 	var options regionCacheOptions

internal/locate/store_cache.go

@@ -54,6 +54,14 @@ const (
 type testingKnobs interface {
 	getMockRequestLiveness() livenessFunc
 	setMockRequestLiveness(f livenessFunc)
+	getStoreLiveness() *storeLiveness
+}
+
+// A policy owns its in-flight probes so connectors with different credentials
+// never share a health result for the same address.
+type storeLiveness struct {
+	security *config.Security
+	flight   singleflight.Group
 }
 
 type storeRegistry interface {
@@ -75,6 +83,7 @@ type storeCache interface {
 	markTiflashComputeStoresNeedReload()
 	markStoreNeedCheck(store *Store)
 	getCheckStoreEvents() <-chan struct{}
+	setStoreLivenessSecurity(security config.Security)
 }
 
 func newStoreCache(pdClient pd.Client) *storeCacheImpl {
@@ -88,6 +97,7 @@ func newStoreCache(pdClient pd.Client) *storeCacheImpl {
 
 type storeCacheImpl struct {
 	pdClient pd.Client
+	liveness atomic.Pointer[storeLiveness]
 
 	testingKnobs struct {
 		// Replace the requestLiveness function for test purpose. Note that in unit tests, if this is not set,
@@ -108,6 +118,22 @@ type storeCacheImpl struct {
 	}
 }
 
+func (c *storeCacheImpl) getStoreLiveness() *storeLiveness {
+	if state := c.liveness.Load(); state != nil {
+		return state
+	}
+	state := &storeLiveness{}
+	if c.liveness.CompareAndSwap(nil, state) {
+		return state
+	}
+	return c.liveness.Load()
+}
+
+func (c *storeCacheImpl) setStoreLivenessSecurity(security config.Security) {
+	security.ClusterVerifyCN = append([]string(nil), security.ClusterVerifyCN...)
+	c.liveness.Store(&storeLiveness{security: &security})
+}
+
 func (c *storeCacheImpl) getMockRequestLiveness() livenessFunc {
 	f := c.testingKnobs.mockRequestLiveness.Load()
 	if f == nil {
@@ -601,7 +627,7 @@ func (l livenessState) injectConstantLiveness(tk testingKnobs) {
 }
 
 var (
-	livenessSf singleflight.Group
+	defaultStoreLiveness storeLiveness
 	// storeLivenessTimeout is the max duration of resolving liveness of a TiKV instance.
 	storeLivenessTimeout = time.Second
 )
@@ -744,8 +770,12 @@ func requestLiveness(ctx context.Context, s *Store, tk testingKnobs) (l liveness
 		return
 	}
 	addr := s.GetAddr()
-	rsCh := livenessSf.DoChan(addr, func() (interface{}, error) {
-		return invokeKVStatusAPI(addr, storeLivenessTimeout), nil
+	state := &defaultStoreLiveness
+	if tk != nil {
+		state = tk.getStoreLiveness()
+	}
+	rsCh := state.flight.DoChan(addr, func() (interface{}, error) {
+		return invokeKVStatusAPI(addr, storeLivenessTimeout, state.security), nil
 	})
 	select {
 	case rs := <-rsCh:
@@ -757,7 +787,7 @@ func requestLiveness(ctx context.Context, s *Store, tk testingKnobs) (l liveness
 	return
 }
 
-func invokeKVStatusAPI(addr string, timeout time.Duration) (l livenessState) {
+func invokeKVStatusAPI(addr string, timeout time.Duration, security *config.Security) (l livenessState) {
 	start := time.Now()
 	defer func() {
 		if l == reachable {
@@ -770,7 +800,7 @@ func invokeKVStatusAPI(addr string, timeout time.Duration) (l livenessState) {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 
-	conn, cli, err := createKVHealthClient(ctx, addr)
+	conn, cli, err := createKVHealthClient(ctx, addr, security)
 	if err != nil {
 		logutil.BgLogger().Info("[health check] create grpc connection failed", zap.String("store", addr), zap.Error(err))
 		l = unreachable
@@ -808,17 +838,17 @@ func invokeKVStatusAPI(addr string, timeout time.Duration) (l livenessState) {
 	return
 }
 
-func createKVHealthClient(ctx context.Context, addr string) (*grpc.ClientConn, healthpb.HealthClient, error) {
-	// Temporarily directly load the config from the global config, however it's not a good idea to let RegionCache to
-	// access it.
-	// TODO: Pass the config in a better way, or use the connPool inner the client directly rather than creating new
-	// connection.
-
+func createKVHealthClient(ctx context.Context, addr string, security *config.Security) (*grpc.ClientConn, healthpb.HealthClient, error) {
 	cfg := config.GetGlobalConfig()
+	// Retain the legacy default only when no connector policy was supplied.
+	// An explicitly empty policy selects plaintext for this connector.
+	if security == nil {
+		security = &cfg.Security
+	}
 
 	opt := grpc.WithTransportCredentials(insecure.NewCredentials())
-	if len(cfg.Security.ClusterSSLCA) != 0 {
-		tlsConfig, err := cfg.Security.ToTLSConfig()
+	if len(security.ClusterSSLCA) != 0 {
+		tlsConfig, err := security.ToTLSConfig()
 		if err != nil {
 			return nil, nil, errors.WithStack(err)
 		}

internal/locate/store_cache_tls_test.go

@@ -0,0 +1,190 @@
+// Copyright 2026 TiKV Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+
+package locate
+
+import (
+	"context"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/tls"
+	"crypto/x509"
+	"crypto/x509/pkix"
+	"encoding/pem"
+	"math/big"
+	"net"
+	"os"
+	"path/filepath"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"github.com/tikv/client-go/v2/config"
+	"github.com/tikv/client-go/v2/config/retry"
+	"github.com/tikv/client-go/v2/tikvrpc"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/health"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
+)
+
+type healthTestPKI struct {
+	security config.Security
+	server   *tls.Config
+}
+
+func newHealthTestPKI(t *testing.T) healthTestPKI {
+	t.Helper()
+	dir := t.TempDir()
+	caKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	require.NoError(t, err)
+	ca := &x509.Certificate{
+		SerialNumber: big.NewInt(1), Subject: pkix.Name{CommonName: "test-health-ca"},
+		NotBefore: time.Now().Add(-time.Hour), NotAfter: time.Now().Add(time.Hour),
+		IsCA: true, BasicConstraintsValid: true,
+		KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature,
+	}
+	caDER, err := x509.CreateCertificate(rand.Reader, ca, ca, &caKey.PublicKey, caKey)
+	require.NoError(t, err)
+	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	require.NoError(t, err)
+	leaf := &x509.Certificate{
+		SerialNumber: big.NewInt(2), Subject: pkix.Name{CommonName: "test-health"},
+		NotBefore: ca.NotBefore, NotAfter: ca.NotAfter,
+		IPAddresses: []net.IP{net.ParseIP("127.0.0.1")},
+		KeyUsage:    x509.KeyUsageDigitalSignature,
+		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
+	}
+	leafDER, err := x509.CreateCertificate(rand.Reader, leaf, ca, &key.PublicKey, caKey)
+	require.NoError(t, err)
+	keyDER, err := x509.MarshalPKCS8PrivateKey(key)
+	require.NoError(t, err)
+	write := func(name, kind string, der []byte) string {
+		path := filepath.Join(dir, name)
+		require.NoError(t, os.WriteFile(path, pem.EncodeToMemory(&pem.Block{Type: kind, Bytes: der}), 0600))
+		return path
+	}
+	security := config.NewSecurity(write("ca.pem", "CERTIFICATE", caDER), write("cert.pem", "CERTIFICATE", leafDER), write("key.pem", "PRIVATE KEY", keyDER), nil)
+	server, err := security.ToTLSConfig()
+	require.NoError(t, err)
+	server.MinVersion = tls.VersionTLS12
+	server.ClientAuth = tls.RequireAndVerifyClientCert
+	return healthTestPKI{security: security, server: server}
+}
+
+func startHealthTestServer(t *testing.T, pki *healthTestPKI, interceptor grpc.UnaryServerInterceptor) (string, *health.Server) {
+	t.Helper()
+	var listenConfig net.ListenConfig
+	listener, err := listenConfig.Listen(t.Context(), "tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	var options []grpc.ServerOption
+	if pki != nil {
+		options = append(options, grpc.Creds(credentials.NewTLS(pki.server)))
+	}
+	if interceptor != nil {
+		options = append(options, grpc.UnaryInterceptor(interceptor))
+	}
+	server := grpc.NewServer(options...)
+	service := health.NewServer()
+	service.SetServingStatus("", healthpb.HealthCheckResponse_SERVING)
+	healthpb.RegisterHealthServer(server, service)
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		_ = server.Serve(listener)
+	}()
+	t.Cleanup(func() { server.Stop(); <-done })
+	return listener.Addr().String(), service
+}
+
+func setHealthTestGlobalSecurity(t *testing.T, security config.Security) {
+	t.Helper()
+	previous := config.GetGlobalConfig()
+	config.UpdateGlobal(func(c *config.Config) { c.Security = security })
+	t.Cleanup(func() { config.StoreGlobalConfig(previous) })
+}
+
+func TestStoreLivenessConnectorMTLS(t *testing.T) {
+	setHealthTestGlobalSecurity(t, config.Security{})
+	a, b := newHealthTestPKI(t), newHealthTestPKI(t)
+	addrA, _ := startHealthTestServer(t, &a, nil)
+	addrB, _ := startHealthTestServer(t, &b, nil)
+	require.Equal(t, unreachable, invokeKVStatusAPI(addrA, 250*time.Millisecond, nil), "legacy empty globals cannot authenticate")
+	require.Equal(t, reachable, invokeKVStatusAPI(addrA, time.Second, &a.security))
+	require.Equal(t, reachable, invokeKVStatusAPI(addrB, time.Second, &b.security))
+	require.Equal(t, unreachable, invokeKVStatusAPI(addrA, 250*time.Millisecond, &b.security), "wrong CA must never fall back")
+	withoutClient := config.NewSecurity(a.security.ClusterSSLCA, "", "", nil)
+	require.Equal(t, unreachable, invokeKVStatusAPI(addrA, 250*time.Millisecond, &withoutClient), "mTLS requires client identity")
+	require.Equal(t, config.Security{}, config.GetGlobalConfig().Security)
+}
+
+func TestStoreLivenessExplicitPlaintext(t *testing.T) {
+	pki := newHealthTestPKI(t)
+	setHealthTestGlobalSecurity(t, pki.security)
+	addr, _ := startHealthTestServer(t, nil, nil)
+	empty := config.Security{}
+	require.Equal(t, reachable, invokeKVStatusAPI(addr, time.Second, &empty))
+	require.Equal(t, unreachable, invokeKVStatusAPI(addr, 250*time.Millisecond, nil), "omitted policy retains global TLS")
+}
+
+func TestStoreLivenessDoesNotShareAnotherConnectorsResult(t *testing.T) {
+	a, b := newHealthTestPKI(t), newHealthTestPKI(t)
+	entered, release := make(chan struct{}, 1), make(chan struct{})
+	var once sync.Once
+	resume := func() { once.Do(func() { close(release) }) }
+	t.Cleanup(resume)
+	addr, _ := startHealthTestServer(t, &a, func(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+		select {
+		case entered <- struct{}{}:
+		default:
+		}
+		select {
+		case <-release:
+			return handler(ctx, req)
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	})
+	good, bad := &storeCacheImpl{}, &storeCacheImpl{}
+	good.setStoreLivenessSecurity(a.security)
+	bad.setStoreLivenessSecurity(b.security)
+	store := newStore(1, addr, "", "", tikvrpc.TiKV, resolved, nil)
+	result := make(chan livenessState, 1)
+	go func() { result <- requestLiveness(context.Background(), store, good) }()
+	select {
+	case <-entered:
+	case <-time.After(time.Second):
+		t.Fatal("authenticated health request did not enter server")
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+	require.Equal(t, unreachable, requestLiveness(ctx, store, bad), "untrusted connector must not inherit an authenticated in-flight probe")
+	good.setStoreLivenessSecurity(b.security)
+	require.Equal(t, unreachable, requestLiveness(ctx, store, good), "new security must not inherit the previous policy's in-flight probe")
+	resume()
+	require.Equal(t, reachable, <-result)
+}
+
+func TestStoreLivenessMTLSBackgroundRecovery(t *testing.T) {
+	setHealthTestGlobalSecurity(t, config.Security{})
+	pki := newHealthTestPKI(t)
+	addr, service := startHealthTestServer(t, &pki, nil)
+	cache := &storeCacheImpl{}
+	regionCache := &RegionCache{stores: cache}
+	regionCache.SetStoreLivenessSecurity(pki.security)
+	store := newStore(1, addr, "", "", tikvrpc.TiKV, resolved, nil)
+	scheduler := newBackgroundRunner(context.Background())
+	t.Cleanup(func() { scheduler.shutdown(true) })
+	service.SetServingStatus("", healthpb.HealthCheckResponse_NOT_SERVING)
+	require.Equal(t, unreachable, store.requestLivenessAndStartHealthCheckLoopIfNeeded(retry.NewBackofferWithVars(context.Background(), 1000, nil), scheduler, cache))
+	require.Equal(t, unreachable, store.getLivenessState())
+	service.SetServingStatus("", healthpb.HealthCheckResponse_SERVING)
+	require.Eventually(t, func() bool { return store.getLivenessState() == reachable }, 3*time.Second, 10*time.Millisecond)
+}

tikv/kv.go

@@ -247,6 +247,16 @@ func (s *KVStore) CheckVisibility(startTS uint64) error {
 // Option is the option for pool.
 type Option func(*KVStore)
 
+// WithStoreLivenessSecurity configures store recovery probes with the same
+// connector-specific TLS policy supplied to the RPC client. An empty policy
+// explicitly selects plaintext; omitting this option preserves global defaults.
+func WithStoreLivenessSecurity(security config.Security) Option {
+	security.ClusterVerifyCN = append([]string(nil), security.ClusterVerifyCN...)
+	return func(store *KVStore) {
+		store.regionCache.SetStoreLivenessSecurity(security)
+	}
+}
+
 // WithPool set the pool
 func WithPool(gp Pool) Option {
 	return func(o *KVStore) {