From a218ef53493e45cce7a480a9adc45253c5f93167 Mon Sep 17 00:00:00 2001
From: Edmund Rhudy <erhudy@users.noreply.github.com>
Date: Fri, 23 Jun 2023 15:20:15 -0400
Subject: [PATCH] fix: adds WebSocket ping to interactive terminal (#14191)

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 5052e38d92c1c..f72688e2a5d3b 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{