feat: support adaptive traffic units and precise quotas

This commit is contained in:
sagitchu
2026-09-24 15:33:53 +08:00
parent 2e845030de
commit 5d60c4fbe1
21 changed files with 579 additions and 173 deletions
+8 -6
View File
@@ -669,7 +669,7 @@ func (r *Repository) GetUserPackageTunnels(userID int64) ([]model.UserTunnelDeta
}
var items []model.UserTunnelDetail
err := r.db.Model(&model.UserTunnel{}).
Select("user_tunnel.id, user_tunnel.user_id, user_tunnel.tunnel_id, tunnel.name AS tunnel_name, user_tunnel.status, tunnel.flow AS tunnel_flow, user_tunnel.flow, user_tunnel.in_flow, user_tunnel.out_flow, user_tunnel.num, user_tunnel.flow_reset_time, user_tunnel.exp_time, user_tunnel.speed_id, speed_limit.name AS speed_limit, speed_limit.speed").
Select("user_tunnel.id, user_tunnel.user_id, user_tunnel.tunnel_id, tunnel.name AS tunnel_name, user_tunnel.status, tunnel.flow AS tunnel_flow, user_tunnel.flow, user_tunnel.flow_mib, user_tunnel.in_flow, user_tunnel.out_flow, user_tunnel.num, user_tunnel.flow_reset_time, user_tunnel.exp_time, user_tunnel.speed_id, speed_limit.name AS speed_limit, speed_limit.speed").
Joins("LEFT JOIN tunnel ON tunnel.id = user_tunnel.tunnel_id").
Joins("LEFT JOIN speed_limit ON speed_limit.id = user_tunnel.speed_id").
Where("user_tunnel.user_id = ?", userID).
@@ -924,7 +924,7 @@ func (r *Repository) ListUsers() ([]map[string]interface{}, error) {
item := map[string]interface{}{
"id": u.ID, "user": u.User, "name": u.User,
"roleId": u.RoleID, "status": u.Status,
"flow": u.Flow, "num": u.Num, "expTime": u.ExpTime,
"flow": u.Flow, "flowMiB": u.FlowMiB, "num": u.Num, "expTime": u.ExpTime,
"flowResetTime": u.FlowResetTime, "createdTime": u.CreatedTime,
"updatedTime": nullableInt64(u.UpdatedTime),
"inFlow": u.InFlow, "outFlow": u.OutFlow,
@@ -2094,7 +2094,7 @@ func (r *Repository) exportUsers() ([]model.UserBackup, error) {
for _, u := range users {
b := model.UserBackup{
ID: u.ID, User: u.User, Pwd: u.Pwd, RoleID: u.RoleID,
ExpTime: u.ExpTime, Flow: u.Flow, InFlow: u.InFlow, OutFlow: u.OutFlow,
ExpTime: u.ExpTime, Flow: u.Flow, FlowMiB: u.FlowMiB, InFlow: u.InFlow, OutFlow: u.OutFlow,
FlowResetTime: u.FlowResetTime, Num: u.Num,
CreatedTime: u.CreatedTime, Status: u.Status,
}
@@ -2270,7 +2270,7 @@ func (r *Repository) exportUserTunnels() ([]model.UserTunnelBackup, error) {
for _, ut := range uts {
b := model.UserTunnelBackup{
ID: ut.ID, UserID: ut.UserID, TunnelID: ut.TunnelID,
Num: ut.Num, Flow: ut.Flow, InFlow: ut.InFlow, OutFlow: ut.OutFlow,
Num: ut.Num, Flow: ut.Flow, FlowMiB: ut.FlowMiB, InFlow: ut.InFlow, OutFlow: ut.OutFlow,
FlowResetTime: ut.FlowResetTime, ExpTime: ut.ExpTime, Status: ut.Status,
}
if ut.SpeedID.Valid {
@@ -2467,6 +2467,7 @@ func importUsers(tx *gorm.DB, users []model.UserBackup, now int64) (int, error)
RoleID: u.RoleID,
ExpTime: u.ExpTime,
Flow: u.Flow,
FlowMiB: u.FlowMiB,
InFlow: u.InFlow,
OutFlow: u.OutFlow,
FlowResetTime: u.FlowResetTime,
@@ -2479,7 +2480,7 @@ func importUsers(tx *gorm.DB, users []model.UserBackup, now int64) (int, error)
err = tx.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "id"}},
DoUpdates: clause.AssignmentColumns([]string{
"user", "pwd", "role_id", "exp_time", "flow", "in_flow", "out_flow",
"user", "pwd", "role_id", "exp_time", "flow", "flow_mib", "in_flow", "out_flow",
"flow_reset_time", "num", "updated_time", "status", "password_changed_at",
}),
}).Create(&item).Error
@@ -2717,6 +2718,7 @@ func importUserTunnels(tx *gorm.DB, userTunnels []model.UserTunnelBackup, _ int6
SpeedID: sql.NullInt64{Int64: ut.SpeedID, Valid: ut.SpeedID > 0},
Num: ut.Num,
Flow: ut.Flow,
FlowMiB: ut.FlowMiB,
InFlow: ut.InFlow,
OutFlow: ut.OutFlow,
FlowResetTime: ut.FlowResetTime,
@@ -2726,7 +2728,7 @@ func importUserTunnels(tx *gorm.DB, userTunnels []model.UserTunnelBackup, _ int6
err := tx.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "id"}},
DoUpdates: clause.AssignmentColumns([]string{
"user_id", "tunnel_id", "speed_id", "num", "flow", "in_flow", "out_flow",
"user_id", "tunnel_id", "speed_id", "num", "flow", "flow_mib", "in_flow", "out_flow",
"flow_reset_time", "exp_time", "status",
}),
}).Create(&item).Error
@@ -37,7 +37,14 @@ func (r *Repository) UserExistsExcluding(username string, excludeID int64) (bool
return cnt > 0, err
}
func (r *Repository) CreateUser(username, pwdHash string, roleID int, expTime, flow, flowResetTime int64, num, status, maxConn int, now int64) (int64, error) {
func optionalFlowMiB(values []int64) int64 {
if len(values) > 0 {
return values[0]
}
return 0
}
func (r *Repository) CreateUser(username, pwdHash string, roleID int, expTime, flow, flowResetTime int64, num, status, maxConn int, now int64, flowMiB ...int64) (int64, error) {
if r == nil || r.db == nil {
return 0, errors.New("repository not initialized")
}
@@ -47,6 +54,7 @@ func (r *Repository) CreateUser(username, pwdHash string, roleID int, expTime, f
RoleID: roleID,
ExpTime: expTime,
Flow: flow,
FlowMiB: optionalFlowMiB(flowMiB),
InFlow: 0,
OutFlow: 0,
FlowResetTime: flowResetTime,
@@ -75,7 +83,7 @@ func (r *Repository) GetUserRoleID(userID int64) (int, error) {
return user.RoleID, nil
}
func (r *Repository) UpdateUserWithPassword(id int64, username, pwdHash string, flow int64, num int, expTime, flowResetTime int64, status, maxConn int, now int64) error {
func (r *Repository) UpdateUserWithPassword(id int64, username, pwdHash string, flow int64, num int, expTime, flowResetTime int64, status, maxConn int, now int64, flowMiB ...int64) error {
if r == nil || r.db == nil {
return errors.New("repository not initialized")
}
@@ -85,6 +93,7 @@ func (r *Repository) UpdateUserWithPassword(id int64, username, pwdHash string,
"user": username,
"pwd": pwdHash,
"flow": flow,
"flow_mib": optionalFlowMiB(flowMiB),
"num": num,
"exp_time": expTime,
"flow_reset_time": flowResetTime,
@@ -95,7 +104,7 @@ func (r *Repository) UpdateUserWithPassword(id int64, username, pwdHash string,
}).Error
}
func (r *Repository) UpdateUserWithoutPassword(id int64, username string, flow int64, num int, expTime, flowResetTime int64, status, maxConn int, now int64) error {
func (r *Repository) UpdateUserWithoutPassword(id int64, username string, flow int64, num int, expTime, flowResetTime int64, status, maxConn int, now int64, flowMiB ...int64) error {
if r == nil || r.db == nil {
return errors.New("repository not initialized")
}
@@ -104,6 +113,7 @@ func (r *Repository) UpdateUserWithoutPassword(id int64, username string, flow i
Updates(map[string]interface{}{
"user": username,
"flow": flow,
"flow_mib": optionalFlowMiB(flowMiB),
"num": num,
"exp_time": expTime,
"flow_reset_time": flowResetTime,
@@ -126,7 +136,7 @@ func (r *Repository) UpdateUserPassword(userID int64, pwdHash string, now int64)
}).Error
}
func (r *Repository) PropagateUserFlowToTunnels(userID int64, flow int64, num int, expTime, flowResetTime int64) {
func (r *Repository) PropagateUserFlowToTunnels(userID int64, flow int64, num int, expTime, flowResetTime int64, flowMiB ...int64) {
if r == nil || r.db == nil {
return
}
@@ -134,6 +144,7 @@ func (r *Repository) PropagateUserFlowToTunnels(userID int64, flow int64, num in
Where("user_id = ?", userID).
Updates(map[string]interface{}{
"flow": flow,
"flow_mib": optionalFlowMiB(flowMiB),
"num": num,
"exp_time": expTime,
"flow_reset_time": flowResetTime,
@@ -650,7 +661,7 @@ func (r *Repository) DeleteUserTunnel(id int64) error {
return r.db.Where("id = ?", id).Delete(&model.UserTunnel{}).Error
}
func (r *Repository) UpdateUserTunnel(id int64, flow int64, num int, expTime, flowResetTime int64, speedID interface{}, status int) error {
func (r *Repository) UpdateUserTunnel(id int64, flow int64, num int, expTime, flowResetTime int64, speedID interface{}, status int, flowMiB ...int64) error {
if r == nil || r.db == nil {
return errors.New("repository not initialized")
}
@@ -658,6 +669,7 @@ func (r *Repository) UpdateUserTunnel(id int64, flow int64, num int, expTime, fl
Where("id = ?", id).
Updates(map[string]interface{}{
"flow": flow,
"flow_mib": optionalFlowMiB(flowMiB),
"num": num,
"exp_time": expTime,
"flow_reset_time": flowResetTime,
@@ -692,7 +704,7 @@ func (r *Repository) GetExistingUserTunnel(userID, tunnelID int64) (id int64, fl
return ut.ID, ut.Flow, int64(ut.Num), ut.ExpTime, ut.FlowResetTime, ut.SpeedID, ut.Status, nil
}
func (r *Repository) InsertUserTunnel(userID, tunnelID int64, speedID interface{}, num int, flow, flowResetTime, expTime int64, status int) error {
func (r *Repository) InsertUserTunnel(userID, tunnelID int64, speedID interface{}, num int, flow, flowResetTime, expTime int64, status int, flowMiB ...int64) error {
if r == nil || r.db == nil {
return errors.New("repository not initialized")
}
@@ -702,6 +714,7 @@ func (r *Repository) InsertUserTunnel(userID, tunnelID int64, speedID interface{
SpeedID: nullInt64FromInterface(speedID),
Num: num,
Flow: flow,
FlowMiB: optionalFlowMiB(flowMiB),
InFlow: 0,
OutFlow: 0,
FlowResetTime: flowResetTime,
@@ -711,7 +724,7 @@ func (r *Repository) InsertUserTunnel(userID, tunnelID int64, speedID interface{
return r.db.Create(&ut).Error
}
func (r *Repository) UpdateUserTunnelFields(id int64, speedID interface{}, flow int64, num int, expTime, flowResetTime int64, status int) error {
func (r *Repository) UpdateUserTunnelFields(id int64, speedID interface{}, flow int64, num int, expTime, flowResetTime int64, status int, flowMiB ...int64) error {
if r == nil || r.db == nil {
return errors.New("repository not initialized")
}
@@ -720,6 +733,7 @@ func (r *Repository) UpdateUserTunnelFields(id int64, speedID interface{}, flow
Updates(map[string]interface{}{
"speed_id": nullInt64FromInterface(speedID),
"flow": flow,
"flow_mib": optionalFlowMiB(flowMiB),
"num": num,
"exp_time": expTime,
"flow_reset_time": flowResetTime,
@@ -1293,7 +1307,7 @@ func (r *Repository) EnsureUserTunnelGrant(userID, tunnelID int64) (int64, bool,
return 0, false, err
}
var user model.User
if err := r.db.Select("flow, num, exp_time, flow_reset_time").Where("id = ?", userID).First(&user).Error; err != nil {
if err := r.db.Select("flow, flow_mib, num, exp_time, flow_reset_time").Where("id = ?", userID).First(&user).Error; err != nil {
return 0, false, err
}
flow := user.Flow
@@ -1305,6 +1319,7 @@ func (r *Repository) EnsureUserTunnelGrant(userID, tunnelID int64) (int64, bool,
TunnelID: tunnelID,
Num: num,
Flow: flow,
FlowMiB: user.FlowMiB,
InFlow: 0,
OutFlow: 0,
FlowResetTime: flowReset,
@@ -0,0 +1,67 @@
package repo
import (
"path/filepath"
"testing"
"time"
"go-backend/internal/store/model"
)
func TestTrafficLimitMiBSurvivesBackupRestore(t *testing.T) {
source, err := Open(filepath.Join(t.TempDir(), "source.db"))
if err != nil {
t.Fatal(err)
}
defer source.Close()
now := time.Now().UnixMilli()
userID, err := source.CreateUser("mib-user", "hash", 1, now+86400000, 1, 1, 10, 1, 0, now, 500)
if err != nil {
t.Fatal(err)
}
tunnel := model.Tunnel{Name: "mib-tunnel", TrafficRatio: 1, Type: 1, Protocol: "tls", Flow: 1, CreatedTime: now, UpdatedTime: now, Status: 1, Inx: 1}
if err := source.DB().Create(&tunnel).Error; err != nil {
t.Fatal(err)
}
if _, _, err := source.EnsureUserTunnelGrant(userID, tunnel.ID); err != nil {
t.Fatal(err)
}
grants, err := source.GetUserPackageTunnels(userID)
if err != nil || len(grants) != 1 || grants[0].FlowMiB != 500 {
t.Fatalf("inherited tunnel quota = %+v, err = %v", grants, err)
}
backup, err := source.ExportAll()
if err != nil {
t.Fatal(err)
}
found := false
for _, user := range backup.Users {
if user.User == "mib-user" {
found = user.Flow == 1 && user.FlowMiB == 500
}
}
if !found {
t.Fatal("500 MiB user quota missing from backup")
}
if len(backup.UserTunnels) != 1 || backup.UserTunnels[0].FlowMiB != 500 {
t.Fatalf("tunnel quota missing from backup: %+v", backup.UserTunnels)
}
dest, err := Open(filepath.Join(t.TempDir(), "dest.db"))
if err != nil {
t.Fatal(err)
}
defer dest.Close()
if _, err := dest.Import(backup, []string{"users", "tunnels", "userTunnels"}); err != nil {
t.Fatal(err)
}
user, err := dest.GetUserByUsername("mib-user")
if err != nil || user == nil || user.Flow != 1 || user.FlowMiB != 500 {
t.Fatalf("restored quota = %+v, err = %v", user, err)
}
grants, err = dest.GetUserPackageTunnels(user.ID)
if err != nil || len(grants) != 1 || grants[0].FlowMiB != 500 {
t.Fatalf("restored tunnel quota = %+v, err = %v", grants, err)
}
}