HY2XS_flamy/apps/service/cron.go

package service

import (
	"github.com/sirupsen/logrus"
	"gorm.io/gorm"
	"hy2xs-admin/dao"
	"hy2xs-admin/model/constant"
	"hy2xs-admin/model/entity"
	"hy2xs-admin/proxy"
	"hy2xs-admin/util"
	"sync"
	"time"
)

var trafficMutex sync.Mutex
var kickMutex sync.Mutex

func CronHandleAccount() {
	go func() {
		if !Hysteria2IsRunning() {
			return
		}

		apiPort, err := GetHysteria2ApiPort()
		if err != nil {
			return
		}

		trafficSecretConfig, err := dao.GetConfig("key = ?", constant.Hysteria2TrafficStatsSecret)
		if err != nil {
			return
		}

		// Сохранение данных трафика
		go saveAccountTraffic(apiPort, *trafficSecretConfig.Value)

		// Принудительное отключение
		go kickAccount(apiPort, *trafficSecretConfig.Value)
	}()
}

func CronResetTraffic() {
	peers, err := dao.ListPeer("1=1")
	if err != nil {
		return
	}
	var ids []int64
	for _, item := range peers {
		ids = append(ids, *item.Id)
	}
	idsList := util.SplitArr(ids, 100)
	for _, item := range idsList {
		if err := dao.UpdatePeer(item, map[string]interface{}{"download_bytes": 0, "upload_bytes": 0}); err != nil {
			continue
		}
	}
}

func saveAccountTraffic(apiPort int64, trafficStatsSecret string) {
	if !trafficMutex.TryLock() {
		return
	}
	defer trafficMutex.Unlock()

	users, err := proxy.NewHysteria2Api(apiPort).ListUsers(true, trafficStatsSecret)
	if err != nil {
		return
	}
	if len(users) == 0 {
		return
	}

	nowMs := time.Now().UnixMilli()
	hourStart := nowMs - (nowMs % int64(time.Hour/time.Millisecond))
	dayStart := nowMs - (nowMs % int64(24*time.Hour/time.Millisecond))

	for key, traffic := range users {
		rxBytes := traffic.Rx
		txBytes := traffic.Tx
		if rxBytes == 0 && txBytes == 0 {
			continue
		}

		peer, peerErr := dao.GetPeer("auth_id = ?", key)
		if peerErr != nil {
			continue
		}
		if peer.Id == nil {
			continue
		}

		authId := key
		if peer.AuthId != nil && *peer.AuthId != "" {
			authId = *peer.AuthId
		}

		sample := entity.TrafficSample{
			PeerId:    peer.Id,
			AuthId:    &authId,
			RxBytes:   &rxBytes,
			TxBytes:   &txBytes,
			SampledAt: &nowMs,
		}
		if err = dao.SaveTrafficSample(sample); err != nil {
			logrus.Errorf("save traffic_sample failed: %v", err)
			continue
		}

		if err = dao.UpdatePeer([]int64{*peer.Id}, map[string]interface{}{
			"download_bytes": gorm.Expr("download_bytes + ?", rxBytes),
			"upload_bytes":   gorm.Expr("upload_bytes + ?", txBytes),
		}); err != nil {
			logrus.Errorf("update peer traffic failed: %v", err)
			continue
		}

		_ = dao.UpsertTrafficAggregateHourly(*peer.Id, hourStart, rxBytes, txBytes)
		_ = dao.UpsertTrafficAggregateDaily(*peer.Id, dayStart, rxBytes, txBytes)
	}
}
func kickAccount(apiPort int64, trafficStatsSecret string) {
	if !kickMutex.TryLock() {
		return
	}
	defer kickMutex.Unlock()

	users, err := proxy.NewHysteria2Api(apiPort).OnlineUsers(trafficStatsSecret)
	if err != nil {
		return
	}
	if len(users) > 0 {
		i := 0
		authIDs := make([]string, len(users))
		for k := range users {
			authIDs[i] = k
			i++
		}
		authIDLists := util.SplitArr(authIDs, 10)
		var wg sync.WaitGroup
		for _, authIDList := range authIDLists {
			wg.Add(1)
			go func(authIDList []string) {
				defer wg.Done()
				now := time.Now().UnixMilli()
				peers, err := dao.ListPeer(`auth_id in ? and (
					disabled = 1
					or (quota_bytes > 0 and quota_bytes < download_bytes + upload_bytes)
					or (expires_at > 0 and ? > expires_at)
					or ? < banned_until
				)`, authIDList, now, now)
				if err != nil {
					return
				}
				kickAuthIDs := make([]string, len(peers))
				j := 0
				for _, item := range peers {
					kickAuthIDs[j] = *item.AuthId
					j++
				}
				if err = proxy.NewHysteria2Api(apiPort).KickUsers(kickAuthIDs, trafficStatsSecret); err != nil {
					return
				}
			}(authIDList)
		}
		wg.Wait()
	}
}