backend: add context to pass parameters need by handler (#144)

pingcap · Dec 6, 2022 · 3cc9581 · 3cc9581
1 parent ed9650d
commit 3cc9581
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 14 deletions.
diff --git a/pkg/proxy/backend/authenticator.go b/pkg/proxy/backend/authenticator.go
@@ -15,6 +15,7 @@
 package backend
 
 import (
+	"context"
 	"crypto/tls"
 	"encoding/binary"
 	"fmt"
@@ -145,15 +146,16 @@ func (auth *Authenticator) handshakeFirstTime(logger *zap.Logger, clientIO *pnet
 	auth.capability = commonCaps.Uint32()
 
 	resp := pnet.ParseHandshakeResponse(pkt)
-	if err = handshakeHandler.HandleHandshakeResp(resp, clientIO.SourceAddr().String()); err != nil {
+	ctx := context.WithValue(context.Background(), ContextKeyClientAddr, clientIO.SourceAddr().String())
+	if err = handshakeHandler.HandleHandshakeResp(ctx, resp); err != nil {
 		return err
 	}
 	auth.user = resp.User
 	auth.dbname = resp.DB
 	auth.collation = resp.Collation
 	auth.attrs = resp.Attrs
 
-	backendIO, err := getBackend(auth, resp)
+	backendIO, err := getBackend(ctx, auth, resp)
 	if err != nil {
 		return err
 	}

diff --git a/pkg/proxy/backend/backend_conn_mgr.go b/pkg/proxy/backend/backend_conn_mgr.go
@@ -54,7 +54,7 @@ type redirectResult struct {
 	to   string
 }
 
-type backendIOGetter func(auth *Authenticator, resp *pnet.HandshakeResp) (*pnet.PacketIO, error)
+type backendIOGetter func(ctx context.Context, auth *Authenticator, resp *pnet.HandshakeResp) (*pnet.PacketIO, error)
 
 // BackendConnManager migrates a session from one BackendConnection to another.
 //
@@ -105,8 +105,8 @@ func NewBackendConnManager(logger *zap.Logger, handshakeHandler HandshakeHandler
 		signalReceived: make(chan struct{}, 1),
 		redirectResCh:  make(chan *redirectResult, 1),
 	}
-	mgr.getBackendIO = func(auth *Authenticator, resp *pnet.HandshakeResp) (*pnet.PacketIO, error) {
-		router, err := handshakeHandler.GetRouter(resp)
+	mgr.getBackendIO = func(ctx context.Context, auth *Authenticator, resp *pnet.HandshakeResp) (*pnet.PacketIO, error) {
+		router, err := handshakeHandler.GetRouter(ctx, resp)
 		if err != nil {
 			return nil, err
 		}

diff --git a/pkg/proxy/backend/backend_conn_mgr_test.go b/pkg/proxy/backend/backend_conn_mgr_test.go
@@ -116,7 +116,7 @@ func newBackendMgrTester(t *testing.T, cfg ...cfgOverrider) *backendMgrTester {
 	return tester
 }
 
-func (ts *backendMgrTester) getBackendIO(auth *Authenticator, _ *pnet.HandshakeResp) (*pnet.PacketIO, error) {
+func (ts *backendMgrTester) getBackendIO(ctx context.Context, auth *Authenticator, _ *pnet.HandshakeResp) (*pnet.PacketIO, error) {
 	addr := ts.tc.backendListener.Addr().String()
 	ts.mp.backendConn = NewBackendConnection(addr)
 	if err := ts.mp.backendConn.Connect(); err != nil {

diff --git a/pkg/proxy/backend/handshake_handler.go b/pkg/proxy/backend/handshake_handler.go
@@ -15,18 +15,31 @@
 package backend
 
 import (
+	"context"
+
 	"github.com/pingcap/TiProxy/lib/util/errors"
 	"github.com/pingcap/TiProxy/pkg/manager/namespace"
 	"github.com/pingcap/TiProxy/pkg/manager/router"
 	pnet "github.com/pingcap/TiProxy/pkg/proxy/net"
 )
 
+type contextKey string
+
+func (k contextKey) String() string {
+	return "handler context key " + string(k)
+}
+
+// Context keys.
+var (
+	ContextKeyClientAddr = contextKey("client_addr")
+)
+
 var _ HandshakeHandler = (*DefaultHandshakeHandler)(nil)
 
 type HandshakeHandler interface {
-	HandleHandshakeResp(resp *pnet.HandshakeResp, sourceAddr string) error
+	HandleHandshakeResp(ctx context.Context, resp *pnet.HandshakeResp) error
 	GetCapability() pnet.Capability
-	GetRouter(resp *pnet.HandshakeResp) (router.Router, error)
+	GetRouter(ctx context.Context, resp *pnet.HandshakeResp) (router.Router, error)
 }
 
 type DefaultHandshakeHandler struct {
@@ -39,15 +52,15 @@ func NewDefaultHandshakeHandler(nsManager *namespace.NamespaceManager) *DefaultH
 	}
 }
 
-func (handler *DefaultHandshakeHandler) HandleHandshakeResp(*pnet.HandshakeResp, string) error {
+func (handler *DefaultHandshakeHandler) HandleHandshakeResp(context.Context, *pnet.HandshakeResp) error {
 	return nil
 }
 
 func (handler *DefaultHandshakeHandler) GetCapability() pnet.Capability {
 	return SupportedServerCapabilities
 }
 
-func (handler *DefaultHandshakeHandler) GetRouter(resp *pnet.HandshakeResp) (router.Router, error) {
+func (handler *DefaultHandshakeHandler) GetRouter(ctx context.Context, resp *pnet.HandshakeResp) (router.Router, error) {
 	ns, ok := handler.nsManager.GetNamespaceByUser(resp.User)
 	if !ok {
 		ns, ok = handler.nsManager.GetNamespace("default")

diff --git a/pkg/proxy/backend/mock_proxy_test.go b/pkg/proxy/backend/mock_proxy_test.go
@@ -15,6 +15,7 @@
 package backend
 
 import (
+	"context"
 	"crypto/tls"
 	"testing"
 
@@ -64,7 +65,7 @@ func newMockProxy(t *testing.T, cfg *proxyConfig) *mockProxy {
 }
 
 func (mp *mockProxy) authenticateFirstTime(clientIO, backendIO *pnet.PacketIO) error {
-	if err := mp.authenticator.handshakeFirstTime(mp.logger, clientIO, mp.handshakeHandler, func(*Authenticator, *pnet.HandshakeResp) (*pnet.PacketIO, error) {
+	if err := mp.authenticator.handshakeFirstTime(mp.logger, clientIO, mp.handshakeHandler, func(context.Context, *Authenticator, *pnet.HandshakeResp) (*pnet.PacketIO, error) {
 		return backendIO, nil
 	}, mp.frontendTLSConfig, mp.backendTLSConfig); err != nil {
 		return err
@@ -107,14 +108,14 @@ type CustomHandshakeHandler struct {
 	outAttrs      map[string]string
 }
 
-func (handler *CustomHandshakeHandler) GetRouter(resp *pnet.HandshakeResp) (router.Router, error) {
+func (handler *CustomHandshakeHandler) GetRouter(ctx context.Context, resp *pnet.HandshakeResp) (router.Router, error) {
 	return nil, nil
 }
 
-func (handler *CustomHandshakeHandler) HandleHandshakeResp(resp *pnet.HandshakeResp, addr string) error {
+func (handler *CustomHandshakeHandler) HandleHandshakeResp(ctx context.Context, resp *pnet.HandshakeResp) error {
 	handler.inUsername = resp.User
 	resp.User = handler.outUsername
-	handler.inAddr = addr
+	handler.inAddr = ctx.Value(ContextKeyClientAddr).(string)
 	resp.Attrs = handler.outAttrs
 	return nil
 }