2.md

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
// 可以运行    webp 支持  移除referer限制,直接写死限制网站访问
// 改变referer头就不可以
// $ curl -i -H "Referer: https://blog.775866.xyz" "https://extract-main-color.775866.xyz/api?img=https://bu.dusays.com/2024/07/04/668675187383f.webp"
// % Total % Received % Xferd Average Speed Time Time Time Current
// Dload Upload Total Spent Left Speed
// 100 18 100 18 0 0 4 0 0:00:04 0:00:03 0:00:01 4HTTP/1.1 200 OK
// Access-Control-Allow-Headers: Content-Type, Referer
// Access-Control-Allow-Methods: GET, OPTIONS
// Access-Control-Allow-Origin: *
// Age: 0
// Cache-Control: public, max-age=0, must-revalidate
// Connection: keep-alive
// Content-Length: 18
// Content-Type: application/json
// Date: Thu, 04 Jul 2024 13:23:21 GMT
// Server: Vercel
// Strict-Transport-Security: max-age=63072000
// X-Vercel-Cache: MISS
// X-Vercel-Id: fra1::iad1::b9s8b-1720099401103-4c92dd2451bc

// {"RGB":"#373a3b"}



package handler

import (
"crypto/md5"
"encoding/base64"
"encoding/json"
"fmt"
"image"
"log"
"net/http"
"os"
"path/filepath"
"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"
"golang.org/x/image/webp"
)

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

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")

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")
}
}

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()

contentType := resp.Header.Get("Content-Type")

var img image.Image

switch contentType {
case "image/webp":
img, err = webp.Decode(resp.Body)
if err != nil {
return "", fmt.Errorf("解码 WebP 图片时出错: %v", err)
}
default:
img, err = imaging.Decode(resp.Body)
if err != nil {
return "", fmt.Errorf("解码图片时出错: %v", 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 isRefererAllowed(referer string) bool {
allowedReferer := "https://blog.775866.xyz"

if strings.HasPrefix(referer, allowedReferer) {
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)
}
}