From b5536dfc7767dc47c9d4bd699404301cf8afaf77 Mon Sep 17 00:00:00 2001
From: Dane Everitt <dane@daneeveritt.com>
Date: Sat, 22 Jan 2022 14:33:03 -0500
Subject: [PATCH] Prevent excessive memory usage when large lines are sent over
 the console

---
 system/utils.go | 58 ++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 43 insertions(+), 15 deletions(-)

diff --git a/system/utils.go b/system/utils.go
index 8bf7781..58aa09d 100644
--- a/system/utils.go
+++ b/system/utils.go
@@ -19,6 +19,11 @@ var (
 	crr = []byte("\r\n")
 )
 
+// The maximum size of the buffer used to send output over the console to
+// clients. Once this length is reached, the line will be truncated and sent
+// as is.
+const maxBufferSize = 64 * 1024
+
 // FirstNotEmpty returns the first string passed in that is not an empty value.
 func FirstNotEmpty(v ...string) string {
 	for _, val := range v {
@@ -37,12 +42,15 @@ func MustInt(v string) int {
 	return i
 }
 
+// ScanReader reads up to 64KB of line from the reader and emits that value
+// over the websocket. If a line exceeds that size, it is truncated and only that
+// amount is sent over.
 func ScanReader(r io.Reader, callback func(line []byte)) error {
 	br := bufio.NewReader(r)
 	// Avoid constantly re-allocating memory when we're flooding lines through this
 	// function by using the same buffer for the duration of the call and just truncating
 	// the value back to 0 every loop.
-	buf := &bytes.Buffer{}
+	var buf bytes.Buffer
 	for {
 		buf.Reset()
 		var err error
@@ -52,32 +60,52 @@ func ScanReader(r io.Reader, callback func(line []byte)) error {
 		for {
 			// Read the line and write it to the buffer.
 			line, isPrefix, err = br.ReadLine()
+
 			// Certain games like Minecraft output absolutely random carriage returns in the output seemingly
 			// in line with that it thinks is the terminal size. Those returns break a lot of output handling,
 			// so we'll just replace them with proper new-lines and then split it later and send each line as
 			// its own event in the response.
-			buf.Write(bytes.Replace(line, cr, crr, -1))
-			// Finish this loop and begin outputting the line if there is no prefix (the line fit into
-			// the default buffer), or if we hit the end of the line.
+			line = bytes.Replace(line, cr, crr, -1)
+			ns := buf.Len() + len(line)
+
+			// If the length of the line value and the current value in the buffer will
+			// exceed the maximum buffer size, chop it down to hit the maximum size and
+			// then send that data over the socket before ending this loop.
+			//
+			// This ensures that we send as much data as possible, without allowing very
+			// long lines to grow the buffer size excessively and potentially DOS the Wings
+			// instance. If the line is not too long, just store the whole value into the
+			// buffer. This is kind of a re-implementation of the bufio.Scanner.Scan() logic
+			// without triggering an error when you exceed this buffer size.
+			if ns > maxBufferSize {
+				buf.Write(line[:len(line)-(ns-maxBufferSize)])
+				break
+			} else {
+				buf.Write(line)
+			}
+			// Finish this loop and begin outputting the line if there is no prefix
+			// (the line fit into the default buffer), or if we hit the end of the line.
 			if !isPrefix || err == io.EOF {
 				break
 			}
-			// If we encountered an error with something in ReadLine that was not an EOF just abort
-			// the entire process here.
+			// If we encountered an error with something in ReadLine that was not an
+			// EOF just abort the entire process here.
 			if err != nil {
 				return err
 			}
 		}
-		// Ensure that the scanner is always able to read the last line.
-		_, _ = buf.Write([]byte("\r\n"))
-		// Publish the line for this loop. Break on new-line characters so every line is sent as a single
-		// output event, otherwise you get funky handling in the browser console.
-		s := bufio.NewScanner(buf)
-		for s.Scan() {
-			callback(s.Bytes())
+
+		// Send the full buffer length over to the event handler to be emitted in
+		// the websocket. The front-end can handle the linebreaks in the middle of
+		// the output, it simply expects that the end of the event emit is a newline.
+		if buf.Len() > 0 {
+			c := make([]byte, buf.Len())
+			copy(c, buf.Bytes())
+			callback(c)
 		}
-		// If the error we got previously that lead to the line being output is an io.EOF we want to
-		// exit the entire looping process.
+
+		// If the error we got previously that lead to the line being output is
+		// an io.EOF we want to exit the entire looping process.
 		if err == io.EOF {
 			break
 		}