From f42c3ac24ac029280d4d409f84fc3492a9aa9676 Mon Sep 17 00:00:00 2001
From: lysu <sulifx@gmail.com>
Date: Mon, 3 Feb 2020 12:34:46 +0800
Subject: [PATCH] plugin: Support logging rejected connection attempts in audit
 log

---
 plugin/audit.go                           |  9 +++++++++
 plugin/conn_ip_example/conn_ip_example.go | 11 +++++++++++
 plugin/conn_ip_example/manifest.toml      |  3 ++-
 server/server.go                          | 12 ++++++++++++
 4 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/plugin/audit.go b/plugin/audit.go
index 644ea591b349e..2150cd83c03aa 100644
--- a/plugin/audit.go
+++ b/plugin/audit.go
@@ -45,6 +45,8 @@ const (
 	ChangeUser
 	// PreAuth presents event before start auth.
 	PreAuth
+	// Reject presents event reject connection event.
+	Reject
 )
 
 func (c ConnectionEvent) String() string {
@@ -57,6 +59,8 @@ func (c ConnectionEvent) String() string {
 		return "ChangeUser"
 	case PreAuth:
 		return "PreAuth"
+	case Reject:
+		return "Reject"
 	}
 	return ""
 }
@@ -85,6 +89,11 @@ type AuditManifest struct {
 	OnParseEvent func(ctx context.Context, sctx *variable.SessionVars, event ParseEvent) error
 }
 
+type (
+	// RejectReasonCtxValue will be used in OnConnectionEvent to pass RejectReason to plugin.
+	RejectReasonCtxValue struct{}
+)
+
 type execStartTimeCtxKeyType struct{}
 
 // ExecStartTimeCtxKey indicates stmt start execution time.
diff --git a/plugin/conn_ip_example/conn_ip_example.go b/plugin/conn_ip_example/conn_ip_example.go
index 5aca75690d317..3a40178bccd05 100644
--- a/plugin/conn_ip_example/conn_ip_example.go
+++ b/plugin/conn_ip_example/conn_ip_example.go
@@ -46,3 +46,14 @@ func OnGeneralEvent(ctx context.Context, sctx *variable.SessionVars, event plugi
 	fmt.Println("variable test: ", variable.GetSysVar("conn_ip_example_test_variable").Value)
 	fmt.Printf("new connection by %s\n", ctx.Value("ip"))
 }
+
+// OnConnectionEvent implements TiDB Audit plugin's OnConnectionEvent SPI.
+func OnConnectionEvent(ctx context.Context, event plugin.ConnectionEvent, info *variable.ConnectionInfo) error {
+	var reason string
+	if r := ctx.Value(plugin.RejectReasonCtxValue{}); r != nil {
+		reason = r.(string)
+	}
+	fmt.Println("conn_ip_example onConnect called")
+	fmt.Printf("conenct event: %s, reason: %s\n", event, reason)
+	return nil
+}
diff --git a/plugin/conn_ip_example/manifest.toml b/plugin/conn_ip_example/manifest.toml
index b57badaf689f0..2cbebb6a47f98 100644
--- a/plugin/conn_ip_example/manifest.toml
+++ b/plugin/conn_ip_example/manifest.toml
@@ -11,5 +11,6 @@ validate = "Validate"
 onInit = "OnInit"
 onShutdown = "OnShutdown"
 export = [
-    {extPoint="OnGeneralEvent", impl="OnGeneralEvent"}
+    {extPoint="OnGeneralEvent", impl="OnGeneralEvent"},
+    {extPoint="OnConnectionEvent", impl="OnConnectionEvent"}
 ]
diff --git a/server/server.go b/server/server.go
index 2aada653dba27..879a995ffdcdc 100644
--- a/server/server.go
+++ b/server/server.go
@@ -413,6 +413,18 @@ func (s *Server) Close() {
 func (s *Server) onConn(conn *clientConn) {
 	ctx := logutil.WithConnID(context.Background(), conn.connectionID)
 	if err := conn.handshake(ctx); err != nil {
+		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 {
+				pluginCtx := context.WithValue(context.Background(), plugin.RejectReasonCtxValue{}, err.Error())
+				return authPlugin.OnConnectionEvent(pluginCtx, plugin.Reject, conn.ctx.GetSessionVars().ConnectionInfo)
+			}
+			return nil
+		})
+		terror.Log(err)
 		// Some keep alive services will send request to TiDB and disconnect immediately.
 		// So we only record metrics.
 		metrics.HandShakeErrorCounter.Inc()