1package captcha
  2
  3import (
  4	"container/list"
  5	"sync"
  6	"time"
  7)
  8
  9// An object implementing Store interface can be registered with SetCustomStore
 10// function to handle storage and retrieval of captcha ids and solutions for
 11// them, replacing the default memory store.
 12type Store interface {
 13	// Set sets the digits for the captcha id.
 14	Set(id string, digits []byte)
 15
 16	// Get returns stored digits for the captcha id. Clear indicates
 17	// whether the captcha must be deleted from the store.
 18	Get(id string, clear bool) (digits []byte)
 19
 20	// Collect deletes expired captchas from the store.  For custom stores
 21	// this method is not called automatically, it is only wired to the
 22	// package's Collect function.  Custom stores must implement their own
 23	// procedure for calling Collect, for example, in Set method.
 24	Collect()
 25}
 26
 27// expValue stores timestamp and id of captchas. It is used in the list inside
 28// memoryStore for indexing generated captchas by timestamp to enable garbage
 29// collection of expired captchas.
 30type expValue struct {
 31	timestamp int64
 32	id        string
 33}
 34
 35// memoryStore is an internal store for captcha ids and their values.
 36type memoryStore struct {
 37	mu  sync.RWMutex
 38	ids map[string][]byte
 39	exp *list.List
 40	// Number of items stored after last collection.
 41	numStored int
 42	// Number of saved items that triggers collection.
 43	collectNum int
 44	// Expiration time of captchas.
 45	expiration int64
 46}
 47
 48// NewMemoryStore returns a new standard memory store for captchas with the
 49// given collection threshold and expiration time in seconds. The returned
 50// store must be registered with SetCustomStore to replace the default one.
 51func NewMemoryStore(collectNum int, expiration int64) Store {
 52	s := new(memoryStore)
 53	s.ids = make(map[string][]byte)
 54	s.exp = list.New()
 55	s.collectNum = collectNum
 56	s.expiration = expiration
 57	return s
 58}
 59
 60func (s *memoryStore) Set(id string, digits []byte) {
 61	s.mu.Lock()
 62	s.ids[id] = digits
 63	s.exp.PushBack(expValue{time.Seconds(), id})
 64	s.numStored++
 65	s.mu.Unlock()
 66	if s.numStored > s.collectNum {
 67		go s.Collect()
 68	}
 69}
 70
 71func (s *memoryStore) Get(id string, clear bool) (digits []byte) {
 72	if !clear {
 73		// When we don't need to clear captcha, acquire read lock.
 74		s.mu.RLock()
 75		defer s.mu.RUnlock()
 76	} else {
 77		s.mu.Lock()
 78		defer s.mu.Unlock()
 79	}
 80	digits, ok := s.ids[id]
 81	if !ok {
 82		return
 83	}
 84	if clear {
 85		s.ids[id] = nil, false
 86		// XXX(dchest) Index (s.exp) will be cleaned when collecting expired
 87		// captchas.  Can't clean it here, because we don't store reference to
 88		// expValue in the map. Maybe store it?
 89	}
 90	return
 91}
 92
 93func (s *memoryStore) Collect() {
 94	now := time.Seconds()
 95	s.mu.Lock()
 96	defer s.mu.Unlock()
 97	s.numStored = 0
 98	for e := s.exp.Front(); e != nil; {
 99		ev, ok := e.Value.(expValue)
100		if !ok {
101			return
102		}
103		if ev.timestamp+s.expiration < now {
104			s.ids[ev.id] = nil, false
105			next := e.Next()
106			s.exp.Remove(e)
107			e = next
108		} else {
109			return
110		}
111	}
112}