perf: 添加session 生命周期

jumpserver · Feb 20, 2024 · 0c00248 · 0c00248
1 parent 31012e9
commit 0c00248
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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)
+	}
+}