Merge pull request #178 from jumpserver/dev

v3.10.8
jumpserver · Apr 18, 2024 · 3b52a49 · 3b52a49
2 parents 82180c8 + 8e92a4d
commit 3b52a49
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
diff --git 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
@@ -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{