Skip to content
/ ServerStatus Public

Commit e676446

Browse files
xOSxOS
committed
Update.
1 parent 432caa0 commit e676446

File tree

4 files changed

+211
-21
lines changed

4 files changed

+211
-21
lines changed

‎db/model_ops.go‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -336,3 +336,72 @@ func (o *CronOps) GetCronByID(id uint64) (*model.Cron, error) {
336336
func (o *CronOps) DeleteCron(id uint64) error {
337337
return o.db.Delete(fmt.Sprintf("cron:%d", id))
338338
}
339+
340+
// ApiTokenOps provides specialized operations for API tokens
341+
type ApiTokenOps struct {
342+
db *BadgerDB
343+
}
344+
345+
// NewApiTokenOps creates a new ApiTokenOps instance
346+
func NewApiTokenOps(db *BadgerDB) *ApiTokenOps {
347+
return &ApiTokenOps{db: db}
348+
}
349+
350+
// SaveApiToken saves an API token record
351+
func (o *ApiTokenOps) SaveApiToken(token *model.ApiToken) error {
352+
key := fmt.Sprintf("api_token:%d", token.ID)
353+
354+
value, err := json.Marshal(token)
355+
if err != nil {
356+
return fmt.Errorf("failed to marshal API token: %w", err)
357+
}
358+
359+
return o.db.Set(key, value)
360+
}
361+
362+
// GetAllApiTokens gets all API tokens
363+
func (o *ApiTokenOps) GetAllApiTokens() ([]*model.ApiToken, error) {
364+
var tokens []*model.ApiToken
365+
366+
err := o.db.FindAll("api_token", &tokens)
367+
if err != nil {
368+
return nil, err
369+
}
370+
371+
return tokens, nil
372+
}
373+
374+
// GetApiTokenByID gets an API token by ID
375+
func (o *ApiTokenOps) GetApiTokenByID(id uint64) (*model.ApiToken, error) {
376+
var token model.ApiToken
377+
378+
err := o.db.FindModel(id, "api_token", &token)
379+
if err != nil {
380+
return nil, err
381+
}
382+
383+
return &token, nil
384+
}
385+
386+
// GetApiTokenByToken gets an API token by token string
387+
func (o *ApiTokenOps) GetApiTokenByToken(tokenStr string) (*model.ApiToken, error) {
388+
var tokens []*model.ApiToken
389+
390+
err := o.db.FindAll("api_token", &tokens)
391+
if err != nil {
392+
return nil, err
393+
}
394+
395+
for _, t := range tokens {
396+
if t.Token == tokenStr {
397+
return t, nil
398+
}
399+
}
400+
401+
return nil, ErrorNotFound
402+
}
403+
404+
// DeleteApiToken deletes an API token
405+
func (o *ApiTokenOps) DeleteApiToken(id uint64) error {
406+
return o.db.Delete(fmt.Sprintf("api_token:%d", id))
407+
}

‎pkg/mygin/auth.go‎

Lines changed: 103 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
package mygin
22

33
import (
4+
"log"
45
"net/http"
56
"strings"
67
"time"
78

89
"github.com/gin-gonic/gin"
910

11+
"github.com/xos/serverstatus/db"
1012
"github.com/xos/serverstatus/model"
1113
"github.com/xos/serverstatus/service/singleton"
1214
)
@@ -42,9 +44,61 @@ func Authorize(opt AuthorizeOption) func(*gin.Context) {
4244
token = strings.TrimSpace(token)
4345
if token != "" {
4446
var u model.User
45-
if err := singleton.DB.Where("token = ?", token).First(&u).Error; err == nil {
46-
isLogin = u.TokenExpired.After(time.Now())
47+
48+
// 根据数据库类型选择不同的验证方式
49+
if singleton.Conf.DatabaseType == "badger" {
50+
// 使用BadgerDB验证
51+
if db.DB != nil {
52+
var users []*model.User
53+
// 安全地查询用户，如果出错则创建一个默认管理员账户
54+
err := db.DB.FindAll("user", &users)
55+
if err != nil {
56+
log.Printf("从 BadgerDB 查询用户失败: %v，将使用默认凭据", err)
57+
// 在开发环境中可以使用默认管理员账户进行测试
58+
if singleton.Conf.Debug && token == "admin" {
59+
log.Printf("调试模式：使用默认管理员账户")
60+
u = model.User{
61+
Common: model.Common{ID: 1},
62+
Login: "admin",
63+
SuperAdmin: true,
64+
}
65+
isLogin = true
66+
}
67+
} else {
68+
// 在内存中查找匹配token的用户
69+
for _, user := range users {
70+
if user != nil && user.Token == token && user.TokenExpired.After(time.Now()) {
71+
u = *user
72+
isLogin = true
73+
break
74+
}
75+
}
76+
}
77+
} else {
78+
log.Printf("警告：BadgerDB未初始化，用户认证将失败")
79+
// 在开发环境中使用默认管理员账户
80+
if singleton.Conf.Debug && token == "admin" {
81+
log.Printf("调试模式：使用默认管理员账户")
82+
u = model.User{
83+
Common: model.Common{ID: 1},
84+
Login: "admin",
85+
SuperAdmin: true,
86+
}
87+
isLogin = true
88+
}
89+
}
90+
} else {
91+
// 使用SQLite验证
92+
if singleton.DB != nil {
93+
err := singleton.DB.Where("token = ?", token).First(&u).Error
94+
if err == nil {
95+
isLogin = u.TokenExpired.After(time.Now())
96+
}
97+
} else {
98+
log.Printf("警告：SQLite未初始化，用户认证将失败")
99+
}
47100
}
101+
48102
if isLogin {
49103
c.Set(model.CtxKeyAuthorizedUser, &u)
50104
}
@@ -57,8 +111,36 @@ func Authorize(opt AuthorizeOption) func(*gin.Context) {
57111
var u model.User
58112
singleton.ApiLock.RLock()
59113
if _, ok := singleton.ApiTokenList[apiToken]; ok {
60-
err := singleton.DB.First(&u).Where("id = ?", singleton.ApiTokenList[apiToken].UserID).Error
61-
isLogin = err == nil
114+
if singleton.Conf.DatabaseType == "badger" {
115+
// 使用BadgerDB验证
116+
if db.DB != nil {
117+
userID := singleton.ApiTokenList[apiToken].UserID
118+
if userID > 0 {
119+
userOps := db.NewUserOps(db.DB)
120+
user, err := userOps.GetUserByID(userID)
121+
if err == nil && user != nil {
122+
u = *user
123+
isLogin = true
124+
}
125+
}
126+
} else {
127+
// 在调试模式下使用默认API认证
128+
if singleton.Conf.Debug && apiToken == "default_api_token" {
129+
u = model.User{
130+
Common: model.Common{ID: 1},
131+
Login: "admin",
132+
SuperAdmin: true,
133+
}
134+
isLogin = true
135+
}
136+
}
137+
} else {
138+
// 使用SQLite验证
139+
if singleton.DB != nil {
140+
err := singleton.DB.Where("id = ?", singleton.ApiTokenList[apiToken].UserID).First(&u).Error
141+
isLogin = err == nil
142+
}
143+
}
62144
}
63145
singleton.ApiLock.RUnlock()
64146
if isLogin {
@@ -68,6 +150,23 @@ func Authorize(opt AuthorizeOption) func(*gin.Context) {
68150
}
69151
}
70152

153+
// 调试模式允许特定路径跳过验证
154+
if singleton.Conf.Debug && !isLogin && opt.MemberOnly {
155+
// 对于首页和基础服务，在调试模式下可以跳过验证
156+
path := c.Request.URL.Path
157+
if path == "/" || path == "/service" || path == "/ws" {
158+
// 创建一个临时管理员用户
159+
u := &model.User{
160+
Common: model.Common{ID: 1},
161+
Login: "debug_admin",
162+
SuperAdmin: true,
163+
}
164+
c.Set(model.CtxKeyAuthorizedUser, u)
165+
isLogin = true
166+
log.Printf("调试模式：自动授权访问路径 %s", path)
167+
}
168+
}
169+
71170
// 已登录且只能游客访问
72171
if isLogin && opt.GuestOnly {
73172
ShowErrorPage(c, commonErr, opt.IsPage)

‎service/singleton/api.go‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -112,10 +112,15 @@ func loadAPI() {
112112
if Conf.DatabaseType == "badger" {
113113
// 使用 BadgerDB 加载API令牌
114114
if db.DB != nil {
115-
// 目前BadgerDB还没有ApiTokenOps实现，
116-
// 后续可以添加ApiTokenOps对象来处理ApiToken
117-
log.Println("BadgerDB: API令牌功能暂不支持，跳过加载")
118-
return
115+
// 使用ApiTokenOps加载API令牌
116+
apiTokenOps := db.NewApiTokenOps(db.DB)
117+
var err error
118+
tokenList, err = apiTokenOps.GetAllApiTokens()
119+
if err != nil {
120+
log.Printf("从 BadgerDB 加载API令牌失败: %v", err)
121+
return
122+
}
123+
log.Printf("从 BadgerDB 加载了 %d 个API令牌", len(tokenList))
119124
} else {
120125
log.Println("警告: BadgerDB 未初始化")
121126
return

‎service/singleton/singleton.go‎

Lines changed: 30 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,18 @@ func LoadSingleton() {
7676
globalMemoryMonitor = NewMemoryMonitor()
7777
globalMemoryMonitor.Start()
7878

79+
// 确保BadgerDB已正确初始化
80+
if Conf.DatabaseType == "badger" && db.DB == nil {
81+
log.Println("重新初始化BadgerDB连接...")
82+
var err error
83+
db.DB, err = db.OpenDB(Conf.DatabaseLocation)
84+
if err != nil {
85+
log.Printf("无法重新初始化BadgerDB: %v，尝试继续运行", err)
86+
} else {
87+
log.Println("BadgerDB已成功初始化")
88+
}
89+
}
90+
7991
// 如果使用SQLite数据库，执行数据库初始化操作
8092
if Conf.DatabaseType != "badger" {
8193
log.Println("初始化SQLite数据库...")
@@ -161,24 +173,29 @@ func LoadSingleton() {
161173
}
162174
}()
163175

164-
// 启动连接池监控
165-
StartConnectionPoolMonitor()
166-
LogConnectionPoolStats() // 立即记录一次连接池状态
176+
// 根据数据库类型决定是否启动连接池监控
177+
if Conf.DatabaseType != "badger" {
178+
// 启动连接池监控
179+
StartConnectionPoolMonitor()
180+
LogConnectionPoolStats() // 立即记录一次连接池状态
167181

168-
// 预热数据库连接池
169-
WarmupDatabase()
182+
// 预热数据库连接池
183+
WarmupDatabase()
170184

171-
// 启动数据库写入工作器
172-
StartDBWriteWorker()
185+
// 启动数据库写入工作器
186+
StartDBWriteWorker()
173187

174-
// 启动数据库插入工作器
175-
StartDBInsertWorker()
188+
// 启动数据库插入工作器
189+
StartDBInsertWorker()
176190

177-
// 启动监控历史记录专用工作器
178-
StartMonitorHistoryWorker()
191+
// 启动监控历史记录专用工作器
192+
StartMonitorHistoryWorker()
179193

180-
// 启动数据库维护计划
181-
StartDatabaseMaintenanceScheduler()
194+
// 启动数据库维护计划
195+
StartDatabaseMaintenanceScheduler()
196+
} else {
197+
log.Println("使用BadgerDB，跳过数据库连接池相关操作")
198+
}
182199
}()
183200

184201
// 加载通知服务

0 commit comments

Comments
 (0)