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.
 12//
 13// It is the responsibility of an object to delete expired and used captchas
 14// when necessary (for example, the default memory store collects them in Set
 15// method after the certain amount of captchas has been stored.)
 16type Store interface {
 17	// Set sets the digits for the captcha id.
 18	Set(id string, digits []byte)
 19
 20	// Get returns stored digits for the captcha id. Clear indicates
 21	// whether the captcha must be deleted from the store.
 22	Get(id string, clear bool) (digits []byte)
 23}
 24
 25// expValue stores timestamp and id of captchas. It is used in the list inside
 26// memoryStore for indexing generated captchas by timestamp to enable garbage
 27// collection of expired captchas.
 28type idByTimeValue struct {
 29	timestamp int64
 30	id        string
 31}
 32
 33// memoryStore is an internal store for captcha ids and their values.
 34type memoryStore struct {
 35	mu         sync.RWMutex
 36	digitsById map[string][]byte
 37	idByTime   *list.List
 38	// Number of items stored since last collection.
 39	numStored int
 40	// Number of saved items that triggers collection.
 41	collectNum int
 42	// Expiration time of captchas.
 43	expiration int64
 44}
 45
 46// NewMemoryStore returns a new standard memory store for captchas with the
 47// given collection threshold and expiration time in seconds. The returned
 48// store must be registered with SetCustomStore to replace the default one.
 49func NewMemoryStore(collectNum int, expiration int64) Store {
 50	s := new(memoryStore)
 51	s.digitsById = make(map[string][]byte)
 52	s.idByTime = list.New()
 53	s.collectNum = collectNum
 54	s.expiration = expiration
 55	return s
 56}
 57
 58func (s *memoryStore) Set(id string, digits []byte) {
 59	s.mu.Lock()
 60	s.digitsById[id] = digits
 61	s.idByTime.PushBack(idByTimeValue{time.Seconds(), id})
 62	s.numStored++
 63	if s.numStored <= s.collectNum {
 64		s.mu.Unlock()
 65		return
 66	}
 67	s.mu.Unlock()
 68	go s.collect()
 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}