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 idByTimeValue 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	digitsById map[string][]byte
 39	idByTime   *list.List
 40	// Number of items stored since 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.digitsById = make(map[string][]byte)
 54	s.idByTime = 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.digitsById[id] = digits
 63	s.idByTime.PushBack(idByTimeValue{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.digitsById[id]
 81	if !ok {
 82		return
 83	}
 84	if clear {
 85		s.digitsById[id] = nil, false
 86		// XXX(dchest) Index (s.idByTime) will be cleaned when
 87		// collecting expired captchas.  Can't clean it here, because
 88		// we don't store reference to expValue in the map.
 89		// Maybe store it?
 90	}
 91	return
 92}
 93
 94func (s *memoryStore) Collect() {
 95	now := time.Seconds()
 96	s.mu.Lock()
 97	defer s.mu.Unlock()
 98	s.numStored = 0
 99	for e := s.idByTime.Front(); e != nil; {
100		ev, ok := e.Value.(idByTimeValue)
101		if !ok {
102			return
103		}
104		if ev.timestamp+s.expiration < now {
105			s.digitsById[ev.id] = nil, false
106			next := e.Next()
107			s.idByTime.Remove(e)
108			e = next
109		} else {
110			return
111		}
112	}
113}