From f865967de078997951b74c2daa36a0401073dfc4 Mon Sep 17 00:00:00 2001
From: lysu <sulifx@gmail.com>
Date: Tue, 2 Jul 2019 11:08:22 +0800
Subject: [PATCH] plugin: setup connection info in session when audit plugin be
 enabled (#10923) (#11013)

---
 plugin/plugin.go               | 15 +++++++++++++++
 server/conn.go                 |  6 +++++-
 server/server.go               | 13 ++++++++-----
 sessionctx/variable/session.go |  3 +++
 4 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/plugin/plugin.go b/plugin/plugin.go
index ee5a7d68c6de6..a7c61cbd5819c 100644
--- a/plugin/plugin.go
+++ b/plugin/plugin.go
@@ -377,6 +377,21 @@ func ForeachPlugin(kind Kind, fn func(plugin *Plugin) error) error {
 	return nil
 }
 
+// IsEnable checks plugin's enable state.
+func IsEnable(kind Kind) bool {
+	plugins := pluginGlobal.plugins()
+	if plugins == nil {
+		return false
+	}
+	for i := range plugins.plugins[kind] {
+		p := &plugins.plugins[kind][i]
+		if p.State == Ready {
+			return true
+		}
+	}
+	return false
+}
+
 // GetAll finds and returns all plugins.
 func GetAll() map[Kind][]Plugin {
 	plugins := pluginGlobal.plugins()
diff --git a/server/conn.go b/server/conn.go
index 8f2b99598a30c..7a56f1f799f8f 100644
--- a/server/conn.go
+++ b/server/conn.go
@@ -1434,10 +1434,14 @@ func (cc *clientConn) handleChangeUser(ctx context.Context, data []byte) error {
 		return err
 	}
 
+	if plugin.IsEnable(plugin.Audit) {
+		cc.ctx.GetSessionVars().ConnectionInfo = cc.connectInfo()
+	}
+
 	err = plugin.ForeachPlugin(plugin.Audit, func(p *plugin.Plugin) error {
 		authPlugin := plugin.DeclareAuditManifest(p.Manifest)
 		if authPlugin.OnConnectionEvent != nil {
-			connInfo := cc.connectInfo()
+			connInfo := cc.ctx.GetSessionVars().ConnectionInfo
 			err = authPlugin.OnConnectionEvent(context.Background(), &auth.UserIdentity{Hostname: connInfo.Host}, plugin.ChangeUser, connInfo)
 			if err != nil {
 				return err
diff --git a/server/server.go b/server/server.go
index c258f074d0ed0..dfc608682023a 100644
--- a/server/server.go
+++ b/server/server.go
@@ -422,11 +422,14 @@ func (s *Server) onConn(conn *clientConn) {
 	s.rwlock.Unlock()
 	metrics.ConnGauge.Set(float64(connections))
 
+	if plugin.IsEnable(plugin.Audit) {
+		conn.ctx.GetSessionVars().ConnectionInfo = conn.connectInfo()
+	}
 	err := plugin.ForeachPlugin(plugin.Audit, func(p *plugin.Plugin) error {
 		authPlugin := plugin.DeclareAuditManifest(p.Manifest)
 		if authPlugin.OnConnectionEvent != nil {
-			connInfo := conn.connectInfo()
-			return authPlugin.OnConnectionEvent(context.Background(), conn.ctx.GetSessionVars().User, plugin.Connected, connInfo)
+			sessionVars := conn.ctx.GetSessionVars()
+			return authPlugin.OnConnectionEvent(context.Background(), sessionVars.User, plugin.Connected, sessionVars.ConnectionInfo)
 		}
 		return nil
 	})
@@ -440,9 +443,9 @@ func (s *Server) onConn(conn *clientConn) {
 	err = plugin.ForeachPlugin(plugin.Audit, func(p *plugin.Plugin) error {
 		authPlugin := plugin.DeclareAuditManifest(p.Manifest)
 		if authPlugin.OnConnectionEvent != nil {
-			connInfo := conn.connectInfo()
-			connInfo.Duration = float64(time.Since(connectedTime)) / float64(time.Millisecond)
-			err := authPlugin.OnConnectionEvent(context.Background(), conn.ctx.GetSessionVars().User, plugin.Disconnect, connInfo)
+			sessionVars := conn.ctx.GetSessionVars()
+			sessionVars.ConnectionInfo.Duration = float64(time.Since(connectedTime)) / float64(time.Millisecond)
+			err := authPlugin.OnConnectionEvent(context.Background(), sessionVars.User, plugin.Disconnect, sessionVars.ConnectionInfo)
 			if err != nil {
 				logutil.Logger(context.Background()).Warn("do connection event failed", zap.String("plugin", authPlugin.Name), zap.Error(err))
 			}
diff --git a/sessionctx/variable/session.go b/sessionctx/variable/session.go
index 83d060a814f93..045445b9300aa 100644
--- a/sessionctx/variable/session.go
+++ b/sessionctx/variable/session.go
@@ -382,6 +382,9 @@ type SessionVars struct {
 
 	// Killed is a flag to indicate that this query is killed.
 	Killed uint32
+
+	// ConnectionInfo indicates current connection info used by current session, only be lazy assigned by plugin.
+	ConnectionInfo *ConnectionInfo
 }
 
 // ConnectionInfo present connection used by audit.