Skip to content

Commit

Permalink
perf: 记录 sftp 和 ssh 命令会话连接记录
Browse files Browse the repository at this point in the history
  • Loading branch information
@LeeEirc
LeeEirc committed Feb 6, 2024
1 parent 4823bd2 commit d68231d
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 3 deletions.
18 changes: 17 additions & 1 deletion pkg/handler/server_ssh.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -332,6 +332,13 @@ func IsScpCommand(rawStr string) bool {
return false
}

func (s *Server) recordSessionLifecycle(sid string, event model.LifecycleEvent, reason string) {
logObj := model.SessionLifecycleLog{Reason: reason}
if err2 := s.jmsService.RecordSessionLifecycleLog(sid, event, logObj); err2 != nil {
logger.Errorf("Record session %s lifecycle %s failed: %s", sid, event, err2)
}
}

func (s *Server) proxyAssetCommand(sess ssh.Session, sshClient *srvconn.SSHClient,
tokeInfo *model.ConnectToken) {
rawStr := sess.RawCommand()
Expand Down Expand Up @@ -371,6 +378,7 @@ func (s *Server) proxyAssetCommand(sess ssh.Session, sshClient *srvconn.SSHClien
}
ctx, cancel := context.WithCancel(sess.Context())
defer cancel()

traceSession := session.NewSession(&respSession, func(task *model.TerminalTask) error {
switch task.Name {
case model.TaskKillSession:
Expand All @@ -393,6 +401,7 @@ func (s *Server) proxyAssetCommand(sess ssh.Session, sshClient *srvconn.SSHClien
logger.Errorf("Get SSH session failed: %s", err)
return
}
s.recordSessionLifecycle(respSession.ID, model.AssetConnectSuccess, "")
defer goSess.Close()
defer sshClient.ReleaseSession(goSess)
go func() {
Expand Down Expand Up @@ -455,6 +464,8 @@ func (s *Server) proxyAssetCommand(sess ssh.Session, sshClient *srvconn.SSHClien
logger.Errorf("User %s Run command %s failed: %s",
tokeInfo.User.String(), rawStr, err)
}
reason := string(model.ReasonErrConnectDisconnect)
s.recordSessionLifecycle(respSession.ID, model.AssetConnectFinished, reason)
}

func (s *Server) proxyVscodeShell(sess ssh.Session, vsReq *vscodeReq, sshClient *srvconn.SSHClient,
Expand Down Expand Up @@ -490,7 +501,7 @@ func (s *Server) proxyVscodeShell(sess ssh.Session, vsReq *vscodeReq, sshClient
logger.Errorf("Get SSH session failed: %s", err)
return err
}

s.recordSessionLifecycle(respSession.ID, model.AssetConnectSuccess, "")
defer goSess.Close()
defer sshClient.ReleaseSession(goSess)
stdOut, err := goSess.StdoutPipe()
Expand All @@ -506,6 +517,7 @@ func (s *Server) proxyVscodeShell(sess ssh.Session, vsReq *vscodeReq, sshClient
err = goSess.Shell()
if err != nil {
logger.Errorf("Get SSH session shell failed: %s", err)
s.recordSessionLifecycle(respSession.ID, model.AssetConnectFinished, err.Error())
return err
}
logger.Infof("User %s start vscode request to %s", vsReq.user, sshClient)
Expand All @@ -525,11 +537,15 @@ func (s *Server) proxyVscodeShell(sess ssh.Session, vsReq *vscodeReq, sshClient
case <-ctx.Done():
logger.Infof("SSH conn[%s] User %s end vscode request %s as session done",
vsReq.reqId, vsReq.user, sshClient)
reason := string(model.ReasonErrConnectDisconnect)
s.recordSessionLifecycle(respSession.ID, model.AssetConnectFinished, reason)
return nil
case now := <-ticker.C:
if vsReq.expireInfo.IsExpired(now) {
logger.Infof("SSH conn[%s] User %s end vscode request %s as permission has expired",
vsReq.reqId, vsReq.user, sshClient)
reason := string(model.ReasonErrPermissionExpired)
s.recordSessionLifecycle(respSession.ID, model.AssetConnectFinished, reason)
return nil
}
logger.Debugf("SSH conn[%s] user %s vscode request still alive", vsReq.reqId, vsReq.user)
Expand Down
18 changes: 16 additions & 2 deletions pkg/srvconn/sftp_asset.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -535,6 +535,7 @@ func (ad *AssetDir) GetSFTPAndRealPath(su *model.PermAccount, path string) (conn
traceSession := session.NewSession(&respSession, terminalFunc)
session.AddSession(traceSession)
ad.sftpTraceSessions[su.String()] = traceSession
ad.recordSessionLifecycle(traceSession.ID, model.AssetConnectSuccess, "")
}
if conn.rootDirPath == "" {
platform := conn.token.Platform
Expand Down Expand Up @@ -581,7 +582,7 @@ func (ad *AssetDir) GetSftpClient(su *model.PermAccount) (conn *SftpConn, err er
if err != nil {
return nil, fmt.Errorf("get connect token account err: %s", err2)
}
return ad.getNewSftpConn(&connectToken)
return ad.getNewSftpConn(&connectToken, su)
}

func (ad *AssetDir) createConnectToken(su *model.PermAccount) (model.ConnectToken, error) {
Expand Down Expand Up @@ -609,7 +610,8 @@ func (ad *AssetDir) createConnectToken(su *model.PermAccount) (model.ConnectToke
return ad.jmsService.GetConnectTokenInfo(tokenInfo.ID)
}

func (ad *AssetDir) getNewSftpConn(connectToken *model.ConnectToken) (conn *SftpConn, err error) {
func (ad *AssetDir) getNewSftpConn(connectToken *model.ConnectToken,
su *model.PermAccount) (conn *SftpConn, err error) {
if ad.detailAsset == nil {
return nil, errNoSelectAsset
}
Expand Down Expand Up @@ -680,6 +682,11 @@ func (ad *AssetDir) getNewSftpConn(connectToken *model.ConnectToken) (conn *Sftp
sshClient.ReleaseSession(sess)
_ = sshClient.Close()
logger.Infof("User %s SSH client(%s) for SFTP release", user.String(), sshClient)
if sftpSession, ok := ad.sftpTraceSessions[su.String()]; ok {
sid := sftpSession.ID
reason := string(model.ReasonErrConnectDisconnect)
ad.recordSessionLifecycle(sid, model.AssetConnectFinished, reason)
}
}()
homeDirPath, err := sftpClient.Getwd()
if err != nil {
Expand Down Expand Up @@ -741,6 +748,13 @@ func (ad *AssetDir) CreateFTPLog(su *model.PermAccount, operate, filename string
return &data
}

func (ad *AssetDir) recordSessionLifecycle(sid string, event model.LifecycleEvent, reason string) {
logObj := model.SessionLifecycleLog{Reason: reason}
if err := ad.jmsService.RecordSessionLifecycleLog(sid, event, logObj); err != nil {
logger.Errorf("Update session %s lifecycle %s failed: %s", sid, event, err)
}
}

func IsExistPath(client *sftp.Client, path string) bool {
_, err := client.Stat(path)
return err == nil
Expand Down

0 comments on commit d68231d

Please sign in to comment.