userID in ai chat

Author: ppolariss

backend/apis/message/api.go

@@ -142,8 +142,8 @@ func ListChatRecords(c *fiber.Ctx) (err error) {
 
 // ListMyChatRecords @ListMyChatRecords
 // @Router /api/chats/{id}/records/me [get]
-// @Summary list my records
-// @Description list my records
+// @Summary list my records in a chat
+// @Description list my records in a chat
 // @Tags ChatRecord
 // @Accept json
 // @Produce json
@@ -259,6 +259,11 @@ func ListMyRecords(c *fiber.Ctx) (err error) {
 // @param Authorization header string true "Bearer和token空格拼接"
 func MossChat(c *websocket.Conn) {
 	var err error
+	userID, err := LoadUserFromWs(c)
+	if err != nil {
+		_ = c.WriteJSON(common.Unauthorized("Unauthorized"))
+		return
+	}
 	defer func() {
 		if err != nil {
 			utils.Logger.Error(
@@ -289,7 +294,7 @@ func MossChat(c *websocket.Conn) {
 	log.Println("Received message from client:", requestMessage)
 	go addRecord(AddRecordsRequest{
 		CreatedAt: time.Now(),
-		UserID:    0,
+		UserID:    userID,
 		RoomID:    "moss",
 		Type:      "toMoss",
 		ToID:      0,
@@ -339,7 +344,7 @@ func MossChat(c *websocket.Conn) {
 					UserID:    0,
 					RoomID:    "moss",
 					Type:      "fromMoss",
-					ToID:      0,
+					ToID:      userID,
 					Message:   mossResponse.Output,
 				})
 				return

backend/models/auth.go

@@ -160,15 +160,23 @@ func GetUserIDFromWs(c *websocket.Conn) (int, error) {
 	if !ok {
 		return 0, common.Unauthorized()
 	}
+	userID := int(id.(float64))
+	var userJwtSecret UserJwtSecret
+	err = DB.Take(&userJwtSecret, userID).Error
+	if err != nil {
+		return 0, common.Unauthorized("Unauthorized")
+	}
+	CheckJWTToken(token, userJwtSecret.Secret)
 	return int(id.(float64)), nil
 }
 
-func LoadUserFromWs(c *websocket.Conn) (*User, error) {
+func LoadUserFromWs(c *websocket.Conn) (int, error) {
 	userID, err := GetUserIDFromWs(c)
 	if err != nil {
-		return nil, err
+		return 0, err
 	}
-	return LoadUserByID(userID)
+	return userID, nil
+	//LoadUserByID(userID)
 }
 
 // parseJWT extracts and parse token