使用数据库存储管理员和用户登录SESSION

2025-12-29 17:26:35 +08:00 · 2023-02-04 15:17:02 +08:00
parent 2e5568dd39
commit e27046c75a
9 changed files with 410 additions and 7 deletions
--- a/internal/db/models/login_session_dao.go
+++ b/internal/db/models/login_session_dao.go
@@ -0,0 +1,208 @@
+package models
+
+import (
+	"encoding/json"
+	"errors"
+	_ "github.com/go-sql-driver/mysql"
+	"github.com/iwind/TeaGo/Tea"
+	"github.com/iwind/TeaGo/dbs"
+	"github.com/iwind/TeaGo/maps"
+	"github.com/iwind/TeaGo/types"
+	"time"
+)
+
+// TODO 定时清理过期的SESSION
+
+type LoginSessionDAO dbs.DAO
+
+func NewLoginSessionDAO() *LoginSessionDAO {
+	return dbs.NewDAO(&LoginSessionDAO{
+		DAOObject: dbs.DAOObject{
+			DB:     Tea.Env,
+			Table:  "edgeLoginSessions",
+			Model:  new(LoginSession),
+			PkName: "id",
+		},
+	}).(*LoginSessionDAO)
+}
+
+var SharedLoginSessionDAO *LoginSessionDAO
+
+func init() {
+	dbs.OnReady(func() {
+		SharedLoginSessionDAO = NewLoginSessionDAO()
+	})
+}
+
+// CreateSession 创建SESSION
+func (this *LoginSessionDAO) CreateSession(tx *dbs.Tx, sid string, ip string, expiresAt int64) (int64, error) {
+	if len(sid) == 0 || len(sid) > 64 {
+		return 0, errors.New("invalid 'sid'")
+	}
+
+	// 是否已存在
+	oldSessionId, err := this.Query(tx).
+		Attr("sid", sid).
+		ResultPk().
+		FindInt64Col(0)
+	if err != nil {
+		return 0, err
+	}
+
+	var op = NewLoginSessionOperator()
+	if oldSessionId > 0 {
+		op.Id = oldSessionId
+	}
+
+	op.Sid = sid
+	op.Ip = ip
+	op.Values = "{}"
+	op.ExpiresAt = expiresAt
+	op.CreatedAt = time.Now().Unix()
+
+	if oldSessionId > 0 {
+		err := this.Save(tx, op)
+		if err != nil {
+			return 0, err
+		}
+		return oldSessionId, nil
+	}
+	return this.SaveInt64(tx, op)
+}
+
+// WriteSessionValue 向SESSION中写入数据
+func (this *LoginSessionDAO) WriteSessionValue(tx *dbs.Tx, sid string, key string, value any) error {
+	if len(sid) == 0 || len(sid) > 64 {
+		return errors.New("invalid 'sid'")
+	}
+
+	// 是否存在
+	sessionOne, err := this.Query(tx).
+		Attr("sid", sid).
+		Find()
+	if err != nil {
+		return err
+	}
+	var sessionId int64
+	var isNewSession = false
+	var valueMap = maps.Map{}
+	if sessionOne != nil {
+		var session = sessionOne.(*LoginSession)
+		if session.IsAvailable() {
+			sessionId = int64(session.Id)
+
+			if !IsNull(session.Values) {
+				err = json.Unmarshal(session.Values, &valueMap)
+				if err != nil {
+					return err
+				}
+			}
+		} else {
+			// 不可用则删除之
+			err = this.Query(tx).
+				Pk(session.Id).
+				DeleteQuickly()
+			if err != nil {
+				return err
+			}
+		}
+	}
+	if sessionId == 0 {
+		// 不存在，则创建之
+		sessionId, err = this.CreateSession(tx, sid, "", time.Now().Unix()+30*86400 /** 默认30天**/)
+		if err != nil {
+			return err
+		}
+		isNewSession = true
+	}
+
+	var sessionOp = NewLoginSessionOperator()
+	sessionOp.Id = sessionId
+
+	// 获取用户ID
+	var adminId int64
+	var userId int64
+
+	switch key {
+	case "adminId":
+		adminId = types.Int64(value)
+	case "userId":
+		userId = types.Int64(value)
+	}
+
+	if adminId > 0 || userId > 0 {
+		sessionOp.AdminId = adminId
+		sessionOp.UserId = userId
+
+		if isNewSession {
+			// 删除此用户之前创建的SESSION，防止单个用户SESSION过多
+			// TODO 将来改成按照活跃时间排序
+			const maxSessionsPerUser = 10
+			oldOnes, err := this.Query(tx).
+				ResultPk().
+				Attr("adminId", adminId).
+				Attr("userId", userId).
+				Asc("createdAt").
+				FindAll()
+			if err != nil {
+				return err
+			}
+			var countOldOnes = len(oldOnes)
+			if countOldOnes > maxSessionsPerUser {
+				var countDeleted int
+				for _, oldOne := range oldOnes {
+					var oldSessionId = int64(oldOne.(*LoginSession).Id)
+					if oldSessionId == sessionId {
+						continue
+					}
+
+					if countDeleted < countOldOnes-maxSessionsPerUser {
+						err = this.Query(tx).
+							Pk(oldSessionId).
+							DeleteQuickly()
+						if err != nil {
+							return err
+						}
+						countDeleted++
+					} else {
+						break
+					}
+				}
+			}
+		}
+	}
+
+	// 写入数据
+	valueMap[key] = value
+	sessionOp.Values = valueMap.AsJSON()
+	return this.Save(tx, sessionOp)
+}
+
+// DeleteSession 删除SESSION
+func (this *LoginSessionDAO) DeleteSession(tx *dbs.Tx, sid string) error {
+	return this.Query(tx).
+		Attr("sid", sid).
+		DeleteQuickly()
+}
+
+// FindSession 查询SESSION
+func (this *LoginSessionDAO) FindSession(tx *dbs.Tx, sid string) (*LoginSession, error) {
+	one, err := this.Query(tx).
+		Attr("sid", sid).
+		Find()
+	if err != nil || one == nil {
+		return nil, err
+	}
+	var session = one.(*LoginSession)
+
+	// 不可用则删除
+	if !session.IsAvailable() {
+		err = this.Query(tx).
+			Pk(session.Id).
+			DeleteQuickly()
+		if err != nil {
+			return nil, err
+		}
+	}
+	return session, nil
+}
--- a/internal/db/models/login_session_dao_test.go
+++ b/internal/db/models/login_session_dao_test.go
@@ -0,0 +1,47 @@
+package models_test
+
+import (
+	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
+	_ "github.com/go-sql-driver/mysql"
+	_ "github.com/iwind/TeaGo/bootstrap"
+	"github.com/iwind/TeaGo/dbs"
+	"testing"
+	"time"
+)
+
+func TestLoginSessionDAO_CreateSession(t *testing.T) {
+	var dao = models.NewLoginSessionDAO()
+	var tx *dbs.Tx
+	sessionId, err := dao.CreateSession(tx, "123456", "192.168.2.40", time.Now().Unix()+3600)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Log("sessionId:", sessionId)
+}
+
+func TestLoginSessionDAO_WriteSessionValue_Admin(t *testing.T) {
+	var dao = models.NewLoginSessionDAO()
+	var tx *dbs.Tx
+	err := dao.WriteSessionValue(tx, "123456", "adminId", 123)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestLoginSessionDAO_WriteSessionValue_User(t *testing.T) {
+	var dao = models.NewLoginSessionDAO()
+	var tx *dbs.Tx
+	err := dao.WriteSessionValue(tx, "123456", "userId", 123)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestLoginSessionDAO_WriteSessionValue(t *testing.T) {
+	var dao = models.NewLoginSessionDAO()
+	var tx *dbs.Tx
+	err := dao.WriteSessionValue(tx, "123456", "key1", "value1")
+	if err != nil {
+		t.Fatal(err)
+	}
+}
--- a/internal/db/models/login_session_model.go
+++ b/internal/db/models/login_session_model.go
@@ -0,0 +1,30 @@
+package models
+
+import "github.com/iwind/TeaGo/dbs"
+
+// LoginSession 登录Session
+type LoginSession struct {
+	Id        uint64   `field:"id"`        // ID
+	AdminId   uint64   `field:"adminId"`   // 管理员ID
+	UserId    uint64   `field:"userId"`    // 用户ID
+	Sid       string   `field:"sid"`       // 令牌
+	Values    dbs.JSON `field:"values"`    // 数据
+	Ip        string   `field:"ip"`        // 登录IP
+	CreatedAt uint64   `field:"createdAt"` // 创建时间
+	ExpiresAt uint64   `field:"expiresAt"` // 过期时间
+}
+
+type LoginSessionOperator struct {
+	Id        any // ID
+	AdminId   any // 管理员ID
+	UserId    any // 用户ID
+	Sid       any // 令牌
+	Values    any // 数据
+	Ip        any // 登录IP
+	CreatedAt any // 创建时间
+	ExpiresAt any // 过期时间
+}
+
+func NewLoginSessionOperator() *LoginSessionOperator {
+	return &LoginSessionOperator{}
+}
--- a/internal/db/models/login_session_model_ext.go
+++ b/internal/db/models/login_session_model_ext.go
@@ -0,0 +1,7 @@
+package models
+
+import "time"
+
+func (this *LoginSession) IsAvailable() bool {
+	return this.ExpiresAt == 0 || int64(this.ExpiresAt) > time.Now().Unix()
+}
--- a/internal/db/models/server_dao_test.go
+++ b/internal/db/models/server_dao_test.go
@@ -45,7 +45,7 @@ func TestServerDAO_CreateManyServers(t *testing.T) {
 func TestServerDAO_ComposeServerConfig(t *testing.T) {
 	dbs.NotifyReady()
 	var tx *dbs.Tx
-	config, err := models.SharedServerDAO.ComposeServerConfigWithServerId(tx, 1, false)
+	config, err := models.SharedServerDAO.ComposeServerConfigWithServerId(tx, 1, false, false)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -55,7 +55,7 @@ func TestServerDAO_ComposeServerConfig(t *testing.T) {
 func TestServerDAO_ComposeServerConfig_AliasServerNames(t *testing.T) {
 	dbs.NotifyReady()
 	var tx *dbs.Tx
-	config, err := models.SharedServerDAO.ComposeServerConfigWithServerId(tx, 14, false)
+	config, err := models.SharedServerDAO.ComposeServerConfigWithServerId(tx, 14, false, false)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -65,7 +65,7 @@ func TestServerDAO_ComposeServerConfig_AliasServerNames(t *testing.T) {
 func TestServerDAO_UpdateServerConfig(t *testing.T) {
 	dbs.NotifyReady()
 	var tx *dbs.Tx
-	config, err := models.SharedServerDAO.ComposeServerConfigWithServerId(tx, 1, false)
+	config, err := models.SharedServerDAO.ComposeServerConfigWithServerId(tx, 1, false, false)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -190,7 +190,7 @@ func TestServerDAO_FindAllEnabledServersWithNode_Cache(t *testing.T) {
 			t.Fatal(err)
 		}
 		for _, server := range servers {
-			_, _ = models.SharedServerDAO.ComposeServerConfig(nil, server, cacheMap, true, false)
+			_, _ = models.SharedServerDAO.ComposeServerConfig(nil, server, false, cacheMap, true, false)
 		}
 	}

@@ -201,7 +201,7 @@ func TestServerDAO_FindAllEnabledServersWithNode_Cache(t *testing.T) {
 			t.Fatal(err)
 		}
 		for _, server := range servers {
-			_, _ = models.SharedServerDAO.ComposeServerConfig(nil, server, cacheMap, true, false)
+			_, _ = models.SharedServerDAO.ComposeServerConfig(nil, server, false, cacheMap, true, false)
 		}
 	}
 	t.Log(time.Since(before).Seconds()*1000, "ms")