修复 bug 替换静态参数出错

xiecat · Sep 10, 2021 · cfee000 · cfee000
1 parent c2f0d60
commit cfee000
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 18 deletions.
diff --git a/internal/plugin/plugin.go b/internal/plugin/plugin.go
@@ -120,14 +120,13 @@ func LoadPlugin(pathFile string) (*BasePlugin, error) {
 	plugin := &BasePlugin{}
 
 	err = yaml.Unmarshal(content, plugin)
-	plugin.setInitConfig()
 	if err != nil {
 		return nil, fmt.Errorf("[%s] yaml format err:%s ", pathFile, err.Error())
 	}
 	return plugin, nil
 }
 
-func (base *BasePlugin) setInitConfig() {
+func (base *BasePlugin) SetInitConfig() {
 	for _, rule := range base.Rule {
 		// InjectJs 检查
 		if rule.InjectJs != nil {
@@ -165,24 +164,24 @@ func (base *BasePlugin) setInitConfig() {
 		if rule.Replace != nil {
 			for _, rp := range rule.Replace {
 				if rp.Response != nil {
-					if rp.Response.Body != nil {
-						if rp.Response.Header != nil {
-							if location, ok := rp.Response.Header["Location"]; ok {
-								// 替换并且检查替换的变量
-								tmpl, err := template.New("test1").Parse(location)
-								if err != nil {
-									log.Fatal(err.Error())
-								}
-								var tpl bytes.Buffer
-								err = tmpl.Execute(&tpl, PluginVariable)
-								if err != nil {
-									log.Fatal(err.Error())
-								}
-								log.Info("[plugin] Parse:Response.Header.Location: %s ==> %s ", location, tpl.String())
-								rp.Response.Header["Location"] = tpl.String()
+					if rp.Response.Header != nil {
+						if location, ok := rp.Response.Header["Location"]; ok {
+							// 替换并且检查替换的变量
+							tmpl, err := template.New("test1").Parse(location)
+							if err != nil {
+								log.Fatal(err.Error())
 							}
+							var tpl bytes.Buffer
+							err = tmpl.Execute(&tpl, PluginVariable)
+							if err != nil {
+								log.Fatal(err.Error())
+							}
+							log.Info("[plugin] Parse:Response.Header.Location: %s ==> %s ", location, tpl.String())
+							rp.Response.Header["Location"] = tpl.String()
 						}
-						// BodyFile
+					}
+					// BodyFile
+					if rp.Response.Body != nil {
 						if rp.Response.Body.File != "" {
 							fname := rp.Response.Body.File
 							if _, ok := replace.BodyFiles[fname]; ok {

diff --git a/internal/reverse/server.go b/internal/reverse/server.go
@@ -183,6 +183,8 @@ func initReverse(options *options.Options) (revMap map[string]struct {
 				log.Fatal("plugin err please check: %s", v.Rules)
 
 			}
+			// 初始化插件配置
+			rule.SetInitConfig()
 			plugin.Plugins[host] = rule
 			plugin.Plugins[hAddr] = rule