M captcha.go → captcha.go

@@ -1,3 +1,5 @@
+// Package captcha implements generation and verification of image and audio
+// CAPTCHAs.
 package captcha
 
 import (
@@ -24,7 +26,12 @@
 var ErrNotFound = os.NewError("captcha with the given id not found")
 
 // globalStore is a shared storage for captchas, generated by New function.
-var globalStore = newStore(StdCollectNum, StdExpiration)
+var globalStore = NewMemoryStore(StdCollectNum, StdExpiration)
+
+// SetCustomStore sets custom storage for captchas.
+func SetCustomStore(s Store) {
+	globalStore = s
+}
 
 // RandomDigits returns a byte slice of the given length containing random
 // digits in range 0-9.
@@ -43,7 +50,7 @@ // New creates a new captcha of the given length, saves it in the internal
 // storage, and returns its id.
 func New(length int) (id string) {
 	id = uniuri.New()
-	globalStore.saveCaptcha(id, RandomDigits(length))
+	globalStore.Set(id, RandomDigits(length))
 	return
 }
 
@@ -54,18 +61,18 @@ // After calling this function, the image or audio presented to a user must be
 // refreshed to show the new captcha representation (WriteImage and WriteAudio
 // will write the new one).
 func Reload(id string) bool {
-	old := globalStore.getDigits(id)
+	old := globalStore.Get(id, false)
 	if old == nil {
 		return false
 	}
-	globalStore.saveCaptcha(id, RandomDigits(len(old)))
+	globalStore.Set(id, RandomDigits(len(old)))
 	return true
 }
 
 // WriteImage writes PNG-encoded image representation of the captcha with the
 // given id. The image will have the given width and height.
 func WriteImage(w io.Writer, id string, width, height int) os.Error {
-	d := globalStore.getDigits(id)
+	d := globalStore.Get(id, false)
 	if d == nil {
 		return ErrNotFound
 	}
@@ -76,7 +83,7 @@
 // WriteAudio writes WAV-encoded audio representation of the captcha with the
 // given id.
 func WriteAudio(w io.Writer, id string) os.Error {
-	d := globalStore.getDigits(id)
+	d := globalStore.Get(id, false)
 	if d == nil {
 		return ErrNotFound
 	}
@@ -93,7 +100,7 @@ func Verify(id string, digits []byte) bool {
 	if digits == nil || len(digits) == 0 {
 		return false
 	}
-	reald := globalStore.getDigitsClear(id)
+	reald := globalStore.Get(id, true)
 	if reald == nil {
 		return false
 	}
@@ -128,7 +135,7 @@ // but still exported to enable freeing memory manually if needed.
 //
 // Collection is launched in a new goroutine.
 func Collect() {
-	go globalStore.collect()
+	go globalStore.Collect()
 }
 
 type captchaHandler struct {
@@ -171,7 +178,7 @@ }
 		//err = WriteAudio(buf, id)
 		//XXX(dchest) Workaround for Chrome: it wants content-length,
 		//or else will start playing NOT from the beginning.
-		d := globalStore.getDigits(id)
+		d := globalStore.Get(id, false)
 		if d == nil {
 			err = ErrNotFound
 		} else {

M captcha_test.go → captcha_test.go

@@ -18,7 +18,7 @@ if Verify(id, []byte{0, 0}) {
 		t.Errorf("verified wrong captcha")
 	}
 	id = New(StdLength)
-	d := globalStore.getDigits(id) // cheating
+	d := globalStore.Get(id, false) // cheating
 	if !Verify(id, d) {
 		t.Errorf("proper captcha not verified")
 	}
@@ -26,9 +26,9 @@ }
 
 func TestReload(t *testing.T) {
 	id := New(StdLength)
-	d1 := globalStore.getDigits(id) // cheating
+	d1 := globalStore.Get(id, false) // cheating
 	Reload(id)
-	d2 := globalStore.getDigits(id) // cheating again
+	d2 := globalStore.Get(id, false) // cheating again
 	if bytes.Equal(d1, d2) {
 		t.Errorf("reload didn't work: %v = %v", d1, d2)
 	}

M store.go → store.go

@@ -6,16 +6,34 @@ "sync"
 	"time"
 )
 
-// expValue stores timestamp and id of captchas. It is used in a list inside
-// store for indexing generated captchas by timestamp to enable garbage
+// An object implementing Store interface can be registered with SetCustomStore
+// function to handle storage and retrieval of captcha ids and solutions for
+// them, replacing the default memory store.
+type Store interface {
+	// Set sets the digits for the captcha id.
+	Set(id string, digits []byte)
+
+	// Get returns stored digits for the captcha id. Clear indicates
+	// whether the captcha must be deleted from the store.
+	Get(id string, clear bool) (digits []byte)
+
+	// Collect deletes expired captchas from the store.  For custom stores
+	// this method is not called automatically, it is only wired to the
+	// package's Collect function.  Custom stores must implement their own
+	// procedure for calling Collect, for example, in Set method.
+	Collect()
+}
+
+// expValue stores timestamp and id of captchas. It is used in the list inside
+// memoryStore for indexing generated captchas by timestamp to enable garbage
 // collection of expired captchas.
 type expValue struct {
 	timestamp int64
 	id        string
 }
 
-// store is an internal store for captcha ids and their values.
-type store struct {
+// memoryStore is an internal store for captcha ids and their values.
+type memoryStore struct {
 	mu  sync.RWMutex
 	ids map[string][]byte
 	exp *list.List
@@ -27,9 +45,11 @@ // Expiration time of captchas.
 	expiration int64
 }
 
-// newStore initializes and returns a new store.
-func newStore(collectNum int, expiration int64) *store {
-	s := new(store)
+// NewMemoryStore returns a new standard memory store for captchas with the
+// given collection threshold and expiration time in seconds. The returned
+// store must be registered with SetCustomStore to replace the default one.
+func NewMemoryStore(collectNum int, expiration int64) Store {
+	s := new(memoryStore)
 	s.ids = make(map[string][]byte)
 	s.exp = list.New()
 	s.collectNum = collectNum
@@ -37,44 +57,40 @@ s.expiration = expiration
 	return s
 }
 
-// saveCaptcha saves the captcha id and the corresponding digits.
-func (s *store) saveCaptcha(id string, digits []byte) {
+func (s *memoryStore) Set(id string, digits []byte) {
 	s.mu.Lock()
 	s.ids[id] = digits
 	s.exp.PushBack(expValue{time.Seconds(), id})
 	s.numStored++
 	s.mu.Unlock()
 	if s.numStored > s.collectNum {
-		go s.collect()
+		go s.Collect()
 	}
 }
 
-// getDigits returns the digits for the given id.
-func (s *store) getDigits(id string) (digits []byte) {
-	s.mu.RLock()
-	defer s.mu.RUnlock()
-	digits, _ = s.ids[id]
-	return
-}
-
-// getDigitsClear returns the digits for the given id, and removes them from
-// the store.
-func (s *store) getDigitsClear(id string) (digits []byte) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
+func (s *memoryStore) Get(id string, clear bool) (digits []byte) {
+	if !clear {
+		// When we don't need to clear captcha, acquire read lock.
+		s.mu.RLock()
+		defer s.mu.RUnlock()
+	} else {
+		s.mu.Lock()
+		defer s.mu.Unlock()
+	}
 	digits, ok := s.ids[id]
 	if !ok {
 		return
 	}
-	s.ids[id] = nil, false
-	// XXX(dchest) Index (s.exp) will be cleaned when collecting expired
-	// captchas.  Can't clean it here, because we don't store reference to
-	// expValue in the map. Maybe store it?
+	if clear {
+		s.ids[id] = nil, false
+		// XXX(dchest) Index (s.exp) will be cleaned when collecting expired
+		// captchas.  Can't clean it here, because we don't store reference to
+		// expValue in the map. Maybe store it?
+	}
 	return
 }
 
-// collect deletes expired captchas from the store.
-func (s *store) collect() {
+func (s *memoryStore) Collect() {
 	now := time.Seconds()
 	s.mu.Lock()
 	defer s.mu.Unlock()

M store_test.go → store_test.go

@@ -6,27 +6,27 @@ "github.com/dchest/uniuri"
 	"testing"
 )
 
-func TestSaveAndGetDigits(t *testing.T) {
-	s := newStore(StdCollectNum, StdExpiration)
+func TestSetGet(t *testing.T) {
+	s := NewMemoryStore(StdCollectNum, StdExpiration)
 	id := "captcha id"
 	d := RandomDigits(10)
-	s.saveCaptcha(id, d)
-	d2 := s.getDigits(id)
+	s.Set(id, d)
+	d2 := s.Get(id, false)
 	if d2 == nil || !bytes.Equal(d, d2) {
 		t.Errorf("saved %v, getDigits returned got %v", d, d2)
 	}
 }
 
-func TestGetDigitsClear(t *testing.T) {
-	s := newStore(StdCollectNum, StdExpiration)
+func TestGetClear(t *testing.T) {
+	s := NewMemoryStore(StdCollectNum, StdExpiration)
 	id := "captcha id"
 	d := RandomDigits(10)
-	s.saveCaptcha(id, d)
-	d2 := s.getDigitsClear(id)
+	s.Set(id, d)
+	d2 := s.Get(id, true)
 	if d2 == nil || !bytes.Equal(d, d2) {
 		t.Errorf("saved %v, getDigitsClear returned got %v", d, d2)
 	}
-	d2 = s.getDigits(id)
+	d2 = s.Get(id, false)
 	if d2 != nil {
 		t.Errorf("getDigitClear didn't clear (%q=%v)", id, d2)
 	}
@@ -35,19 +35,19 @@
 func TestCollect(t *testing.T) {
 	//TODO(dchest): can't test automatic collection when saving, because
 	//it's currently launched in a different goroutine.
-	s := newStore(10, -1)
+	s := NewMemoryStore(10, -1)
 	// create 10 ids
 	ids := make([]string, 10)
 	d := RandomDigits(10)
 	for i := range ids {
 		ids[i] = uniuri.New()
-		s.saveCaptcha(ids[i], d)
+		s.Set(ids[i], d)
 	}
-	s.collect()
+	s.Collect()
 	// Must be already collected
 	nc := 0
 	for i := range ids {
-		d2 := s.getDigits(ids[i])
+		d2 := s.Get(ids[i], false)
 		if d2 != nil {
 			t.Errorf("%d: not collected", i)
 			nc++