api/img2color.go

package handler

import (
	"crypto/md5"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"image"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/go-redis/redis/v8"
	"github.com/joho/godotenv"
	"github.com/lucasb-eyer/go-colorful"
	"github.com/nfnt/resize"
	"github.com/disintegration/imaging"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"golang.org/x/net/context"
)

var redisClient *redis.Client
var mongoClient *mongo.Client
var cacheEnabled bool
var useMongoDB bool
var redisDB int
var mongoDB string
var ctx = context.Background()
var colorsCollection *mongo.Collection
var allowedReferers []string

func init() {
	currentDir, err := os.Getwd()
	if err != nil {
		fmt.Printf("获取当前工作目录路径时出错：%v\n", err)
		return
	}

	envFile := filepath.Join(currentDir, ".env")

	err = godotenv.Load(envFile)
	if err != nil {
		fmt.Printf("加载 .env 文件时出错：%v\n", err)
		return
	}

	redisAddr := os.Getenv("REDIS_ADDRESS")
	redisPassword := os.Getenv("REDIS_PASSWORD")
	cacheEnabledStr := os.Getenv("USE_REDIS_CACHE")
	redisDBStr := os.Getenv("REDIS_DB")
	mongoDB = os.Getenv("MONGO_DB")
	mongoURI := os.Getenv("MONGO_URI")
	referers := os.Getenv("ALLOWED_REFERERS")

	redisDB, err = strconv.Atoi(redisDBStr)
	if err != nil {
		redisDB = 0
	}

	redisClient = redis.NewClient(&redis.Options{
		Addr:     redisAddr,
		Password: redisPassword,
		DB:       redisDB,
	})

	cacheEnabled = cacheEnabledStr == "true"

	useMongoDBStr := os.Getenv("USE_MONGODB")
	useMongoDB = useMongoDBStr == "true"
	if useMongoDB {
		log.Println("连接到MongoDB...")
		clientOptions := options.Client().ApplyURI(mongoURI)
		mongoClient, err = mongo.Connect(ctx, clientOptions)
		if err != nil {
			log.Fatalf("连接到MongoDB时出错：%v", err)
		}
		log.Println("已连接到MongoDB！")

		colorsCollection = mongoClient.Database(mongoDB).Collection("colors")
	}

	allowedReferers = parseReferers(referers)
}

func calculateMD5Hash(data []byte) string {
	hash := md5.Sum(data)
	return base64.StdEncoding.EncodeToString(hash[:])
}

func extractMainColor(imgURL string) (string, error) {
	md5Hash := calculateMD5Hash([]byte(imgURL))

	if cacheEnabled && redisClient != nil {
		cachedColor, err := redisClient.Get(ctx, md5Hash).Result()
		if err == nil && cachedColor != "" {
			return cachedColor, nil
		}
	}

	req, err := http.NewRequest("GET", imgURL, nil)
	if err != nil {
		return "", err
	}

	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36 Edg/115.0.1901.253")

	client := http.DefaultClient
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	var img image.Image

	img, err = imaging.Decode(resp.Body)

	if err != nil {
		return "", err
	}

	img = resize.Resize(50, 0, img, resize.Lanczos3)

	bounds := img.Bounds()
	var r, g, b uint32
	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			c := img.At(x, y)
			r0, g0, b0, _ := c.RGBA()
			r += r0
			g += g0
			b += b0
		}
	}

	totalPixels := uint32(bounds.Dx() * bounds.Dy())
	averageR := r / totalPixels
	averageG := g / totalPixels
	averageB := b / totalPixels

	mainColor := colorful.Color{R: float64(averageR) / 0xFFFF, G: float64(averageG) / 0xFFFF, B: float64(averageB) / 0xFFFF}

	colorHex := mainColor.Hex()

	if cacheEnabled && redisClient != nil {
		_, err := redisClient.Set(ctx, md5Hash, colorHex, 0).Result()
		if err != nil {
			log.Printf("将结果存储在缓存中时出错：%v\n", err)
		}
	}

	if useMongoDB && colorsCollection != nil {
		_, err := colorsCollection.InsertOne(ctx, bson.M{
			"url":   imgURL,
			"color": colorHex,
		})
		if err != nil {
			log.Printf("将结果存储在MongoDB中时出错：%v\n", err)
		}
	}

	return colorHex, nil
}

func handleImageColor(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", "GET, OPTIONS")
	w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Referer")

	if r.Method == http.MethodOptions {
		w.WriteHeader(http.StatusOK)
		return
	}

	referer := r.Header.Get("Referer")
	if !isRefererAllowed(referer) {
		http.Error(w, "禁止访问", http.StatusForbidden)
		return
	}

	imgURL := r.URL.Query().Get("img")
	if imgURL == "" {
		http.Error(w, "缺少img参数", http.StatusBadRequest)
		return
	}

	color, err := extractMainColor(imgURL)
	if err != nil {
		http.Error(w, fmt.Sprintf("提取主色调失败：%v", err), http.StatusInternalServerError)
		return
	}

	data := map[string]string{
		"RGB": color,
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(data)
}

func Handler(w http.ResponseWriter, r *http.Request) {
	handleImageColor(w, r)
}

func parseReferers(referers string) []string {
	refererList := strings.Split(referers, ",")
	for i, referer := range refererList {
		refererList[i] = strings.TrimSpace(referer)
	}
	return refererList
}

func isRefererAllowed(referer string) bool {
	if len(allowedReferers) == 0 {
		return true
	}

	for _, allowedReferer := range allowedReferers {
		allowedReferer = strings.ReplaceAll(allowedReferer, ".", "\\.")
		allowedReferer = strings.ReplaceAll(allowedReferer, "*", ".*")
		match, _ := regexp.MatchString(allowedReferer, referer)
		if match {
			return true
		}
	}

	return false
}

func main() {
	http.HandleFunc("/api", Handler)

	port := os.Getenv("PORT")
	if port == "" {
		port = "3000"
	}

	log.Printf("服务器监听在：%s...\n", port)
	err := http.ListenAndServe(":"+port, nil)
	if err != nil {
		log.Fatalf("启动服务器时出错：%v\n", err)
	}
}