THRIFT-3805 Golang server susceptible to memory spike from malformed message
Client: Go
Patch: Michael Scott Leuthaeuser <scott.leuthaeuser@gmail.com>
diff --git a/lib/go/thrift/binary_protocol.go b/lib/go/thrift/binary_protocol.go
index e1b4056..690d341 100644
--- a/lib/go/thrift/binary_protocol.go
+++ b/lib/go/thrift/binary_protocol.go
@@ -20,6 +20,7 @@
package thrift
import (
+ "bytes"
"encoding/binary"
"errors"
"fmt"
@@ -473,6 +474,8 @@
return NewTProtocolException(err)
}
+const readLimit = 32768
+
func (p *TBinaryProtocol) readStringBody(size int32) (value string, err error) {
if size < 0 {
return "", nil
@@ -480,12 +483,32 @@
if uint64(size) > p.trans.RemainingBytes() {
return "", invalidDataLength
}
- var buf []byte
- if int(size) <= len(p.buffer) {
- buf = p.buffer[0:size]
- } else {
- buf = make([]byte, size)
+
+ var (
+ buf bytes.Buffer
+ e error
+ b []byte
+ )
+
+ switch {
+ case int(size) <= len(p.buffer):
+ b = p.buffer[:size] // avoids allocation for small reads
+ case int(size) < readLimit:
+ b = make([]byte, size)
+ default:
+ b = make([]byte, readLimit)
}
- _, e := io.ReadFull(p.trans, buf)
- return string(buf), NewTProtocolException(e)
+
+ for size > 0 {
+ _, e = io.ReadFull(p.trans, b)
+ buf.Write(b)
+ if e != nil {
+ break
+ }
+ size -= readLimit
+ if size < readLimit && size > 0 {
+ b = b[:size]
+ }
+ }
+ return buf.String(), NewTProtocolException(e)
}