From 0c002483a640ee6f30c72de801d8d6731d6e9edb Mon Sep 17 00:00:00 2001
From: feng <1304903146@qq.com>
Date: Tue, 20 Feb 2024 17:28:55 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E6=B7=BB=E5=8A=A0session=20=E7=94=9F?=
 =?UTF-8?q?=E5=91=BD=E5=91=A8=E6=9C=9F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 Dockerfile             |  2 +-
 go.mod                 |  2 +-
 go.sum                 |  4 ++--
 pkg/jms/jms_session.go | 11 +++++++++++
 pkg/jms/session.go     | 19 +++++++++++++++++++
 5 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index b9387d9..28275fb 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -81,7 +81,7 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=kael \
 
 WORKDIR /opt
 
-ARG WISP_VERSION=v0.1.17
+ARG WISP_VERSION=v0.1.19
 RUN set -ex \
     && wget https://github.com/jumpserver/wisp/releases/download/${WISP_VERSION}/wisp-${WISP_VERSION}-linux-${TARGETARCH}.tar.gz \
     && tar -xf wisp-${WISP_VERSION}-linux-${TARGETARCH}.tar.gz -C /usr/local/bin/ --strip-components=1 \
diff --git a/go.mod b/go.mod
index 6d11c38..94decde 100644
--- a/go.mod
+++ b/go.mod
@@ -7,7 +7,7 @@ require (
 	github.com/gin-gonic/gin v1.9.1
 	github.com/google/uuid v1.4.0
 	github.com/gorilla/websocket v1.5.1
-	github.com/jumpserver/wisp v0.1.17
+	github.com/jumpserver/wisp v0.1.19
 	github.com/lestrrat-go/file-rotatelogs v2.4.0+incompatible
 	github.com/sashabaranov/go-openai v1.17.9
 	github.com/spf13/viper v1.17.0
diff --git a/go.sum b/go.sum
index 67cd74c..a5902ff 100644
--- a/go.sum
+++ b/go.sum
@@ -49,8 +49,8 @@ github.com/jonboulle/clockwork v0.4.0 h1:p4Cf1aMWXnXAUh8lVfewRBx1zaTSYKrKMF2g3ST
 github.com/jonboulle/clockwork v0.4.0/go.mod h1:xgRqUGwRcjKCO1vbZUEtSLrqKoPSsUpK7fnezOII0kc=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
-github.com/jumpserver/wisp v0.1.17 h1:i41iDwyz3MpCIIMEaNI43pKP5+lpzfqUSrVo3OnOq50=
-github.com/jumpserver/wisp v0.1.17/go.mod h1:6EHU3fk0PJoDE+Qmj89YJrDb1K8HPcj/sowltmYbKVg=
+github.com/jumpserver/wisp v0.1.19 h1:EEQLjgJApA6Le8CIAZF0ufuBQOnkDUybPAlYdz1dKxU=
+github.com/jumpserver/wisp v0.1.19/go.mod h1:6EHU3fk0PJoDE+Qmj89YJrDb1K8HPcj/sowltmYbKVg=
 github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
 github.com/klauspost/cpuid/v2 v2.2.5 h1:0E5MSMDEoAulmXNFquVs//DdoomxaoTY1kUhbc/qbZg=
 github.com/klauspost/cpuid/v2 v2.2.5/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws=
diff --git a/pkg/jms/jms_session.go b/pkg/jms/jms_session.go
index 3afb6a7..62c2050 100644
--- a/pkg/jms/jms_session.go
+++ b/pkg/jms/jms_session.go
@@ -36,6 +36,11 @@ func (jmss *JMSSession) ActiveSession() {
 	jmss.CommandHandler = NewCommandHandler(
 		jmss.Websocket, jmss.Session, jmss.CommandACLs, jmss.JMSState,
 	)
+	go jmss.SessionHandler.recordSessionLife(
+		jmss.Session,
+		protobuf.SessionLifecycleLogRequest_AssetConnectSuccess,
+		"",
+	)
 	go jmss.MaximumIdleTimeDetection()
 	go jmss.MaxSessionTimeDetection()
 }
@@ -79,11 +84,17 @@ func (jmss *JMSSession) MaxSessionTimeDetection() {
 
 func (jmss *JMSSession) Close(reason string) {
 	jmss.CurrentAskInterrupt = true
+	go jmss.SessionHandler.recordSessionLife(
+		jmss.Session,
+		protobuf.SessionLifecycleLogRequest_AssetConnectFinished,
+		"",
+	)
 	time.Sleep(1 * time.Second)
 	jmss.ReplayHandler.Upload()
 	jmss.SessionHandler.closeSession(jmss.Session)
 	GlobalSessionManager.UnregisterSession(jmss)
 	jmss.NotifyToClose(reason)
+
 }
 
 func (jmss *JMSSession) NotifyToClose(reason string) {
diff --git a/pkg/jms/session.go b/pkg/jms/session.go
index 7cec6f3..e1d0c8e 100644
--- a/pkg/jms/session.go
+++ b/pkg/jms/session.go
@@ -88,3 +88,22 @@ func (sh *SessionHandler) closeSession(session *protobuf.Session) {
 		logger.GlobalLogger.Error(errorMessage)
 	}
 }
+
+func (sh *SessionHandler) recordSessionLife(
+	session *protobuf.Session,
+	event protobuf.SessionLifecycleLogRequest_EventType,
+	reason string,
+) {
+	ctx := context.Background()
+	req := &protobuf.SessionLifecycleLogRequest{
+		SessionId: session.Id,
+		Event:     event,
+		Reason:    reason,
+	}
+
+	resp, _ := grpc.GlobalGrpcClient.Client.RecordSessionLifecycleLog(ctx, req)
+	if !resp.Status.Ok {
+		errorMessage := fmt.Sprintf("Failed to record session: %s", resp.Status.Err)
+		logger.GlobalLogger.Error(errorMessage)
+	}
+}