From 311b76b75b1aeec29708f94c432ad81e30c9cf9c Mon Sep 17 00:00:00 2001
From: Alex Vanin <alexey@nspcc.ru>
Date: Fri, 14 Aug 2020 15:29:48 +0300
Subject: [PATCH] Add proto marshal helper for string

Signed-off-by: Alex Vanin <alexey@nspcc.ru>
---
 util/proto/marshal.go      |   8 ++++
 util/proto/marshal_test.go |  88 ++++++++++++++++++++++++++-----------
 util/proto/test/test.pb.go | Bin 7825 -> 9049 bytes
 util/proto/test/test.proto |   1 +
 4 files changed, 72 insertions(+), 25 deletions(-)

diff --git a/util/proto/marshal.go b/util/proto/marshal.go
index cd9c493..2d5f9e6 100644
--- a/util/proto/marshal.go
+++ b/util/proto/marshal.go
@@ -38,6 +38,14 @@ func BytesSize(field int, v []byte) int {
 	return VarUIntSize(uint64(prefix)) + VarUIntSize(uint64(ln)) + ln
 }
 
+func StringMarshal(field int, buf []byte, v string) (int, error) {
+	return BytesMarshal(field, buf, []byte(v))
+}
+
+func StringSize(field int, v string) int {
+	return BytesSize(field, []byte(v))
+}
+
 // varUIntSize returns length of varint byte sequence for uint64 value 'x'.
 func VarUIntSize(x uint64) int {
 	return (bits.Len64(x|1) + 6) / 7
diff --git a/util/proto/marshal_test.go b/util/proto/marshal_test.go
index ccc1cca..66a8b62 100644
--- a/util/proto/marshal_test.go
+++ b/util/proto/marshal_test.go
@@ -11,9 +11,10 @@ import (
 
 type stablePrimitives struct {
 	FieldA []byte
+	FieldB string
 }
 
-func (s *stablePrimitives) stableMarshal(buf []byte) ([]byte, error) {
+func (s *stablePrimitives) stableMarshal(buf []byte, wrongField bool) ([]byte, error) {
 	if s == nil {
 		return []byte{}, nil
 	}
@@ -23,32 +24,24 @@ func (s *stablePrimitives) stableMarshal(buf []byte) ([]byte, error) {
 	}
 
 	var (
-		i, offset int
+		i, offset, fieldNum int
 	)
 
-	offset, err := proto.BytesMarshal(1, buf, s.FieldA)
+	fieldNum = 1
+	if wrongField {
+		fieldNum++
+	}
+	offset, err := proto.BytesMarshal(fieldNum, buf, s.FieldA)
 	if err != nil {
 		return nil, errors.Wrap(err, "can't marshal field a")
 	}
 	i += offset
 
-	return buf, nil
-}
-
-func (s *stablePrimitives) stableMarshalWrongFieldNum(buf []byte) ([]byte, error) {
-	if s == nil {
-		return []byte{}, nil
+	fieldNum = 2
+	if wrongField {
+		fieldNum++
 	}
-
-	if buf == nil {
-		buf = make([]byte, s.stableSize())
-	}
-
-	var (
-		i, offset int
-	)
-
-	offset, err := proto.BytesMarshal(1+1, buf, s.FieldA)
+	offset, err = proto.StringMarshal(fieldNum, buf, s.FieldB)
 	if err != nil {
 		return nil, errors.Wrap(err, "can't marshal field a")
 	}
@@ -58,7 +51,8 @@ func (s *stablePrimitives) stableMarshalWrongFieldNum(buf []byte) ([]byte, error
 }
 
 func (s *stablePrimitives) stableSize() int {
-	return proto.BytesSize(1, s.FieldA)
+	return proto.BytesSize(1, s.FieldA) +
+		proto.StringSize(2, s.FieldB)
 }
 
 func TestBytesMarshal(t *testing.T) {
@@ -86,11 +80,7 @@ func testBytesMarshal(t *testing.T, data []byte, wrongField bool) {
 		transport = test.Primitives{FieldA: data}
 	)
 
-	if !wrongField {
-		wire, err = custom.stableMarshal(nil)
-	} else {
-		wire, err = custom.stableMarshalWrongFieldNum(nil)
-	}
+	wire, err = custom.stableMarshal(nil, wrongField)
 	require.NoError(t, err)
 
 	wireGen, err := transport.Marshal()
@@ -116,3 +106,51 @@ func testBytesMarshal(t *testing.T, data []byte, wrongField bool) {
 		require.Len(t, result.FieldA, 0)
 	}
 }
+
+func TestStringMarshal(t *testing.T) {
+	t.Run("not empty", func(t *testing.T) {
+		data := "Hello World"
+		testStringMarshal(t, data, false)
+		testStringMarshal(t, data, true)
+	})
+
+	t.Run("empty", func(t *testing.T) {
+		testStringMarshal(t, "", false)
+	})
+}
+
+func testStringMarshal(t *testing.T, s string, wrongField bool) {
+	var (
+		wire []byte
+		err  error
+
+		custom    = stablePrimitives{FieldB: s}
+		transport = test.Primitives{FieldB: s}
+	)
+
+	wire, err = custom.stableMarshal(nil, wrongField)
+	require.NoError(t, err)
+
+	wireGen, err := transport.Marshal()
+	require.NoError(t, err)
+
+	if !wrongField {
+		// we can check equality because single field cannot be unstable marshalled
+		require.Equal(t, wireGen, wire)
+	} else {
+		require.NotEqual(t, wireGen, wire)
+	}
+
+	result := new(test.Primitives)
+	err = result.Unmarshal(wire)
+	require.NoError(t, err)
+
+	if !wrongField {
+		require.Len(t, result.FieldB, len(s))
+		if len(s) > 0 {
+			require.Equal(t, s, result.FieldB)
+		}
+	} else {
+		require.Len(t, result.FieldB, 0)
+	}
+}
diff --git a/util/proto/test/test.pb.go b/util/proto/test/test.pb.go
index 5db45e5e93083bb3f11eb3bfa1ed9caf1e76b93d..9a1afd3504b96959d5752ffcf80b9a17af8575d4 100644
GIT binary patch
delta 755
zcmY*X&ui2`6lQm|-B~SdwXNVD90P@oWMz^~lijT=G+>LxdeMuNRr6yak!=!^S<NC+
zLBX@kU*K5~Atw+12kPCEH;>|<VCK!LHHVk?zW06K%gfB0(YM_Xn~j6rh)x0r|NrDH
zi{eKW_At#7mUz=}sCY9*bIs6_lxeYhL`R`eKJc{1c@oRQ;kCwk%_}c7-ksvxXIIYc
zQ6~L6Dwdx7^`f%uw9xYw$|##=F;*1*aq&mp(95^Y)3v5$YM2zh&EKQX$A%#a1E9gl
zyf1Jaa2@jAmOy%at%XS8>d*#(NkuLrkrp;|0rv+W1O65uHrUCL$39R)fG82vfesUN
zUlg<_VjprzwP6AQ!~|~(atp+Bejro?X3isY3B)jj0o4nWcA>k1EBCQrMN>2i+b{Re
z9Mo+`#n&;x&sqpYxI4lV8h4Ja%73m_pr53(PO!sv4x`~%!+%=*>yp`aR8>%-Znuh(
zGsdIJTsO*BOQgOssj8EP<%{J{3k$xRQ*7}2=kHx{6b(O(GJ1EGQeh=d6`?0?Hb3lJ
zG{(~?W|qmnU)<-j)eU}a%|r4gR89VK^%DQmR=Y8~12;|xh9g`_I)6!E9EAduB>rn{
k5B;mQ*I!<CsxOs-`Lefz8^th9a&|KgXfamH&5cR(H;xY0lK=n!

delta 267
zcmccVHqmy2D)Z)P%yU^L_jBB3GO*m7$H~gbk!Yf$U{H~qIJuhZOuca;h+|+0A}y0a
zq`3u%G=eBHM~I{vg4G!6aB&(`7=o3i8iRx^z@!O8sS#KNENcuFH%tYIBqxDLW3Ukr
zqm#hwKnpNsl7f&nFaz6Q0Cs?Z0Z0Nurlo<{5Hi&i#7+d;X$Cgi5G-qEyg7l#mua&b
Rzbw<{<$|4Tn^~npH~?+rMOFX+

diff --git a/util/proto/test/test.proto b/util/proto/test/test.proto
index cf51a44..81a1058 100644
--- a/util/proto/test/test.proto
+++ b/util/proto/test/test.proto
@@ -4,4 +4,5 @@ package test;
 
 message Primitives {
     bytes field_a = 1;
+    string field_b = 2;
 }
\ No newline at end of file