github.com/sagernet/quic-go@v0.43.1-beta.1/token_store.go (about) 1 package quic 2 3 import ( 4 "sync" 5 6 "github.com/sagernet/quic-go/internal/utils" 7 list "github.com/sagernet/quic-go/internal/utils/linkedlist" 8 ) 9 10 type singleOriginTokenStore struct { 11 tokens []*ClientToken 12 len int 13 p int 14 } 15 16 func newSingleOriginTokenStore(size int) *singleOriginTokenStore { 17 return &singleOriginTokenStore{tokens: make([]*ClientToken, size)} 18 } 19 20 func (s *singleOriginTokenStore) Add(token *ClientToken) { 21 s.tokens[s.p] = token 22 s.p = s.index(s.p + 1) 23 s.len = utils.Min(s.len+1, len(s.tokens)) 24 } 25 26 func (s *singleOriginTokenStore) Pop() *ClientToken { 27 s.p = s.index(s.p - 1) 28 token := s.tokens[s.p] 29 s.tokens[s.p] = nil 30 s.len = utils.Max(s.len-1, 0) 31 return token 32 } 33 34 func (s *singleOriginTokenStore) Len() int { 35 return s.len 36 } 37 38 func (s *singleOriginTokenStore) index(i int) int { 39 mod := len(s.tokens) 40 return (i + mod) % mod 41 } 42 43 type lruTokenStoreEntry struct { 44 key string 45 cache *singleOriginTokenStore 46 } 47 48 type lruTokenStore struct { 49 mutex sync.Mutex 50 51 m map[string]*list.Element[*lruTokenStoreEntry] 52 q *list.List[*lruTokenStoreEntry] 53 capacity int 54 singleOriginSize int 55 } 56 57 var _ TokenStore = &lruTokenStore{} 58 59 // NewLRUTokenStore creates a new LRU cache for tokens received by the client. 60 // maxOrigins specifies how many origins this cache is saving tokens for. 61 // tokensPerOrigin specifies the maximum number of tokens per origin. 62 func NewLRUTokenStore(maxOrigins, tokensPerOrigin int) TokenStore { 63 return &lruTokenStore{ 64 m: make(map[string]*list.Element[*lruTokenStoreEntry]), 65 q: list.New[*lruTokenStoreEntry](), 66 capacity: maxOrigins, 67 singleOriginSize: tokensPerOrigin, 68 } 69 } 70 71 func (s *lruTokenStore) Put(key string, token *ClientToken) { 72 s.mutex.Lock() 73 defer s.mutex.Unlock() 74 75 if el, ok := s.m[key]; ok { 76 entry := el.Value 77 entry.cache.Add(token) 78 s.q.MoveToFront(el) 79 return 80 } 81 82 if s.q.Len() < s.capacity { 83 entry := &lruTokenStoreEntry{ 84 key: key, 85 cache: newSingleOriginTokenStore(s.singleOriginSize), 86 } 87 entry.cache.Add(token) 88 s.m[key] = s.q.PushFront(entry) 89 return 90 } 91 92 elem := s.q.Back() 93 entry := elem.Value 94 delete(s.m, entry.key) 95 entry.key = key 96 entry.cache = newSingleOriginTokenStore(s.singleOriginSize) 97 entry.cache.Add(token) 98 s.q.MoveToFront(elem) 99 s.m[key] = elem 100 } 101 102 func (s *lruTokenStore) Pop(key string) *ClientToken { 103 s.mutex.Lock() 104 defer s.mutex.Unlock() 105 106 var token *ClientToken 107 if el, ok := s.m[key]; ok { 108 s.q.MoveToFront(el) 109 cache := el.Value.cache 110 token = cache.Pop() 111 if cache.Len() == 0 { 112 s.q.Remove(el) 113 delete(s.m, key) 114 } 115 } 116 return token 117 }