0
0
Fork 0
mirror of https://github.com/schollz/croc.git synced 2025-10-11 13:21:00 +02:00

bug fix: prevent crazy number of bytes getting into comm

This commit is contained in:
Zack Scholl 2020-03-27 09:32:41 -07:00
parent 44c3d43fa0
commit b60a841044
2 changed files with 26 additions and 3 deletions

View file

@ -8,8 +8,11 @@ import (
"time"
"github.com/pkg/errors"
"github.com/schollz/logger"
)
const MAXBYTES = 1000000
// Comm is some basic TCP communication
type Comm struct {
connection net.Conn
@ -89,12 +92,18 @@ func (c *Comm) Read() (buf []byte, numBytes int, bs []byte, err error) {
rbuf := bytes.NewReader(header)
err = binary.Read(rbuf, binary.LittleEndian, &numBytesUint32)
if err != nil {
fmt.Println("binary.Read failed:", err)
err = fmt.Errorf("binary.Read failed: %s", err.Error())
return
}
numBytes = int(numBytesUint32)
if numBytes > MAXBYTES {
err = fmt.Errorf("too many bytes: %d", numBytes)
logger.Error(err)
return
}
buf = make([]byte, 0)
for {
// log.Debugf("bytes: %d/%d",len(buf),numBytes)
// log.Debugf("bytes: %d/%d", len(buf), numBytes)
tmp := make([]byte, numBytes-len(buf))
n, errRead := c.connection.Read(tmp)
if errRead != nil {