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
 	}