From 9e17cdfc7647935103a536521b9ad4be77237c5c Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Fri, 26 Aug 2022 14:18:07 +0300
Subject: [PATCH] [#413] util/proto: Marshal repeated []byte fields correctly

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
---
 util/proto/marshal.go      | 16 +++++++++++-----
 util/proto/marshal_test.go |  8 ++++++++
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/util/proto/marshal.go b/util/proto/marshal.go
index 589c0be2..a82478b1 100644
--- a/util/proto/marshal.go
+++ b/util/proto/marshal.go
@@ -24,7 +24,10 @@ func BytesMarshal(field int, buf, v []byte) int {
 	if len(v) == 0 {
 		return 0
 	}
+	return bytesMarshal(field, buf, v)
+}
 
+func bytesMarshal(field int, buf, v []byte) int {
 	prefix := field<<3 | 0x2
 
 	// buf length check can prevent panic at PutUvarint, but it will make
@@ -41,10 +44,13 @@ func BytesSize(field int, v []byte) int {
 	if ln == 0 {
 		return 0
 	}
+	return bytesSize(field, v)
+}
 
+func bytesSize(field int, v []byte) int {
 	prefix := field<<3 | 0x2
 
-	return VarUIntSize(uint64(prefix)) + VarUIntSize(uint64(ln)) + ln
+	return VarUIntSize(uint64(prefix)) + VarUIntSize(uint64(len(v))) + len(v)
 }
 
 func StringMarshal(field int, buf []byte, v string) int {
@@ -141,7 +147,7 @@ func RepeatedBytesMarshal(field int, buf []byte, v [][]byte) int {
 	var offset int
 
 	for i := range v {
-		offset += BytesMarshal(field, buf[offset:], v[i])
+		offset += bytesMarshal(field, buf[offset:], v[i])
 	}
 
 	return offset
@@ -149,7 +155,7 @@ func RepeatedBytesMarshal(field int, buf []byte, v [][]byte) int {
 
 func RepeatedBytesSize(field int, v [][]byte) (size int) {
 	for i := range v {
-		size += BytesSize(field, v[i])
+		size += bytesSize(field, v[i])
 	}
 
 	return size
@@ -159,7 +165,7 @@ func RepeatedStringMarshal(field int, buf []byte, v []string) int {
 	var offset int
 
 	for i := range v {
-		offset += StringMarshal(field, buf[offset:], v[i])
+		offset += bytesMarshal(field, buf[offset:], []byte(v[i]))
 	}
 
 	return offset
@@ -167,7 +173,7 @@ func RepeatedStringMarshal(field int, buf []byte, v []string) int {
 
 func RepeatedStringSize(field int, v []string) (size int) {
 	for i := range v {
-		size += StringSize(field, v[i])
+		size += bytesSize(field, []byte(v[i]))
 	}
 
 	return size
diff --git a/util/proto/marshal_test.go b/util/proto/marshal_test.go
index e19f4ee4..cda1440d 100644
--- a/util/proto/marshal_test.go
+++ b/util/proto/marshal_test.go
@@ -313,6 +313,10 @@ func TestRepeatedBytesMarshal(t *testing.T) {
 		testRepeatedBytesMarshal(t, [][]byte{}, false)
 	})
 
+	t.Run("empty element", func(t *testing.T) {
+		testRepeatedBytesMarshal(t, [][]byte{{1}, {}}, false)
+	})
+
 	t.Run("nil", func(t *testing.T) {
 		testRepeatedBytesMarshal(t, nil, false)
 	})
@@ -325,6 +329,10 @@ func TestRepeatedStringMarshal(t *testing.T) {
 		testRepeatedStringMarshal(t, data, true)
 	})
 
+	t.Run("empty element", func(t *testing.T) {
+		testRepeatedStringMarshal(t, []string{""}, false)
+	})
+
 	t.Run("empty", func(t *testing.T) {
 		testRepeatedStringMarshal(t, []string{}, false)
 	})