doucan/app/internal/captcha/captcha.go

package captcha

import (
	"bytes"
	"context"
	"doudizhu-server/internal/redis"
	"encoding/base64"
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"image/png"
	"math/rand"
	"time"
)

const (
	charset    = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
	captchaLen = 4
	width      = 120
	height     = 40
	CaptchaTTL = 5 * time.Minute
)

type Manager struct {
	rdb *redis.Client
}

func NewManager(rdb *redis.Client) *Manager {
	return &Manager{rdb: rdb}
}

func (m *Manager) Generate(ctx context.Context) (captchaID string, imageBase64 string, err error) {
	code := generateCode(captchaLen)
	captchaID = generateID()

	img := createImage(code)

	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		return "", "", fmt.Errorf("failed to encode captcha image: %w", err)
	}

	imageBase64 = "data:image/png;base64," + base64.StdEncoding.EncodeToString(buf.Bytes())

	if err := m.rdb.Set(ctx, redis.CaptchaKey(captchaID), code, CaptchaTTL); err != nil {
		return "", "", fmt.Errorf("failed to store captcha: %w", err)
	}

	return captchaID, imageBase64, nil
}

func (m *Manager) Verify(ctx context.Context, captchaID, code string) bool {
	var storedCode string
	if err := m.rdb.Get(ctx, redis.CaptchaKey(captchaID), &storedCode); err != nil {
		return false
	}

	m.rdb.Delete(ctx, redis.CaptchaKey(captchaID))

	return storedCode == code
}

func generateCode(length int) string {
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	b := make([]byte, length)
	for i := range b {
		b[i] = charset[r.Intn(len(charset))]
	}
	return string(b)
}

func generateID() string {
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	const idCharset = "abcdefghijklmnopqrstuvwxyz0123456789"
	b := make([]byte, 16)
	for i := range b {
		b[i] = idCharset[r.Intn(len(idCharset))]
	}
	return string(b)
}

func createImage(code string) image.Image {
	r := rand.New(rand.NewSource(time.Now().UnixNano()))

	img := image.NewRGBA(image.Rect(0, 0, width, height))

	bgColor := color.RGBA{
		R: uint8(200 + r.Intn(56)),
		G: uint8(200 + r.Intn(56)),
		B: uint8(200 + r.Intn(56)),
		A: 255,
	}
	draw.Draw(img, img.Bounds(), &image.Uniform{bgColor}, image.Point{}, draw.Src)

	for i := 0; i < 50; i++ {
		x := r.Intn(width)
		y := r.Intn(height)
		noiseColor := color.RGBA{
			R: uint8(r.Intn(256)),
			G: uint8(r.Intn(256)),
			B: uint8(r.Intn(256)),
			A: 255,
		}
		img.Set(x, y, noiseColor)
	}

	for i := 0; i < 3; i++ {
		x1 := r.Intn(width)
		y1 := r.Intn(height)
		x2 := r.Intn(width)
		y2 := r.Intn(height)
		lineColor := color.RGBA{
			R: uint8(r.Intn(200)),
			G: uint8(r.Intn(200)),
			B: uint8(r.Intn(200)),
			A: 255,
		}
		drawLine(img, x1, y1, x2, y2, lineColor)
	}

	charWidth := width / captchaLen
	for i, c := range code {
		x := i*charWidth + charWidth/4
		y := 8 + r.Intn(10)

		textColor := color.RGBA{
			R: uint8(r.Intn(150)),
			G: uint8(r.Intn(150)),
			B: uint8(r.Intn(150)),
			A: 255,
		}
		drawChar(img, x, y, string(c), textColor)
	}

	return img
}

func drawLine(img *image.RGBA, x1, y1, x2, y2 int, c color.Color) {
	dx := abs(x2 - x1)
	dy := abs(y2 - y1)
	sx, sy := 1, 1
	if x1 >= x2 {
		sx = -1
	}
	if y1 >= y2 {
		sy = -1
	}
	err := dx - dy

	for {
		img.Set(x1, y1, c)
		if x1 == x2 && y1 == y2 {
			break
		}
		e2 := err * 2
		if e2 > -dy {
			err -= dy
			x1 += sx
		}
		if e2 < dx {
			err += dx
			y1 += sy
		}
	}
}

func abs(x int) int {
	if x < 0 {
		return -x
	}
	return x
}

func drawChar(img *image.RGBA, x, y int, char string, c color.Color) {
	patterns := map[string][]string{
		"A": {"  ##  ", " #  # ", "######", "#    #", "#    #"},
		"B": {"##### ", "#    #", "##### ", "#    #", "##### "},
		"C": {" #####", "#     ", "#     ", "#     ", " #####"},
		"D": {"##### ", "#    #", "#    #", "#    #", "##### "},
		"E": {"######", "#     ", "####  ", "#     ", "######"},
		"F": {"######", "#     ", "####  ", "#     ", "#     "},
		"G": {" #####", "#     ", "#  ###", "#    #", " #####"},
		"H": {"#    #", "#    #", "######", "#    #", "#    #"},
		"J": {"  ####", "    # ", "    # ", "#   # ", " ###  "},
		"K": {"#   # ", "#  #  ", "###   ", "#  #  ", "#   # "},
		"L": {"#     ", "#     ", "#     ", "#     ", "######"},
		"M": {"#    #", "##  ##", "# ## #", "#    #", "#    #"},
		"N": {"#    #", "##   #", "# #  #", "#  # #", "#   ##"},
		"P": {"##### ", "#    #", "##### ", "#     ", "#     "},
		"Q": {" #####", "#    #", "#  # #", "#   # ", " #### #"},
		"R": {"##### ", "#    #", "##### ", "#  #  ", "#   # "},
		"S": {" #####", "#     ", " #####", "     #", "##### "},
		"T": {"######", "  ##  ", "  ##  ", "  ##  ", "  ##  "},
		"U": {"#    #", "#    #", "#    #", "#    #", " #### "},
		"V": {"#    #", "#    #", " #  # ", " #  # ", "  ##  "},
		"W": {"#    #", "#    #", "# ## #", "##  ##", "#    #"},
		"X": {"#    #", " #  # ", "  ##  ", " #  # ", "#    #"},
		"Y": {"#    #", " #  # ", "  ##  ", "  ##  ", "  ##  "},
		"Z": {"######", "   ## ", "  ##  ", " ##   ", "######"},
		"2": {" #####", "#    #", "   ## ", "  #   ", "######"},
		"3": {" #####", "     #", "  ####", "     #", " #####"},
		"4": {"#    #", "#    #", "######", "     #", "     #"},
		"5": {"######", "#     ", "##### ", "     #", "##### "},
		"6": {" #####", "#     ", "##### ", "#    #", " #####"},
		"7": {"######", "    # ", "   #  ", "  #   ", "  #   "},
		"8": {" #####", "#    #", " #####", "#    #", " #####"},
		"9": {" #####", "#    #", " ######", "     #", " #####"},
	}

	pattern, ok := patterns[char]
	if !ok {
		return
	}

	for row, line := range pattern {
		for col, ch := range line {
			if ch == '#' {
				px := x + col*2
				py := y + row*2
				if px >= 0 && px < width-1 && py >= 0 && py < height-1 {
					img.Set(px, py, c)
					img.Set(px+1, py, c)
					img.Set(px, py+1, c)
					img.Set(px+1, py+1, c)
				}
			}
		}
	}
}