diff --git a/internal/web/job/periodic_traffic_reset_client_test.go b/internal/web/job/periodic_traffic_reset_client_test.go index 66a872306..cfb776d95 100644 --- a/internal/web/job/periodic_traffic_reset_client_test.go +++ b/internal/web/job/periodic_traffic_reset_client_test.go @@ -26,6 +26,7 @@ type seededClient struct { recordEnable bool quotaEnable bool total int64 + inboundCycle string } // seedClientOnCycle creates an inbound that never resets on its own, a client @@ -43,9 +44,13 @@ func seedClientOnCycle(t *testing.T, port int, c seededClient) { if err != nil { t.Fatalf("marshal settings: %v", err) } + inboundCycle := c.inboundCycle + if inboundCycle == "" { + inboundCycle = "never" + } ib := model.Inbound{ UserId: 1, Enable: true, Port: port, Protocol: model.VLESS, - Tag: "inbound-" + c.email, TrafficReset: "never", Settings: string(settings), + Tag: "inbound-" + c.email, TrafficReset: inboundCycle, Settings: string(settings), } if err := db.Create(&ib).Error; err != nil { t.Fatalf("create inbound: %v", err) @@ -185,6 +190,46 @@ func TestPeriodicTrafficResetClients(t *testing.T) { }) } +func TestPeriodicTrafficResetInboundRestoresClients(t *testing.T) { + t.Run("restores a client the quota switched off", func(t *testing.T) { + initResetJobDB(t) + seedClientOnCycle(t, 41101, seededClient{ + email: "depleted-ib@example.com", cycle: "never", day: 1, + recordEnable: false, quotaEnable: false, total: 1000, inboundCycle: "daily", + }) + + NewPeriodicTrafficResetJob("daily", time.UTC).Run() + + row := trafficFor(t, "depleted-ib@example.com") + if row.Up != 0 || row.Down != 0 || !row.Enable { + t.Fatalf("traffic not reset: up=%d down=%d enable=%v", row.Up, row.Down, row.Enable) + } + if rec := recordFor(t, "depleted-ib@example.com"); !rec.Enable { + t.Fatal("clients.enable still false after the inbound's reset: the client stays out of the core") + } + if enabled := settingsEnableOf(t, 41101); !enabled { + t.Fatal("the inbound settings JSON still has the client disabled") + } + }) + + t.Run("leaves a client the operator switched off", func(t *testing.T) { + initResetJobDB(t) + seedClientOnCycle(t, 41102, seededClient{ + email: "banned-ib@example.com", cycle: "never", day: 1, + recordEnable: false, quotaEnable: true, total: 100000, inboundCycle: "daily", + }) + + NewPeriodicTrafficResetJob("daily", time.UTC).Run() + + if rec := recordFor(t, "banned-ib@example.com"); rec.Enable { + t.Fatal("an operator-disabled client was switched back on by the inbound's reset") + } + if enabled := settingsEnableOf(t, 41102); enabled { + t.Fatal("an operator-disabled client was enabled in the inbound settings JSON") + } + }) +} + func settingsEnableOf(t *testing.T, port int) bool { t.Helper() var stored model.Inbound diff --git a/internal/web/job/periodic_traffic_reset_job.go b/internal/web/job/periodic_traffic_reset_job.go index ffe1e2802..bc43b7d90 100644 --- a/internal/web/job/periodic_traffic_reset_job.go +++ b/internal/web/job/periodic_traffic_reset_job.go @@ -113,10 +113,13 @@ func (j *PeriodicTrafficResetJob) resetInboundsOnSchedule() { } } - resetClientErr := j.clientService.ResetAllClientTraffics(&j.inboundService, inbound.Id) + nr, resetClientErr := j.clientService.ResetAllClientTraffics(&j.inboundService, inbound.Id) if resetClientErr != nil { logger.Warning("Failed to reset traffic for all users of inbound", inbound.Id, ":", resetClientErr) } + if nr { + needRestart.Store(true) + } if resetInboundErr == nil && resetClientErr == nil { resetCount.Add(1) diff --git a/internal/web/service/client_bulk_reset_reenable_test.go b/internal/web/service/client_bulk_reset_reenable_test.go index 7293d386e..db0860b7c 100644 --- a/internal/web/service/client_bulk_reset_reenable_test.go +++ b/internal/web/service/client_bulk_reset_reenable_test.go @@ -27,3 +27,24 @@ func TestBulkResetTraffic_ReenablesDisabledClient(t *testing.T) { t.Fatalf("%s: traffic after reset up=%d down=%d, want 0/0", email, tr.Up, tr.Down) } } + +func TestResetAllTraffics_ReenablesDepletedClientEverywhere(t *testing.T) { + setupBulkDB(t) + svc := &ClientService{} + inboundSvc := &InboundService{} + + depleted := seedLocalDisabledClient(t, svc, 52011, "", "spent-all@x", 1000, 0, 600, 500) + banned := seedLocalDisabledClient(t, svc, 52012, "", "banned-all@x", 100000, 0, 600, 500) + + if _, err := svc.ResetAllTraffics(); err != nil { + t.Fatalf("ResetAllTraffics: %v", err) + } + + assertEnableEverywhere(t, svc, inboundSvc, depleted.Id, "spent-all@x", true) + if recordEnableOf(t, svc, "banned-all@x") { + t.Fatal("banned-all@x: an operator-disabled client was switched back on by Reset All Client Traffic") + } + if jsonClientEnable(t, inboundSvc, banned.Id, "banned-all@x") { + t.Fatal("banned-all@x: an operator-disabled client was enabled in the inbound JSON") + } +} diff --git a/internal/web/service/client_crud.go b/internal/web/service/client_crud.go index 2ff74b41f..30a3ab19b 100644 --- a/internal/web/service/client_crud.go +++ b/internal/web/service/client_crud.go @@ -128,13 +128,17 @@ func (c ClientResetCycle) Depleted() bool { return c.Total > 0 && c.Used >= c.Total } +func clientResetCycles(tx *gorm.DB) *gorm.DB { + return tx.Table("clients c"). + Select("c.email, c.traffic_reset_day, c.enable, COALESCE(ct.total, 0) AS total, COALESCE(ct.up, 0) + COALESCE(ct.down, 0) AS used"). + Joins("LEFT JOIN client_traffics ct ON ct.email = c.email") +} + // GetClientsByTrafficReset returns the clients whose own reset cycle matches the // period, independent of the cycle configured on the inbounds they belong to. func (s *ClientService) GetClientsByTrafficReset(period string) ([]ClientResetCycle, error) { var cycles []ClientResetCycle - err := database.GetDB().Table("clients c"). - Select("c.email, c.traffic_reset_day, c.enable, COALESCE(ct.total, 0) AS total, COALESCE(ct.up, 0) + COALESCE(ct.down, 0) AS used"). - Joins("LEFT JOIN client_traffics ct ON ct.email = c.email"). + err := clientResetCycles(database.GetDB()). Where("c.traffic_reset = ?", period). Scan(&cycles).Error if err != nil { @@ -143,6 +147,30 @@ func (s *ClientService) GetClientsByTrafficReset(period string) ([]ClientResetCy return cycles, nil } +// quotaDisabledEmails lists the disabled clients whose usage reached their quota, +// only those in within unless it is nil. Read it before zeroing the counters, +// which hides why a client is off. +func quotaDisabledEmails(tx *gorm.DB, within []string) ([]string, error) { + var cycles []ClientResetCycle + if err := clientResetCycles(tx).Where("c.enable = ?", false).Scan(&cycles).Error; err != nil { + return nil, err + } + var in map[string]bool + if within != nil { + in = make(map[string]bool, len(within)) + for _, e := range within { + in[e] = true + } + } + var emails []string + for _, c := range cycles { + if c.Depleted() && (in == nil || in[c.Email]) { + emails = append(emails, c.Email) + } + } + return emails, nil +} + // Create applies the client to every requested inbound: one failing inbound no // longer aborts the others, so the error can name several and needRestart holds. func (s *ClientService) Create(inboundSvc *InboundService, payload *ClientCreatePayload) (bool, error) { diff --git a/internal/web/service/client_traffic.go b/internal/web/service/client_traffic.go index 93016eeac..fb05aa5f4 100644 --- a/internal/web/service/client_traffic.go +++ b/internal/web/service/client_traffic.go @@ -124,26 +124,34 @@ func (s *ClientService) BulkResetTraffic(inboundSvc *InboundService, emails []st return affected, nil } -func (s *ClientService) ResetAllClientTraffics(inboundSvc *InboundService, id int) error { +func (s *ClientService) ResetAllClientTraffics(inboundSvc *InboundService, id int) (bool, error) { var resetNodes []int - var resetEmails []string + var resetEmails, restore []string err := submitTrafficWrite(func() error { var inner error - resetEmails, resetNodes, inner = s.resetAllClientTrafficsLocked(id) + resetEmails, restore, resetNodes, inner = s.resetAllClientTrafficsLocked(id) return inner }) - if err == nil { - inboundSvc.resetMtprotoClientQuotas(resetEmails) - inboundSvc.deliverNodeResetsNow(resetNodes) + if err != nil { + return false, err } - return err + inboundSvc.resetMtprotoClientQuotas(resetEmails) + inboundSvc.deliverNodeResetsNow(resetNodes) + needRestart := false + if len(restore) > 0 { + var enableErr error + if _, needRestart, enableErr = s.BulkSetEnable(inboundSvc, restore, true); enableErr != nil { + logger.Warning("Failed to re-enable depleted clients after resetting inbound traffic:", enableErr) + } + } + return needRestart, nil } -func (s *ClientService) resetAllClientTrafficsLocked(id int) ([]string, []int, error) { +func (s *ClientService) resetAllClientTrafficsLocked(id int) ([]string, []string, []int, error) { db := database.GetDB() now := time.Now().Unix() * 1000 var resetNodes []int - var reset []string + var reset, restore []string if err := db.Transaction(func(tx *gorm.DB) error { // client_traffics.inbound_id is stale: it reflects the inbound the row was @@ -168,6 +176,11 @@ func (s *ClientService) resetAllClientTrafficsLocked(id int) ([]string, []int, e } reset = resetEmails + var qErr error + if restore, qErr = quotaDisabledEmails(tx, resetEmails); qErr != nil { + return qErr + } + if err := adjustGroupBaselinesForRemovedTraffic(tx, resetEmails); err != nil { return err } @@ -189,7 +202,6 @@ func (s *ClientService) resetAllClientTrafficsLocked(id int) ([]string, []int, e return err } } - var qErr error if resetNodes, qErr = queueNodeResets(tx, resetEmails); qErr != nil { return qErr } @@ -207,16 +219,21 @@ func (s *ClientService) resetAllClientTrafficsLocked(id int) ([]string, []int, e return result.Error }); err != nil { - return nil, nil, err + return nil, nil, nil, err } - return reset, resetNodes, nil + return reset, restore, resetNodes, nil } func (s *ClientService) ResetAllTraffics() (bool, error) { var affected int64 var resetNodes []int + var restore []string err := submitTrafficWrite(func() error { return database.GetDB().Transaction(func(tx *gorm.DB) error { + var qErr error + if restore, qErr = quotaDisabledEmails(tx, nil); qErr != nil { + return qErr + } res := tx.Model(&xray.ClientTraffic{}). Where("1 = 1"). Updates(map[string]any{"enable": true, "up": 0, "down": 0}) @@ -230,7 +247,6 @@ func (s *ClientService) ResetAllTraffics() (bool, error) { if err := tx.Where("1 = 1").Delete(&model.NodeClientTraffic{}).Error; err != nil { return err } - var qErr error resetNodes, qErr = queueNodeResets(tx, nil) return qErr }) @@ -241,5 +257,12 @@ func (s *ClientService) ResetAllTraffics() (bool, error) { inbounds := &InboundService{} inbounds.resetAllMtprotoQuotas() inbounds.deliverNodeResetsNow(resetNodes) - return affected > 0, nil + needRestart := false + if len(restore) > 0 { + var enableErr error + if _, needRestart, enableErr = s.BulkSetEnable(inbounds, restore, true); enableErr != nil { + logger.Warning("Failed to re-enable depleted clients after resetting all traffic:", enableErr) + } + } + return affected > 0 || needRestart, nil } diff --git a/internal/web/service/mtproto_quota_reset_test.go b/internal/web/service/mtproto_quota_reset_test.go index efcfec0ef..5b87109bf 100644 --- a/internal/web/service/mtproto_quota_reset_test.go +++ b/internal/web/service/mtproto_quota_reset_test.go @@ -80,7 +80,7 @@ func TestPanelResetsZeroSidecarQuota(t *testing.T) { }) t.Run("inbound clients", func(t *testing.T) { ib, logPath := startQuotaSidecar(t, 46207, model.Client{Enable: true}) - if err := (&ClientService{}).ResetAllClientTraffics(&InboundService{}, ib.Id); err != nil { + if _, err := (&ClientService{}).ResetAllClientTraffics(&InboundService{}, ib.Id); err != nil { t.Fatalf("ResetAllClientTraffics: %v", err) } if got := quotaResets(t, logPath); !slices.Equal(got, []string{"mtga", "mtgb"}) { @@ -128,7 +128,7 @@ func TestInboundClientResetKeepsOtherSidecarQuotas(t *testing.T) { if err := (&ClientService{}).SyncInbound(nil, other.Id, []model.Client{{Email: "vless-only", ID: "11111111-1111-1111-1111-1111111111ab", Enable: true}}); err != nil { t.Fatalf("SyncInbound: %v", err) } - if err := (&ClientService{}).ResetAllClientTraffics(&InboundService{}, other.Id); err != nil { + if _, err := (&ClientService{}).ResetAllClientTraffics(&InboundService{}, other.Id); err != nil { t.Fatalf("ResetAllClientTraffics: %v", err) } if got := quotaResets(t, logPath); len(got) != 0 { diff --git a/internal/web/service/node_reset_undelivered_test.go b/internal/web/service/node_reset_undelivered_test.go index 988fd5b20..535157b37 100644 --- a/internal/web/service/node_reset_undelivered_test.go +++ b/internal/web/service/node_reset_undelivered_test.go @@ -56,7 +56,8 @@ func TestNodeResetNotDeliveredDoesNotRedisableClient(t *testing.T) { return err }}, {"inbound", func(svc *InboundService, ib *model.Inbound) error { - return (&ClientService{}).ResetAllClientTraffics(svc, ib.Id) + _, err := (&ClientService{}).ResetAllClientTraffics(svc, ib.Id) + return err }}, {"all", func(*InboundService, *model.Inbound) error { _, err := (&ClientService{}).ResetAllTraffics() @@ -145,7 +146,8 @@ func TestNodeResetDeliveredRightAway(t *testing.T) { return err }}, {"inbound", func(svc *InboundService, ib *model.Inbound) error { - return (&ClientService{}).ResetAllClientTraffics(svc, ib.Id) + _, err := (&ClientService{}).ResetAllClientTraffics(svc, ib.Id) + return err }}, {"all", func(*InboundService, *model.Inbound) error { _, err := (&ClientService{}).ResetAllTraffics()