add some basic test

Signed-off-by: Yang Keao <[email protected]>

Author: YangKeao

privilege/privileges/ldap/BUILD.bazel

@@ -21,8 +21,14 @@ go_library(
 go_test(
     name = "ldap_test",
     timeout = "short",
-    srcs = ["ldap_common_test.go"],
+    srcs = [
+        "ldap_common_test.go",
+        "mock_ldap_server_test.go",
+    ],
     embed = [":ldap"],
     flaky = True,
-    deps = ["@com_github_stretchr_testify//require"],
+    deps = [
+        "@com_github_go_ldap_ldap_v3//:ldap",
+        "@com_github_stretchr_testify//require",
+    ],
 )

privilege/privileges/ldap/ldap_common.go

@@ -143,6 +143,8 @@ func (impl *ldapAuthImpl) getConnection() (*ldap.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
+
+	// FIXME: if the connection in the pool is killed or timeout, re-initialize the connection.
 	return conn.(*ldap.Conn), nil
 }
 

privilege/privileges/ldap/ldap_common_test.go

@@ -15,8 +15,14 @@
 package ldap
 
 import (
-	"github.com/stretchr/testify/require"
+	"context"
+	"math/rand"
+	"net"
 	"testing"
+	"time"
+
+	"github.com/go-ldap/ldap/v3"
+	"github.com/stretchr/testify/require"
 )
 
 const listenPortRangeStart = 11310
@@ -29,3 +35,58 @@ func TestCanonicalizeDN(t *testing.T) {
 	require.Equal(t, impl.canonicalizeDN("yka", "cn=y,dc=ping,dc=cap"), "cn=y,dc=ping,dc=cap")
 	require.Equal(t, impl.canonicalizeDN("yka", "+dc=ping,dc=cap"), "cn=yka,dc=ping,dc=cap")
 }
+
+func getConnectionsWithinTimeout(t *testing.T, ch chan net.Conn, timeout time.Duration, count int) []net.Conn {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	result := []net.Conn{}
+	for count > 0 {
+		count--
+
+		select {
+		case conn := <-ch:
+			result = append(result, conn)
+		case <-ctx.Done():
+			require.Fail(t, "fail to get connections")
+		}
+	}
+
+	return result
+}
+
+func TestLDAPConnectionPool(t *testing.T) {
+	// allocate a random port between the port range
+	port := rand.Int()%listenPortRangeLength + listenPortRangeStart
+	ldapServer := NewMockLDAPServer(t)
+	ldapServer.Listen(port)
+	defer ldapServer.Close()
+	conns := ldapServer.GetConnections()
+
+	impl := &ldapAuthImpl{ldapServerHost: "localhost", ldapServerPort: port}
+	impl.SetInitCapacity(256)
+	impl.SetMaxCapacity(1024)
+	conn, err := impl.getConnection()
+	require.NoError(t, err)
+	impl.putConnection(conn)
+
+	getConnectionsWithinTimeout(t, conns, time.Second, 1)
+
+	// test allocating 255 more connections
+	var clientConnections []*ldap.Conn
+	for i := 0; i < 256; i++ {
+		conn, err := impl.getConnection()
+		require.NoError(t, err)
+
+		clientConnections = append(clientConnections, conn)
+	}
+	getConnectionsWithinTimeout(t, conns, time.Second, 255)
+	for _, conn := range clientConnections {
+		impl.putConnection(conn)
+	}
+
+	clientConnections = clientConnections[:]
+
+	// now, the max capacity is somehow meaningless
+	// TODO: auto scalling the capacity of LDAP connection pool
+}

privilege/privileges/ldap/mock_ldap_server_test.go

@@ -0,0 +1,85 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ldap
+
+import (
+	"fmt"
+	"net"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// MockLDAPServer is a mock LDAP server to help testing the LDAP authentication
+type MockLDAPServer struct {
+	t  *testing.T
+	wg sync.WaitGroup
+
+	listener net.Listener
+	closed   atomic.Bool
+
+	connections chan net.Conn
+}
+
+// NewMockLDAPServer creates the MockLDAPServer
+func NewMockLDAPServer(t *testing.T) *MockLDAPServer {
+	return &MockLDAPServer{
+		t:           t,
+		connections: make(chan net.Conn),
+	}
+}
+
+// Listen listens on the specific port
+func (s *MockLDAPServer) Listen(port int) {
+	l, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", port))
+	require.NoError(s.t, err)
+
+	s.wg.Add(1)
+	go func() {
+		for {
+			conn, err := l.Accept()
+			if err != nil {
+				if s.closed.Load() {
+					break
+				}
+
+				require.NoError(s.t, err)
+			}
+			go func() {
+				s.connections <- conn
+			}()
+		}
+		s.wg.Done()
+	}()
+
+	s.listener = l
+	return
+}
+
+// Close closes the listener
+func (s *MockLDAPServer) Close() {
+	s.closed.Store(true)
+	err := s.listener.Close()
+	require.NoError(s.t, err)
+
+	s.wg.Wait()
+}
+
+// GetConnections returns all live connections
+func (s *MockLDAPServer) GetConnections() chan net.Conn {
+	return s.connections
+}

server/conn.go

@@ -203,7 +203,7 @@ func (cc *clientConn) authSwitchRequest(ctx context.Context, plugin string) ([]b
 		clientPlugin = mysql.AuthMySQLClearPassword
 	}
 	failpoint.Inject("FakeAuthSwitch", func() {
-		failpoint.Return([]byte(plugin), nil)
+		failpoint.Return([]byte(clientPlugin), nil)
 	})
 	enclen := 1 + len(clientPlugin) + 1 + len(cc.salt) + 1
 	data := cc.alloc.AllocWithLen(4, enclen)

server/conn_test.go

@@ -1630,13 +1630,13 @@ func TestAuthSessionTokenPlugin(t *testing.T) {
 	// create a token without TLS
 	tk1 := testkit.NewTestKitWithSession(t, store, tc.Session)
 	tc.Session.GetSessionVars().ConnectionInfo = cc.connectInfo()
-	tk1.Session().Auth(&auth.UserIdentity{Username: "auth_session_token", Hostname: "localhost"}, nil, nil)
+	tk1.Session().Auth(&auth.UserIdentity{Username: "auth_session_token", Hostname: "localhost"}, nil, nil, nil)
 	tk1.MustQuery("show session_states")
 
 	// create a token with TLS
 	cc.tlsConn = &tls.Conn{}
 	tc.Session.GetSessionVars().ConnectionInfo = cc.connectInfo()
-	tk1.Session().Auth(&auth.UserIdentity{Username: "auth_session_token", Hostname: "localhost"}, nil, nil)
+	tk1.Session().Auth(&auth.UserIdentity{Username: "auth_session_token", Hostname: "localhost"}, nil, nil, nil)
 	tk1.MustQuery("show session_states")
 
 	// create a token with UnixSocket
@@ -1986,3 +1986,42 @@ func TestProcessInfoForExecuteCommand(t *testing.T) {
 		0x0A, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0}))
 	require.Equal(t, cc.ctx.Session.ShowProcess().Info, "select sum(col1) from t where col1 < ? and col1 > 100")
 }
+
+func TestLDAPAuthSwitch(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	cfg := newTestConfig()
+	cfg.Port = 0
+	cfg.Status.StatusPort = 0
+	drv := NewTiDBDriver(store)
+	srv, err := NewServer(cfg, drv)
+	require.NoError(t, err)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("CREATE USER test_simple_ldap IDENTIFIED WITH authentication_ldap_simple AS 'uid=test_simple_ldap,dc=example,dc=com'")
+
+	cc := &clientConn{
+		connectionID: 1,
+		alloc:        arena.NewAllocator(1024),
+		chunkAlloc:   chunk.NewAllocator(),
+		pkt: &packetIO{
+			bufWriter: bufio.NewWriter(bytes.NewBuffer(nil)),
+		},
+		server: srv,
+		user:   "test_simple_ldap",
+	}
+	se, _ := session.CreateSession4Test(store)
+	tc := &TiDBContext{
+		Session: se,
+		stmts:   make(map[int]*TiDBStatement),
+	}
+	cc.setCtx(tc)
+	cc.isUnixSocket = true
+
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/server/FakeAuthSwitch", "return(1)"))
+	respAuthSwitch, err := cc.checkAuthPlugin(context.Background(), &handshakeResponse41{
+		Capability: mysql.ClientProtocol41 | mysql.ClientPluginAuth,
+		User:       "test_simple_ldap",
+	})
+	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/server/FakeAuthSwitch"))
+	require.NoError(t, err)
+	require.Equal(t, []byte(mysql.AuthMySQLClearPassword), respAuthSwitch)
+}