From b509fd71170643bc748408dd5534ae148bfd8c77 Mon Sep 17 00:00:00 2001
From: luke-lombardi <33990301+luke-lombardi@users.noreply.github.com>
Date: Tue, 31 Dec 2024 18:42:23 -0500
Subject: [PATCH] refactor buffer size

---
 pkg/abstractions/shell/http.go  | 24 ++++++++++++++++++++----
 pkg/abstractions/shell/shell.go |  1 +
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/pkg/abstractions/shell/http.go b/pkg/abstractions/shell/http.go
index 9b8bb8438..ca148ff04 100644
--- a/pkg/abstractions/shell/http.go
+++ b/pkg/abstractions/shell/http.go
@@ -81,21 +81,37 @@ func (g *shellGroup) ShellConnect(ctx echo.Context) error {
 	// TODO: confirm disconnects happen when python client exits
 	defer log.Println("disconnected")
 
+	// Create a context that will be canceled when the client disconnects
+	clientCtx, clientCancel := context.WithCancel(ctx.Request().Context())
+	defer clientCancel()
+
 	var wg sync.WaitGroup
 	wg.Add(2)
 
 	// Copy from client -> container
 	go func() {
 		defer wg.Done()
-		buf := make([]byte, 32*1024) // 32KB buffer
-		io.CopyBuffer(containerConn, conn, buf)
+		buf := make([]byte, shellProxyBufferSizeKb)
+		select {
+		case <-clientCtx.Done():
+			log.Println("client disconnected")
+			return
+		default:
+			io.CopyBuffer(containerConn, conn, buf)
+		}
 	}()
 
 	// Copy from container -> client
 	go func() {
 		defer wg.Done()
-		buf := make([]byte, 32*1024) // 32KB buffer
-		io.CopyBuffer(conn, containerConn, buf)
+		buf := make([]byte, shellProxyBufferSizeKb)
+		select {
+		case <-clientCtx.Done():
+			log.Println("client disconnected")
+			return
+		default:
+			io.CopyBuffer(conn, containerConn, buf)
+		}
 	}()
 
 	wg.Wait()
diff --git a/pkg/abstractions/shell/shell.go b/pkg/abstractions/shell/shell.go
index c63e29e01..391796dde 100644
--- a/pkg/abstractions/shell/shell.go
+++ b/pkg/abstractions/shell/shell.go
@@ -25,6 +25,7 @@ import (
 const (
 	shellRoutePrefix              string        = "/shell"
 	shellContainerPrefix          string        = "shell"
+	shellProxyBufferSizeKb        int           = 32 * 1024
 	defaultContainerCpu           int64         = 100
 	defaultContainerMemory        int64         = 128
 	containerDialTimeoutDurationS time.Duration = 300 * time.Second