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() } }