From df7e1b3d988fb39f03b0c2dc5d61bfb7cb8068fb Mon Sep 17 00:00:00 2001
From: feng <1304903146@qq.com>
Date: Tue, 12 Mar 2024 15:44:21 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=E8=8E=B7=E5=8F=96=20?=
 =?UTF-8?q?remoteIP=20=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/httpd/router/chat.go | 3 ++-
 pkg/jms/session.go       | 8 ++------
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/pkg/httpd/router/chat.go b/pkg/httpd/router/chat.go
index 72eb69d..e043143 100644
--- a/pkg/httpd/router/chat.go
+++ b/pkg/httpd/router/chat.go
@@ -19,6 +19,7 @@ var ChatApi = new(_ChatApi)
 type _ChatApi struct{}
 
 func (s *_ChatApi) ChatHandler(ctx *gin.Context) {
+	remoteIP := ctx.ClientIP()
 	conn, err := ws.UpgradeWsConn(ctx)
 	if err != nil {
 		ctx.JSON(http.StatusInternalServerError, gin.H{"error": "websocket upgrade failed"})
@@ -35,7 +36,7 @@ func (s *_ChatApi) ChatHandler(ctx *gin.Context) {
 
 	currentJMSS := make([]*jms.JMSSession, 0)
 	tokenHandler := jms.NewTokenHandler()
-	sessionHandler := jms.NewSessionHandler(conn)
+	sessionHandler := jms.NewSessionHandler(conn, remoteIP)
 	authInfo, err := tokenHandler.GetTokenAuthInfo(token)
 	if err != nil {
 		ctx.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
diff --git a/pkg/jms/session.go b/pkg/jms/session.go
index e1d0c8e..340a3c4 100644
--- a/pkg/jms/session.go
+++ b/pkg/jms/session.go
@@ -16,17 +16,13 @@ type SessionHandler struct {
 	RemoteAddress string
 }
 
-func NewSessionHandler(websocket *websocket.Conn) *SessionHandler {
+func NewSessionHandler(websocket *websocket.Conn, remoteIP string) *SessionHandler {
 	return &SessionHandler{
 		Websocket:     websocket,
-		RemoteAddress: getRemoteAddress(websocket),
+		RemoteAddress: remoteIP,
 	}
 }
 
-func getRemoteAddress(websocket *websocket.Conn) string {
-	return websocket.RemoteAddr().String()
-}
-
 func (sh *SessionHandler) CreateNewSession(authInfo *protobuf.TokenAuthInfo, prompt string) *JMSSession {
 	session := sh.createSession(authInfo)
 	return &JMSSession{