提交代码

controller/yidun/yidun_callback_controller.go

@@ -4,6 +4,7 @@ import (
 	dataengineService "cid/service/dataengine"
 	"context"
 	"fmt"
+	"strings"
 
 	"gitea.redpowerfuture.com/red-future/common/beans"
 	"github.com/gogf/gf/v2/frame/g"
@@ -40,8 +41,14 @@ type PollResult struct {
 // 易盾回调格式: POST /yidun/callback/receiveImage
 // Body: callbackData={"antispam":{...}}
 func (c *YidunCallbackController) ReceiveImageCallback(r *ghttp.Request) {
+	// IP 白名单校验
+	if !checkCallbackIP(r) {
+		r.Response.WriteJson(CallbackResult{Code: 403, Msg: "IP not allowed"})
+		return
+	}
+
 	ctx := r.Context()
-	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "yidun_callback"})
+	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "yidun_callback", TenantId: 1})
 
 	// 易盾推送的数据在请求体中
 	var callbackData string
@@ -86,8 +93,14 @@ func (c *YidunCallbackController) ReceiveImageCallback(r *ghttp.Request) {
 // 易盾回调格式: POST /yidun/callback/receiveVideo
 // Body: callbackData={"antispam":{...}}
 func (c *YidunCallbackController) ReceiveVideoCallback(r *ghttp.Request) {
+	// IP 白名单校验
+	if !checkCallbackIP(r) {
+		r.Response.WriteJson(CallbackResult{Code: 403, Msg: "IP not allowed"})
+		return
+	}
+
 	ctx := r.Context()
-	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "yidun_callback"})
+	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "yidun_callback", TenantId: 1})
 
 	// 易盾推送的数据在请求体中
 	var callbackData string
@@ -136,7 +149,7 @@ func (c *YidunCallbackController) ReceiveVideoCallback(r *ghttp.Request) {
 // 格式: POST /yidun/callback/poll
 func (c *YidunCallbackController) PollAllResults(r *ghttp.Request) {
 	ctx := r.Context()
-	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin"})
+	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin", TenantId: 1})
 
 	g.Log().Info(ctx, "开始轮询所有待查询的检测结果...")
 
@@ -172,7 +185,7 @@ func (c *YidunCallbackController) PollAllResults(r *ghttp.Request) {
 // 格式: POST /yidun/callback/pollImage
 func (c *YidunCallbackController) PollImageResults(r *ghttp.Request) {
 	ctx := r.Context()
-	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin"})
+	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin", TenantId: 1})
 
 	g.Log().Info(ctx, "开始轮询图片待查询的检测结果...")
 
@@ -198,7 +211,7 @@ func (c *YidunCallbackController) PollImageResults(r *ghttp.Request) {
 // 格式: POST /yidun/callback/pollVideo
 func (c *YidunCallbackController) PollVideoResults(r *ghttp.Request) {
 	ctx := r.Context()
-	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin"})
+	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin", TenantId: 1})
 
 	g.Log().Info(ctx, "开始轮询视频待查询的检测结果...")
 
@@ -224,7 +237,7 @@ func (c *YidunCallbackController) PollVideoResults(r *ghttp.Request) {
 // 格式: POST /yidun/callback/pollTask
 func (c *YidunCallbackController) PollByTaskID(r *ghttp.Request) {
 	ctx := r.Context()
-	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin"})
+	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin", TenantId: 1})
 
 	taskID := r.Get("taskId", "").String()
 	taskType := r.Get("type", "").String() // image 或 video
@@ -262,7 +275,7 @@ func (c *YidunCallbackController) PollByTaskID(r *ghttp.Request) {
 // 格式: GET /yidun/callback/pendingCount
 func (c *YidunCallbackController) GetPendingCount(r *ghttp.Request) {
 	ctx := r.Context()
-	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin"})
+	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin", TenantId: 1})
 
 	count, err := dataengineService.MaterialVerify.GetPendingResultsCount(ctx)
 	if err != nil {
@@ -287,7 +300,7 @@ func (c *YidunCallbackController) GetPendingCount(r *ghttp.Request) {
 // 格式: POST /yidun/callback/processImage
 func (c *YidunCallbackController) ProcessImageCallback(r *ghttp.Request) {
 	ctx := r.Context()
-	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin"})
+	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin", TenantId: 1})
 
 	var req struct {
 		CallbackData string `json:"callbackData" v:"required#回调数据不能为空"`
@@ -311,7 +324,7 @@ func (c *YidunCallbackController) ProcessImageCallback(r *ghttp.Request) {
 // 格式: POST /yidun/callback/processVideo
 func (c *YidunCallbackController) ProcessVideoCallback(r *ghttp.Request) {
 	ctx := r.Context()
-	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin"})
+	ctx = context.WithValue(ctx, "user", &beans.User{UserName: "admin", TenantId: 1})
 
 	var req struct {
 		CallbackData string `json:"callbackData" v:"required#回调数据不能为空"`
@@ -338,3 +351,19 @@ func toString(v interface{}) string {
 	}
 	return ""
 }
+
+// checkCallbackIP 校验回调请求IP是否在白名单内
+func checkCallbackIP(r *ghttp.Request) bool {
+	allowedIPs := g.Cfg().MustGet(r.Context(), "yidun.callback_allowed_ips", "").String()
+	if allowedIPs == "" {
+		return true // 未配置白名单，跳过校验
+	}
+	clientIP := r.GetClientIp()
+	for _, ip := range strings.Split(allowedIPs, ",") {
+		if strings.TrimSpace(ip) == clientIP {
+			return true
+		}
+	}
+	g.Log().Warningf(r.Context(), "回调IP不在白名单中, clientIP=%s, allowedIPs=%s", clientIP, allowedIPs)
+	return false
+}