From 96c13ec6b20651ced20af849ca9f62766aee166f Mon Sep 17 00:00:00 2001
From: "Alexander \"Arav\" Andreev" <me@arav.top>
Date: Mon, 27 Jun 2022 01:20:36 +0400
Subject: [PATCH] Added support for different styles for CAPTCHA.

---
 internal/captcha/dwelling_captcha.go |  6 +++---
 internal/captcha/instance.go         |  4 ++--
 internal/handlers/handlers.go        |  3 ++-
 pkg/captcha/captcha.go               | 10 +++++-----
 pkg/captcha/db.go                    |  4 ++--
 5 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/internal/captcha/dwelling_captcha.go b/internal/captcha/dwelling_captcha.go
index 5877f93..dc209e8 100644
--- a/internal/captcha/dwelling_captcha.go
+++ b/internal/captcha/dwelling_captcha.go
@@ -25,7 +25,7 @@ func NewDwellingCaptcha(expiry time.Duration) *DwellingCaptcha {
 			ExpireIn: captcha.ExpiryDate(expiry)}}
 }
 
-func (c *DwellingCaptcha) generateImage() *image.Image {
+func (c *DwellingCaptcha) generateImage(style string) *image.Image {
 	ctx := gg.NewContext(dwImageWidth, dwImageHeight)
 
 	// fill background
@@ -65,9 +65,9 @@ func (c *DwellingCaptcha) generateImage() *image.Image {
 	return &c.Image
 }
 
-func (c *DwellingCaptcha) GetImage() *image.Image {
+func (c *DwellingCaptcha) GetImage(style string) *image.Image {
 	if c.Image == nil {
-		return c.generateImage()
+		return c.generateImage(style)
 	}
 
 	return &c.Image
diff --git a/internal/captcha/instance.go b/internal/captcha/instance.go
index 28e6a94..238f75a 100644
--- a/internal/captcha/instance.go
+++ b/internal/captcha/instance.go
@@ -17,8 +17,8 @@ func New(data string, captcha captcha.ICaptcha) (captcha.ICaptcha, captcha.ID) {
 	return captchaDb.New(data, captcha)
 }
 
-func Image(id captcha.ID) (*image.Image, error) {
-	return captchaDb.Image(id)
+func Image(id captcha.ID, style string) (*image.Image, error) {
+	return captchaDb.Image(id, style)
 }
 
 func Solve(id captcha.ID, answer captcha.Answer) (bool, error) {
diff --git a/internal/handlers/handlers.go b/internal/handlers/handlers.go
index 3e8a3f4..567dd9a 100644
--- a/internal/handlers/handlers.go
+++ b/internal/handlers/handlers.go
@@ -24,8 +24,9 @@ func (h *CaptchaHandlers) New(w http.ResponseWriter, r *http.Request) {
 
 func (h *CaptchaHandlers) Image(w http.ResponseWriter, r *http.Request) {
 	captchaID := pcaptcha.ID(server.GetURLParam(r, "captcha"))
+	captchaStyle := r.URL.Query().Get("style")
 
-	captchaImage, err := captcha.Image(captchaID)
+	captchaImage, err := captcha.Image(captchaID, captchaStyle)
 	if err != nil {
 		w.WriteHeader(http.StatusNotFound)
 		return
diff --git a/pkg/captcha/captcha.go b/pkg/captcha/captcha.go
index 141fd6d..4e16375 100644
--- a/pkg/captcha/captcha.go
+++ b/pkg/captcha/captcha.go
@@ -19,10 +19,10 @@ type Answer string
 // ICaptcha interface that should be implemented by captcha.
 type ICaptcha interface {
 	// generateImage is used to create a captcha image.
-	generateImage() *image.Image
+	generateImage(style string) *image.Image
 	// GetImage() returns an image of captcha. Calls
 	// generateImage() method if image doesn't exist.
-	GetImage() *image.Image
+	GetImage(style string) *image.Image
 	// GetAnswer returns a pregenerated answer.
 	GetAnswer() Answer
 	// Solve compares a stored answer with a passed one.
@@ -50,13 +50,13 @@ type Captcha struct {
 	ExpireIn time.Time
 }
 
-func (c *Captcha) generateImage() *image.Image {
+func (c *Captcha) generateImage(style string) *image.Image {
 	return nil
 }
 
-func (c *Captcha) GetImage() *image.Image {
+func (c *Captcha) GetImage(style string) *image.Image {
 	if c.Image == nil {
-		return c.generateImage()
+		return c.generateImage(style)
 	}
 
 	return &c.Image
diff --git a/pkg/captcha/db.go b/pkg/captcha/db.go
index e8963b4..992cf07 100644
--- a/pkg/captcha/db.go
+++ b/pkg/captcha/db.go
@@ -78,9 +78,9 @@ func (cdb *CaptchaDB) New(data string, captcha ICaptcha) (ICaptcha, ID) {
 }
 
 // Image returns image for a captcha.
-func (cdb *CaptchaDB) Image(id ID) (*image.Image, error) {
+func (cdb *CaptchaDB) Image(id ID, style string) (*image.Image, error) {
 	if c, ok := cdb.DB[id]; ok {
-		return c.GetImage(), nil
+		return c.GetImage(style), nil
 	}
 	return nil, errorNotFound
 }