justcaptcha/pkg/captcha/db.go

package captcha

import (
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"image"
	"strconv"
	"sync"
	"time"
)

var errorNotFound = errors.New("captcha not found")

var expiredScanInterval = 60 * time.Second

type ID string

// NewID generates an ID as a sha256 hash of additionalData, current time
// and answer encoded with base64 in raw URL variant.
func NewID(additionalData string, answer Answer) ID {
	idHash := sha256.New()

	idHash.Write([]byte(additionalData))
	idHash.Write([]byte(strconv.FormatInt(time.Now().UnixMicro(), 16)))
	idHash.Write([]byte(answer))

	return ID(base64.RawURLEncoding.EncodeToString(idHash.Sum(nil)))
}

type CaptchaDB interface {
	New(data string) (Captcha, ID)
	SetExpiry(expiry time.Duration)
	GetExpiry() time.Duration
	Image(id ID) (*image.Image, error)
	Solve(id ID, answer Answer) (bool, error)
	IsSolved(id ID) bool
}

type InMemoryCaptchaDB struct {
	sync.Mutex

	DB       map[ID]Captcha
	ExpireIn time.Duration
}

// New accepts an Captcha instance, generates an ID and store it in a database.
// `data` string is an additional random data used to generate an ID,
// e.g. IP-address.
func (cdb *InMemoryCaptchaDB) New(data string, captcha Captcha) (Captcha, ID) {
	id := NewID(data, captcha.GetAnswer())

	cdb.Lock()
	cdb.DB[id] = captcha
	cdb.Unlock()

	return captcha, id
}

// SetExpiry stores expire value and starts a goroutine
// that checks for expired CAPTCHAs.
func (cdb *InMemoryCaptchaDB) SetExpiry(expire time.Duration) {
	cdb.ExpireIn = expire
	if expire < expiredScanInterval {
		expiredScanInterval = expire
	}

	go func() {
		for {
			sleepFor := expiredScanInterval - (time.Duration(time.Now().Second()) % expiredScanInterval)
			time.Sleep(sleepFor)

			cdb.Lock()
			for id, captcha := range cdb.DB {
				if time.Since(captcha.Expiry()) >= cdb.ExpireIn {
					delete(cdb.DB, id)
				}
			}
			cdb.Unlock()
		}
	}()
}

// GetExpiry returns time for how long captcha will last.
func (cdb *InMemoryCaptchaDB) GetExpiry() time.Duration {
	return cdb.ExpireIn
}

// Image returns image for a captcha.
func (cdb *InMemoryCaptchaDB) Image(id ID, style string) (*image.Image, error) {
	cdb.Lock()
	defer cdb.Unlock()
	if c, ok := cdb.DB[id]; ok {
		return c.Image(style), nil
	}
	return nil, errorNotFound
}

// Solve compares given answer with a stored one and if failed
// deletes a captcha from database.
func (cdb *InMemoryCaptchaDB) Solve(id ID, answer Answer) (bool, error) {
	cdb.Lock()
	defer cdb.Unlock()
	if c, ok := cdb.DB[id]; ok {
		ok = c.Solve(answer)
		if !ok {
			delete(cdb.DB, id)
		}
		return ok, nil
	}
	return false, errorNotFound
}

// IsSolved checks if captcha was solved and removes it
// from a database.
func (cdb *InMemoryCaptchaDB) IsSolved(id ID) (bool, error) {
	cdb.Lock()
	defer cdb.Unlock()
	if c, ok := cdb.DB[id]; ok {
		ok = c.IsSolved()
		delete(cdb.DB, id)
		return ok, nil
	}
	return false, errorNotFound
}