Kaynağa Gözat

Fix: gRPC multi accepting empty bytes (#411)

Jim Han 4 yıl önce
ebeveyn
işleme
0470381fe2

+ 1 - 1
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{

+ 1 - 1
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

+ 16 - 9
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
 	}