160 lines
3.4 KiB
Go
160 lines
3.4 KiB
Go
package common
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"math/rand"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/eryajf/go-ldap-admin/config"
|
|
|
|
ldap "github.com/go-ldap/ldap/v3"
|
|
)
|
|
|
|
var ldapPool *LdapConnPool
|
|
var ldapInit = false
|
|
var ldapInitOne sync.Once
|
|
|
|
// Init 初始化连接
|
|
func InitLDAP() {
|
|
if ldapInit {
|
|
return
|
|
}
|
|
|
|
ldapInitOne.Do(func() {
|
|
ldapInit = true
|
|
})
|
|
|
|
// Dail有两个参数 network, address, 返回 (*Conn, error)
|
|
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))
|
|
}
|
|
err = ldapConn.Bind(config.Conf.Ldap.AdminDN, config.Conf.Ldap.AdminPass)
|
|
if err != nil {
|
|
Log.Panicf("绑定admin账号异常: %v", err)
|
|
panic(fmt.Errorf("绑定admin账号异常: %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)
|
|
|
|
// 隐藏密码
|
|
showDsn := fmt.Sprintf(
|
|
"%s:******@tcp(%s)",
|
|
config.Conf.Ldap.AdminDN,
|
|
config.Conf.Ldap.Url,
|
|
)
|
|
|
|
Log.Info("初始化ldap完成! dsn: ", showDsn)
|
|
}
|
|
|
|
// GetLDAPConn 获取 LDAP 连接
|
|
func GetLDAPConn() (*ldap.Conn, error) {
|
|
return ldapPool.GetConnection()
|
|
}
|
|
|
|
// PutLDAPConn 放回 LDAP 连接
|
|
func PutLADPConn(conn *ldap.Conn) {
|
|
ldapPool.PutConnection(conn)
|
|
}
|
|
|
|
type LdapConnPool struct {
|
|
mu sync.Mutex
|
|
conns []*ldap.Conn
|
|
reqConns map[uint64]chan *ldap.Conn
|
|
openConn int
|
|
maxOpen int
|
|
}
|
|
|
|
// 获取一个 ladp Conn
|
|
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()
|
|
// 发现连接已经 close 重新获取连接
|
|
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()
|
|
}
|
|
}
|
|
|
|
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
|
|
} else {
|
|
lcp.openConn--
|
|
if !conn.IsClosing() {
|
|
lcp.conns = append(lcp.conns, conn)
|
|
}
|
|
}
|
|
}
|
|
|
|
// 获取下一个请求令牌
|
|
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
|
|
}
|