ldap-1-backend/public/common/ldap.go

/*
Package common 提供 LDAP 连接池管理功能

LDAP 连接池是系统与 OpenLDAP 服务器交互的核心组件，负责：
- 管理 LDAP 连接的创建、复用和销毁
- 提供线程安全的连接获取和归还机制
- 控制最大连接数，避免资源耗尽
- 支持连接的健康检查和自动重连

连接池的设计目标：
1. 提高 LDAP 操作的性能，避免频繁建立连接
2. 控制并发连接数，保护 LDAP 服务器资源
3. 提供稳定可靠的连接管理机制
4. 支持优雅的连接回收和错误处理
*/
package common

import (
	"fmt"
	"log"
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/eryajf/go-ldap-admin/config"

	ldap "github.com/go-ldap/ldap/v3"
)

// LDAP 连接池相关全局变量
var (
	ldapPool    *LdapConnPool // LDAP 连接池实例
	ldapInit    = false       // LDAP 初始化标志
	ldapInitOne sync.Once     // 确保 LDAP 只初始化一次
)

// InitLDAP 初始化 LDAP 连接池
// 该函数在系统启动时调用，负责：
// 1. 建立与 LDAP 服务器的初始连接
// 2. 验证管理员账户凭据
// 3. 初始化连接池结构
// 4. 设置连接池参数
func InitLDAP() {
	// 防止重复初始化
	if ldapInit {
		return
	}

	// 使用 sync.Once 确保只初始化一次
	ldapInitOne.Do(func() {
		ldapInit = true
	})

	// ==================== 建立初始 LDAP 连接 ====================

	// 创建 LDAP 连接，设置 5 秒连接超时
	ldapConn, err := ldap.DialURL(config.Conf.Ldap.Url, ldap.DialWithDialer(&net.Dialer{Timeout: 5 * time.Second}))
	if err != nil {
		Log.Panicf("初始化 LDAP 连接异常: %v", err)
		panic(fmt.Errorf("初始化 LDAP 连接异常: %v", err))
	}

	// 使用管理员账户绑定 LDAP 连接
	err = ldapConn.Bind(config.Conf.Ldap.AdminDN, config.Conf.Ldap.AdminPass)
	if err != nil {
		Log.Panicf("绑定 LDAP 管理员账号异常: %v", err)
		panic(fmt.Errorf("绑定 LDAP 管理员账号异常: %v", err))
	}

	// ==================== 初始化连接池 ====================

	// 创建连接池实例
	ldapPool = &LdapConnPool{
		conns:    make([]*ldap.Conn, 0),                    // 可用连接切片
		reqConns: make(map[uint64]chan *ldap.Conn),         // 等待连接的请求映射
		openConn: 0,                                        // 当前打开的连接数
		maxOpen:  config.Conf.Ldap.MaxConn,                 // 最大连接数
	}

	// 将初始连接放入连接池
	PutLADPConn(ldapConn)

	// ==================== 输出初始化信息 ====================

	// 构建显示用的 DSN（隐藏密码）
	showDsn := fmt.Sprintf(
		"%s:******@tcp(%s)",
		config.Conf.Ldap.AdminDN,
		config.Conf.Ldap.Url,
	)

	Log.Info("🔗 LDAP 连接池初始化完成! DSN: ", showDsn)
}

// GetLDAPConn 从连接池获取 LDAP 连接
// 返回一个可用的 LDAP 连接，使用完毕后需要调用 PutLADPConn 归还
func GetLDAPConn() (*ldap.Conn, error) {
	return ldapPool.GetConnection()
}

// PutLADPConn 将 LDAP 连接归还到连接池
// 连接使用完毕后必须调用此方法归还，以便其他请求复用
func PutLADPConn(conn *ldap.Conn) {
	ldapPool.PutConnection(conn)
}

// LdapConnPool LDAP 连接池结构体
// 实现了线程安全的连接池管理机制
type LdapConnPool struct {
	mu       sync.Mutex                    // 互斥锁，保证线程安全
	conns    []*ldap.Conn                  // 可用连接切片
	reqConns map[uint64]chan *ldap.Conn    // 等待连接的请求映射表
	openConn int                           // 当前打开的连接数
	maxOpen  int                           // 最大允许的连接数
}

// GetConnection 从连接池获取一个可用的 LDAP 连接
// 该方法实现了连接池的核心逻辑：
// 1. 优先从池中获取现有连接
// 2. 检查连接健康状态
// 3. 控制最大连接数限制
// 4. 支持连接等待队列
func (lcp *LdapConnPool) GetConnection() (*ldap.Conn, error) {
	lcp.mu.Lock()

	// ==================== 从连接池获取现有连接 ====================

	// 检查连接池中是否有可用连接
	connNum := len(lcp.conns)
	if connNum > 0 {
		lcp.openConn++                                    // 增加打开连接计数
		conn := lcp.conns[0]                             // 获取第一个连接
		copy(lcp.conns, lcp.conns[1:])                   // 移除已获取的连接
		lcp.conns = lcp.conns[:connNum-1]                // 调整切片长度

		lcp.mu.Unlock()

		// 检查连接是否已关闭，如果已关闭则创建新连接
		if conn.IsClosing() {
			return initLDAPConn()
		}
		return conn, nil
	}

	// ==================== 连接数限制处理 ====================

	// 当连接池为空且已达到最大连接数限制时
	if lcp.maxOpen != 0 && lcp.openConn >= lcp.maxOpen {
		// 创建等待队列，当有连接归还时会通知等待的请求
		req := make(chan *ldap.Conn, 1)
		reqKey := lcp.nextRequestKeyLocked()             // 生成唯一请求键
		lcp.reqConns[reqKey] = req                       // 将请求加入等待队列
		lcp.mu.Unlock()

		// 阻塞等待连接归还
		return <-req, nil
	} else {
		// 未达到连接数限制，创建新连接
		lcp.openConn++
		lcp.mu.Unlock()
		return initLDAPConn()
	}
}

// PutConnection 将 LDAP 连接归还到连接池
// 该方法负责连接的回收和分配：
// 1. 优先满足等待队列中的请求
// 2. 将健康的连接放回连接池
// 3. 丢弃已关闭的连接
func (lcp *LdapConnPool) PutConnection(conn *ldap.Conn) {
	log.Println("🔄 归还一个 LDAP 连接到连接池")
	lcp.mu.Lock()
	defer lcp.mu.Unlock()

	// ==================== 优先处理等待队列 ====================

	// 检查是否有等待连接的请求
	if num := len(lcp.reqConns); num > 0 {
		var req chan *ldap.Conn
		var reqKey uint64

		// 获取第一个等待请求
		for reqKey, req = range lcp.reqConns {
			break
		}

		// 从等待队列中移除该请求
		delete(lcp.reqConns, reqKey)

		// 将连接发送给等待的请求
		req <- conn
		return
	}

	// ==================== 连接回收处理 ====================

	// 没有等待请求时，将连接放回连接池
	lcp.openConn--                                      // 减少打开连接计数

	// 只有健康的连接才放回连接池
	if !conn.IsClosing() {
		lcp.conns = append(lcp.conns, conn)
	}
}

// nextRequestKeyLocked 生成下一个唯一的请求键
// 用于标识等待队列中的请求，确保每个等待请求都有唯一标识
// 注意：调用此方法时必须已经持有互斥锁
func (lcp *LdapConnPool) nextRequestKeyLocked() uint64 {
	for {
		reqKey := rand.Uint64()                         // 生成随机数作为请求键
		if _, ok := lcp.reqConns[reqKey]; !ok {         // 确保键的唯一性
			return reqKey
		}
	}
}

// 获取 ladp 连接
func initLDAPConn() (*ldap.Conn, error) {
	ldap, err := ldap.DialURL(config.Conf.Ldap.Url, ldap.DialWithDialer(&net.Dialer{Timeout: 5 * time.Second}))
	if err != nil {
		return nil, err
	}
	err = ldap.Bind(config.Conf.Ldap.AdminDN, config.Conf.Ldap.AdminPass)
	if err != nil {
		return nil, err
	}
	return ldap, err
}