Prevent excessive memory usage when large lines are sent over the console
This commit is contained in:
parent
45418c86dd
commit
b5536dfc77
|
@ -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
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue
Block a user