picgo/handler/captcha.go

package handler

import (
	"bytes"
	"image"
	"image/png"
	"net"
	"net/http"
	"picgo/corelib"
	"picgo/corelib/captcha"
	"strings"
)

func CaptchaHandler(w http.ResponseWriter, r *http.Request) {
	var (
		res []byte
		err error
	)
	captchaId := getCaptchaId(r)
	corelib.Logger.Info("captchaHandler cid: ", captchaId)
	if res, err = generateCaptcha(captchaId); err != nil {
		corelib.WriteJsonResponse(w, 1030, "生产验证码失败", nil)
		return
	}
	// 设置响应头
	w.Header().Set("Content-Type", "image/png")
	// 将图片文件内容写入响应
	_, err = w.Write(res)
	if err != nil {
		http.Error(w, "Failed to serve image.", http.StatusInternalServerError)
	}
}

// user agent + client ip  md5 生成验证码key
func getCaptchaId(r *http.Request) string {
	clientIP := GetRealIP(r)
	userAgent := r.Header.Get("User-Agent")
	return corelib.Md5Hash(clientIP + userAgent)
}

// 生成验证码图片
func generateCaptcha(captchaId string) ([]byte, error) {
	var (
		b   bytes.Buffer
		img *image.RGBA
		err error
	)
	if img, err = captcha.GenerateCaptcha(captchaId); err != nil {
		corelib.Logger.Fatalln("生成验证码错误", err)
		return nil, err
		//return nil, corelib.CAPTCHA_GENERATE_ERROR
	}
	if err = png.Encode(&b, img); err != nil {
		corelib.Logger.Fatalln("验证码encode错误", err)
		return nil, err
		//return nil, corelib.CAPTCHA_GENERATE_ERROR
	}
	return b.Bytes(), nil
}

// GetRealIP 获取用户的真实IP地址
func GetRealIP(r *http.Request) string {
	// 检查 X-Forwarded-For 头
	xff := r.Header.Get("X-Forwarded-For")
	if xff != "" {
		// X-Forwarded-For 头可以包含多个IP地址，取第一个
		ips := strings.Split(xff, ",")
		if len(ips) > 0 {
			return strings.TrimSpace(ips[0])
		}
	}

	// 检查 X-Real-IP 头
	xri := r.Header.Get("X-Real-IP")
	if xri != "" {
		return xri
	}

	// 如果没有这些头，则使用远程地址
	ip, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}

	return ip
}