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	if s.numStored <= s.collectNum {
 66		s.mu.Unlock()
 67		return
 68	}
 69	s.mu.Unlock()
 70	go s.Collect()
 71}
 72
 73func (s *memoryStore) Get(id string, clear bool) (digits []byte) {
 74	if !clear {
 75		// When we don't need to clear captcha, acquire read lock.
 76		s.mu.RLock()
 77		defer s.mu.RUnlock()
 78	} else {
 79		s.mu.Lock()
 80		defer s.mu.Unlock()
 81	}
 82	digits, ok := s.digitsById[id]
 83	if !ok {
 84		return
 85	}
 86	if clear {
 87		s.digitsById[id] = nil, false
 88		// XXX(dchest) Index (s.idByTime) will be cleaned when
 89		// collecting expired captchas.  Can't clean it here, because
 90		// we don't store reference to expValue in the map.
 91		// Maybe store it?
 92	}
 93	return
 94}
 95
 96func (s *memoryStore) Collect() {
 97	now := time.Seconds()
 98	s.mu.Lock()
 99	defer s.mu.Unlock()
100	s.numStored = 0
101	for e := s.idByTime.Front(); e != nil; {
102		ev, ok := e.Value.(idByTimeValue)
103		if !ok {
104			return
105		}
106		if ev.timestamp+s.expiration < now {
107			s.digitsById[ev.id] = nil, false
108			next := e.Next()
109			s.idByTime.Remove(e)
110			e = next
111		} else {
112			return
113		}
114	}
115}