github.com/MerlinKodo/quic-go@v0.39.2/token_store.go (about)

     1  package quic
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/MerlinKodo/quic-go/internal/utils"
     7  	list "github.com/MerlinKodo/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  }