Fix: use bufio.Reader on grpc to avoid panic

This commit is contained in:
Dreamacro 2021-04-14 00:16:59 +08:00
parent c9943fb857
commit a9cbd9ec98
1 changed files with 11 additions and 44 deletions

View File

@ -64,6 +64,7 @@ func (g *Conn) initRequest() {
if !g.close.Load() { if !g.close.Load() {
g.response = response g.response = response
g.br = bufio.NewReader(response.Body)
} else { } else {
response.Body.Close() response.Body.Close()
} }
@ -75,24 +76,13 @@ func (g *Conn) Read(b []byte) (n int, err error) {
return 0, g.err return 0, g.err
} }
if g.br != nil { if g.remain > 0 {
remain := g.br.Buffered()
if len(b) < remain {
remain = len(b)
}
n, err = g.br.Read(b[:remain])
if g.br.Buffered() == 0 {
g.br = nil
}
return
} else if g.remain > 0 {
size := g.remain size := g.remain
if len(b) < size { if len(b) < size {
size = len(b) size = len(b)
} }
n, err = g.response.Body.Read(b[:size]) n, err = io.ReadFull(g.br, b[:size])
g.remain -= n g.remain -= n
return return
} else if g.response == nil { } else if g.response == nil {
@ -100,55 +90,32 @@ func (g *Conn) Read(b []byte) (n int, err error) {
} }
// 0x00 grpclength(uint32) 0x0A uleb128 payload // 0x00 grpclength(uint32) 0x0A uleb128 payload
buf := make([]byte, 6) _, err = g.br.Discard(6)
_, err = io.ReadFull(g.response.Body, buf)
if err != nil { if err != nil {
return 0, err return 0, err
} }
br := bufio.NewReaderSize(g.response.Body, 16) protobufPayloadLen, err := binary.ReadUvarint(g.br)
protobufPayloadLen, err := binary.ReadUvarint(br)
if err != nil { if err != nil {
return 0, ErrInvalidLength return 0, ErrInvalidLength
} }
bufferedSize := br.Buffered() size := int(protobufPayloadLen)
if len(b) < bufferedSize { if len(b) < size {
n, err = br.Read(b) size = len(b)
g.br = br
remain := int(protobufPayloadLen) - n - g.br.Buffered()
if remain < 0 {
return 0, ErrInvalidLength
}
g.remain = remain
return
} }
_, err = br.Read(b[:bufferedSize]) n, err = io.ReadFull(g.br, b[:size])
if err != nil { if err != nil {
return return
} }
offset := int(protobufPayloadLen) remain := int(protobufPayloadLen) - n
if len(b) < int(protobufPayloadLen) {
offset = len(b)
}
if offset == bufferedSize {
return bufferedSize, nil
}
n, err = io.ReadFull(g.response.Body, b[bufferedSize:offset])
if err != nil {
return
}
remain := int(protobufPayloadLen) - offset
if remain > 0 { if remain > 0 {
g.remain = remain g.remain = remain
} }
return offset, nil return n, nil
} }
func (g *Conn) Write(b []byte) (n int, err error) { func (g *Conn) Write(b []byte) (n int, err error) {