From 44e4c59493c233f11f3685c69f0b50dbbc49c3fb Mon Sep 17 00:00:00 2001
From: skiffer-git <44203734@qq.com>
Date: Mon, 15 Aug 2022 18:34:03 +0800
Subject: [PATCH] Restrict user login with IP

---
 internal/demo/register/ip_limit.go     | 6 +++---
 internal/demo/register/set_password.go | 2 ++
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/internal/demo/register/ip_limit.go b/internal/demo/register/ip_limit.go
index bcd84d489..c2bdd6d07 100644
--- a/internal/demo/register/ip_limit.go
+++ b/internal/demo/register/ip_limit.go
@@ -34,7 +34,7 @@ func QueryIP(c *gin.Context) {
 	ips, err := imdb.QueryUserIPLimits(req.IP)
 	if err != nil {
 		log.NewError(req.OperationID, "GetInvitationCode failed", req.IP)
-		c.JSON(http.StatusInternalServerError, gin.H{"errCode": constant.ErrDB, "errMsg": "QueryUserIPLimits error!"})
+		c.JSON(http.StatusInternalServerError, gin.H{"errCode": constant.ErrDB.ErrCode, "errMsg": "QueryUserIPLimits error!"})
 		return
 	}
 	resp.IP = req.IP
@@ -85,7 +85,7 @@ func AddIPLimit(c *gin.Context) {
 		LimitTime:     time.Time{},
 	}); err != nil {
 		log.NewError(req.OperationID, utils.GetSelfFuncName(), err.Error(), req.IP, req.LimitTime)
-		c.JSON(http.StatusInternalServerError, gin.H{"errCode": constant.ErrDB, "errMsg": "InsertOneIntoIpLimits error!"})
+		c.JSON(http.StatusInternalServerError, gin.H{"errCode": constant.ErrDB.ErrCode, "errMsg": "InsertOneIntoIpLimits error!"})
 		return
 	}
 	c.JSON(http.StatusOK, gin.H{"errCode": 0, "errMsg": ""})
@@ -122,7 +122,7 @@ func QueryUserIPLimitLogin(c *gin.Context) {
 	resp, err := imdb.GetIpLimitsLoginByUserID(req.UserID)
 	if err != nil {
 		log.NewError(req.OperationID, utils.GetSelfFuncName(), err.Error(), req.UserID)
-		c.JSON(http.StatusInternalServerError, gin.H{"errCode": constant.ErrDB, "errMsg": "GetIpLimitsByUserID error!"})
+		c.JSON(http.StatusInternalServerError, gin.H{"errCode": constant.ErrDB.ErrCode, "errMsg": "GetIpLimitsByUserID error!"})
 		return
 	}
 	log.NewInfo(req.OperationID, utils.GetSelfFuncName(), "resp:", resp)
diff --git a/internal/demo/register/set_password.go b/internal/demo/register/set_password.go
index 632254383..def5b8c16 100644
--- a/internal/demo/register/set_password.go
+++ b/internal/demo/register/set_password.go
@@ -47,9 +47,11 @@ func SetPassword(c *gin.Context) {
 	if LimitError != nil {
 		log.Error(params.OperationID, utils.GetSelfFuncName(), LimitError, ip)
 		c.JSON(http.StatusBadRequest, gin.H{"errCode": constant.FormattingError, "errMsg": LimitError.Error()})
+		return
 	}
 	if Limited {
 		c.JSON(http.StatusBadRequest, gin.H{"errCode": constant.FormattingError, "errMsg": "limited"})
+		return
 	}
 
 	var account string