From 7857c3ba61dfabf6f9cce24499b93cc851f14f6c Mon Sep 17 00:00:00 2001
From: Henry Barreto <henry.barreto@ossystems.com.br>
Date: Fri, 29 Nov 2024 11:18:12 -0300
Subject: [PATCH] refactor(ssh): graceful disconnect session record websocket

---
 ssh/server/channels/utils.go |  2 +-
 ssh/session/camera.go        | 37 ++++++++++++++++++++++++++++++++++++
 ssh/session/session.go       |  5 ++---
 3 files changed, 40 insertions(+), 4 deletions(-)
 create mode 100644 ssh/session/camera.go

diff --git a/ssh/server/channels/utils.go b/ssh/server/channels/utils.go
index a4cd0fec31d..d4fff0d4edd 100644
--- a/ssh/server/channels/utils.go
+++ b/ssh/server/channels/utils.go
@@ -89,7 +89,7 @@ func pipe(ctx gliderssh.Context, sess *session.Session, client gossh.Channel, ag
 				}
 
 				if recording {
-					if err := camera.WriteJSON(&models.SessionRecorded{ //nolint:errcheck
+					if err := camera.WriteFrame(&models.SessionRecorded{ //nolint:errcheck
 						UID:       sess.UID,
 						Namespace: sess.Lookup["domain"],
 						Message:   string(buffer[:read]),
diff --git a/ssh/session/camera.go b/ssh/session/camera.go
new file mode 100644
index 00000000000..fee1505d162
--- /dev/null
+++ b/ssh/session/camera.go
@@ -0,0 +1,37 @@
+package session
+
+import (
+	"time"
+
+	"github.com/gorilla/websocket"
+	"github.com/shellhub-io/shellhub/pkg/models"
+)
+
+type Camera struct {
+	conn *websocket.Conn
+}
+
+// Close closes the camera's WebSocket connections normally.
+func (c *Camera) Close() error {
+	if err := c.conn.WriteControl(
+		websocket.CloseMessage,
+		websocket.FormatCloseMessage(websocket.CloseNormalClosure, "session record connection done"),
+		time.Now().Add(time.Minute),
+	); err != nil {
+		return err
+	}
+
+	return c.conn.Close()
+}
+
+// WriteFrame writes a session's frame into the WebSocket connection.
+func (c *Camera) WriteFrame(frame *models.SessionRecorded) error {
+	return c.conn.WriteJSON(frame)
+}
+
+// NewCamera creates a new camera, using a WebSocket connections.
+func NewCamera(conn *websocket.Conn) *Camera {
+	return &Camera{
+		conn: conn,
+	}
+}
diff --git a/ssh/session/session.go b/ssh/session/session.go
index 510eeddc2aa..6cafef10467 100644
--- a/ssh/session/session.go
+++ b/ssh/session/session.go
@@ -11,7 +11,6 @@ import (
 	"time"
 
 	gliderssh "github.com/gliderlabs/ssh"
-	"github.com/gorilla/websocket"
 	"github.com/shellhub-io/shellhub/pkg/api/internalclient"
 	"github.com/shellhub-io/shellhub/pkg/api/requests"
 	"github.com/shellhub-io/shellhub/pkg/cache"
@@ -425,7 +424,7 @@ func (s *Session) Auth(ctx gliderssh.Context, auth Auth) error {
 	return nil
 }
 
-func (s *Session) Record(ctx context.Context, url string) (*websocket.Conn, error) {
+func (s *Session) Record(ctx context.Context, url string) (*Camera, error) {
 	conn, err := s.api.RecordSession(ctx, s.UID, url)
 	if err != nil {
 		log.WithError(err).Error("failed to start the record session process")
@@ -433,7 +432,7 @@ func (s *Session) Record(ctx context.Context, url string) (*websocket.Conn, erro
 		return nil, err
 	}
 
-	return conn, nil
+	return NewCamera(conn), nil
 }
 
 func Event[D any](sess *Session, t string, data []byte) {