From d9aba868dbe482e061d4c04e1266b29b54801075 Mon Sep 17 00:00:00 2001
From: "Alexander \"Arav\" Andreev" <me@arav.top>
Date: Thu, 20 Oct 2022 22:55:37 +0400
Subject: [PATCH] db.go. InMemoryCaptchaDB was moved out. ErrorNotFound is now
 public. And DefaultExpiredScanInterval is now a public const. For more
 security a string of random data added to the end of a hash.

---
 pkg/captcha/db.go | 126 +++++-----------------------------------------
 1 file changed, 12 insertions(+), 114 deletions(-)

diff --git a/pkg/captcha/db.go b/pkg/captcha/db.go
index 9416c1d..cab8d6d 100644
--- a/pkg/captcha/db.go
+++ b/pkg/captcha/db.go
@@ -1,147 +1,45 @@
 package captcha
 
 import (
+	"crypto/rand"
 	"crypto/sha256"
 	"encoding/base64"
 	"errors"
 	"image"
 	"strconv"
-	"sync"
 	"time"
 )
 
-var errorNotFound = errors.New("captcha not found")
+var ErrorNotFound = errors.New("captcha not found")
 
-var defaultExpiredScanInterval = 60 * time.Second
+const DefaultExpiredScanInterval = 60 * time.Second
 
+// ID is a CAPTCHA identifier.
 type ID string
 
-// NewID generates an ID as a sha256 hash of additionalData, current time
-// and answer encoded with base64 in raw URL variant.
+// NewID generates an ID as a sha256 hash of additionalData (usually IP-address),
+// current time, answer and more, it adds a set of random bytes and encodes all
+// of it 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))
+	randData := make([]byte, 32)
+	rand.Read(randData)
+	idHash.Write(randData)
 
 	return ID(base64.RawURLEncoding.EncodeToString(idHash.Sum(nil)))
 }
 
+// CaptchaDB interface with all necessary methods.
 type CaptchaDB interface {
 	New(data string, captcha Captcha) (Captcha, ID)
 	GetExpiry() time.Duration
+	SetExpiry(expiry time.Duration)
 	Image(id ID, style string) (*image.Image, error)
 	Solve(id ID, answer Answer) (bool, error)
 	IsSolved(id ID) (bool, error)
 	Remove(id ID) 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.Answer())
-
-	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 {
-		delete(cdb.db, id)
-		return c.IsSolved(), nil
-	}
-	return false, errorNotFound
-}
-
-// Remove a CAPTCHA from a database.
-func (cdb *InMemoryCaptchaDB) Remove(id ID) error {
-	cdb.Lock()
-	defer cdb.Unlock()
-	if _, ok := cdb.db[id]; ok {
-		delete(cdb.db, id)
-		return nil
-	}
-	return errorNotFound
 }