域名限制

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
 }

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

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

helper/func_test.go

@@ -512,8 +512,9 @@ func TestSliceToMap(t *testing.T) {
 		v  string
 	}
 	type args struct {
-		arr []ss
+		arr         []ss
-		fn  func(ss) int
+		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)
 			}
 		})

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

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{

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 {