Merge pull request #182 from komari-monitor/dev

fix: 修复2FA验证逻辑

api/client/report.go

@@ -65,7 +65,7 @@ func WebSocketReport(c *gin.Context) {
 	}
 	upgrader := websocket.Upgrader{
 		CheckOrigin: func(r *http.Request) bool {
-			return true
+			return true // 被控
 		},
 	}
 	// Upgrade the HTTP connection to a WebSocket connection

api/client/terminal.go

@@ -22,7 +22,7 @@ func EstablishConnection(c *gin.Context) {
 	}
 	upgrader := websocket.Upgrader{
 		CheckOrigin: func(r *http.Request) bool {
-			return true // TODO: 检查源站
+			return true // 被控
 		},
 	}
 	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)

api/login.go

@@ -48,13 +48,15 @@ func Login(c *gin.Context) {
 	}
 	// 2FA
 	user, _ := accounts.GetUserByUUID(uuid)
-	if user.TwoFactor != "" && data.TwoFa == "" {
-		RespondError(c, http.StatusUnauthorized, "2FA code is required")
-		return
-	}
-	if ok, err := accounts.Verify2Fa(uuid, data.TwoFa); err != nil && ok {
-		RespondError(c, http.StatusUnauthorized, "Invalid 2FA code")
-		return
+	if user.TwoFactor != "" { // 开启了2FA
+		if data.TwoFa == "" {
+			RespondError(c, http.StatusUnauthorized, "2FA code is required")
+			return
+		}
+		if ok, err := accounts.Verify2Fa(uuid, data.TwoFa); err != nil || !ok {
+			RespondError(c, http.StatusUnauthorized, "Invalid 2FA code")
+			return
+		}
 	}
 	// Create session
 	session, err := accounts.CreateSession(uuid, 2592000, c.Request.UserAgent(), c.ClientIP(), "password")

api/terminal.go

@@ -30,9 +30,7 @@ func RequestTerminal(c *gin.Context) {
 		return
 	}
 	upgrader := websocket.Upgrader{
-		CheckOrigin: func(r *http.Request) bool {
-			return true
-		},
+		CheckOrigin: ws.CheckOrigin,
 	}
 	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
 	if err != nil {

ws/clients.go

@@ -16,9 +16,7 @@ func GetClients(c *gin.Context) {
 		return
 	}
 	upgrader := websocket.Upgrader{
-		CheckOrigin: func(r *http.Request) bool {
-			return true
-		},
+		CheckOrigin: CheckOrigin,
 	}
 	// Upgrade the HTTP connection to a WebSocket connection
 	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)

ws/origin.go

@@ -0,0 +1,19 @@
+package ws
+
+import (
+	"net/http"
+	"net/url"
+)
+
+func CheckOrigin(r *http.Request) bool {
+	origin := r.Header.Get("Origin")
+	if origin == "" {
+		return false
+	}
+	host := r.Host
+	originUrl, err := url.Parse(origin)
+	if err != nil {
+		return false
+	}
+	return originUrl.Host == host
+}