From c2ba3e4d3d502670f561c03d532a30a6f3695bb8 Mon Sep 17 00:00:00 2001
From: wood chen <wood@czl.net>
Date: Sun, 17 Nov 2024 08:15:29 +0800
Subject: [PATCH] feat(config): add UnmarshalJSON method and update
 NewProxyHandler parameter type

---
 internal/config/types.go  | 63 +++++++++++++++++++++++++--------------
 internal/handler/proxy.go | 11 ++-----
 2 files changed, 44 insertions(+), 30 deletions(-)

diff --git a/internal/config/types.go b/internal/config/types.go
index f3ed8ca..56f5047 100644
--- a/internal/config/types.go
+++ b/internal/config/types.go
@@ -1,5 +1,7 @@
 package config
 
+import "encoding/json"
+
 type Config struct {
 	MAP         map[string]PathConfig `json:"MAP"` // 改为使用PathConfig
 	Compression CompressionConfig     `json:"Compression"`
@@ -28,27 +30,44 @@ type FixedPathConfig struct {
 }
 
 // 添加一个辅助方法来处理字符串到 PathConfig 的转换
-func NewPathConfig(target interface{}) PathConfig {
-	switch v := target.(type) {
-	case string:
-		// 简单字符串格式
-		return PathConfig{
-			DefaultTarget: v,
-		}
-	case map[string]interface{}:
-		// 完整配置格式
-		config := PathConfig{
-			DefaultTarget: v["DefaultTarget"].(string),
-		}
-		if extMap, ok := v["ExtensionMap"].(map[string]interface{}); ok {
-			config.ExtensionMap = make(map[string]string)
-			for ext, target := range extMap {
-				config.ExtensionMap[ext] = target.(string)
-			}
-		}
-		return config
-	default:
-		// 处理异常情况
-		return PathConfig{}
+func (c *Config) UnmarshalJSON(data []byte) error {
+	// 创建一个临时结构来解析原始JSON
+	type TempConfig struct {
+		MAP         map[string]json.RawMessage `json:"MAP"`
+		Compression CompressionConfig          `json:"Compression"`
+		FixedPaths  []FixedPathConfig          `json:"FixedPaths"`
 	}
+
+	var temp TempConfig
+	if err := json.Unmarshal(data, &temp); err != nil {
+		return err
+	}
+
+	// 初始化 MAP
+	c.MAP = make(map[string]PathConfig)
+
+	// 处理每个路径配置
+	for key, raw := range temp.MAP {
+		// 尝试作为字符串解析
+		var strValue string
+		if err := json.Unmarshal(raw, &strValue); err == nil {
+			c.MAP[key] = PathConfig{
+				DefaultTarget: strValue,
+			}
+			continue
+		}
+
+		// 如果不是字符串，尝试作为PathConfig解析
+		var pathConfig PathConfig
+		if err := json.Unmarshal(raw, &pathConfig); err != nil {
+			return err
+		}
+		c.MAP[key] = pathConfig
+	}
+
+	// 复制其他字段
+	c.Compression = temp.Compression
+	c.FixedPaths = temp.FixedPaths
+
+	return nil
 }
diff --git a/internal/handler/proxy.go b/internal/handler/proxy.go
index d869ce3..9f7b83f 100644
--- a/internal/handler/proxy.go
+++ b/internal/handler/proxy.go
@@ -21,15 +21,10 @@ type ProxyHandler struct {
 	pathMap map[string]config.PathConfig
 }
 
-func NewProxyHandler(pathMap map[string]interface{}) *ProxyHandler {
-	convertedMap := make(map[string]config.PathConfig)
-
-	for path, target := range pathMap {
-		convertedMap[path] = config.NewPathConfig(target)
-	}
-
+// 修改参数类型
+func NewProxyHandler(pathMap map[string]config.PathConfig) *ProxyHandler {
 	return &ProxyHandler{
-		pathMap: convertedMap,
+		pathMap: pathMap,
 	}
 }