mirror of https://github.com/v2ray/v2ray-core
				
				
				
			Generate user hash on demand
							parent
							
								
									46abe9de5e
								
							
						
					
					
						commit
						ab723fa0b5
					
				
							
								
								
									
										68
									
								
								userset.go
								
								
								
								
							
							
						
						
									
										68
									
								
								userset.go
								
								
								
								
							| 
						 | 
					@ -1,9 +1,11 @@
 | 
				
			||||||
package core
 | 
					package core
 | 
				
			||||||
 | 
					
 | 
				
			||||||
import (
 | 
					import (
 | 
				
			||||||
 | 
						"container/heap"
 | 
				
			||||||
	"time"
 | 
						"time"
 | 
				
			||||||
 | 
					
 | 
				
			||||||
	v2hash "github.com/v2ray/v2ray-core/hash"
 | 
						v2hash "github.com/v2ray/v2ray-core/hash"
 | 
				
			||||||
 | 
						"github.com/v2ray/v2ray-core/log"
 | 
				
			||||||
)
 | 
					)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
const (
 | 
					const (
 | 
				
			||||||
| 
						 | 
					@ -19,6 +21,7 @@ type UserSet interface {
 | 
				
			||||||
type TimedUserSet struct {
 | 
					type TimedUserSet struct {
 | 
				
			||||||
	validUserIds []ID
 | 
						validUserIds []ID
 | 
				
			||||||
	userHashes   map[string]indexTimePair
 | 
						userHashes   map[string]indexTimePair
 | 
				
			||||||
 | 
						hash2Remove  hashEntrySet
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
type indexTimePair struct {
 | 
					type indexTimePair struct {
 | 
				
			||||||
| 
						 | 
					@ -31,22 +34,62 @@ type hashEntry struct {
 | 
				
			||||||
	timeSec int64
 | 
						timeSec int64
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					type hashEntrySet []*hashEntry
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					func (set hashEntrySet) Len() int {
 | 
				
			||||||
 | 
						return len(set)
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					func (set hashEntrySet) Less(i, j int) bool {
 | 
				
			||||||
 | 
						return set[i].timeSec < set[j].timeSec
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					func (set hashEntrySet) Swap(i, j int) {
 | 
				
			||||||
 | 
						tmp := set[i]
 | 
				
			||||||
 | 
						set[i] = set[j]
 | 
				
			||||||
 | 
						set[j] = tmp
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					func (set *hashEntrySet) Push(value interface{}) {
 | 
				
			||||||
 | 
						entry := value.(*hashEntry)
 | 
				
			||||||
 | 
						*set = append(*set, entry)
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					func (set *hashEntrySet) Pop() interface{} {
 | 
				
			||||||
 | 
						old := *set
 | 
				
			||||||
 | 
						n := len(old)
 | 
				
			||||||
 | 
						v := old[n-1]
 | 
				
			||||||
 | 
						*set = old[:n-1]
 | 
				
			||||||
 | 
						return v
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
func NewTimedUserSet() UserSet {
 | 
					func NewTimedUserSet() UserSet {
 | 
				
			||||||
	vuSet := new(TimedUserSet)
 | 
						vuSet := new(TimedUserSet)
 | 
				
			||||||
	vuSet.validUserIds = make([]ID, 0, 16)
 | 
						vuSet.validUserIds = make([]ID, 0, 16)
 | 
				
			||||||
	vuSet.userHashes = make(map[string]indexTimePair)
 | 
						vuSet.userHashes = make(map[string]indexTimePair)
 | 
				
			||||||
 | 
						vuSet.hash2Remove = make(hashEntrySet, 0, cacheDurationSec*10)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
	go vuSet.updateUserHash(time.Tick(updateIntervalSec * time.Second))
 | 
						go vuSet.updateUserHash(time.Tick(updateIntervalSec * time.Second))
 | 
				
			||||||
	return vuSet
 | 
						return vuSet
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					func (us *TimedUserSet) generateNewHashes(lastSec, nowSec int64, idx int, id ID) {
 | 
				
			||||||
 | 
						idHash := v2hash.NewTimeHash(v2hash.HMACHash{})
 | 
				
			||||||
 | 
						for lastSec < nowSec+cacheDurationSec {
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
							idHash := idHash.Hash(id.Bytes, lastSec)
 | 
				
			||||||
 | 
							log.Debug("Valid User Hash: %v", idHash)
 | 
				
			||||||
 | 
							heap.Push(&us.hash2Remove, &hashEntry{string(idHash), lastSec})
 | 
				
			||||||
 | 
							us.userHashes[string(idHash)] = indexTimePair{idx, lastSec}
 | 
				
			||||||
 | 
							lastSec++
 | 
				
			||||||
 | 
						}
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
func (us *TimedUserSet) updateUserHash(tick <-chan time.Time) {
 | 
					func (us *TimedUserSet) updateUserHash(tick <-chan time.Time) {
 | 
				
			||||||
	now := time.Now().UTC()
 | 
						now := time.Now().UTC()
 | 
				
			||||||
	lastSec := now.Unix() - cacheDurationSec
 | 
						lastSec := now.Unix()
 | 
				
			||||||
 | 
					 | 
				
			||||||
	hash2Remove := make(chan hashEntry, cacheDurationSec*3*len(us.validUserIds))
 | 
					 | 
				
			||||||
	lastSec2Remove := now.Unix()
 | 
						lastSec2Remove := now.Unix()
 | 
				
			||||||
	idHash := v2hash.NewTimeHash(v2hash.HMACHash{})
 | 
					
 | 
				
			||||||
	for {
 | 
						for {
 | 
				
			||||||
		now := <-tick
 | 
							now := <-tick
 | 
				
			||||||
		nowSec := now.UTC().Unix()
 | 
							nowSec := now.UTC().Unix()
 | 
				
			||||||
| 
						 | 
					@ -54,26 +97,27 @@ func (us *TimedUserSet) updateUserHash(tick <-chan time.Time) {
 | 
				
			||||||
		remove2Sec := nowSec - cacheDurationSec
 | 
							remove2Sec := nowSec - cacheDurationSec
 | 
				
			||||||
		if remove2Sec > lastSec2Remove {
 | 
							if remove2Sec > lastSec2Remove {
 | 
				
			||||||
			for lastSec2Remove+1 < remove2Sec {
 | 
								for lastSec2Remove+1 < remove2Sec {
 | 
				
			||||||
				entry := <-hash2Remove
 | 
									front := heap.Pop(&us.hash2Remove)
 | 
				
			||||||
 | 
									entry := front.(*hashEntry)
 | 
				
			||||||
				lastSec2Remove = entry.timeSec
 | 
									lastSec2Remove = entry.timeSec
 | 
				
			||||||
				delete(us.userHashes, entry.hash)
 | 
									delete(us.userHashes, entry.hash)
 | 
				
			||||||
			}
 | 
								}
 | 
				
			||||||
		}
 | 
							}
 | 
				
			||||||
 | 
					 | 
				
			||||||
		for lastSec < nowSec+cacheDurationSec {
 | 
					 | 
				
			||||||
		for idx, id := range us.validUserIds {
 | 
							for idx, id := range us.validUserIds {
 | 
				
			||||||
				idHash := idHash.Hash(id.Bytes, lastSec)
 | 
								us.generateNewHashes(lastSec, nowSec, idx, id)
 | 
				
			||||||
				hash2Remove <- hashEntry{string(idHash), lastSec}
 | 
					 | 
				
			||||||
				us.userHashes[string(idHash)] = indexTimePair{idx, lastSec}
 | 
					 | 
				
			||||||
			}
 | 
					 | 
				
			||||||
			lastSec++
 | 
					 | 
				
			||||||
		}
 | 
							}
 | 
				
			||||||
	}
 | 
						}
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
func (us *TimedUserSet) AddUser(user User) error {
 | 
					func (us *TimedUserSet) AddUser(user User) error {
 | 
				
			||||||
	id := user.Id
 | 
						id := user.Id
 | 
				
			||||||
 | 
						idx := len(us.validUserIds)
 | 
				
			||||||
	us.validUserIds = append(us.validUserIds, id)
 | 
						us.validUserIds = append(us.validUserIds, id)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
						nowSec := time.Now().UTC().Unix()
 | 
				
			||||||
 | 
						lastSec := nowSec - cacheDurationSec
 | 
				
			||||||
 | 
						us.generateNewHashes(lastSec, nowSec, idx, id)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
	return nil
 | 
						return nil
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
		Loading…
	
		Reference in New Issue