schollz/croc
1
0
Fork 0
mirror of https://github.com/schollz/croc.git synced 2024-11-23 15:44:16 +03:00
Code Issues Projects Releases Wiki Activity

bug fix: prevent crazy number of bytes getting into comm

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
main.go
View File

@ -13,7 +13,21 @@ import (
func main() {
// "github.com/pkg/profile"
// defer profile.Start(profile.CPUProfile).Stop()
// go func() {
// for {
// f, err := os.Create("croc.pprof")
// if err != nil {
// panic(err)
// }
// runtime.GC() // get up-to-date statistics
// if err := pprof.WriteHeapProfile(f); err != nil {
// panic(err)
// }
// f.Close()
// time.Sleep(3 * time.Second)
// fmt.Println("wrote profile")
// }
// }()
if err := cli.Run(); err != nil {
fmt.Println(err)
}

13
src/comm/comm.go
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 {