go ldap连接池

简介

场景：写了个定时任务，每天下午5点根据HR数据对全量LDAP用户作更新，因为LDAP协议是一种可以快速读但修改操作慢的协议，因此定时任务执行时间会很长。于是决定并发地对全量用户进行更新，假设并发数是20，那么一次性会开始创建20个LDAP连接，因为LDAP连接速度较慢，会报错LDAP Result Code 200 \"Network Error\": ldap: connection is in startls phase。

通过思考和调研，决定仿照MySQL、Redis连接池那样在主进程启动后先初始化LDAP连接池。找到并改写5年前一个国外的老仓库，并测试通过，完美地解决了问题。现将源码和使用方式给出，后面也会作优化~

源码： RandolphCYG/ldapPool

使用方式

// 拉下包
go get "github.com/RandolphCYG/ldapPool"
// 导入
import ldappool "github.com/RandolphCYG/ldapPool"

示例

在自己的项目中写一个Init函数，主进程中解析出来ldap服务器连接信息后初始化ldap连接池，得到一个全局的连接池对象LdapPool，用到的地方用LdapPool.Get()获取连接即可。后续开发若有并发修改ldap信息的需求，即可快速从连接池取连接，作并发操作。

package main

import (
	"crypto/tls"
	"fmt"
	"time"
	"unicode/utf16"

	ldappool "github.com/RandolphCYG/ldapPool"
	"github.com/go-ldap/ldap/v3"
	"github.com/sirupsen/logrus"
	"gorm.io/gorm"
)

func main() {
	c := &ldapCfg{
		ConnUrl:       "ldap://192.168.x.xx:389",
		SslEncryption: true,
		Timeout:       5 * time.Second,
		BaseDn:        "DC=xxx,DC=com",
		AdminAccount:  "CN=Admin,CN=Users,DC=xxx,DC=com",
		Password:      "xxxxx",
	}
	// 初始化配置ldap连接池
	Init(c)

	// 待查询用户
	ldapUser := &LdapAttributes{
		Num:         "工号",
		DisplayName: "姓名",
	}

	entry, err := FetchUser(ldapUser)
	if err != nil {
		logrus.Error("fail to fetch ldap user", err)
	}
	fmt.Println(entry.DN)
}

// LdapConnCfg LDAP服务器连接配置
type ldapCfg struct {
	gorm.Model
	// 连接地址
	ConnUrl string `json:"conn_url" gorm:"type:varchar(255);unique_index;not null;comment:连接地址 逻辑外键"`
	// SSL加密方式
	SslEncryption bool `json:"ssl_encryption" gorm:"type:tinyint;length:1;comment:SSL加密方式"`
	// 超时设置
	Timeout time.Duration `json:"timeout" gorm:"type:int;comment:超时设置"`
	// 根目录
	BaseDn string `json:"base_dn" gorm:"type:varchar(255);not null;comment:根目录"`
	// 用户名
	AdminAccount string `json:"admin_account" gorm:"type:varchar(255);not null;comment:用户名"`
	// 密码
	Password string `json:"password" gorm:"type:varchar(255);not null;comment:密码"`
}

var LdapCfg *ldapCfg
var LdapPool ldappool.Pool

// Init 初始化连接池
func Init(c *ldapCfg) (err error) {
	LdapCfg = &ldapCfg{
		ConnUrl:       c.ConnUrl,
		SslEncryption: c.SslEncryption,
		Timeout:       c.Timeout,
		BaseDn:        c.BaseDn,
		AdminAccount:  c.AdminAccount,
		Password:      c.Password,
	}
	// 初始化ldap连接池 TODO 待确认参数
	n := utf16.Encode([]rune("3"))
	LdapPool, err = ldappool.NewChannelPool(50, 1000, "test",
		func(s string) (ldap.Client, error) {
			conn, err := ldap.DialURL(LdapCfg.ConnUrl)
			if err != nil {
				logrus.Error("Fail to dial ldap url, err: ", err)
			}

			// 重新连接TLS
			if err = conn.StartTLS(&tls.Config{InsecureSkipVerify: true}); err != nil {
				logrus.Error("Fail to start tls, err: ", err)
			}

			// 与只读用户绑定
			if err = conn.Bind(LdapCfg.AdminAccount, LdapCfg.Password); err != nil {
				logrus.Error("admin user auth failed, err: ", err)
			}
			return conn, nil
		}, n)
	if err != nil {
		logrus.Error(err.Error())
	}
	return
}

type LdapAttributes struct {
	// ldap字段
	Num         string `json:"employeeNumber" gorm:"type:varchar(100);unique_index"`    // 工号
	Sam         string `json:"sAMAccountName" gorm:"type:varchar(128);unique_index"`    // SAM账号
	Dn          string `json:"distinguishedName" gorm:"type:varchar(100);unique_index"` // dn
	AccountCtl  string `json:"UserAccountControl" gorm:"type:varchar(100)"`             // 用户账户控制
	Expire      int64  `json:"accountExpires" gorm:"type:int(30)"`                      // 账户过期时间
	PwdLastSet  string `json:"pwdLastSet" gorm:"type:varchar(100)"`                     // 用户下次登录必须修改密码
	WhenCreated string `json:"whenCreated" gorm:"type:varchar(100)"`                    // 创建时间
	WhenChanged string `json:"whenChanged" gorm:"type:varchar(100)"`                    // 修改时间
	DisplayName string `json:"displayName" gorm:"type:varchar(32)"`                     // 真实姓名
	Sn          string `json:"sn" gorm:"type:varchar(100)"`                             // 姓
	Name        string `json:"name" gorm:"type:varchar(100)"`                           // 姓名
	GivenName   string `json:"givenName" gorm:"type:varchar(100)"`                      // 名
	Email       string `json:"mail" gorm:"type:varchar(128);unique_index"`              // 邮箱
	Phone       string `json:"mobile" gorm:"type:varchar(32);unique_index"`             // 移动电话
	Company     string `json:"company" gorm:"type:varchar(128)"`                        // 公司
	Depart      string `json:"department" gorm:"type:varchar(128)"`                     // 部门
	Title       string `json:"title" gorm:"type:varchar(100)"`                          // 职务
}

var attrs = []string{
	"employeeNumber",     // 工号
	"sAMAccountName",     // SAM账号
	"distinguishedName",  // dn
	"UserAccountControl", // 用户账户控制
	"accountExpires",     // 账户过期时间
	"pwdLastSet",         // 用户下次登录必须修改密码
	"whenCreated",        // 创建时间
	"whenChanged",        // 修改时间
	"displayName",        // 显示名
	"sn",                 // 姓
	"name",
	"givenName",  // 名
	"mail",       // 邮箱
	"mobile",     // 手机号
	"company",    // 公司
	"department", // 部门
	"title",      // 职务
	"cn",         // common name
}

/* 根据cn查询用户 注意: cn查询不到则会返回管理员用户
 * 这里的查询条件必须保证每个用户必须有
 * 根据cn查询用户 [sam登录名字段也出现了不同的版本 邮箱\手机号都可能更换掉 真实姓名存在重复可能]
 */
func FetchUser(user *LdapAttributes) (result *ldap.Entry, err error) {
	// 获取连接
	LdapConn, err := LdapPool.Get()
	if err != nil {
		logrus.Error("Fail to get ldap connection, err: ", err)
		return
	}
	defer LdapConn.Close()

	ldapFilterCn := "(cn=" + user.DisplayName + user.Num + ")"
	searchFilter := "(objectClass=organizationalPerson)"

	if user.DisplayName != "" && user.Num != "" {
		searchFilter += ldapFilterCn
	}
	searchFilter = "(&" + searchFilter + ")"

	searchRequest := ldap.NewSearchRequest(
		LdapCfg.BaseDn,
		ldap.ScopeWholeSubtree, ldap.NeverDerefAliases, 0, 0, false,
		searchFilter,
		attrs,
		nil,
	)

	// 这里LdapConn 为nil
	sr, err := LdapConn.Search(searchRequest)
	if err != nil {
		logrus.Error("Fail to fetch user, err: ", err)
		return
	}
	if len(sr.Entries) > 0 && len(sr.Entries[0].Attributes) > 0 {
		result = sr.Entries[0]
	}
	return
}