From f7f4feda486c0e287bdcbb187a5ffb4b0a430219 Mon Sep 17 00:00:00 2001 From: idy Date: Fri, 16 Oct 2026 11:21:12 +0800 Subject: [PATCH 1/2] locate: preserve connector TLS for store liveness probes Signed-off-by: idy --- README.md | 11 ++ internal/locate/region_cache.go | 7 + internal/locate/store_cache.go | 56 +++++-- internal/locate/store_cache_tls_test.go | 189 ++++++++++++++++++++++++ tikv/kv.go | 10 ++ 5 files changed, 260 insertions(+), 13 deletions(-) create mode 100644 internal/locate/store_cache_tls_test.go diff --git a/README.md b/README.md index de124849c5..91e4356b8a 100644 --- a/README.md +++ b/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). diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go index 2002ad72f3..e99bb8f3e8 100644 --- a/internal/locate/region_cache.go +++ b/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 diff --git a/internal/locate/store_cache.go b/internal/locate/store_cache.go index 38f36b5b8e..55f5adc030 100644 --- a/internal/locate/store_cache.go +++ b/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 { @@ -593,7 +619,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 ) @@ -736,8 +762,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: @@ -749,7 +779,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 { @@ -762,7 +792,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 @@ -800,17 +830,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) } diff --git a/internal/locate/store_cache_tls_test.go b/internal/locate/store_cache_tls_test.go new file mode 100644 index 0000000000..a093e5f774 --- /dev/null +++ b/internal/locate/store_cache_tls_test.go @@ -0,0 +1,189 @@ +// 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() + listener, err := net.Listen("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) +} diff --git a/tikv/kv.go b/tikv/kv.go index 66cee5da1d..024a8127d7 100644 --- a/tikv/kv.go +++ b/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) { From 5a308bf7545de2e7408539a840e905f45b1d7a99 Mon Sep 17 00:00:00 2001 From: idy Date: Fri, 16 Oct 2026 12:07:15 +0800 Subject: [PATCH 2/2] locate: use test context for TLS listener creation Signed-off-by: idy --- internal/locate/store_cache_tls_test.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/internal/locate/store_cache_tls_test.go b/internal/locate/store_cache_tls_test.go index a093e5f774..ea16575888 100644 --- a/internal/locate/store_cache_tls_test.go +++ b/internal/locate/store_cache_tls_test.go @@ -81,7 +81,8 @@ func newHealthTestPKI(t *testing.T) healthTestPKI { func startHealthTestServer(t *testing.T, pki *healthTestPKI, interceptor grpc.UnaryServerInterceptor) (string, *health.Server) { t.Helper() - listener, err := net.Listen("tcp", "127.0.0.1:0") + 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 {