From 98344cf34ec61f51064c768a04b0151b920f7336 Mon Sep 17 00:00:00 2001
From: Edmund Rhudy <erhudy@users.noreply.github.com>
Date: Fri, 7 Jul 2023 10:59:08 -0400
Subject: [PATCH] fix: adds WebSocket ping to interactive terminal (#14191)
 (#14192)

This adds a WebSocket ping message on a 5-second interval, sent
from the server to the client. This ensures that the interactive
terminal will remain open and won't be closed by load balancers
that are reaping idle connections.

Signed-off-by: Edmund Rhudy <erhudy@users.noreply.github.com>
---
 server/application/terminal.go  |  5 +++++
 server/application/websocket.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/server/application/terminal.go b/server/application/terminal.go
index 667ff529ae076..6424c89e97670 100644
--- a/server/application/terminal.go
+++ b/server/application/terminal.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"io"
 	"net/http"
+	"time"
 
 	"github.com/argoproj/gitops-engine/pkg/utils/kube"
 	log "github.com/sirupsen/logrus"
@@ -228,6 +229,10 @@ func (s *terminalHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 	defer session.Done()
 
+	// send pings across the WebSocket channel at regular intervals to keep it alive through
+	// load balancers which may close an idle connection after some period of time
+	go session.StartKeepalives(time.Second * 5)
+
 	if isValidShell(s.allowedShells, shell) {
 		cmd := []string{shell}
 		err = startProcess(kubeClientset, config, namespace, podName, container, cmd, session)
diff --git a/server/application/websocket.go b/server/application/websocket.go
index ff72aa28644da..fdac5a76c592b 100644
--- a/server/application/websocket.go
+++ b/server/application/websocket.go
@@ -51,6 +51,23 @@ func (t *terminalSession) Done() {
 	close(t.doneChan)
 }
 
+func (t *terminalSession) StartKeepalives(dur time.Duration) {
+	ticker := time.NewTicker(dur)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			err := t.Ping()
+			if err != nil {
+				log.Errorf("ping error: %v", err)
+				return
+			}
+		case <-t.doneChan:
+			return
+		}
+	}
+}
+
 // Next called in a loop from remotecommand as long as the process is running
 func (t *terminalSession) Next() *remotecommand.TerminalSize {
 	select {
@@ -86,6 +103,17 @@ func (t *terminalSession) Read(p []byte) (int, error) {
 	}
 }
 
+// Ping called periodically to ensure connection stays alive through load balancers
+func (t *terminalSession) Ping() error {
+	t.writeLock.Lock()
+	err := t.wsConn.WriteMessage(websocket.PingMessage, []byte("ping"))
+	t.writeLock.Unlock()
+	if err != nil {
+		log.Errorf("ping message err: %v", err)
+	}
+	return err
+}
+
 // Write called from remotecommand whenever there is any output
 func (t *terminalSession) Write(p []byte) (int, error) {
 	msg, err := json.Marshal(TerminalMessage{