1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-10-13 06:23:48 -04:00
v2fly/userset.go
2015-09-19 21:56:00 +02:00

131 lines
2.7 KiB
Go

package core
import (
"container/heap"
"time"
v2hash "github.com/v2ray/v2ray-core/hash"
"github.com/v2ray/v2ray-core/log"
)
const (
updateIntervalSec = 10
cacheDurationSec = 120
)
type UserSet interface {
AddUser(user User) error
GetUser(timeHash []byte) (*ID, int64, bool)
}
type TimedUserSet struct {
validUserIds []ID
userHashes map[string]indexTimePair
hash2Remove hashEntrySet
}
type indexTimePair struct {
index int
timeSec int64
}
type hashEntry struct {
hash string
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 {
vuSet := new(TimedUserSet)
vuSet.validUserIds = make([]ID, 0, 16)
vuSet.userHashes = make(map[string]indexTimePair)
vuSet.hash2Remove = make(hashEntrySet, 0, cacheDurationSec*10)
go vuSet.updateUserHash(time.Tick(updateIntervalSec * time.Second))
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) {
now := time.Now().UTC()
lastSec := now.Unix()
lastSec2Remove := now.Unix()
for {
now := <-tick
nowSec := now.UTC().Unix()
remove2Sec := nowSec - cacheDurationSec
if remove2Sec > lastSec2Remove {
for lastSec2Remove+1 < remove2Sec {
front := heap.Pop(&us.hash2Remove)
entry := front.(*hashEntry)
lastSec2Remove = entry.timeSec
delete(us.userHashes, entry.hash)
}
}
for idx, id := range us.validUserIds {
us.generateNewHashes(lastSec, nowSec, idx, id)
}
}
}
func (us *TimedUserSet) AddUser(user User) error {
id := user.Id
idx := len(us.validUserIds)
us.validUserIds = append(us.validUserIds, id)
nowSec := time.Now().UTC().Unix()
lastSec := nowSec - cacheDurationSec
us.generateNewHashes(lastSec, nowSec, idx, id)
return nil
}
func (us TimedUserSet) GetUser(userHash []byte) (*ID, int64, bool) {
pair, found := us.userHashes[string(userHash)]
if found {
return &us.validUserIds[pair.index], pair.timeSec, true
}
return nil, 0, false
}