From d214def59eadad988ba2db52ca62a88e703f7464 Mon Sep 17 00:00:00 2001
From: LeonardWang <wangdeyu1994@163.com>
Date: Thu, 25 Mar 2021 23:34:26 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E4=BB=A3=E7=A0=81=E7=94=9F?=
 =?UTF-8?q?=E6=88=90=E4=B8=AD=E5=AF=B9gorm=E7=9A=84=E9=87=8D=E5=A4=8D?=
 =?UTF-8?q?=E6=A3=80=E6=9F=A5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/utils/injectionCode.go | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/server/utils/injectionCode.go b/server/utils/injectionCode.go
index a9f2ba5a..575a3705 100644
--- a/server/utils/injectionCode.go
+++ b/server/utils/injectionCode.go
@@ -65,9 +65,9 @@ func AutoInjectionCode(filepath string, funcName string, codeData string) error
 	}
 
 	// 在指定函数名，且函数中startComment和endComment都存在时，进行区间查重
-	if (codeStartPos != -1 && codeEndPos != srcDataLen) && (startCommentPos != -1 && endCommentPos != srcDataLen) && expectedFunction != nil {
+	if (codeStartPos != -1 && codeEndPos <= srcDataLen) && (startCommentPos != -1 && endCommentPos != srcDataLen) && expectedFunction != nil {
 		if exist := checkExist(&srcData, startCommentPos, endCommentPos, expectedFunction.Body, codeData); exist {
-			fmt.Println("已存在")
+			fmt.Printf("文件 %s 待插入数据 %s 已存在\n", filepath, codeData)
 			return nil // 这里不需要返回错误？
 		}
 	}
@@ -122,6 +122,21 @@ func checkExist(srcData *[]byte, startPos int, endPos int, blockStmt *ast.BlockS
 			if checkExist(srcData, startPos, endPos, stmt, target) {
 				return true
 			}
+		case *ast.AssignStmt:
+			// 为 model 中的代码进行检查
+			if len(stmt.Rhs) > 0 {
+				if callExpr, ok := stmt.Rhs[0].(*ast.CallExpr); ok {
+					for _, arg := range callExpr.Args {
+						if int(arg.Pos()) > startPos && int(arg.End()) < endPos {
+							text := string((*srcData)[int(arg.Pos()-1):int(arg.End())])
+							key := strings.TrimSpace(text)
+							if key == target {
+								return true
+							}
+						}
+					}
+				}
+			}
 		}
 	}
 	return false