From 90120056d4080d38ecf1da14ecd2dccef16bd1ee Mon Sep 17 00:00:00 2001
From: xing <fthvgb1>
Date: Mon, 10 Oct 2022 21:07:39 +0800
Subject: [PATCH] =?UTF-8?q?=E5=9F=9F=E5=90=8D=E9=99=90=E5=88=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 actions/common/comments.go       |  2 +-
 config.example.yaml              |  4 ++++
 helper/func.go                   |  7 ++++++-
 helper/func_test.go              | 20 ++++++++++++++++----
 middleware/validateservername.go | 23 +++++++++++++++++++++++
 route/route.go                   |  8 +++++++-
 vars/config.go                   |  1 +
 7 files changed, 58 insertions(+), 7 deletions(-)
 create mode 100644 middleware/validateservername.go

diff --git a/actions/common/comments.go b/actions/common/comments.go
index 680419d..9c43df4 100644
--- a/actions/common/comments.go
+++ b/actions/common/comments.go
@@ -60,5 +60,5 @@ func getCommentByIds(args ...any) (map[uint64]models.WpComments, error) {
 	}
 	return helper.SliceToMap(r, func(t models.WpComments) uint64 {
 		return t.CommentId
-	}), err
+	}, true), err
 }
diff --git a/config.example.yaml b/config.example.yaml
index ec9d9f5..fcb9709 100644
--- a/config.example.yaml
+++ b/config.example.yaml
@@ -57,3 +57,7 @@ commentsCacheTime: 24h
 gzip: false
 # 提交评论url
 postCommentUrl: http://wp.test/wp-comments-post.php
+# TrustIps
+trustIps: []
+# trust servername 信任的域名
+trustServerNames: ["xy.test","blog.xy.test"]
diff --git a/helper/func.go b/helper/func.go
index 284e103..99607e1 100644
--- a/helper/func.go
+++ b/helper/func.go
@@ -231,10 +231,15 @@ func SliceSelfReverse[T any](arr []T) []T {
 	return arr
 }
 
-func SliceToMap[K comparable, V any](arr []V, fn func(V) K) map[K]V {
+func SliceToMap[K comparable, V any](arr []V, fn func(V) K, isCoverPrev bool) map[K]V {
 	m := make(map[K]V)
 	for _, v := range arr {
 		k := fn(v)
+		if !isCoverPrev {
+			if _, ok := m[k]; ok {
+				continue
+			}
+		}
 		m[k] = v
 	}
 	return m
diff --git a/helper/func_test.go b/helper/func_test.go
index 1e3766b..539fe8a 100644
--- a/helper/func_test.go
+++ b/helper/func_test.go
@@ -512,8 +512,9 @@ func TestSliceToMap(t *testing.T) {
 		v  string
 	}
 	type args struct {
-		arr []ss
-		fn  func(ss) int
+		arr         []ss
+		fn          func(ss) int
+		isCoverPrev bool
 	}
 	tests := []struct {
 		name string
@@ -523,17 +524,28 @@ func TestSliceToMap(t *testing.T) {
 		{
 			name: "t1",
 			args: args{
-				arr: []ss{{1, "k1"}, {2, "v2"}},
+				arr: []ss{{1, "k1"}, {2, "v2"}, {2, "v3"}},
 				fn: func(s ss) int {
 					return s.id
 				},
+				isCoverPrev: true,
+			},
+			want: map[int]ss{1: {1, "k1"}, 2: {2, "v3"}},
+		}, {
+			name: "t2",
+			args: args{
+				arr: []ss{{1, "k1"}, {2, "v2"}, {2, "v3"}},
+				fn: func(s ss) int {
+					return s.id
+				},
+				isCoverPrev: false,
 			},
 			want: map[int]ss{1: {1, "k1"}, 2: {2, "v2"}},
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if got := SliceToMap(tt.args.arr, tt.args.fn); !reflect.DeepEqual(got, tt.want) {
+			if got := SliceToMap(tt.args.arr, tt.args.fn, tt.args.isCoverPrev); !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("SliceToMap() = %v, want %v", got, tt.want)
 			}
 		})
diff --git a/middleware/validateservername.go b/middleware/validateservername.go
new file mode 100644
index 0000000..57febb2
--- /dev/null
+++ b/middleware/validateservername.go
@@ -0,0 +1,23 @@
+package middleware
+
+import (
+	"github.com/gin-gonic/gin"
+	"github/fthvgb1/wp-go/helper"
+	"github/fthvgb1/wp-go/vars"
+	"net/http"
+	"strings"
+)
+
+func ValidateServerNames() func(ctx *gin.Context) {
+	serverName := helper.SliceToMap(vars.Conf.TrustServerNames, func(v string) string {
+		return v
+	}, true)
+	return func(c *gin.Context) {
+		if len(serverName) > 0 {
+			if _, ok := serverName[strings.Split(c.Request.Host, ":")[0]]; !ok {
+				c.Status(http.StatusForbidden)
+				c.Abort()
+			}
+		}
+	}
+}
diff --git a/route/route.go b/route/route.go
index c295474..09cf0ef 100644
--- a/route/route.go
+++ b/route/route.go
@@ -36,7 +36,13 @@ func SetupRouter() *gin.Engine {
 			return t.Format("2006年 01月 02日")
 		},
 	}).SetTemplate()
-	r.Use(gin.Logger(), middleware.FlowLimit(vars.Conf.MaxRequestSleepNum, vars.Conf.MaxRequestNum, vars.Conf.SingleIpSearchNum, vars.Conf.SleepTime), gin.Recovery(), middleware.SetStaticFileCache)
+	r.Use(
+		middleware.ValidateServerNames(),
+		gin.Logger(),
+		middleware.FlowLimit(vars.Conf.MaxRequestSleepNum, vars.Conf.MaxRequestNum, vars.Conf.SingleIpSearchNum, vars.Conf.SleepTime),
+		gin.Recovery(),
+		middleware.SetStaticFileCache,
+	)
 	//gzip 因为一般会用nginx做反代时自动使用gzip,所以go这边本身可以不用
 	if vars.Conf.Gzip {
 		r.Use(gzip.Gzip(gzip.DefaultCompression, gzip.WithExcludedPaths([]string{
diff --git a/vars/config.go b/vars/config.go
index bd19dd4..f9e7050 100644
--- a/vars/config.go
+++ b/vars/config.go
@@ -34,6 +34,7 @@ type Config struct {
 	Gzip                    bool            `yaml:"gzip"`
 	PostCommentUrl          string          `yaml:"postCommentUrl"`
 	TrustIps                []string        `yaml:"trustIps"`
+	TrustServerNames        []string        `yaml:"trustServerNames"`
 }
 
 type Mysql struct {