From 0ff706259f20b81f1a96c4acffd9dde2a0435095 Mon Sep 17 00:00:00 2001
From: Scott Leggett <scott@sl.id.au>
Date: Fri, 23 Dec 2022 14:57:59 +0800
Subject: [PATCH] fix: handle terminal resizing correctly

---
 internal/k8s/exec.go                 | 11 +++++---
 internal/k8s/termsizequeue.go        | 40 ++++++++++++++++++++++++++++
 internal/sshserver/sessionhandler.go |  6 ++---
 3 files changed, 50 insertions(+), 7 deletions(-)
 create mode 100644 internal/k8s/termsizequeue.go

diff --git a/internal/k8s/exec.go b/internal/k8s/exec.go
index 2c9e6a28..56041975 100644
--- a/internal/k8s/exec.go
+++ b/internal/k8s/exec.go
@@ -8,6 +8,7 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/gliderlabs/ssh"
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -191,7 +192,7 @@ func (c *Client) getExecutor(ctx context.Context, namespace, deployment,
 // shell, running in a pod inside the deployment.
 func (c *Client) Exec(ctx context.Context, namespace, deployment,
 	container string, command []string, stdio io.ReadWriter, stderr io.Writer,
-	tty bool) error {
+	tty bool, winch <-chan ssh.Window) error {
 	exec, err := c.getExecutor(ctx, namespace, deployment, container, command,
 		stderr, tty)
 	if err != nil {
@@ -199,8 +200,10 @@ func (c *Client) Exec(ctx context.Context, namespace, deployment,
 	}
 	// execute the command
 	return exec.StreamWithContext(ctx, remotecommand.StreamOptions{
-		Stdin:  stdio,
-		Stdout: stdio,
-		Stderr: stderr,
+		Stdin:             stdio,
+		Stdout:            stdio,
+		Stderr:            stderr,
+		Tty:               tty,
+		TerminalSizeQueue: newTermSizeQueue(ctx, winch),
 	})
 }
diff --git a/internal/k8s/termsizequeue.go b/internal/k8s/termsizequeue.go
new file mode 100644
index 00000000..1ff619ac
--- /dev/null
+++ b/internal/k8s/termsizequeue.go
@@ -0,0 +1,40 @@
+package k8s
+
+import (
+	"context"
+
+	"github.com/gliderlabs/ssh"
+	"k8s.io/client-go/tools/remotecommand"
+)
+
+type termSizeQueue struct {
+	send chan *remotecommand.TerminalSize
+}
+
+// newTermSizeQueue returns a termSizeQueue which implements the
+// remotecommand.TerminalSizeQueue interface. It starts a goroutine which exits
+// when the given context is done.
+func newTermSizeQueue(ctx context.Context, winch <-chan ssh.Window) *termSizeQueue {
+	tsq := termSizeQueue{
+		send: make(chan *remotecommand.TerminalSize, 1),
+	}
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				close(tsq.send)
+				return
+			case window := <-winch:
+				tsq.send <- &remotecommand.TerminalSize{
+					Width:  uint16(window.Width),
+					Height: uint16(window.Height),
+				}
+			}
+		}
+	}()
+	return &tsq
+}
+
+func (t *termSizeQueue) Next() *remotecommand.TerminalSize {
+	return <-t.send
+}
diff --git a/internal/sshserver/sessionhandler.go b/internal/sshserver/sessionhandler.go
index a4921539..1b040459 100644
--- a/internal/sshserver/sessionhandler.go
+++ b/internal/sshserver/sessionhandler.go
@@ -99,8 +99,8 @@ func sessionHandler(log *zap.Logger, c *k8s.Client, sftp bool) ssh.Handler {
 			}
 			return
 		}
-		// check if a pty was requested
-		_, _, pty := s.Pty()
+		// check if a pty was requested, and get the window size channel
+		_, winch, pty := s.Pty()
 		// extract info passed through the context by the authhandler
 		ctx := s.Context()
 		eid, ok := ctx.Value(environmentIDKey).(int)
@@ -137,7 +137,7 @@ func sessionHandler(log *zap.Logger, c *k8s.Client, sftp bool) ssh.Handler {
 			zap.Strings("command", cmd),
 		)
 		err = c.Exec(s.Context(), s.User(), deployment, container, cmd, s,
-			s.Stderr(), pty)
+			s.Stderr(), pty, winch)
 		if err != nil {
 			log.Warn("couldn't execute command",
 				zap.String("sessionID", sid),