From 17f9692ce65004217aa67beb912529c397f5e0b0 Mon Sep 17 00:00:00 2001
From: wood chen <wood@czl.net>
Date: Sun, 17 Nov 2024 08:24:57 +0800
Subject: [PATCH] refactor(config): optimize extension mapping handling and add
 utility methods

---
 data/config.json          |  4 ++--
 internal/config/types.go  | 44 +++++++++++++++++++++++++++++++++++----
 internal/handler/proxy.go |  4 +---
 3 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/data/config.json b/data/config.json
index 3fee8d0..1bc047e 100644
--- a/data/config.json
+++ b/data/config.json
@@ -3,8 +3,8 @@
     "/path1": {
       "DefaultTarget": "https://path1.com/path/path/path",
       "ExtensionMap": {
-        "jpg": "https://path1-img.com/path/path/path",
-        "png": "https://path1-img.com/path/path/path"
+        "jpg,png,avif": "https://path1-img.com/path/path/path",
+        "mp4,webm": "https://path1-video.com/path/path/path"
       }
     },
     "/path2": "https://path2.com",
diff --git a/internal/config/types.go b/internal/config/types.go
index 56f5047..b2a1b9a 100644
--- a/internal/config/types.go
+++ b/internal/config/types.go
@@ -1,6 +1,9 @@
 package config
 
-import "encoding/json"
+import (
+	"encoding/json"
+	"strings"
+)
 
 type Config struct {
 	MAP         map[string]PathConfig `json:"MAP"` // 改为使用PathConfig
@@ -9,8 +12,9 @@ type Config struct {
 }
 
 type PathConfig struct {
-	DefaultTarget string            `json:"DefaultTarget"` // 默认回源地址
-	ExtensionMap  map[string]string `json:"ExtensionMap"`  // 特定后缀的回源地址
+	DefaultTarget   string            `json:"DefaultTarget"` // 默认回源地址
+	ExtensionMap    map[string]string `json:"ExtensionMap"`  // 特定后缀的回源地址
+	processedExtMap map[string]string // 内部使用，存储拆分后的映射
 }
 
 type CompressionConfig struct {
@@ -51,9 +55,11 @@ func (c *Config) UnmarshalJSON(data []byte) error {
 		// 尝试作为字符串解析
 		var strValue string
 		if err := json.Unmarshal(raw, &strValue); err == nil {
-			c.MAP[key] = PathConfig{
+			pathConfig := PathConfig{
 				DefaultTarget: strValue,
 			}
+			pathConfig.ProcessExtensionMap() // 处理扩展名映射
+			c.MAP[key] = pathConfig
 			continue
 		}
 
@@ -62,6 +68,7 @@ func (c *Config) UnmarshalJSON(data []byte) error {
 		if err := json.Unmarshal(raw, &pathConfig); err != nil {
 			return err
 		}
+		pathConfig.ProcessExtensionMap() // 处理扩展名映射
 		c.MAP[key] = pathConfig
 	}
 
@@ -71,3 +78,32 @@ func (c *Config) UnmarshalJSON(data []byte) error {
 
 	return nil
 }
+
+// 添加处理扩展名映射的方法
+func (p *PathConfig) ProcessExtensionMap() {
+	if p.ExtensionMap == nil {
+		return
+	}
+
+	p.processedExtMap = make(map[string]string)
+	for exts, target := range p.ExtensionMap {
+		// 分割扩展名
+		for _, ext := range strings.Split(exts, ",") {
+			ext = strings.TrimSpace(ext) // 移除可能的空格
+			if ext != "" {
+				p.processedExtMap[ext] = target
+			}
+		}
+	}
+}
+
+// 添加获取目标URL的方法
+func (p *PathConfig) GetTargetForExt(ext string) string {
+	if p.processedExtMap == nil {
+		p.ProcessExtensionMap()
+	}
+	if target, exists := p.processedExtMap[ext]; exists {
+		return target
+	}
+	return p.DefaultTarget
+}
diff --git a/internal/handler/proxy.go b/internal/handler/proxy.go
index 9f7b83f..d7f7086 100644
--- a/internal/handler/proxy.go
+++ b/internal/handler/proxy.go
@@ -79,9 +79,7 @@ func (h *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		ext := strings.ToLower(path.Ext(decodedPath))
 		if ext != "" {
 			ext = ext[1:] // 移除开头的点
-			if alternativeTarget, exists := pathConfig.ExtensionMap[ext]; exists {
-				targetBase = alternativeTarget
-			}
+			targetBase = pathConfig.GetTargetForExt(ext)
 		}
 	}