From 0470381fe2ec092a15ee6449b3b235ebe4bc8f8f Mon Sep 17 00:00:00 2001 From: Jim Han <50871214+JimhHan@users.noreply.github.com> Date: Sun, 21 Mar 2021 17:16:52 +0800 Subject: [PATCH] Fix: gRPC multi accepting empty bytes (#411) --- common/buf/buffer.go | 2 +- transport/internet/grpc/encoding/hunkconn.go | 2 +- transport/internet/grpc/encoding/multiconn.go | 25 ++++++++++++------- 3 files changed, 18 insertions(+), 11 deletions(-) diff --git a/common/buf/buffer.go b/common/buf/buffer.go index ec0c32bd..9eef9997 100644 --- a/common/buf/buffer.go +++ b/common/buf/buffer.go @@ -38,7 +38,7 @@ func NewExisted(b []byte) *Buffer { oLen := len(b) if oLen < Size { - b = append(b, make([]byte, Size-oLen)...) + b = b[:Size] } return &Buffer{ diff --git a/transport/internet/grpc/encoding/hunkconn.go b/transport/internet/grpc/encoding/hunkconn.go index 6cfde6e1..a15a9577 100644 --- a/transport/internet/grpc/encoding/hunkconn.go +++ b/transport/internet/grpc/encoding/hunkconn.go @@ -100,7 +100,7 @@ func (h *HunkReaderWriter) ReadMultiBuffer() (buf.MultiBuffer, error) { } } - if cap(h.buf) == buf.Size { + if cap(h.buf) >= buf.Size { b := h.buf h.index = len(h.buf) return buf.MultiBuffer{buf.NewExisted(b)}, nil diff --git a/transport/internet/grpc/encoding/multiconn.go b/transport/internet/grpc/encoding/multiconn.go index d247b3eb..31e8bb4b 100644 --- a/transport/internet/grpc/encoding/multiconn.go +++ b/transport/internet/grpc/encoding/multiconn.go @@ -79,16 +79,20 @@ func (h *MultiHunkReaderWriter) ReadMultiBuffer() (buf.MultiBuffer, error) { var mb = make(buf.MultiBuffer, 0, len(h.buf)) for _, b := range h.buf { - if cap(b) >= buf.Size { - mb = append(mb, buf.NewExisted(b)) + if len(b) == 0 { continue } - nb := buf.New() - nb.Extend(int32(len(b))) - copy(nb.Bytes(), b) + if cap(b) >= buf.Size { + mb = append(mb, buf.NewExisted(b)) + } else { + nb := buf.New() + nb.Extend(int32(len(b))) + copy(nb.Bytes(), b) + + mb = append(mb, nb) + } - mb = append(mb, nb) } return mb, nil } @@ -99,12 +103,15 @@ func (h *MultiHunkReaderWriter) WriteMultiBuffer(mb buf.MultiBuffer) error { return io.ErrClosedPipe } - hunk := &MultiHunk{Data: make([][]byte, len(mb))} + hunks := make([][]byte, 0, len(mb)) + for _, b := range mb { - hunk.Data = append(hunk.Data, b.Bytes()) + if b.Len() > 0 { + hunks = append(hunks, b.Bytes()) + } } - err := h.hc.Send(hunk) + err := h.hc.Send(&MultiHunk{Data: hunks}) if err != nil { return err }