DragonCoreSSH-NewWEB/quota.go

package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/time/rate"
)

const (
	quotaActionBlock    = "block"
	quotaActionThrottle = "throttle"
)

var errDataQuotaExceeded = errors.New("data quota exceeded")

func normalizeQuotaAction(v string) string {
	if strings.EqualFold(strings.TrimSpace(v), quotaActionThrottle) {
		return quotaActionThrottle
	}
	return quotaActionBlock
}

func quotaThrottleMbpsOrDefault(v int) int {
	if v <= 0 {
		return 1
	}
	return v
}

type sshTrafficDelta struct {
	Uplink   int64
	Downlink int64
}

var sshTrafficPersistenceMu sync.Mutex

func (s *Store) AddSSHUserTrafficBatch(ctx context.Context, deltas map[string]sshTrafficDelta) error {
	if s == nil || len(deltas) == 0 {
		return nil
	}
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	stmt, err := tx.PrepareContext(ctx, `
        UPDATE ssh_users SET
            total_uplink_bytes   = GREATEST(total_uplink_bytes + GREATEST($2::BIGINT, 0), 0),
            total_downlink_bytes = GREATEST(total_downlink_bytes + GREATEST($3::BIGINT, 0), 0)
        WHERE username = $1`)
	if err != nil {
		_ = tx.Rollback()
		return err
	}
	defer stmt.Close()
	for username, d := range deltas {
		if strings.TrimSpace(username) == "" || (d.Uplink == 0 && d.Downlink == 0) {
			continue
		}
		if _, err := stmt.ExecContext(ctx, username, d.Uplink, d.Downlink); err != nil {
			_ = tx.Rollback()
			return err
		}
	}
	return tx.Commit()
}

func (s *Store) ResetSSHUserTraffic(ctx context.Context, username string) error {
	if s == nil || strings.TrimSpace(username) == "" {
		return nil
	}
	_, err := s.db.ExecContext(ctx, `
        UPDATE ssh_users
        SET total_uplink_bytes = 0, total_downlink_bytes = 0
        WHERE username = $1`, username)
	return err
}

func initSSHRuntimeUsage(u *UserState, uplink, downlink int64) {
	if u == nil {
		return
	}
	if uplink < 0 {
		uplink = 0
	}
	if downlink < 0 {
		downlink = 0
	}
	atomic.StoreInt64(&u.TotalUplinkBytes, uplink)
	atomic.StoreInt64(&u.TotalDownlinkBytes, downlink)
	atomic.StoreInt64(&u.totalBytes, uplink+downlink)
	atomic.StoreInt64(&u.pendingUplinkBytes, 0)
	atomic.StoreInt64(&u.pendingDownlinkBytes, 0)
}

func resetSSHRuntimeUsageLocked(u *UserState) {
	if u == nil {
		return
	}
	initSSHRuntimeUsage(u, 0, 0)
	u.mu.Lock()
	u.quotaLimiter = nil
	u.quotaLimiterMbps = 0
	u.mu.Unlock()
}

func resetSSHRuntimeUsage(username string) {
	u, ok := userMgr.Get(username)
	if !ok || u == nil {
		return
	}
	u.trafficMu.Lock()
	resetSSHRuntimeUsageLocked(u)
	u.trafficMu.Unlock()
}

func resetSSHUserTrafficAccounting(ctx context.Context, store *Store, username string) error {
	u, _ := userMgr.Get(username)
	if u != nil {
		u.trafficMu.Lock()
		defer u.trafficMu.Unlock()
	}
	sshTrafficPersistenceMu.Lock()
	defer sshTrafficPersistenceMu.Unlock()
	if err := store.ResetSSHUserTraffic(ctx, username); err != nil {
		return err
	}
	if u != nil {
		resetSSHRuntimeUsageLocked(u)
	}
	return nil
}

func sshUserQuotaBlocked(u *UserState) bool {
	if u == nil {
		return false
	}
	u.mu.Lock()
	quota := u.Cfg.DataQuotaBytes
	action := normalizeQuotaAction(u.Cfg.QuotaAction)
	u.mu.Unlock()
	return quota > 0 && action == quotaActionBlock && atomic.LoadInt64(&u.totalBytes) >= quota
}

func sshQuotaLimiter(u *UserState, mbps int) *rate.Limiter {
	mbps = quotaThrottleMbpsOrDefault(mbps)
	u.mu.Lock()
	defer u.mu.Unlock()
	if u.quotaLimiter == nil || u.quotaLimiterMbps != mbps {
		bps := mbpsToBytesPerSec(mbps)
		burst := int(bps)
		if burst < copyBufSize {
			burst = copyBufSize
		}
		u.quotaLimiter = rate.NewLimiter(rate.Limit(bps), burst)
		u.quotaLimiterMbps = mbps
	}
	return u.quotaLimiter
}

func reserveSSHUserBytes(u *UserState, requested int) (allowed int, throttle *rate.Limiter, stopAfter bool) {
	if u == nil || requested <= 0 {
		return 0, nil, false
	}
	u.mu.Lock()
	quota := u.Cfg.DataQuotaBytes
	action := normalizeQuotaAction(u.Cfg.QuotaAction)
	throttleMbps := u.Cfg.QuotaThrottleMbps
	u.mu.Unlock()

	n := int64(requested)
	if quota <= 0 {
		atomic.AddInt64(&u.totalBytes, n)
		return requested, nil, false
	}

	if action == quotaActionThrottle {
		previous := atomic.AddInt64(&u.totalBytes, n) - n
		if previous+n > quota {
			return requested, sshQuotaLimiter(u, throttleMbps), false
		}
		return requested, nil, false
	}

	for {
		used := atomic.LoadInt64(&u.totalBytes)
		remaining := quota - used
		if remaining <= 0 {
			return 0, nil, true
		}
		take := n
		if take > remaining {
			take = remaining
		}
		if atomic.CompareAndSwapInt64(&u.totalBytes, used, used+take) {
			return int(take), nil, take < n || used+take >= quota
		}
	}
}

func finishSSHUserReservation(u *UserState, uplink bool, reserved, written int) {
	if u == nil || reserved <= 0 {
		return
	}
	if written < 0 {
		written = 0
	}
	if written > reserved {
		written = reserved
	}
	if written < reserved {
		atomic.AddInt64(&u.totalBytes, -int64(reserved-written))
	}
	if written == 0 {
		return
	}
	if uplink {
		atomic.AddInt64(&u.TotalUplinkBytes, int64(written))
		atomic.AddInt64(&u.pendingUplinkBytes, int64(written))
	} else {
		atomic.AddInt64(&u.TotalDownlinkBytes, int64(written))
		atomic.AddInt64(&u.pendingDownlinkBytes, int64(written))
	}
}

type sshQuotaWriter struct {
	w      io.Writer
	user   *UserState
	uplink bool
}

func (qw sshQuotaWriter) Write(p []byte) (int, error) {
	if qw.user != nil {
		qw.user.trafficMu.RLock()
		defer qw.user.trafficMu.RUnlock()
	}
	allowed, quotaLimiter, stopAfter := reserveSSHUserBytes(qw.user, len(p))
	if allowed <= 0 {
		return 0, errDataQuotaExceeded
	}
	if quotaLimiter != nil {
		if err := quotaLimiter.WaitN(context.Background(), allowed); err != nil {
			finishSSHUserReservation(qw.user, qw.uplink, allowed, 0)
			return 0, err
		}
	}
	n, err := qw.w.Write(p[:allowed])
	finishSSHUserReservation(qw.user, qw.uplink, allowed, n)
	if err != nil {
		return n, err
	}
	if stopAfter || allowed < len(p) {
		return n, errDataQuotaExceeded
	}
	return n, nil
}

func startSSHUserTrafficFlusher(store *Store) {
	if store == nil {
		return
	}
	go func() {
		ticker := time.NewTicker(5 * time.Second)
		defer ticker.Stop()
		for range ticker.C {
			flushSSHUserTraffic(store)
		}
	}()
}

func flushSSHUserTraffic(store *Store) {
	if store == nil {
		return
	}
	sshTrafficPersistenceMu.Lock()
	defer sshTrafficPersistenceMu.Unlock()
	deltas := make(map[string]sshTrafficDelta)
	states := make(map[string]*UserState)
	for _, u := range userMgr.List() {
		if u == nil || strings.TrimSpace(u.Cfg.Username) == "" {
			continue
		}
		up := atomic.SwapInt64(&u.pendingUplinkBytes, 0)
		down := atomic.SwapInt64(&u.pendingDownlinkBytes, 0)
		if up == 0 && down == 0 {
			continue
		}
		username := u.Cfg.Username
		deltas[username] = sshTrafficDelta{Uplink: up, Downlink: down}
		states[username] = u
	}
	if len(deltas) == 0 {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := store.AddSSHUserTrafficBatch(ctx, deltas); err != nil {
		log.Printf("ssh traffic flush failed: %v", err)
		for username, d := range deltas {
			if u := states[username]; u != nil {
				atomic.AddInt64(&u.pendingUplinkBytes, d.Uplink)
				atomic.AddInt64(&u.pendingDownlinkBytes, d.Downlink)
			}
		}
	}
}

func validateQuotaConfig(quotaBytes int64, action string, throttleMbps int) error {
	if quotaBytes < 0 {
		return fmt.Errorf("data_quota_bytes must be non-negative")
	}
	action = normalizeQuotaAction(action)
	if quotaBytes > 0 && action == quotaActionThrottle && throttleMbps < 0 {
		return fmt.Errorf("quota_throttle_mbps must be non-negative")
	}
	return nil
}