Skip to content

Commit 4ed7810

Browse files
committed
Update.
1 parent 53cc27e commit 4ed7810

File tree

3 files changed

+228
-14
lines changed

3 files changed

+228
-14
lines changed

db/badger.go

Lines changed: 158 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -217,8 +217,23 @@ func (b *BadgerDB) FindModel(id uint64, modelType string, result interface{}) er
217217
// 反序列化到结果
218218
return json.Unmarshal(ruleJSON, result)
219219
default:
220-
// 其他类型的记录,使用标准处理方式
221-
return json.Unmarshal(data, result)
220+
// 其他类型的记录,也需要进行字段类型转换
221+
var dataMap map[string]interface{}
222+
if err := json.Unmarshal(data, &dataMap); err != nil {
223+
return err
224+
}
225+
226+
// 转换字段类型,确保布尔字段正确
227+
convertDbFieldTypes(&dataMap)
228+
229+
// 重新序列化为 JSON
230+
dataJSON, err := json.Marshal(dataMap)
231+
if err != nil {
232+
return err
233+
}
234+
235+
// 反序列化到结果
236+
return json.Unmarshal(dataJSON, result)
222237
}
223238
}
224239

@@ -495,18 +510,151 @@ func (b *BadgerDB) FindAll(prefix string, result interface{}) error {
495510
}
496511

497512
return nil
513+
case "api_token":
514+
// API令牌记录需要特殊处理布尔字段和Token字段
515+
var tokens []*map[string]interface{}
516+
for _, item := range items {
517+
var data map[string]interface{}
518+
if err := json.Unmarshal(item, &data); err != nil {
519+
continue
520+
}
521+
522+
// 转换字段类型,确保布尔字段正确
523+
convertDbFieldTypes(&data)
524+
tokens = append(tokens, &data)
525+
}
526+
527+
// 重新序列化为 JSON
528+
tokensJSON, err := json.Marshal(tokens)
529+
if err != nil {
530+
return err
531+
}
532+
533+
// 反序列化到结果,然后手动设置Token字段
534+
if err := json.Unmarshal(tokensJSON, result); err != nil {
535+
return err
536+
}
537+
538+
// 手动设置Token字段(因为它可能有json:"-"标签)
539+
if tokenSlice, ok := result.(*[]*model.ApiToken); ok {
540+
for i, token := range *tokenSlice {
541+
if token != nil && i < len(tokens) {
542+
tokenData := tokens[i]
543+
if tokenVal, exists := (*tokenData)["Token"]; exists {
544+
if tokenStr, isStr := tokenVal.(string); isStr {
545+
token.Token = tokenStr
546+
}
547+
}
548+
}
549+
}
550+
}
551+
552+
return nil
553+
case "ddns_profile":
554+
// DDNS配置记录需要特殊处理布尔字段
555+
var profiles []*map[string]interface{}
556+
for _, item := range items {
557+
var data map[string]interface{}
558+
if err := json.Unmarshal(item, &data); err != nil {
559+
continue
560+
}
561+
562+
// 转换字段类型,确保布尔字段正确
563+
convertDbFieldTypes(&data)
564+
profiles = append(profiles, &data)
565+
}
566+
567+
// 重新序列化为 JSON
568+
profilesJSON, err := json.Marshal(profiles)
569+
if err != nil {
570+
return err
571+
}
572+
573+
return json.Unmarshal(profilesJSON, result)
574+
case "notification":
575+
// 通知配置记录需要特殊处理布尔字段
576+
var notifications []*map[string]interface{}
577+
for _, item := range items {
578+
var data map[string]interface{}
579+
if err := json.Unmarshal(item, &data); err != nil {
580+
continue
581+
}
582+
583+
// 转换字段类型,确保布尔字段正确
584+
convertDbFieldTypes(&data)
585+
notifications = append(notifications, &data)
586+
}
587+
588+
// 重新序列化为 JSON
589+
notificationsJSON, err := json.Marshal(notifications)
590+
if err != nil {
591+
return err
592+
}
593+
594+
return json.Unmarshal(notificationsJSON, result)
595+
case "nat":
596+
// NAT配置记录需要特殊处理布尔字段
597+
var nats []*map[string]interface{}
598+
for _, item := range items {
599+
var data map[string]interface{}
600+
if err := json.Unmarshal(item, &data); err != nil {
601+
continue
602+
}
603+
604+
// 转换字段类型,确保布尔字段正确
605+
convertDbFieldTypes(&data)
606+
nats = append(nats, &data)
607+
}
608+
609+
// 重新序列化为 JSON
610+
natsJSON, err := json.Marshal(nats)
611+
if err != nil {
612+
return err
613+
}
614+
615+
return json.Unmarshal(natsJSON, result)
616+
case "cron":
617+
// 定时任务记录需要特殊处理布尔字段
618+
var crons []*map[string]interface{}
619+
for _, item := range items {
620+
var data map[string]interface{}
621+
if err := json.Unmarshal(item, &data); err != nil {
622+
continue
623+
}
624+
625+
// 转换字段类型,确保布尔字段正确
626+
convertDbFieldTypes(&data)
627+
crons = append(crons, &data)
628+
}
629+
630+
// 重新序列化为 JSON
631+
cronsJSON, err := json.Marshal(crons)
632+
if err != nil {
633+
return err
634+
}
635+
636+
return json.Unmarshal(cronsJSON, result)
498637
default:
499-
// 其他类型的记录,使用标准处理方式
500-
itemsJSON := "["
501-
for i, item := range items {
502-
if i > 0 {
503-
itemsJSON += ","
638+
// 其他类型的记录,也需要进行字段类型转换
639+
var others []*map[string]interface{}
640+
for _, item := range items {
641+
var data map[string]interface{}
642+
if err := json.Unmarshal(item, &data); err != nil {
643+
continue
504644
}
505-
itemsJSON += string(item)
645+
646+
// 转换字段类型,确保布尔字段正确
647+
convertDbFieldTypes(&data)
648+
others = append(others, &data)
649+
}
650+
651+
// 重新序列化为 JSON
652+
othersJSON, err := json.Marshal(others)
653+
if err != nil {
654+
return err
506655
}
507-
itemsJSON += "]"
508656

509-
return json.Unmarshal([]byte(itemsJSON), result)
657+
return json.Unmarshal(othersJSON, result)
510658
}
511659
}
512660

model/rule.go

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
package model
22

33
import (
4+
"encoding/json"
45
"slices"
6+
"strconv"
57
"strings"
68
"time"
79

@@ -38,6 +40,56 @@ type Rule struct {
3840
LastCycleStatus map[uint64]interface{} `json:"-"`
3941
}
4042

43+
// MarshalJSON 自定义 JSON 序列化,确保 Ignore 字段正确序列化
44+
func (r Rule) MarshalJSON() ([]byte, error) {
45+
// 创建一个临时结构体,将 Ignore 字段转换为字符串键的 map
46+
type Alias Rule
47+
aux := struct {
48+
Alias
49+
Ignore map[string]bool `json:"ignore,omitempty"`
50+
}{
51+
Alias: Alias(r),
52+
}
53+
54+
// 转换 Ignore 字段
55+
if r.Ignore != nil {
56+
aux.Ignore = make(map[string]bool)
57+
for k, v := range r.Ignore {
58+
aux.Ignore[strconv.FormatUint(k, 10)] = v
59+
}
60+
}
61+
62+
return json.Marshal(aux)
63+
}
64+
65+
// UnmarshalJSON 自定义 JSON 反序列化,确保 Ignore 字段正确反序列化
66+
func (r *Rule) UnmarshalJSON(data []byte) error {
67+
// 创建一个临时结构体,接收字符串键的 Ignore 字段
68+
type Alias Rule
69+
aux := struct {
70+
*Alias
71+
Ignore map[string]bool `json:"ignore,omitempty"`
72+
}{
73+
Alias: (*Alias)(r),
74+
}
75+
76+
if err := json.Unmarshal(data, &aux); err != nil {
77+
return err
78+
}
79+
80+
// 转换 Ignore 字段
81+
if aux.Ignore != nil {
82+
r.Ignore = make(map[uint64]bool)
83+
for k, v := range aux.Ignore {
84+
if id, err := strconv.ParseUint(k, 10, 64); err == nil {
85+
r.Ignore[id] = v
86+
}
87+
}
88+
}
89+
90+
return nil
91+
}
92+
4193
func percentage(used, total uint64) float64 {
4294
if total == 0 {
4395
return 0

service/rpc/server.go

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -566,10 +566,24 @@ func (s *ServerHandler) ReportSystemInfo(c context.Context, r *pb.Host) (*pb.Rec
566566
singleton.ServerList[clientID].PrevTransferOutSnapshot = 0
567567

568568
// 确保从数据库读取最新的累计流量值(只在重启时读取一次)
569-
var server model.Server
570-
if err := singleton.DB.First(&server, clientID).Error; err == nil {
571-
singleton.ServerList[clientID].CumulativeNetInTransfer = server.CumulativeNetInTransfer
572-
singleton.ServerList[clientID].CumulativeNetOutTransfer = server.CumulativeNetOutTransfer
569+
if singleton.Conf.DatabaseType == "badger" {
570+
// 使用BadgerDB读取累计流量
571+
if db.DB != nil {
572+
serverOps := db.NewServerOps(db.DB)
573+
if server, err := serverOps.GetServer(clientID); err == nil && server != nil {
574+
singleton.ServerList[clientID].CumulativeNetInTransfer = server.CumulativeNetInTransfer
575+
singleton.ServerList[clientID].CumulativeNetOutTransfer = server.CumulativeNetOutTransfer
576+
}
577+
}
578+
} else {
579+
// 使用SQLite读取累计流量
580+
if singleton.DB != nil {
581+
var server model.Server
582+
if err := singleton.DB.First(&server, clientID).Error; err == nil {
583+
singleton.ServerList[clientID].CumulativeNetInTransfer = server.CumulativeNetInTransfer
584+
singleton.ServerList[clientID].CumulativeNetOutTransfer = server.CumulativeNetOutTransfer
585+
}
586+
}
573587
}
574588
}
575589

0 commit comments

Comments
 (0)