From 8484945bb22b1236e01f6369bb1e153696f11c59 Mon Sep 17 00:00:00 2001 From: egregors Date: Sun, 31 Aug 2025 20:57:53 +0200 Subject: [PATCH] feat: implement periodic traffic reset job and integrate with cron scheduler --- web/job/periodic_traffic_reset_job.go | 124 +++++ ...odic_traffic_reset_job_integration_test.go | 333 +++++++++++++ web/job/periodic_traffic_reset_job_test.go | 455 ++++++++++++++++++ web/web.go | 3 + 4 files changed, 915 insertions(+) create mode 100644 web/job/periodic_traffic_reset_job.go create mode 100644 web/job/periodic_traffic_reset_job_integration_test.go create mode 100644 web/job/periodic_traffic_reset_job_test.go diff --git a/web/job/periodic_traffic_reset_job.go b/web/job/periodic_traffic_reset_job.go new file mode 100644 index 00000000..c718f177 --- /dev/null +++ b/web/job/periodic_traffic_reset_job.go @@ -0,0 +1,124 @@ +package job + +import ( + "sync" + "time" + + "x-ui/database/model" + "x-ui/logger" + "x-ui/web/service" +) + +const ( + resetTypeNever = "never" + resetTypeDaily = "daily" + resetTypeWeekly = "weekly" + resetTypeMonthly = "monthly" +) + +var validResetTypes = map[string]bool{ + resetTypeNever: true, + resetTypeDaily: true, + resetTypeWeekly: true, + resetTypeMonthly: true, +} + +type PeriodicTrafficResetJob struct { + inboundService service.InboundService + lastResetTimes map[string]time.Time + mu sync.RWMutex +} + +func NewPeriodicTrafficResetJob() *PeriodicTrafficResetJob { + return &PeriodicTrafficResetJob{ + lastResetTimes: make(map[string]time.Time), + mu: sync.RWMutex{}, + } +} + +func (j *PeriodicTrafficResetJob) Run() { + inbounds, err := j.inboundService.GetAllInbounds() + if err != nil { + logger.Warning("Failed to get inbounds for traffic reset:", err) + return + } + + resetCount := 0 + now := time.Now() + + for _, inbound := range inbounds { + if !j.shouldResetTraffic(inbound, now) { + continue + } + + if err := j.inboundService.ResetAllClientTraffics(inbound.Id); err != nil { + logger.Warning("Failed to reset traffic for inbound", inbound.Id, ":", err) + continue + } + + j.updateLastResetTime(inbound, now) + resetCount++ + logger.Infof("Reset traffic for inbound %d (%s)", inbound.Id, inbound.Remark) + } + + if resetCount > 0 { + logger.Infof("Periodic traffic reset completed: %d inbounds reset", resetCount) + } +} + +func (j *PeriodicTrafficResetJob) shouldResetTraffic(inbound *model.Inbound, now time.Time) bool { + if !validResetTypes[inbound.PeriodicTrafficReset] || inbound.PeriodicTrafficReset == resetTypeNever { + return false + } + + resetKey := j.getResetKey(inbound) + lastReset := j.getLastResetTime(resetKey) + + switch inbound.PeriodicTrafficReset { + case resetTypeDaily: + return j.shouldResetDaily(now, lastReset) + case resetTypeWeekly: + return j.shouldResetWeekly(now, lastReset) + case resetTypeMonthly: + return j.shouldResetMonthly(now, lastReset) + default: + return false + } +} + +func (j *PeriodicTrafficResetJob) shouldResetDaily(now, lastReset time.Time) bool { + if lastReset.IsZero() { + return now.Hour() == 0 && now.Minute() < 10 + } + return now.Sub(lastReset) >= 24*time.Hour && now.Hour() == 0 && now.Minute() < 10 +} + +func (j *PeriodicTrafficResetJob) shouldResetWeekly(now, lastReset time.Time) bool { + if lastReset.IsZero() { + return now.Weekday() == time.Sunday && now.Hour() == 0 && now.Minute() < 10 + } + return now.Sub(lastReset) >= 7*24*time.Hour && now.Weekday() == time.Sunday && now.Hour() == 0 && now.Minute() < 10 +} + +func (j *PeriodicTrafficResetJob) shouldResetMonthly(now, lastReset time.Time) bool { + if lastReset.IsZero() { + return now.Day() == 1 && now.Hour() == 0 && now.Minute() < 10 + } + return now.Sub(lastReset) >= 28*24*time.Hour && now.Day() == 1 && now.Hour() == 0 && now.Minute() < 10 +} + +func (j *PeriodicTrafficResetJob) getResetKey(inbound *model.Inbound) string { + return inbound.PeriodicTrafficReset + "_" + inbound.Tag +} + +func (j *PeriodicTrafficResetJob) getLastResetTime(key string) time.Time { + j.mu.RLock() + defer j.mu.RUnlock() + return j.lastResetTimes[key] +} + +func (j *PeriodicTrafficResetJob) updateLastResetTime(inbound *model.Inbound, resetTime time.Time) { + j.mu.Lock() + defer j.mu.Unlock() + j.lastResetTimes[j.getResetKey(inbound)] = resetTime +} diff --git a/web/job/periodic_traffic_reset_job_integration_test.go b/web/job/periodic_traffic_reset_job_integration_test.go new file mode 100644 index 00000000..65190ce0 --- /dev/null +++ b/web/job/periodic_traffic_reset_job_integration_test.go @@ -0,0 +1,333 @@ +package job + +import ( + "fmt" + "testing" + "time" + + "x-ui/database/model" +) + +// Tests to ensure the job handles various real-world scenarios properly +func TestPeriodicTrafficResetJobNewInstance(t *testing.T) { + job := NewPeriodicTrafficResetJob() + + if job == nil { + t.Fatal("Expected job to be created, got nil") + } + + if job.lastResetTimes == nil { + t.Fatal("Expected lastResetTimes map to be initialized") + } + + if len(job.lastResetTimes) != 0 { + t.Error("Expected lastResetTimes map to be empty initially") + } +} + +func TestPeriodicTrafficResetJobBoundaryConditions(t *testing.T) { + job := NewPeriodicTrafficResetJob() + + // Test edge cases for time boundaries + testCases := []struct { + name string + resetType string + now time.Time + lastReset time.Time + expectedResult bool + }{ + { + name: "daily reset at exact midnight boundary", + resetType: "daily", + now: time.Date(2024, 1, 15, 0, 0, 0, 0, time.UTC), + lastReset: time.Time{}, + expectedResult: true, + }, + { + name: "daily reset one second after midnight boundary", + resetType: "daily", + now: time.Date(2024, 1, 15, 0, 0, 1, 0, time.UTC), + lastReset: time.Time{}, + expectedResult: true, + }, + { + name: "weekly reset on Sunday at exact midnight", + resetType: "weekly", + now: time.Date(2024, 1, 14, 0, 0, 0, 0, time.UTC), // Sunday + lastReset: time.Time{}, + expectedResult: true, + }, + { + name: "monthly reset on 1st at exact midnight", + resetType: "monthly", + now: time.Date(2024, 2, 1, 0, 0, 0, 0, time.UTC), + lastReset: time.Time{}, + expectedResult: true, + }, + { + name: "daily reset at 10-minute boundary (should not reset)", + resetType: "daily", + now: time.Date(2024, 1, 15, 0, 10, 0, 0, time.UTC), + lastReset: time.Time{}, + expectedResult: false, + }, + { + name: "weekly reset at 10-minute boundary (should not reset)", + resetType: "weekly", + now: time.Date(2024, 1, 14, 0, 10, 0, 0, time.UTC), // Sunday + lastReset: time.Time{}, + expectedResult: false, + }, + { + name: "monthly reset at 10-minute boundary (should not reset)", + resetType: "monthly", + now: time.Date(2024, 2, 1, 0, 10, 0, 0, time.UTC), + lastReset: time.Time{}, + expectedResult: false, + }, + { + name: "leap year handling - February 29th to March 1st", + resetType: "monthly", + now: time.Date(2024, 3, 1, 0, 5, 0, 0, time.UTC), + lastReset: time.Date(2024, 2, 1, 0, 5, 0, 0, time.UTC), // 29 days ago + expectedResult: true, + }, + { + name: "weekly reset after 6 days and 23 hours (should not reset)", + resetType: "weekly", + now: time.Date(2024, 1, 20, 23, 0, 0, 0, time.UTC), // Saturday + lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), // Previous Sunday + expectedResult: false, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + inbound := &model.Inbound{ + Id: 1, + Tag: "test-inbound", + PeriodicTrafficReset: tc.resetType, + } + + // Set up last reset time if provided + if !tc.lastReset.IsZero() { + job.updateLastResetTime(inbound, tc.lastReset) + } + + result := job.shouldResetTraffic(inbound, tc.now) + if result != tc.expectedResult { + t.Errorf("Expected %v, got %v", tc.expectedResult, result) + } + + // Clean up for next test + job.lastResetTimes = make(map[string]time.Time) + }) + } +} + +func TestPeriodicTrafficResetJobTimezoneHandling(t *testing.T) { + job := NewPeriodicTrafficResetJob() + + inbound := &model.Inbound{ + Id: 1, + Tag: "timezone-test", + PeriodicTrafficReset: "daily", + } + + // Test that UTC midnight behaves consistently + utcMidnight := time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC) + result := job.shouldResetTraffic(inbound, utcMidnight) + + if !result { + t.Error("Expected reset at UTC midnight") + } + + // Test that the same instant in different timezones produces the same absolute behavior + // This tests consistency rather than timezone conversion + utcTime := time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC) + utcTimestamp := utcTime.Unix() + + // Create equivalent time from timestamp + equivalentTime := time.Unix(utcTimestamp, 0).UTC() + + result1 := job.shouldResetTraffic(inbound, utcTime) + result2 := job.shouldResetTraffic(inbound, equivalentTime) + + if result1 != result2 { + t.Errorf("Equivalent times should behave consistently: %v vs %v", result1, result2) + } +} + +func TestPeriodicTrafficResetJobConcurrentSafety(t *testing.T) { + job := NewPeriodicTrafficResetJob() + + inbound := &model.Inbound{ + Id: 1, + Tag: "concurrent-test", + PeriodicTrafficReset: "daily", + } + + // Test concurrent operations + done := make(chan bool, 100) + + // Start multiple goroutines performing different operations + for i := 0; i < 50; i++ { + go func(i int) { + defer func() { done <- true }() + + // Mix of operations + testTime := time.Date(2024, 1, 15+i%5, 0, 5, 0, 0, time.UTC) + job.updateLastResetTime(inbound, testTime) + _ = job.getLastResetTime(job.getResetKey(inbound)) + _ = job.shouldResetTraffic(inbound, testTime) + }(i) + } + + // Wait for all goroutines to complete + for i := 0; i < 50; i++ { + select { + case <-done: + // Good + case <-time.After(5 * time.Second): + t.Fatal("Concurrent test timed out - possible deadlock") + } + } +} + +func TestPeriodicTrafficResetJobValidationRobustness(t *testing.T) { + job := NewPeriodicTrafficResetJob() + + testCases := []struct { + name string + inbound *model.Inbound + now time.Time + expected bool + }{ + { + name: "nil inbound", + inbound: nil, + now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + // This should panic or handle gracefully - depends on implementation + }, + { + name: "empty tag", + inbound: &model.Inbound{ + Id: 1, + Tag: "", + PeriodicTrafficReset: "daily", + }, + now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: true, // Should still work with empty tag + }, + { + name: "very long tag", + inbound: &model.Inbound{ + Id: 1, + Tag: "very-very-very-long-tag-name-that-might-cause-issues-with-memory-or-key-generation-in-the-reset-tracking-system", + PeriodicTrafficReset: "daily", + }, + now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: true, + }, + { + name: "special characters in tag", + inbound: &model.Inbound{ + Id: 1, + Tag: "test-inbound_with.special@chars#123", + PeriodicTrafficReset: "daily", + }, + now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: true, + }, + { + name: "zero time", + inbound: &model.Inbound{ + Id: 1, + Tag: "zero-time-test", + PeriodicTrafficReset: "daily", + }, + now: time.Time{}, // Zero time + expected: true, // Zero time has hour=0, minute=0, so should reset according to logic + }, + { + name: "far future time", + inbound: &model.Inbound{ + Id: 1, + Tag: "future-test", + PeriodicTrafficReset: "daily", + }, + now: time.Date(2099, 12, 31, 0, 5, 0, 0, time.UTC), + expected: true, + }, + { + name: "far past time", + inbound: &model.Inbound{ + Id: 1, + Tag: "past-test", + PeriodicTrafficReset: "daily", + }, + now: time.Date(1970, 1, 1, 0, 5, 0, 0, time.UTC), + expected: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + defer func() { + if r := recover(); r != nil { + if tc.inbound == nil { + // Expected panic for nil inbound + t.Log("Expected panic for nil inbound:", r) + } else { + t.Errorf("Unexpected panic: %v", r) + } + } + }() + + if tc.inbound == nil { + // Skip execution for nil test case + return + } + + result := job.shouldResetTraffic(tc.inbound, tc.now) + if result != tc.expected { + t.Errorf("Expected %v, got %v", tc.expected, result) + } + }) + } +} + +func TestPeriodicTrafficResetJobMemoryManagement(t *testing.T) { + job := NewPeriodicTrafficResetJob() + + // Test with many inbounds to check memory usage + const numInbounds = 1000 + testTime := time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC) + + for i := 0; i < numInbounds; i++ { + inbound := &model.Inbound{ + Id: i, + Tag: fmt.Sprintf("test-inbound-%d", i), + PeriodicTrafficReset: "daily", + } + job.updateLastResetTime(inbound, testTime) + } + + // Verify all entries are stored + if len(job.lastResetTimes) != numInbounds { + t.Errorf("Expected %d entries in lastResetTimes, got %d", numInbounds, len(job.lastResetTimes)) + } + + // Test retrieval of all entries + for i := 0; i < numInbounds; i++ { + inbound := &model.Inbound{ + Id: i, + Tag: fmt.Sprintf("test-inbound-%d", i), + PeriodicTrafficReset: "daily", + } + retrievedTime := job.getLastResetTime(job.getResetKey(inbound)) + if !retrievedTime.Equal(testTime) { + t.Errorf("Expected time %v for inbound %d, got %v", testTime, i, retrievedTime) + } + } +} diff --git a/web/job/periodic_traffic_reset_job_test.go b/web/job/periodic_traffic_reset_job_test.go new file mode 100644 index 00000000..e6116995 --- /dev/null +++ b/web/job/periodic_traffic_reset_job_test.go @@ -0,0 +1,455 @@ +package job + +import ( + "testing" + "time" + + "x-ui/database/model" +) + +func TestShouldResetDaily(t *testing.T) { + job := NewPeriodicTrafficResetJob() + + tests := []struct { + name string + now time.Time + lastReset time.Time + expected bool + }{ + { + name: "first time reset at midnight", + now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + lastReset: time.Time{}, // zero time + expected: true, + }, + { + name: "first time reset not at midnight", + now: time.Date(2024, 1, 15, 12, 0, 0, 0, time.UTC), + lastReset: time.Time{}, // zero time + expected: false, + }, + { + name: "reset after 24 hours at midnight", + now: time.Date(2024, 1, 16, 0, 5, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: true, + }, + { + name: "reset after 24 hours but not at midnight", + now: time.Date(2024, 1, 16, 12, 0, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "reset same day", + now: time.Date(2024, 1, 15, 0, 8, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "reset after 23 hours", + now: time.Date(2024, 1, 15, 23, 5, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "midnight but minute > 10", + now: time.Date(2024, 1, 16, 0, 15, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "edge case: exactly at minute 10", + now: time.Date(2024, 1, 16, 0, 10, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "edge case: minute 9 (last valid minute)", + now: time.Date(2024, 1, 16, 0, 9, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := job.shouldResetDaily(tt.now, tt.lastReset) + if result != tt.expected { + t.Errorf("shouldResetDaily() = %v, expected %v", result, tt.expected) + } + }) + } +} + +func TestShouldResetWeekly(t *testing.T) { + job := NewPeriodicTrafficResetJob() + + tests := []struct { + name string + now time.Time + lastReset time.Time + expected bool + }{ + { + name: "first time reset on Sunday at midnight", + now: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), // Sunday + lastReset: time.Time{}, // zero time + expected: true, + }, + { + name: "first time reset on Monday", + now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), // Monday + lastReset: time.Time{}, // zero time + expected: false, + }, + { + name: "reset after 7 days on Sunday at midnight", + now: time.Date(2024, 1, 21, 0, 5, 0, 0, time.UTC), // Sunday + lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), // Previous Sunday + expected: true, + }, + { + name: "reset after 7 days on Sunday but not at midnight", + now: time.Date(2024, 1, 21, 12, 0, 0, 0, time.UTC), // Sunday + lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "reset after 6 days on Saturday", + now: time.Date(2024, 1, 20, 0, 5, 0, 0, time.UTC), // Saturday + lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "reset same Sunday", + now: time.Date(2024, 1, 14, 0, 8, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "reset after 8 days but on Monday", + now: time.Date(2024, 1, 22, 0, 5, 0, 0, time.UTC), // Monday + lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), // Previous Sunday + expected: false, + }, + { + name: "edge case: Sunday but minute > 10", + now: time.Date(2024, 1, 21, 0, 15, 0, 0, time.UTC), // Sunday + lastReset: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := job.shouldResetWeekly(tt.now, tt.lastReset) + if result != tt.expected { + t.Errorf("shouldResetWeekly() = %v, expected %v", result, tt.expected) + } + }) + } +} + +func TestShouldResetMonthly(t *testing.T) { + job := NewPeriodicTrafficResetJob() + + tests := []struct { + name string + now time.Time + lastReset time.Time + expected bool + }{ + { + name: "first time reset on 1st day at midnight", + now: time.Date(2024, 2, 1, 0, 5, 0, 0, time.UTC), + lastReset: time.Time{}, // zero time + expected: true, + }, + { + name: "first time reset on 15th day", + now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + lastReset: time.Time{}, // zero time + expected: false, + }, + { + name: "reset after 28 days on 1st day at midnight", + now: time.Date(2024, 2, 1, 0, 5, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 4, 0, 5, 0, 0, time.UTC), // 28 days ago + expected: true, + }, + { + name: "reset after 32 days on 1st day at midnight", + now: time.Date(2024, 3, 1, 0, 5, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 29, 0, 5, 0, 0, time.UTC), // 32 days ago + expected: true, + }, + { + name: "reset after 28 days on 1st day but not at midnight", + now: time.Date(2024, 2, 1, 12, 0, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 4, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "reset after 27 days", + now: time.Date(2024, 1, 31, 0, 5, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 4, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "reset same day (1st)", + now: time.Date(2024, 1, 1, 0, 8, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 1, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "reset after 28 days but on 2nd day", + now: time.Date(2024, 2, 2, 0, 5, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 5, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "leap year February edge case", + now: time.Date(2024, 3, 1, 0, 5, 0, 0, time.UTC), // March 1st in leap year + lastReset: time.Date(2024, 2, 1, 0, 5, 0, 0, time.UTC), // February 1st (29 days ago) + expected: true, + }, + { + name: "edge case: 1st day but minute > 10", + now: time.Date(2024, 2, 1, 0, 15, 0, 0, time.UTC), + lastReset: time.Date(2024, 1, 4, 0, 5, 0, 0, time.UTC), + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := job.shouldResetMonthly(tt.now, tt.lastReset) + if result != tt.expected { + t.Errorf("shouldResetMonthly() = %v, expected %v", result, tt.expected) + } + }) + } +} + +func TestShouldResetTraffic(t *testing.T) { + tests := []struct { + name string + inbound *model.Inbound + now time.Time + setup func(*PeriodicTrafficResetJob, *model.Inbound) + expected bool + }{ + { + name: "never reset type", + inbound: &model.Inbound{ + Id: 1, + Tag: "test1", + PeriodicTrafficReset: "never", + }, + now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "invalid reset type", + inbound: &model.Inbound{ + Id: 2, + Tag: "test2", + PeriodicTrafficReset: "invalid", + }, + now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: false, + }, + { + name: "daily reset should trigger", + inbound: &model.Inbound{ + Id: 3, + Tag: "test3", + PeriodicTrafficReset: "daily", + }, + now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + expected: true, + }, + { + name: "weekly reset should trigger", + inbound: &model.Inbound{ + Id: 4, + Tag: "test4", + PeriodicTrafficReset: "weekly", + }, + now: time.Date(2024, 1, 14, 0, 5, 0, 0, time.UTC), // Sunday + expected: true, + }, + { + name: "monthly reset should trigger", + inbound: &model.Inbound{ + Id: 5, + Tag: "test5", + PeriodicTrafficReset: "monthly", + }, + now: time.Date(2024, 2, 1, 0, 5, 0, 0, time.UTC), + expected: true, + }, + { + name: "daily reset with existing reset time - should not trigger", + inbound: &model.Inbound{ + Id: 6, + Tag: "test6", + PeriodicTrafficReset: "daily", + }, + now: time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC), + setup: func(j *PeriodicTrafficResetJob, inbound *model.Inbound) { + j.updateLastResetTime(inbound, time.Date(2024, 1, 15, 0, 3, 0, 0, time.UTC)) + }, + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create new job for each test to avoid state interference + testJob := NewPeriodicTrafficResetJob() + + if tt.setup != nil { + tt.setup(testJob, tt.inbound) + } + + result := testJob.shouldResetTraffic(tt.inbound, tt.now) + if result != tt.expected { + t.Errorf("shouldResetTraffic() = %v, expected %v", result, tt.expected) + } + }) + } +} + +func TestGetResetKey(t *testing.T) { + job := NewPeriodicTrafficResetJob() + + tests := []struct { + name string + inbound *model.Inbound + expected string + }{ + { + name: "daily reset key", + inbound: &model.Inbound{ + Tag: "test-inbound", + PeriodicTrafficReset: "daily", + }, + expected: "daily_test-inbound", + }, + { + name: "weekly reset key", + inbound: &model.Inbound{ + Tag: "proxy-1", + PeriodicTrafficReset: "weekly", + }, + expected: "weekly_proxy-1", + }, + { + name: "monthly reset key", + inbound: &model.Inbound{ + Tag: "vless-tcp", + PeriodicTrafficReset: "monthly", + }, + expected: "monthly_vless-tcp", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := job.getResetKey(tt.inbound) + if result != tt.expected { + t.Errorf("getResetKey() = %v, expected %v", result, tt.expected) + } + }) + } +} + +func TestLastResetTimeOperations(t *testing.T) { + job := NewPeriodicTrafficResetJob() + + inbound := &model.Inbound{ + Id: 1, + Tag: "test-inbound", + PeriodicTrafficReset: "daily", + } + + // Test getting non-existent reset time (should return zero time) + resetTime := job.getLastResetTime(job.getResetKey(inbound)) + if !resetTime.IsZero() { + t.Errorf("Expected zero time for non-existent key, got %v", resetTime) + } + + // Test updating and getting reset time + testTime := time.Date(2024, 1, 15, 0, 5, 0, 0, time.UTC) + job.updateLastResetTime(inbound, testTime) + + retrievedTime := job.getLastResetTime(job.getResetKey(inbound)) + if !retrievedTime.Equal(testTime) { + t.Errorf("Expected %v, got %v", testTime, retrievedTime) + } + + // Test updating existing reset time + newTime := time.Date(2024, 1, 16, 0, 5, 0, 0, time.UTC) + job.updateLastResetTime(inbound, newTime) + + retrievedTime = job.getLastResetTime(job.getResetKey(inbound)) + if !retrievedTime.Equal(newTime) { + t.Errorf("Expected %v, got %v", newTime, retrievedTime) + } +} + +// Test for concurrent access to lastResetTimes map +func TestConcurrentAccess(t *testing.T) { + job := NewPeriodicTrafficResetJob() + + inbound := &model.Inbound{ + Id: 1, + Tag: "concurrent-test", + PeriodicTrafficReset: "daily", + } + + // Run multiple goroutines to test concurrent access + done := make(chan bool, 20) + + // 10 writers + for i := 0; i < 10; i++ { + go func(i int) { + testTime := time.Date(2024, 1, 15+i, 0, 5, 0, 0, time.UTC) + job.updateLastResetTime(inbound, testTime) + done <- true + }(i) + } + + // 10 readers + for i := 0; i < 10; i++ { + go func() { + _ = job.getLastResetTime(job.getResetKey(inbound)) + done <- true + }() + } + + // Wait for all goroutines to complete + for i := 0; i < 20; i++ { + <-done + } + + // If we reach here without panic/deadlock, the concurrent access is safe + t.Log("Concurrent access test passed") +} + +func TestValidResetTypes(t *testing.T) { + expectedTypes := []string{"never", "daily", "weekly", "monthly"} + + for _, resetType := range expectedTypes { + if !validResetTypes[resetType] { + t.Errorf("Expected reset type %s to be valid", resetType) + } + } + + // Test invalid type + if validResetTypes["invalid"] { + t.Error("Expected 'invalid' to not be a valid reset type") + } +} diff --git a/web/web.go b/web/web.go index 35ccec70..e5be536f 100644 --- a/web/web.go +++ b/web/web.go @@ -266,6 +266,9 @@ func (s *Server) startTask() { // check client ips from log file every day s.cron.AddJob("@daily", job.NewClearLogsJob()) + // Check for periodic traffic resets every 10 minutes + s.cron.AddJob("@every 10m", job.NewPeriodicTrafficResetJob()) + // Make a traffic condition every day, 8:30 var entry cron.EntryID isTgbotenabled, err := s.settingService.GetTgbotEnabled()