Prevent excessive memory usage when large lines are sent over the console

This commit is contained in:
Dane Everitt 2022-01-22 14:33:03 -05:00
parent 45418c86dd
commit b5536dfc77

View File

@ -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
}