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 defaultExpiredScanInterval = 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 Captcha) (Captcha, ID)
	GetExpiry() time.Duration
	Image(id ID, style string) (*image.Image, error)
	Solve(id ID, answer Answer) (bool, error)
	IsSolved(id ID) (bool, error)
	cleanExpired()
}

type InMemoryCaptchaDB struct {
	sync.Mutex

	db                 map[ID]Captcha
	expireIn           time.Duration
	expireScanInterval time.Duration
}

func NewInMemoryCaptchaDB(expire time.Duration) *InMemoryCaptchaDB {
	db := &InMemoryCaptchaDB{
		db:       make(map[ID]Captcha),
		expireIn: expire}

	if expire < defaultExpiredScanInterval {
		db.expireScanInterval = expire
	} else {
		db.expireScanInterval = defaultExpiredScanInterval
	}

	db.cleanExpired()

	return db
}

// 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
}

// cleanExpired starts a goroutine that deletes expired CAPTCHAs.
func (cdb *InMemoryCaptchaDB) cleanExpired() {
	go func() {
		for {
			sleepFor := cdb.expireScanInterval - (time.Duration(time.Now().Second()) % cdb.expireScanInterval)
			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
}