go-gin中session实现redis前缀和db库选择+单点登录

分别实现了redigo中自动加前缀和session中自动加前缀

等有空了整理一个demo放到github上，到时候求个小星星

1. 在gin-contrib/sessions/redis库中redis的前缀是被封装起来了，所以自定义前缀没有内部方法
2. 在这里我们自己实现一下NewStoreWithDBPrefix方法
3. 配置文件可以看到引用的redis配置
4. 这样方面多个项目在同一个redis中方便管理
5. 也方便做单点登录

项目目录

app/
├── common/
│   ├── global/
│   ├────  global.go
├── config/
│   ├── application.yaml
├── initialize/
│   ├── config.go
│   ├── initialize.go
│   ├── redis.go
│   ├── store.go
├── middleware/
│   ├── middleware.go
├── router/
│   ├── router.go
├── main.go

下面是代码实现

配置文件：config/application.yaml

redis: &redisConfig
  addr: 127.0.0.1
  port: 6379
  password: 123456
  db: 3
  size: 20
  max-idle: 24
  active: 10
  auth: true
  prefix: demo
  timeout: 180
store:
  size: 10
  redis: *redisConfig
  key-pairs: sessionKey

全局变量：common/global/global.go

package global

import (
	"demo/config"

	"github.com/gin-contrib/sessions"
	"github.com/robfig/cron/v3"

	"github.com/garyburd/redigo/redis"
	ut "github.com/go-playground/universal-translator"
	"github.com/jordan-wright/email"
	"go.uber.org/zap"
	"gorm.io/gorm"
)

var (
	CONFIG    config.Config
	DB        *gorm.DB
	Trans     *ut.Translator
	REDISPoll *redis.Pool
	RootDir   string
	EmailPool *email.Pool
	Store     *sessions.Store
	Logger    *zap.SugaredLogger
	Cron      *cron.Cron
)

配置文件初始化：initialize/config.go

package initialize

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"

	"demo/common/global"
	"demo/config"

	"github.com/joho/godotenv"
	"gopkg.in/yaml.v2"
)

/**
 * 读取配置文件,最先被初始化
 */

func InitilizeConfig() {
	// 获取当前工作目录
	rootDir, err := os.Getwd()
	if err != nil {
		log.Fatalf("Failed to get current working directory: %s", err)
	}

	// 这是单元测试
	if strings.HasSuffix(rootDir, "service") || strings.HasSuffix(rootDir, "command") || strings.HasSuffix(rootDir, "util") || strings.HasSuffix(rootDir, "model") {
		rootDir = filepath.Dir(rootDir)
	}
	if strings.HasSuffix(rootDir, "oss") {
		rootDir = filepath.Dir(filepath.Dir(rootDir))
	}

	// rootDir = "/Users/dupeisheng.vendor/go/src/gitlab.bj.sensetime.com/entry_manage"
	// 设置工作目录
	if err := os.Chdir(rootDir); err != nil {
		log.Fatalf("Failed to change working directory: %s", err)
	}
	fmt.Println("工作目录：", rootDir)

	// 加载.env
	godotenv.Load(".env")

	// 设置当前环境
	mode := os.Getenv("Mode")
	fmt.Println("当前环境：", mode)
	if mode != "" {
		mode = "-" + mode
	}

	// 加载配置文件
	configPath := rootDir + "/config/application" + mode + ".yaml"
	yamlFile, err := ioutil.ReadFile(configPath)
	if err != nil {
		log.Panicf("Failed to read file , cause is %s", err.Error())
	}
	config := config.Config{}
	err = yaml.Unmarshal(yamlFile, &config)
	global.CONFIG = config
	if err != nil {
		fmt.Println(err.Error())
	}
}

session初始化：initialize/store.go

package initialize

import (
	"fmt"
	"demo/common/global"
	"github.com/boj/redistore"
	"github.com/gin-contrib/sessions"
	"github.com/gin-contrib/sessions/redis"
	"github.com/spf13/cast"
)

const SessionMaxAge = 8 * 3600

func InitilizeStore() {
	var Store sessions.Store
	Store, err := NewStoreWithDBPrefix(
		global.CONFIG.Store.Size,
		"tcp",
		global.CONFIG.Store.Redis.Addr+":"+global.CONFIG.Store.Redis.Port,
		global.CONFIG.Store.Redis.Password,
		cast.ToString(global.CONFIG.Store.Redis.Db),
		// []byte(global.CONFIG.Store.Redis.Prefix),
		[]byte(global.CONFIG.Store.KeyPairs),
	)
	if err != nil {
		fmt.Println("创建Session-Redis存储失败" + err.Error())
	}
	Store.Options(sessions.Options{
		MaxAge: SessionMaxAge,
	})
	global.Store = &Store
}

type store struct {
	*redistore.RediStore
}

func NewStoreWithDBPrefix(size int, network, address, password, DB string, keyPairs ...[]byte) (redis.Store, error) {
	s, err := redistore.NewRediStoreWithDB(size, network, address, password, DB, keyPairs...)
	if err != nil {
		return nil, err
	}
	// 这里设置前缀
	s.SetKeyPrefix(global.CONFIG.Store.Redis.Prefix + ":session:")
	return &store{s}, nil
}

func (c *store) Options(options sessions.Options) {
	c.RediStore.Options = options.ToGorillaOptions()
}

redis初始化：initialize/redis.go

package initialize

import (
	"fmt"
	"log"
	"strings"
	"time"

	"demo/common/global"
	"demo/config"

	"github.com/garyburd/redigo/redis"
	"github.com/samber/lo"
	"github.com/spf13/cast"
)

/**
 * 初始化redis，并赋值给全局变量
 */

func InitilizeRedis() {
	// 创建redis连接池
	global.REDISPoll = GetRedisPool(global.CONFIG)
}

func GetRedisPool(config config.Config) *redis.Pool {
	return &redis.Pool{
		MaxIdle:     config.Redis.MaxIdle, // 最大空闲连接数
		MaxActive:   config.Redis.Active,  // 最大连接数
		IdleTimeout: time.Duration(config.Redis.Timeout) * time.Second,
		Wait:        true, // 超过连接数后是否等待
		Dial: func() (redis.Conn, error) {
			redisUri := fmt.Sprintf("%s:%s", config.Redis.Addr, config.Redis.Port)
			var redisConn redis.Conn
			var err error

			if config.Redis.Auth {
				redisConn, err = redis.Dial("tcp", redisUri, redis.DialPassword(config.Redis.Password))
			} else {
				redisConn, err = redis.Dial("tcp", redisUri)
			}

			if err != nil {
				log.Println("获取连接失败:" + err.Error())
				return nil, err
			}

			// 添加 Redis 前缀
			if config.Redis.Prefix != "" {
				redisConn = &PrefixedConn{Conn: redisConn, config: config.Redis}
			}

			return redisConn, nil
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			if time.Since(t) < time.Minute {
				return nil
			}
			_, err := c.Do("PING")
			return err
		},
	}
}

// PrefixedConn 是一个实现了 redis.Conn 接口的自定义结构体，它在键名前添加了前缀
type PrefixedConn struct {
	redis.Conn
	config config.Redis
}

// 允许适配前缀的命令
var commandsWithPrefix = []string{
	"GET", "SET", "EXISTS", "DEL", "TYPE",
	"RPUSH", "LPOP", "RPOP", "LLEN", "LRANGE",
	"SADD", "SREM", "SISMEMBER", "SMEMBERS", "SCARD",
	"HSET", "HMSET", "HGET", "HGETALL",
	"ZADD", "ZRANGE", "ZRANGEBYSCORE", "ZREVRANGEBYSCORE", "ZREM",
	"INCR", "INCRBY",
	"WATCH", "MULTI", "EXEC", "EXPIRE",
}

// Do 实现了 redis.Conn 接口的 Do 方法
func (c *PrefixedConn) Do(command string, args ...any) (any, error) {
	// 执行命令前切换到数据库
	if _, err := c.Conn.Do("SELECT", c.config.Db); err != nil {
		return nil, err
	}
	command = strings.ToUpper(command)
	// 判断是不是单独实现了命令 不用反射吧-性能问题
	switch command {
	case "MGET":
		return c.MGET(args...)
	case "MSET":
		return c.MSET(args...)
	}
	// 加前缀
	if len(args) > 0 {
		key := args[0].(string)
		if lo.IndexOf[string](commandsWithPrefix, command) != -1 && key != "" {
			args[0] = c.config.Prefix + ":" + key
		}
	}
	return c.Conn.Do(command, args...)
}

// MGET 实现了批量获取命令
func (c *PrefixedConn) MGET(args ...any) (interface{}, error) {
	for i := range args {
		args[i] = c.config.Prefix + ":" + cast.ToString(args[i])
	}

	return c.Conn.Do("MGET", args...)
}

// MSET 实现了批量设置命令
func (c *PrefixedConn) MSET(args ...any) (interface{}, error) {
	for i := range args {
		if i%2 == 0 {
			args[i] = c.config.Prefix + ":" + cast.ToString(args[i])
		}
	}
	return c.Conn.Do("MSET", args...)
}

gin路由：router/router.go

package router

import (
	"github.com/gin-contrib/cors"
	"github.com/gin-contrib/sessions"
	"github.com/gin-gonic/gin"
)
// InitRouters 初始化路由
func InitRouters() *gin.Engine {
	// 创建默认带中间件的路由 Logger、Recovery
	r := gin.Default()
	// 允许跨域
	r.Use(cors.Default())
	r.Use(sessions.Sessions("sessionID", *global.Store))
}

下面是单点登录实现逻辑

middleware/middleware.go

package middleware

import (
	"encoding/json"
	"net/http"
	myerror "demo/common/error"
	"demo/common/global"
	"demo/model"
	"demo/util"

	"github.com/gin-contrib/sessions"
	"github.com/gin-gonic/gin"
	"github.com/gomodule/redigo/redis"
	"github.com/spf13/cast"
)

const UserSessionExpireDuration = 8 * 3600
const UserSessionMapPrefix = "session:user:"

type SessionData struct {
	Member model.Member `json:"member"`
}

// 保存登录信息
func GenerateSession(c *gin.Context, sessionData SessionData) error {
	c.Request.Header.Set("Cookie", "")
	redisConn := global.REDISPoll.Get()
	defer redisConn.Close()
	sessionKey := UserSessionMapPrefix + cast.ToString(sessionData.Member.ID)
	// 单点登录删除旧的登录态
	oldSessionID, err := redis.String(redisConn.Do("get", sessionKey))
	if err != nil && err != redis.ErrNil {
		return err
	}
	_, err = redisConn.Do("del", oldSessionID)
	if err != nil {
		return err
	}
	// 保存用户相关
	token := cast.ToString(sessionData.Member.ID) + sessionData.Member.Name + util.Md5(util.GetUUID()+util.RandString(20))
	session := sessions.Default(c)
	by, err := json.Marshal(sessionData)
	if err != nil {
		return err
	}
	session.Set("authenticated", true)
	session.Set("data", by)
	session.Set("member_id", sessionData.Member.ID)
	session.Set("member_name", sessionData.Member.Name)
	session.Set("token", token)
	session.Save()
	sessionVal := "session:" + session.ID()
	_, err = redisConn.Do("set", sessionKey, sessionVal, "EX", UserSessionExpireDuration)
	return err
}

// 身份验证
func Authentication(ctx *gin.Context) {
	session := sessions.Default(ctx)
	if auth, ok := session.Get("authenticated").(bool); !ok || !auth {
		ctx.JSON(http.StatusUnauthorized, gin.H{
			"code": myerror.SERVER_UNAUTHORIZED_ERROR,
			"msg":  myerror.SERVER_UNAUTHORIZED_ERROR.String(),
		})
		ctx.Abort()
		return
	}
	session.Save()
	// 保存用户相关信息到ctx
	memberId := session.Get("member_id")
	memberName := session.Get("member_name")
	data := session.Get("data")
	var sessionData SessionData
	if data != nil {
		by := data.([]byte)
		err := json.Unmarshal(by, &sessionData)
		if err != nil {
			ctx.JSON(http.StatusUnauthorized, gin.H{
				"code": myerror.SERVER_SYSTEM_ERROR,
				"msg":  myerror.SERVER_SYSTEM_ERROR.String(),
			})
			ctx.Abort()
			return
		}
	}
	ctx.Set("member_id", memberId)
	ctx.Set("member_name", memberName)
	ctx.Set("member_info", sessionData.Member)
	ctx.Set("data", data)
	// 登录续时
	redisConn := global.REDISPoll.Get()
	defer redisConn.Close()
	sessionKey := UserSessionMapPrefix + cast.ToString(memberId)
	sessionVal := "session:" + session.ID()
	redisConn.Do("set", sessionKey, sessionVal, "EX", UserSessionExpireDuration)
	ctx.Next()
}