From 7064c26e5ec2bfefd6d1259b44050ce639a88d6e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BE=90=E6=B6=9B?= <midnite@outlook.com>
Date: Fri, 12 Aug 2022 17:02:44 +0800
Subject: [PATCH] =?UTF-8?q?fix(security):=E4=BF=AE=E5=A4=8D=E5=9C=A8?=
 =?UTF-8?q?=E8=8E=B7=E5=8F=96=E7=94=A8=E6=88=B7=E4=BC=9A=E8=AF=9D=E4=BB=A5?=
 =?UTF-8?q?=E5=90=8E=EF=BC=8C=E7=94=A8=E6=88=B7=E4=BC=9A=E8=AF=9D=E4=B8=8D?=
 =?UTF-8?q?=E5=AD=98=E5=9C=A8=E7=9A=84=E6=83=85=E5=86=B5=E4=B8=8B=E4=BE=9D?=
 =?UTF-8?q?=E6=97=A7=E4=BC=9A=E4=BF=9D=E5=AD=98=E8=87=B3=E4=B8=8A=E4=B8=8B?=
 =?UTF-8?q?=E6=96=87=E7=9A=84=E9=97=AE=E9=A2=98=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 router/router.go     |  2 +-
 security/security.go | 25 ++++++++++++++++++++-----
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/router/router.go b/router/router.go
index 896a568..9471b3d 100644
--- a/router/router.go
+++ b/router/router.go
@@ -32,7 +32,7 @@ func Recover(c *gin.Context) {
 			//打印错误堆栈信息
 			log.Printf("panic: %v\n", r)
 			debug.PrintStack()
-			response.NewResult(c).Error(500, "服务器内部错误")
+			// response.NewResult(c).Error(500, "服务器内部错误")
 		}
 	}()
 	//继续后续接口调用
diff --git a/security/security.go b/security/security.go
index 8cff1e3..4f33855 100644
--- a/security/security.go
+++ b/security/security.go
@@ -3,6 +3,7 @@ package security
 import (
 	"electricity_bill_calc/cache"
 	"electricity_bill_calc/model"
+	"log"
 	"net/http"
 	"strings"
 
@@ -18,7 +19,7 @@ func SessionRecovery(c *gin.Context) {
 		token := strings.Fields(auth)[1]
 		session, err := cache.RetreiveSession(token)
 
-		if err == nil {
+		if err == nil && session != nil {
 			c.Set("session", session)
 		}
 	}
@@ -29,7 +30,11 @@ func SessionRecovery(c *gin.Context) {
 // ! 通过该中间件以后，是可以保证上下文中一定具有用户会话信息的。
 func MustAuthenticated(c *gin.Context) {
 	session, exists := c.Get("session")
-	if _, ok := session.(*model.Session); !exists || session == nil || !ok {
+	if !exists || session == nil {
+		c.AbortWithStatus(http.StatusForbidden)
+		return
+	}
+	if _, ok := session.(*model.Session); !ok {
 		c.AbortWithStatus(http.StatusForbidden)
 	}
 	c.Next()
@@ -39,7 +44,10 @@ func MustAuthenticated(c *gin.Context) {
 // ! 通过该中间件以后，是可以保证上下文中一定具有用户会话信息的。
 func EnterpriseAuthorize(c *gin.Context) {
 	session, exists := c.Get("session")
-	if sess, ok := session.(*model.Session); !exists || !ok || sess.Type != 0 {
+	if !exists || session == nil {
+		c.AbortWithStatus(http.StatusForbidden)
+	}
+	if sess, ok := session.(*model.Session); !ok || sess.Type != 0 {
 		c.AbortWithStatus(http.StatusForbidden)
 	}
 	c.Next()
@@ -49,7 +57,10 @@ func EnterpriseAuthorize(c *gin.Context) {
 // ! 通过该中间件以后，是可以保证上下文中一定具有用户会话信息的。
 func ManagementAuthorize(c *gin.Context) {
 	session, exists := c.Get("session")
-	if sess, ok := session.(*model.Session); !exists || !ok || (sess.Type != 1 && sess.Type != 2) {
+	if !exists || session == nil {
+		c.AbortWithStatus(http.StatusForbidden)
+	}
+	if sess, ok := session.(*model.Session); !ok || (sess.Type != 1 && sess.Type != 2) {
 		c.AbortWithStatus(http.StatusForbidden)
 	}
 	c.Next()
@@ -59,7 +70,11 @@ func ManagementAuthorize(c *gin.Context) {
 // ! 通过该中间件以后，是可以保证上下文中一定具有用户会话信息的。
 func OPSAuthorize(c *gin.Context) {
 	session, exists := c.Get("session")
-	if sess, ok := session.(*model.Session); !exists || !ok || sess.Type != 2 {
+	log.Printf("[debug]session exists: %v, %v \n", exists, session)
+	if !exists {
+		c.AbortWithStatus(http.StatusForbidden)
+	}
+	if sess, ok := session.(*model.Session); !ok || sess.Type != 2 {
 		c.AbortWithStatus(http.StatusForbidden)
 	}
 	c.Next()