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