From e70f9fc4b05575b905e2ec8b97283a903fe43ce0 Mon Sep 17 00:00:00 2001
From: Matt Robenolt <matt@ydekproductions.com>
Date: Mon, 10 Jan 2022 16:29:19 -0800
Subject: [PATCH] mysql: Add ConnectionReady callback to Handler interface

This is is useful to determine the intermediate state after
NewConnection and before ConnectionClosed, basically for a handler to
know if authentication, etc, were successful. NewConnection is way too
soon in the process to determine if the connection got fully
established.

Signed-off-by: Matt Robenolt <matt@ydekproductions.com>
---
 go/mysql/conn_flaky_test.go              | 4 ++++
 go/mysql/fakesqldb/server.go             | 5 +++++
 go/mysql/mysql_fuzzer.go                 | 6 ++++++
 go/mysql/server.go                       | 8 ++++++++
 go/mysql/server_flaky_test.go            | 3 +++
 go/vt/vtgate/plugin_mysql_server.go      | 2 ++
 go/vt/vtgate/plugin_mysql_server_test.go | 3 +++
 7 files changed, 31 insertions(+)

diff --git a/go/mysql/conn_flaky_test.go b/go/mysql/conn_flaky_test.go
index 38f6e6fd3bd..be47b9d0318 100644
--- a/go/mysql/conn_flaky_test.go
+++ b/go/mysql/conn_flaky_test.go
@@ -820,6 +820,10 @@ func (t testRun) NewConnection(c *Conn) {
 	panic("implement me")
 }
 
+func (t testRun) ConnectionReady(c *Conn) {
+	panic("implement me")
+}
+
 func (t testRun) ConnectionClosed(c *Conn) {
 	panic("implement me")
 }
diff --git a/go/mysql/fakesqldb/server.go b/go/mysql/fakesqldb/server.go
index a61f40e8cbf..92f19b2ce3e 100644
--- a/go/mysql/fakesqldb/server.go
+++ b/go/mysql/fakesqldb/server.go
@@ -308,6 +308,11 @@ func (db *DB) NewConnection(c *mysql.Conn) {
 	db.connections[c.ConnectionID] = c
 }
 
+// ConnectionReady is part of the mysql.Handler interface.
+func (db *DB) ConnectionReady(c *mysql.Conn) {
+
+}
+
 // ConnectionClosed is part of the mysql.Handler interface.
 func (db *DB) ConnectionClosed(c *mysql.Conn) {
 	db.mu.Lock()
diff --git a/go/mysql/mysql_fuzzer.go b/go/mysql/mysql_fuzzer.go
index d65169f30a2..a5b41d5f41c 100644
--- a/go/mysql/mysql_fuzzer.go
+++ b/go/mysql/mysql_fuzzer.go
@@ -87,6 +87,9 @@ type fuzztestRun struct{}
 func (t fuzztestRun) NewConnection(c *Conn) {
 }
 
+func (t fuzztestRun) ConnectionReady(c *Conn) {
+}
+
 func (t fuzztestRun) ConnectionClosed(c *Conn) {
 }
 
@@ -275,6 +278,9 @@ func (th *fuzzTestHandler) NewConnection(c *Conn) {
 	th.lastConn = c
 }
 
+func (th *fuzzTestHandler) ConnectionReady(_ *Conn) {
+}
+
 func (th *fuzzTestHandler) ConnectionClosed(_ *Conn) {
 }
 
diff --git a/go/mysql/server.go b/go/mysql/server.go
index 31ebb3c0e07..2651e52f06b 100644
--- a/go/mysql/server.go
+++ b/go/mysql/server.go
@@ -94,6 +94,10 @@ type Handler interface {
 	// In particular, ServerStatusAutocommit might be set.
 	NewConnection(c *Conn)
 
+	// ConnectionReady is called after the connection handshake, but
+	// before we begin to process commands.
+	ConnectionReady(c *Conn)
+
 	// ConnectionClosed is called when a connection is closed.
 	ConnectionClosed(c *Conn)
 
@@ -470,6 +474,10 @@ func (l *Listener) handle(conn net.Conn, connectionID uint32, acceptTime time.Ti
 		log.Warningf("Slow connection from %s: %v", c, connectTime)
 	}
 
+	// Tell our handler that we're finished handshake and are ready to
+	// process commands.
+	l.handler.ConnectionReady(c)
+
 	for {
 		kontinue := c.handleNextCommand(l.handler)
 		if !kontinue {
diff --git a/go/mysql/server_flaky_test.go b/go/mysql/server_flaky_test.go
index 7dd5e694446..864b2bc025d 100644
--- a/go/mysql/server_flaky_test.go
+++ b/go/mysql/server_flaky_test.go
@@ -110,6 +110,9 @@ func (th *testHandler) NewConnection(c *Conn) {
 	th.lastConn = c
 }
 
+func (th *testHandler) ConnectionReady(_ *Conn) {
+}
+
 func (th *testHandler) ConnectionClosed(_ *Conn) {
 }
 
diff --git a/go/vt/vtgate/plugin_mysql_server.go b/go/vt/vtgate/plugin_mysql_server.go
index b31e8465e8c..6de23087f21 100644
--- a/go/vt/vtgate/plugin_mysql_server.go
+++ b/go/vt/vtgate/plugin_mysql_server.go
@@ -102,6 +102,8 @@ func (vh *vtgateHandler) NewConnection(c *mysql.Conn) {
 	vh.connections[c] = true
 }
 
+func (vh *vtgateHandler) ConnectionReady(_ *mysql.Conn) {}
+
 func (vh *vtgateHandler) numConnections() int {
 	vh.mu.Lock()
 	defer vh.mu.Unlock()
diff --git a/go/vt/vtgate/plugin_mysql_server_test.go b/go/vt/vtgate/plugin_mysql_server_test.go
index 1caa14346a1..f86740a54e8 100644
--- a/go/vt/vtgate/plugin_mysql_server_test.go
+++ b/go/vt/vtgate/plugin_mysql_server_test.go
@@ -46,6 +46,9 @@ func (th *testHandler) NewConnection(c *mysql.Conn) {
 	th.lastConn = c
 }
 
+func (th *testHandler) ConnectionReady(c *mysql.Conn) {
+}
+
 func (th *testHandler) ConnectionClosed(c *mysql.Conn) {
 }