From 1fc6dd7c1abfeed020a6ced2cc28d10bea302f3a Mon Sep 17 00:00:00 2001
From: Evan Simkowitz <esimkowitz@users.noreply.github.com>
Date: Mon, 21 Oct 2024 14:05:52 -0700
Subject: [PATCH] Fix code scanning alert no. 50: Size computation for
 allocation may overflow (#1088)

Fixes
[https://github.com/wavetermdev/waveterm/security/code-scanning/50](https://github.com/wavetermdev/waveterm/security/code-scanning/50)

To fix the problem, we need to ensure that the size computation for the
allocation does not overflow. This can be achieved by validating the
length of `barr` before performing the arithmetic operation. We will set
a maximum allowable size for `barr` to ensure that the sum of
`oscPrefixLen(oscNum)` and `len(barr)` does not exceed the maximum value
for an `int`.

1. Define a maximum allowable size for `barr` (e.g., 64 MB).
2. Check the length of `barr` against this maximum size before
performing the allocation.
3. If `barr` exceeds the maximum size, return an error.


_Suggested fixes powered by Copilot Autofix. Review carefully before
merging._

---------

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
---
 pkg/blockcontroller/blockcontroller.go |  5 ++++-
 pkg/wshutil/wshrpcio.go                |  8 +++++---
 pkg/wshutil/wshutil.go                 | 20 ++++++++++++++------
 3 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/pkg/blockcontroller/blockcontroller.go b/pkg/blockcontroller/blockcontroller.go
index 4e9bc4886..146df6b38 100644
--- a/pkg/blockcontroller/blockcontroller.go
+++ b/pkg/blockcontroller/blockcontroller.go
@@ -383,7 +383,10 @@ func (bc *BlockController) DoRunShellCommand(rc *RunShellOpts, blockMeta waveobj
 	go func() {
 		// handles outputCh -> shellInputCh
 		for msg := range wshProxy.ToRemoteCh {
-			encodedMsg := wshutil.EncodeWaveOSCBytes(wshutil.WaveServerOSC, msg)
+			encodedMsg, err := wshutil.EncodeWaveOSCBytes(wshutil.WaveServerOSC, msg)
+			if err != nil {
+				log.Printf("error encoding OSC message: %v\n", err)
+			}
 			shellInputCh <- &BlockInputUnion{InputData: encodedMsg}
 		}
 	}()
diff --git a/pkg/wshutil/wshrpcio.go b/pkg/wshutil/wshrpcio.go
index 00d564d71..5d9dd27ba 100644
--- a/pkg/wshutil/wshrpcio.go
+++ b/pkg/wshutil/wshrpcio.go
@@ -79,10 +79,12 @@ func AdaptMsgChToPty(outputCh chan []byte, oscEsc string, output io.Writer) erro
 		panic("oscEsc must be 5 characters")
 	}
 	for msg := range outputCh {
-		barr := EncodeWaveOSCBytes(oscEsc, msg)
-		_, err := output.Write(barr)
+		barr, err := EncodeWaveOSCBytes(oscEsc, msg)
 		if err != nil {
-			return fmt.Errorf("error writing to output: %w", err)
+			return fmt.Errorf("error encoding osc message (AdaptMsgChToPty): %w", err)
+		}
+		if _, err := output.Write(barr); err != nil {
+			return fmt.Errorf("error writing osc message (AdaptMsgChToPty): %w", err)
 		}
 	}
 	return nil
diff --git a/pkg/wshutil/wshutil.go b/pkg/wshutil/wshutil.go
index 926fa8353..79cdc6080 100644
--- a/pkg/wshutil/wshutil.go
+++ b/pkg/wshutil/wshutil.go
@@ -67,9 +67,13 @@ func makeOscPrefix(oscNum string) []byte {
 	return output
 }
 
-func EncodeWaveOSCBytes(oscNum string, barr []byte) []byte {
+func EncodeWaveOSCBytes(oscNum string, barr []byte) ([]byte, error) {
 	if len(oscNum) != 5 {
-		panic("oscNum must be 5 characters")
+		return nil, fmt.Errorf("oscNum must be 5 characters")
+	}
+	const maxSize = 64 * 1024 * 1024 // 64 MB
+	if len(barr) > maxSize {
+		return nil, fmt.Errorf("input data too large")
 	}
 	hasControlChars := false
 	for _, b := range barr {
@@ -85,7 +89,7 @@ func EncodeWaveOSCBytes(oscNum string, barr []byte) []byte {
 		copyOscPrefix(output, oscNum)
 		copy(output[oscPrefixLen(oscNum):], barr)
 		output[len(output)-1] = BEL
-		return output
+		return output, nil
 	}
 
 	var buf bytes.Buffer
@@ -101,7 +105,7 @@ func EncodeWaveOSCBytes(oscNum string, barr []byte) []byte {
 		}
 	}
 	buf.WriteByte(BEL)
-	return buf.Bytes()
+	return buf.Bytes(), nil
 }
 
 func EncodeWaveOSCMessageEx(oscNum string, msg *RpcMessage) ([]byte, error) {
@@ -112,7 +116,7 @@ func EncodeWaveOSCMessageEx(oscNum string, msg *RpcMessage) ([]byte, error) {
 	if err != nil {
 		return nil, fmt.Errorf("error marshalling message to json: %w", err)
 	}
-	return EncodeWaveOSCBytes(oscNum, barr), nil
+	return EncodeWaveOSCBytes(oscNum, barr)
 }
 
 var termModeLock = sync.Mutex{}
@@ -194,7 +198,11 @@ func SetupTerminalRpcClient(serverImpl ServerImpl) (*WshRpc, io.Reader) {
 	rpcClient := MakeWshRpc(messageCh, outputCh, wshrpc.RpcContext{}, serverImpl)
 	go func() {
 		for msg := range outputCh {
-			barr := EncodeWaveOSCBytes(WaveOSC, msg)
+			barr, err := EncodeWaveOSCBytes(WaveOSC, msg)
+			if err != nil {
+				fmt.Fprintf(os.Stderr, "Error encoding OSC message: %v\n", err)
+				continue
+			}
 			os.Stdout.Write(barr)
 		}
 	}()