From 2dd0a291ed092009e9332d76fefbb72cd8342898 Mon Sep 17 00:00:00 2001
From: Pudong Zheng <tianshi8650@gmail.com>
Date: Tue, 28 Jun 2022 08:59:19 +0000
Subject: [PATCH] [ssh-gateway] add track for send closed heartbeat

---
 components/ws-proxy/pkg/sshproxy/forward.go | 24 +++++++++++++++++----
 components/ws-proxy/pkg/sshproxy/server.go  |  2 ++
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/components/ws-proxy/pkg/sshproxy/forward.go b/components/ws-proxy/pkg/sshproxy/forward.go
index 14568e25cae1d2..045d03c8951194 100644
--- a/components/ws-proxy/pkg/sshproxy/forward.go
+++ b/components/ws-proxy/pkg/sshproxy/forward.go
@@ -9,7 +9,9 @@ import (
 	"sync"
 	"time"
 
+	"github.com/gitpod-io/gitpod/common-go/analytics"
 	"github.com/gitpod-io/gitpod/common-go/log"
+	tracker "github.com/gitpod-io/gitpod/ws-proxy/pkg/analytics"
 	"github.com/gitpod-io/golang-crypto/ssh"
 	"golang.org/x/net/context"
 )
@@ -29,7 +31,7 @@ func (s *Server) ChannelForward(ctx context.Context, session *Session, targetCon
 		return
 	}
 	if originChannel.ChannelType() == "session" {
-		originChan = startHeartbeatingChannel(originChan, s.Heartbeater, session.InstanceID)
+		originChan = startHeartbeatingChannel(originChan, s.Heartbeater, session)
 	}
 	defer originChan.Close()
 
@@ -90,7 +92,19 @@ func (s *Server) ChannelForward(ctx context.Context, session *Session, targetCon
 	log.WithFields(log.OWI("", session.WorkspaceID, session.InstanceID)).Debug("session forward stop")
 }
 
-func startHeartbeatingChannel(c ssh.Channel, heartbeat Heartbeat, instanceID string) ssh.Channel {
+func TrackIDECloseSignal(session *Session) {
+	propertics := make(map[string]interface{})
+	propertics["workspaceId"] = session.WorkspaceID
+	propertics["instanceId"] = session.InstanceID
+	propertics["clientKind"] = "ssh"
+	tracker.Track(analytics.TrackMessage{
+		Identity:   analytics.Identity{UserID: session.OwnerUserId},
+		Event:      "ide_close_signal",
+		Properties: propertics,
+	})
+}
+
+func startHeartbeatingChannel(c ssh.Channel, heartbeat Heartbeat, session *Session) ssh.Channel {
 	ctx, cancel := context.WithCancel(context.Background())
 	res := &heartbeatingChannel{
 		Channel: c,
@@ -108,10 +122,12 @@ func startHeartbeatingChannel(c ssh.Channel, heartbeat Heartbeat, instanceID str
 				}
 				res.sawActivity = false
 				res.mux.Unlock()
-				heartbeat.SendHeartbeat(instanceID, false)
+				heartbeat.SendHeartbeat(session.InstanceID, false)
 			case <-ctx.Done():
 				if res.requestedPty {
-					heartbeat.SendHeartbeat(instanceID, true)
+					heartbeat.SendHeartbeat(session.InstanceID, true)
+					TrackIDECloseSignal(session)
+					log.WithField("instanceId", session.InstanceID).Info("send closed heartbeat")
 				}
 				return
 			}
diff --git a/components/ws-proxy/pkg/sshproxy/server.go b/components/ws-proxy/pkg/sshproxy/server.go
index c0da02e3b3e14c..c0daa49c43c48f 100644
--- a/components/ws-proxy/pkg/sshproxy/server.go
+++ b/components/ws-proxy/pkg/sshproxy/server.go
@@ -74,6 +74,7 @@ type Session struct {
 
 	WorkspaceID string
 	InstanceID  string
+	OwnerUserId string
 
 	PublicKey           ssh.PublicKey
 	WorkspacePrivateKey ssh.Signer
@@ -251,6 +252,7 @@ func (s *Server) HandleConn(c net.Conn) {
 		Conn:                clientConn,
 		WorkspaceID:         workspaceId,
 		InstanceID:          wsInfo.InstanceID,
+		OwnerUserId:         wsInfo.OwnerUserId,
 		WorkspacePrivateKey: key,
 	}
 	remoteAddr := wsInfo.IPAddress + ":23001"