balance, backend: support traffic replay when backend signing cert is unavailable (#734)

Author: djshow832

pkg/balance/observer/backend_health.go

@@ -5,6 +5,7 @@ package observer
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/pingcap/tiproxy/lib/config"
 )
@@ -16,6 +17,10 @@ type BackendHealth struct {
 	PingErr error
 	// The backend version that returned to the client during handshake.
 	ServerVersion string
+	// The last time checking the signing cert.
+	lastCheckSigningCertTime time.Time
+	// Whether the backend has set the signing cert. If not, the connection redirection will be disabled.
+	SupportRedirection bool
 	// Whether the backend in the same zone with TiProxy. If TiProxy location is undefined, take all backends as local.
 	Local bool
 }
@@ -37,10 +42,6 @@ func (bh *BackendHealth) setLocal(cfg *config.Config) {
 	bh.Local = false
 }
 
-func (bh *BackendHealth) Equals(health BackendHealth) bool {
-	return bh.Healthy == health.Healthy && bh.ServerVersion == health.ServerVersion && bh.Local == health.Local
-}
-
 func (bh *BackendHealth) String() string {
 	str := "down"
 	if bh.Healthy {
@@ -49,6 +50,9 @@ func (bh *BackendHealth) String() string {
 	if bh.PingErr != nil {
 		str += fmt.Sprintf(", err: %s", bh.PingErr.Error())
 	}
+	if !bh.SupportRedirection {
+		str += ", support redirection: false"
+	}
 	return str
 }
 

pkg/balance/observer/backend_observer.go

@@ -123,8 +123,9 @@ func (bo *DefaultBackendObserver) checkHealth(ctx context.Context, backends map[
 	if !bo.healthCheckConfig.Enable {
 		for addr, backend := range backends {
 			curBackendHealth[addr] = &BackendHealth{
-				BackendInfo: *backend,
-				Healthy:     true,
+				BackendInfo:        *backend,
+				SupportRedirection: true,
+				Healthy:            true,
 			}
 		}
 		return curBackendHealth
@@ -139,7 +140,8 @@ func (bo *DefaultBackendObserver) checkHealth(ctx context.Context, backends map[
 				if ctx.Err() != nil {
 					return
 				}
-				health := bo.hc.Check(ctx, addr, info)
+				lastHealth := bo.curBackends[addr]
+				health := bo.hc.Check(ctx, addr, info, lastHealth)
 				health.setLocal(cfg)
 				lock.Lock()
 				curBackendHealth[addr] = health

pkg/balance/observer/backend_observer_test.go

@@ -284,8 +284,9 @@ func (ts *observerTestSuite) addBackend() (string, BackendInfo) {
 	}
 	ts.fetcher.setBackend(addr, info)
 	ts.hc.setBackend(addr, &BackendHealth{
-		BackendInfo: *info,
-		Healthy:     true,
+		BackendInfo:        *info,
+		SupportRedirection: true,
+		Healthy:            true,
 	})
 	return addr, *info
 }

pkg/balance/observer/health_check.go

@@ -21,11 +21,14 @@ import (
 
 // HealthCheck is used to check the backends of one backend. One can pass a customized health check function to the observer.
 type HealthCheck interface {
-	Check(ctx context.Context, addr string, info *BackendInfo) *BackendHealth
+	Check(ctx context.Context, addr string, info *BackendInfo, lastHealth *BackendHealth) *BackendHealth
 }
 
 const (
 	statusPathSuffix = "/status"
+	configPathSuffix = "/config"
+
+	checkSigningCertInterval = time.Minute
 )
 
 type backendHttpStatusRespBody struct {
@@ -34,6 +37,14 @@ type backendHttpStatusRespBody struct {
 	GitHash     string `json:"git_hash"`
 }
 
+type backendHttpConfigRespBody struct {
+	Security security `json:"security"`
+}
+
+type security struct {
+	SessionTokenSigningCert string `json:"session-token-signing-cert"`
+}
+
 type DefaultHealthCheck struct {
 	cfg     *config.HealthCheck
 	logger  *zap.Logger
@@ -51,7 +62,7 @@ func NewDefaultHealthCheck(httpCli *http.Client, cfg *config.HealthCheck, logger
 	}
 }
 
-func (dhc *DefaultHealthCheck) Check(ctx context.Context, addr string, info *BackendInfo) *BackendHealth {
+func (dhc *DefaultHealthCheck) Check(ctx context.Context, addr string, info *BackendInfo, lastBh *BackendHealth) *BackendHealth {
 	bh := &BackendHealth{
 		BackendInfo: *info,
 		Healthy:     true,
@@ -64,6 +75,10 @@ func (dhc *DefaultHealthCheck) Check(ctx context.Context, addr string, info *Bac
 		return bh
 	}
 	dhc.checkSqlPort(ctx, addr, bh)
+	if !bh.Healthy {
+		return bh
+	}
+	dhc.queryConfig(ctx, info, bh, lastBh)
 	return bh
 }
 
@@ -123,3 +138,48 @@ func (dhc *DefaultHealthCheck) checkStatusPort(ctx context.Context, info *Backen
 		bh.PingErr = errors.Wrapf(err, "connect status port failed")
 	}
 }
+
+func (dhc *DefaultHealthCheck) queryConfig(ctx context.Context, info *BackendInfo, bh *BackendHealth, lastBh *BackendHealth) {
+	if ctx.Err() != nil {
+		return
+	}
+	// Using static backends, no status port.
+	if info == nil || len(info.IP) == 0 {
+		return
+	}
+
+	now := time.Now()
+	if lastBh != nil {
+		bh.SupportRedirection = lastBh.SupportRedirection
+		if lastBh.lastCheckSigningCertTime.Add(checkSigningCertInterval).After(now) {
+			bh.lastCheckSigningCertTime = lastBh.lastCheckSigningCertTime
+			return
+		}
+	} else {
+		// Assume it has the signing cert if reading config fails.
+		bh.SupportRedirection = true
+	}
+	bh.lastCheckSigningCertTime = now
+
+	var err error
+	defer func() {
+		if lastBh == nil || lastBh.SupportRedirection != bh.SupportRedirection {
+			dhc.logger.Info("backend has updated signing cert", zap.Bool("support_redirection", bh.SupportRedirection), zap.Error(err))
+		}
+	}()
+
+	addr := net.JoinHostPort(info.IP, strconv.Itoa(int(info.StatusPort)))
+	b := backoff.WithContext(backoff.WithMaxRetries(backoff.NewConstantBackOff(dhc.cfg.RetryInterval), uint64(dhc.cfg.MaxRetries)), ctx)
+	var resp []byte
+	if resp, err = dhc.httpCli.Get(addr, configPathSuffix, b, dhc.cfg.DialTimeout); err != nil {
+		return
+	}
+	var respBody backendHttpConfigRespBody
+	if err = json.Unmarshal(resp, &respBody); err != nil {
+		dhc.logger.Error("unmarshal body in healthy check failed", zap.String("addr", addr), zap.String("resp body", string(resp)), zap.Error(err))
+		return
+	}
+	if len(respBody.Security.SessionTokenSigningCert) == 0 {
+		bh.SupportRedirection = false
+	}
+}

pkg/balance/observer/health_check_test.go

@@ -25,19 +25,19 @@ func TestReadServerVersion(t *testing.T) {
 	hc := NewDefaultHealthCheck(nil, newHealthCheckConfigForTest(), lg)
 	backend, info := newBackendServer(t)
 	backend.setServerVersion("1.0")
-	health := hc.Check(context.Background(), backend.sqlAddr, info)
+	health := hc.Check(context.Background(), backend.sqlAddr, info, nil)
 	require.Equal(t, "1.0", health.ServerVersion)
 	backend.stopSQLServer()
 	backend.setServerVersion("2.0")
 	backend.startSQLServer()
-	health = hc.Check(context.Background(), backend.sqlAddr, info)
+	health = hc.Check(context.Background(), backend.sqlAddr, info, nil)
 	require.Equal(t, "2.0", health.ServerVersion)
 	backend.stopSQLServer()
 
 	//test for respBody not ok
-	backend.mockHttpHandler.setHTTPRespBody("")
+	backend.mockHttpHandler.setStatusRespBody("")
 	backend.startSQLServer()
-	health = hc.Check(context.Background(), backend.sqlAddr, info)
+	health = hc.Check(context.Background(), backend.sqlAddr, info, nil)
 	require.False(t, health.Healthy)
 	require.NotNil(t, health.PingErr)
 	require.Equal(t, true, strings.Contains(health.PingErr.Error(), "unexpected end of JSON input"))
@@ -52,37 +52,47 @@ func TestHealthCheck(t *testing.T) {
 	hc := NewDefaultHealthCheck(nil, cfg, lg)
 	backend, info := newBackendServer(t)
 	backend.setServerVersion("1.0")
-	health := hc.Check(context.Background(), backend.sqlAddr, info)
+	backend.setHasSigningCert(true)
+	health := hc.Check(context.Background(), backend.sqlAddr, info, nil)
 	require.True(t, health.Healthy)
 
 	backend.stopSQLServer()
-	health = hc.Check(context.Background(), backend.sqlAddr, info)
+	health = hc.Check(context.Background(), backend.sqlAddr, info, nil)
 	require.False(t, health.Healthy)
 	backend.startSQLServer()
-	health = hc.Check(context.Background(), backend.sqlAddr, info)
+	health = hc.Check(context.Background(), backend.sqlAddr, info, nil)
 	require.True(t, health.Healthy)
-
 	backend.setHTTPResp(false)
-	health = hc.Check(context.Background(), backend.sqlAddr, info)
+
+	health = hc.Check(context.Background(), backend.sqlAddr, info, nil)
 	require.False(t, health.Healthy)
 	require.NotContains(t, text.String(), "unmarshal body")
 	backend.setHTTPResp(true)
-	health = hc.Check(context.Background(), backend.sqlAddr, info)
+	health = hc.Check(context.Background(), backend.sqlAddr, info, nil)
 	require.True(t, health.Healthy)
 
 	backend.setHTTPWait(time.Second + cfg.DialTimeout)
-	health = hc.Check(context.Background(), backend.sqlAddr, info)
+	health = hc.Check(context.Background(), backend.sqlAddr, info, nil)
 	require.False(t, health.Healthy)
 	backend.setHTTPWait(time.Duration(0))
-	health = hc.Check(context.Background(), backend.sqlAddr, info)
+	health = hc.Check(context.Background(), backend.sqlAddr, info, nil)
 	require.True(t, health.Healthy)
 
 	backend.setSqlResp(false)
-	health = hc.Check(context.Background(), backend.sqlAddr, info)
+	health = hc.Check(context.Background(), backend.sqlAddr, info, nil)
 	require.False(t, health.Healthy)
 	backend.setSqlResp(true)
-	health = hc.Check(context.Background(), backend.sqlAddr, info)
+	health = hc.Check(context.Background(), backend.sqlAddr, info, nil)
 	require.True(t, health.Healthy)
+
+	require.True(t, health.SupportRedirection)
+	health.lastCheckSigningCertTime = time.Time{}
+	health = hc.Check(context.Background(), backend.sqlAddr, info, health)
+	require.True(t, health.SupportRedirection)
+	backend.setHasSigningCert(false)
+	health.lastCheckSigningCertTime = time.Time{}
+	health = hc.Check(context.Background(), backend.sqlAddr, info, health)
+	require.False(t, health.SupportRedirection)
 	backend.close()
 }
 
@@ -105,7 +115,7 @@ func newBackendServer(t *testing.T) (*backendServer, *BackendInfo) {
 	}
 	backend.startHTTPServer()
 	backend.setHTTPResp(true)
-	backend.setHTTPRespBody("")
+	backend.setStatusRespBody("")
 	backend.setSqlResp(true)
 	backend.startSQLServer()
 	return backend, &BackendInfo{
@@ -120,7 +130,7 @@ func (srv *backendServer) setServerVersion(version string) {
 		GitHash:     "",
 	}
 	body, _ := json.Marshal(resp)
-	srv.mockHttpHandler.setHTTPRespBody(string(body))
+	srv.mockHttpHandler.setStatusRespBody(string(body))
 }
 
 func (srv *backendServer) startHTTPServer() {

pkg/balance/observer/mock_test.go

@@ -5,6 +5,7 @@ package observer
 
 import (
 	"context"
+	"encoding/json"
 	"net/http"
 	"sync"
 	"testing"
@@ -81,7 +82,7 @@ func newMockHealthCheck() *mockHealthCheck {
 	}
 }
 
-func (mhc *mockHealthCheck) Check(_ context.Context, addr string, info *BackendInfo) *BackendHealth {
+func (mhc *mockHealthCheck) Check(_ context.Context, addr string, info *BackendInfo, _ *BackendHealth) *BackendHealth {
 	mhc.Lock()
 	defer mhc.Unlock()
 	mhc.backends[addr].BackendInfo = *info
@@ -109,18 +110,28 @@ func (mhc *mockHealthCheck) removeBackend(addr string) {
 }
 
 type mockHttpHandler struct {
-	t        *testing.T
-	httpOK   atomic.Bool
-	respBody atomic.String
-	wait     atomic.Int64
+	t      *testing.T
+	httpOK atomic.Bool
+	status atomic.String
+	config atomic.String
+	wait   atomic.Int64
 }
 
 func (handler *mockHttpHandler) setHTTPResp(succeed bool) {
 	handler.httpOK.Store(succeed)
 }
 
-func (handler *mockHttpHandler) setHTTPRespBody(body string) {
-	handler.respBody.Store(body)
+func (handler *mockHttpHandler) setStatusRespBody(body string) {
+	handler.status.Store(body)
+}
+
+func (handler *mockHttpHandler) setHasSigningCert(hasSigningCert bool) {
+	var resp backendHttpConfigRespBody
+	if hasSigningCert {
+		resp.Security.SessionTokenSigningCert = "/tmp"
+	}
+	b, _ := json.Marshal(resp)
+	handler.config.Store(string(b))
 }
 
 func (handler *mockHttpHandler) setHTTPWait(wait time.Duration) {
@@ -134,7 +145,11 @@ func (handler *mockHttpHandler) ServeHTTP(w http.ResponseWriter, r *http.Request
 	}
 	if handler.httpOK.Load() {
 		w.WriteHeader(http.StatusOK)
-		_, _ = w.Write([]byte(handler.respBody.Load()))
+		if r.RequestURI == statusPathSuffix {
+			_, _ = w.Write([]byte(handler.status.Load()))
+		} else {
+			_, _ = w.Write([]byte(handler.config.Load()))
+		}
 	} else {
 		w.WriteHeader(http.StatusInternalServerError)
 	}

pkg/balance/router/router.go

@@ -146,11 +146,11 @@ func (b *backendWrapper) GetBackendInfo() observer.BackendInfo {
 	return info
 }
 
-func (b *backendWrapper) Equals(health observer.BackendHealth) bool {
+func (b *backendWrapper) SupportRedirection() bool {
 	b.mu.RLock()
-	equal := b.mu.BackendHealth.Equals(health)
+	supportRedirection := b.mu.SupportRedirection
 	b.mu.RUnlock()
-	return equal
+	return supportRedirection
 }
 
 func (b *backendWrapper) String() string {