package main import ( "encoding/json" "os" "path/filepath" "sync" "time" ) const ( mebibyte = 1024 * 1024 persistInterval = 5 * time.Second ) type usageState struct { Day string `json:"day"` Bytes int64 `json:"bytes"` } type Limiter struct { mu sync.Mutex cfg Config loc *time.Location now func() time.Time statePath string serverLimit ServerLimits serverSet ServerSettings hasServer bool battery *batteryWatch day string used int64 active int lastPersist time.Time } func NewLimiter(cfg Config, statePath string, now func() time.Time) (*Limiter, error) { loc, err := cfg.Location() if err != nil { return nil, err } limiter := &Limiter{cfg: cfg, loc: loc, now: now, statePath: statePath, battery: newBatteryWatch(systemBattery, now)} limiter.load() return limiter, nil } func (l *Limiter) load() { if l.statePath == "" { return } data, err := os.ReadFile(l.statePath) if err != nil || len(data) > 4096 { return } var state usageState if json.Unmarshal(data, &state) != nil || state.Bytes < 0 { return } l.day, l.used = state.Day, state.Bytes } func (l *Limiter) persist(force bool) { if l.statePath == "" { return } moment := l.now() if !force && moment.Sub(l.lastPersist) < persistInterval { return } l.lastPersist = moment data, err := json.Marshal(usageState{Day: l.day, Bytes: l.used}) if err != nil { return } if err := os.MkdirAll(filepath.Dir(l.statePath), 0o700); err != nil { return } _ = writeFileAtomic(l.statePath, data) } func (l *Limiter) Flush() { l.mu.Lock() defer l.mu.Unlock() l.roll() l.persist(true) } func (l *Limiter) ApplyServer(limits *ServerLimits, settings *ServerSettings) { l.mu.Lock() defer l.mu.Unlock() if limits != nil { l.serverLimit = *limits } if settings != nil { l.serverSet = *settings } l.hasServer = limits != nil || settings != nil } func (l *Limiter) roll() { today := l.now().In(l.loc).Format("2006-01-02") if l.day != today { l.day = today l.used = 0 } } func minPositive(values ...int) int { best := 0 for _, value := range values { if value > 0 && (best == 0 || value < best) { best = value } } return best } func (l *Limiter) capBytes() int64 { mb := minPositive(l.cfg.BandwidthMBDay, l.serverSet.BandwidthMBDay, l.serverLimit.MaxDonorMBPerDay) return int64(mb) * mebibyte } func (l *Limiter) concurrencyCap() int { return minPositive(l.cfg.MaxConcurrent, l.serverSet.MaxConcurrent, l.serverLimit.MaxConcurrentPerNode) } func minutesOf(text string) int { var hours, minutes int if len(text) != 5 { return -1 } hours = int(text[0]-'0')*10 + int(text[1]-'0') minutes = int(text[3]-'0')*10 + int(text[4]-'0') return hours*60 + minutes } func InWindow(from, to string, moment time.Time) bool { if from == "" || to == "" || from == to { return true } start, end := minutesOf(from), minutesOf(to) if start < 0 || end < 0 { return true } now := moment.Hour()*60 + moment.Minute() if start < end { return now >= start && now < end } return now >= start || now < end } func (l *Limiter) hoursCode() string { moment := l.now() if !InWindow(l.cfg.ActiveFrom, l.cfg.ActiveTo, moment.In(l.loc)) { return RefuseHours } if l.serverSet.ActiveFrom != "" { zone, err := time.LoadLocation(l.serverSet.Timezone) if err != nil { zone = time.UTC } if !InWindow(l.serverSet.ActiveFrom, l.serverSet.ActiveTo, moment.In(zone)) { return RefuseHours } } return "" } func (l *Limiter) SetBatteryProbe(probe BatteryProbe, debugf func(format string, args ...any)) { l.mu.Lock() defer l.mu.Unlock() l.battery = newBatteryWatch(probe, l.now) l.battery.debugf = debugf } func (l *Limiter) Acquire() string { l.mu.Lock() defer l.mu.Unlock() l.roll() if code := l.hoursCode(); code != "" { return code } if l.serverSet.PauseOnBattery && l.battery.OnBattery() { return RefuseBattery } if l.used >= l.capBytes() { return RefuseBandwidth } if l.active >= l.concurrencyCap() { return RefuseConcurrency } l.active++ return "" } func (l *Limiter) Release() { l.mu.Lock() defer l.mu.Unlock() if l.active > 0 { l.active-- } } func (l *Limiter) Remaining() int64 { l.mu.Lock() defer l.mu.Unlock() l.roll() left := l.capBytes() - l.used if left < 0 { return 0 } return left } func (l *Limiter) AddBytes(n int64) { if n <= 0 { return } l.mu.Lock() defer l.mu.Unlock() l.roll() l.used += n l.persist(false) } func (l *Limiter) Usage() (day string, used, capacity int64) { l.mu.Lock() defer l.mu.Unlock() l.roll() return l.day, l.used, l.capBytes() } func ReadUsage(path string, loc *time.Location, now time.Time) int64 { data, err := os.ReadFile(path) if err != nil || len(data) > 4096 { return 0 } var state usageState if json.Unmarshal(data, &state) != nil || state.Day != now.In(loc).Format("2006-01-02") { return 0 } return state.Bytes } func (l *Limiter) ConcurrencyCap() int { l.mu.Lock() defer l.mu.Unlock() return l.concurrencyCap() }