github.com/apernet/quic-go@v0.43.1-0.20240515053213-5e9e635fd9f0/token_store.go (about)

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