Merge pull request #2319 from seanlaff/customNetListener

Add net.Listener as server config option

Author: Hydrocharged

go.mod

@@ -25,6 +25,7 @@ require (
 	golang.org/x/sys v0.12.0
 	golang.org/x/text v0.6.0
 	golang.org/x/tools v0.13.0
+	google.golang.org/grpc v1.53.0
 	gopkg.in/src-d/go-errors.v1 v1.0.0
 	gopkg.in/yaml.v3 v3.0.1
 )
@@ -38,7 +39,6 @@ require (
 	github.com/tetratelabs/wazero v1.1.0 // indirect
 	golang.org/x/mod v0.12.0 // indirect
 	google.golang.org/genproto v0.0.0-20230110181048-76db0878b65f // indirect
-	google.golang.org/grpc v1.53.0 // indirect
 	google.golang.org/protobuf v1.28.1 // indirect
 	gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b // indirect
 )

server/server.go

@@ -143,18 +143,21 @@ func newServerFromHandler(cfg Config, e *sqle.Engine, sm *SessionManager, handle
 		cfg.MaxConnections = 0
 	}
 
+	l := cfg.Listener
 	var unixSocketInUse error
+	if l == nil {
+		if portInUse(cfg.Address) {
+			unixSocketInUse = fmt.Errorf("Port %s already in use.", cfg.Address)
+		}
 
-	if portInUse(cfg.Address) {
-		unixSocketInUse = fmt.Errorf("Port %s already in use.", cfg.Address)
-	}
-
-	l, err := NewListener(cfg.Protocol, cfg.Address, cfg.Socket)
-	if err != nil {
-		if errors.Is(err, UnixSocketInUseError) {
-			unixSocketInUse = err
-		} else {
-			return nil, err
+		var err error
+		l, err = NewListener(cfg.Protocol, cfg.Address, cfg.Socket)
+		if err != nil {
+			if errors.Is(err, UnixSocketInUseError) {
+				unixSocketInUse = err
+			} else {
+				return nil, err
+			}
 		}
 	}
 

server/server_config.go

@@ -16,6 +16,7 @@ package server
 
 import (
 	"crypto/tls"
+	"net"
 	"time"
 
 	"github.com/dolthub/vitess/go/mysql"
@@ -43,6 +44,9 @@ type Config struct {
 	Protocol string
 	// Address of the server.
 	Address string
+	// Custom listener for the mysql server. Use this if you don't want ports or unix sockets to be opened automatically.
+	// This can be useful in testing by using a pure go net.Conn implementation.
+	Listener net.Listener
 	// Tracer to use in the server. By default, a noop tracer will be used if
 	// no tracer is provided.
 	Tracer trace.Tracer

server/server_test.go

@@ -0,0 +1,94 @@
+package server_test
+
+import (
+	"context"
+	"database/sql"
+	"net"
+	"testing"
+	"time"
+
+	vsql "github.com/dolthub/vitess/go/mysql"
+	"github.com/go-sql-driver/mysql"
+	"github.com/stretchr/testify/require"
+
+	sqle "github.com/dolthub/go-mysql-server"
+	"github.com/dolthub/go-mysql-server/memory"
+	"github.com/dolthub/go-mysql-server/server"
+	gsql "github.com/dolthub/go-mysql-server/sql"
+	"github.com/dolthub/go-mysql-server/sql/mysql_db"
+	"google.golang.org/grpc/test/bufconn"
+)
+
+// TestSeverCustomListener verifies a caller can provide their own net.Conn implementation for the server to use
+func TestSeverCustomListener(t *testing.T) {
+	dbName := "mydb"
+	// create a net.Conn thats based on a golang buffer
+	buffer := 1024
+	listener := bufconn.Listen(buffer)
+
+	// create the memory database
+	memdb := memory.NewDatabase(dbName)
+	pro := memory.NewDBProvider(memdb)
+	engine := sqle.NewDefault(pro)
+
+	// server config with custom listener
+	cfg := server.Config{Listener: listener}
+	// since we're using a memory db, we can't rely on server.DefaultSessionBuilder as it causes panics, so explicitly build a memorySessionBuilder
+	sessionBuilder := func(ctx context.Context, c *vsql.Conn, addr string) (gsql.Session, error) {
+		host := ""
+		user := ""
+		mysqlConnectionUser, ok := c.UserData.(mysql_db.MysqlConnectionUser)
+		if ok {
+			host = mysqlConnectionUser.Host
+			user = mysqlConnectionUser.User
+		}
+		client := gsql.Client{Address: host, User: user, Capabilities: c.Capabilities}
+		return memory.NewSession(gsql.NewBaseSessionWithClientServer(addr, client, c.ConnectionID), pro), nil
+	}
+	s, err := server.NewServer(cfg, engine, sessionBuilder, nil)
+	require.NoError(t, err)
+
+	networkName := "testNetwork"
+	// wire up go-mysql-driver to the listener
+	mysql.RegisterDialContext(networkName, func(ctx context.Context, addr string) (net.Conn, error) {
+		return listener.DialContext(ctx)
+	})
+	driver, err := mysql.NewConnector(&mysql.Config{
+		DBName:               dbName,
+		Addr:                 "bufconn",
+		Net:                  networkName,
+		Passwd:               "",
+		User:                 "root",
+		AllowNativePasswords: true,
+	})
+	require.NoError(t, err)
+
+	// start go-mysql-server
+	go func() {
+		err := s.Start()
+		require.NoError(t, err)
+	}()
+
+	// open the db, ping it, and run some execs/queries
+	db := sql.OpenDB(driver)
+
+	var pingErr error
+	for i := 0; i < 3; i++ {
+		if pingErr = db.Ping(); pingErr == nil {
+			break
+		}
+		time.Sleep(time.Second)
+	}
+	require.NoError(t, pingErr)
+
+	_, err = db.Exec("CREATE TABLE table1 (id int)")
+	require.NoError(t, err)
+
+	row := db.QueryRow("SHOW TABLES")
+	var tableName string
+	err = row.Scan(&tableName)
+	require.NoError(t, err)
+	if tableName != "table1" {
+		t.Fatalf("expected to find table1, but got %s", tableName)
+	}
+}

sql/mysql_db/mysql_db.go

@@ -769,9 +769,14 @@ func (db *MySQLDb) AuthMethod(user, addr string) (string, error) {
 	} else {
 		splitHost, _, err := net.SplitHostPort(addr)
 		if err != nil {
-			return "", err
+			if err.(*net.AddrError).Err == "missing port in address" {
+				host = addr
+			} else {
+				return "", err
+			}
+		} else {
+			host = splitHost
 		}
-		host = splitHost
 	}
 
 	rd := db.Reader()
@@ -801,7 +806,11 @@ func (db *MySQLDb) ValidateHash(salt []byte, user string, authResponse []byte, a
 	} else {
 		host, _, err = net.SplitHostPort(addr.String())
 		if err != nil {
-			return nil, err
+			if err.(*net.AddrError).Err == "missing port in address" {
+				host = addr.String()
+			} else {
+				return nil, err
+			}
 		}
 	}
 
@@ -837,7 +846,11 @@ func (db *MySQLDb) Negotiate(c *mysql.Conn, user string, addr net.Addr) (mysql.G
 	} else {
 		host, _, err = net.SplitHostPort(addr.String())
 		if err != nil {
-			return nil, err
+			if err.(*net.AddrError).Err == "missing port in address" {
+				host = addr.String()
+			} else {
+				return nil, err
+			}
 		}
 	}