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

user io.ReadFull for shorter connection read code

Browse source
This commit is contained in:
fooofei 2020-10-04 10:37:36 +08:00
parent f484635321
commit c877497cc9
1 changed files with 13 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

41
src/comm/comm.go
View file

@ -4,6 +4,7 @@ import (
"bytes"
"encoding/binary"
"fmt"
"io"
"net"
"time"
@ -85,22 +86,15 @@ func (c *Comm) Read() (buf []byte, numBytes int, bs []byte, err error) {
if err := c.connection.SetReadDeadline(time.Now().Add(3 * time.Hour)); err != nil {
log.Warnf("error setting read deadline: %v", err)
}
// must clear the timeout setting
defer c.connection.SetDeadline(time.Time{})
// read until we get 4 bytes for the header
var header []byte
numBytes = 4
for {
tmp := make([]byte, numBytes-len(header))
n, errRead := c.connection.Read(tmp)
if errRead != nil {
err = errRead
log.Debugf("initial read error: %v", err)
return
}
header = append(header, tmp[:n]...)
if numBytes == len(header) {
break
}
header := make([]byte, 4)
_, err = io.ReadFull(c.connection, header)
if err != nil {
log.Debugf("initial read error: %v", err)
return
}
var numBytesUint32 uint32
@ -117,25 +111,16 @@ func (c *Comm) Read() (buf []byte, numBytes int, bs []byte, err error) {
log.Debug(err)
return
}
buf = make([]byte, 0)
// shorten the reading deadline in case getting weird data
if err := c.connection.SetReadDeadline(time.Now().Add(10 * time.Second)); err != nil {
log.Warnf("error setting read deadline: %v", err)
}
for {
// log.Debugf("bytes: %d/%d", len(buf), numBytes)
tmp := make([]byte, numBytes-len(buf))
n, errRead := c.connection.Read(tmp)
if errRead != nil {
err = errRead
log.Debugf("consecutive read error: %v", err)
return
}
buf = append(buf, tmp[:n]...)
if numBytes == len(buf) {
break
}
buf = make([]byte, numBytes)
_, err = io.ReadFull(c.connection, buf)
if err != nil {
log.Debugf("consecutive read error: %v", err)
return
}
return
}