From a6025e6fabb2f503c940b548cdedf376ef3d39f8 Mon Sep 17 00:00:00 2001
From: RockYang <yangjian102621@gmail.com>
Date: Thu, 7 Dec 2023 14:02:13 +0800
Subject: [PATCH] fix: fixed bug for prompt code format, prevent xss attacks

---
 api/handler/chatimpl/azure_handler.go   |  3 ++-
 api/handler/chatimpl/baidu_handler.go   |  3 ++-
 api/handler/chatimpl/chatglm_handler.go |  3 ++-
 api/handler/chatimpl/openai_handler.go  |  3 ++-
 api/handler/chatimpl/xunfei_handler.go  |  3 ++-
 web/src/views/ChatPlus.vue              | 12 ++++--------
 6 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/api/handler/chatimpl/azure_handler.go b/api/handler/chatimpl/azure_handler.go
index a4138616..489ce012 100644
--- a/api/handler/chatimpl/azure_handler.go
+++ b/api/handler/chatimpl/azure_handler.go
@@ -10,6 +10,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"gorm.io/gorm"
+	"html/template"
 	"io"
 	"strings"
 	"time"
@@ -200,7 +201,7 @@ func (h *ChatHandler) sendAzureMessage(
 					RoleId:     role.Id,
 					Type:       types.PromptMsg,
 					Icon:       userVo.Avatar,
-					Content:    prompt,
+					Content:    template.HTMLEscapeString(prompt),
 					Tokens:     promptToken,
 					UseContext: useContext,
 				}
diff --git a/api/handler/chatimpl/baidu_handler.go b/api/handler/chatimpl/baidu_handler.go
index f20756b8..bc92b8a2 100644
--- a/api/handler/chatimpl/baidu_handler.go
+++ b/api/handler/chatimpl/baidu_handler.go
@@ -9,6 +9,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"html/template"
 	"io"
 	"net/http"
 	"strings"
@@ -156,7 +157,7 @@ func (h *ChatHandler) sendBaiduMessage(
 					RoleId:     role.Id,
 					Type:       types.PromptMsg,
 					Icon:       userVo.Avatar,
-					Content:    prompt,
+					Content:    template.HTMLEscapeString(prompt),
 					Tokens:     promptToken,
 					UseContext: true,
 				}
diff --git a/api/handler/chatimpl/chatglm_handler.go b/api/handler/chatimpl/chatglm_handler.go
index 0efa47a3..0a9f74df 100644
--- a/api/handler/chatimpl/chatglm_handler.go
+++ b/api/handler/chatimpl/chatglm_handler.go
@@ -10,6 +10,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/golang-jwt/jwt/v5"
+	"html/template"
 	"io"
 	"strings"
 	"time"
@@ -135,7 +136,7 @@ func (h *ChatHandler) sendChatGLMMessage(
 					RoleId:     role.Id,
 					Type:       types.PromptMsg,
 					Icon:       userVo.Avatar,
-					Content:    prompt,
+					Content:    template.HTMLEscapeString(prompt),
 					Tokens:     promptToken,
 					UseContext: true,
 				}
diff --git a/api/handler/chatimpl/openai_handler.go b/api/handler/chatimpl/openai_handler.go
index cda0836c..d8eaf6be 100644
--- a/api/handler/chatimpl/openai_handler.go
+++ b/api/handler/chatimpl/openai_handler.go
@@ -10,6 +10,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"gorm.io/gorm"
+	"html/template"
 	"io"
 	"strings"
 	"time"
@@ -199,7 +200,7 @@ func (h *ChatHandler) sendOpenAiMessage(
 					RoleId:     role.Id,
 					Type:       types.PromptMsg,
 					Icon:       userVo.Avatar,
-					Content:    prompt,
+					Content:    template.HTMLEscapeString(prompt),
 					Tokens:     promptToken,
 					UseContext: useContext,
 				}
diff --git a/api/handler/chatimpl/xunfei_handler.go b/api/handler/chatimpl/xunfei_handler.go
index 2221d9e9..be39a84b 100644
--- a/api/handler/chatimpl/xunfei_handler.go
+++ b/api/handler/chatimpl/xunfei_handler.go
@@ -12,6 +12,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/gorilla/websocket"
+	"html/template"
 	"io"
 	"net/http"
 	"net/url"
@@ -198,7 +199,7 @@ func (h *ChatHandler) sendXunFeiMessage(
 				RoleId:     role.Id,
 				Type:       types.PromptMsg,
 				Icon:       userVo.Avatar,
-				Content:    prompt,
+				Content:    template.HTMLEscapeString(prompt),
 				Tokens:     promptToken,
 				UseContext: true,
 			}
diff --git a/web/src/views/ChatPlus.vue b/web/src/views/ChatPlus.vue
index 64e6ce8c..22159e19 100644
--- a/web/src/views/ChatPlus.vue
+++ b/web/src/views/ChatPlus.vue
@@ -245,7 +245,7 @@ import {
   VideoPause
 } from '@element-plus/icons-vue'
 import 'highlight.js/styles/a11y-dark.css'
-import {dateFormat, isMobile, randString, removeArrayItem, renderInputText, UUID} from "@/utils/libs";
+import {dateFormat, isMobile, randString, removeArrayItem, UUID} from "@/utils/libs";
 import {ElMessage, ElMessageBox} from "element-plus";
 import hl from "highlight.js";
 import {getSessionId, getUserToken, removeUserToken} from "@/store/session";
@@ -690,13 +690,12 @@ const sendMessage = function () {
   if (prompt.value.trim().length === 0 || canSend.value === false) {
     return false;
   }
-
   // 追加消息
   chatData.value.push({
     type: "prompt",
     id: randString(32),
     icon: loginUser.value.avatar,
-    content: renderInputText(prompt.value),
+    content: md.render(prompt.value),
     created_at: new Date().getTime(),
   });
 
@@ -761,10 +760,7 @@ const loadChatHistory = function (chatId) {
     }
     showHello.value = false
     for (let i = 0; i < data.length; i++) {
-      if (data[i].type === "prompt") {
-        chatData.value.push(data[i]);
-        continue;
-      } else if (data[i].type === "mj") {
+      if (data[i].type === "mj") {
         data[i].content = JSON.parse(data[i].content)
         data[i].content.html = md.render(data[i].content?.content)
         chatData.value.push(data[i]);
@@ -802,7 +798,7 @@ const reGenerate = function () {
     type: "prompt",
     id: randString(32),
     icon: loginUser.value.avatar,
-    content: renderInputText(text)
+    content: md.render(text)
   });
   socket.value.send(text);
 }