fix: copy data from big frame msg

🤦 I forgot to actually copy data from the frame on `len(payload) > bufferPoolingThreshold`.
Also add test to ensure that this works correctly.

Signed-off-by: Dmitriy Matrenichev <[email protected]>

Author: Dmitriy Matrenichev

go.mod

@@ -20,3 +20,5 @@ require (
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20241007155032-5fefd90f89a9 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
+
+retract v0.5.0

go.sum

@@ -19,8 +19,6 @@ golang.org/x/sys v0.26.0 h1:KHjCJyddX0LoSTb3J+vWpupP9p0oznkqVk/IfjymZbo=
 golang.org/x/sys v0.26.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/text v0.19.0 h1:kTxAhCbGbxhK0IwgSKiMO5awPoDQ0RpfiVYBfK860YM=
 golang.org/x/text v0.19.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 h1:e7S5W7MGGLaSu8j3YjdezkZ+m1/Nm0uRVRMEMGk26Xs=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142/go.mod h1:UqMtugtsSgubUsoxbuAoiCXvqvErP7Gf0so0mK9tHxU=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20241007155032-5fefd90f89a9 h1:QCqS/PdaHTSWGvupk2F/ehwHtGc0/GYkT+3GAcR1CCc=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20241007155032-5fefd90f89a9/go.mod h1:GX3210XPVPUjJbTUbvwI8f2IpZDMZuPJWDzDuebbviI=
 google.golang.org/grpc v1.67.1 h1:zWnc1Vrcno+lHZCOofnIMvycFcc0QRGIzm9dhnDX68E=

proxy/codec.go

@@ -57,6 +57,8 @@ func (c *rawCodec) Marshal(v any) (data mem.BufferSlice, err error) {
 	} else {
 		pool := mem.DefaultBufferPool()
 		buf := pool.Get(len(f.payload))
+		copy(*buf, f.payload)
+
 		data = append(data, mem.NewBuffer(buf, pool))
 	}
 

proxy/codec_test.go

@@ -1,36 +1,73 @@
 package proxy_test
 
 import (
+	"bytes"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/mem"
 
 	"github.com/siderolabs/grpc-proxy/proxy"
+	talos_testproto "github.com/siderolabs/grpc-proxy/testservice"
 )
 
 func TestCodec_ReadYourWrites(t *testing.T) {
-	framePtr := proxy.NewFrame(nil)
-	data := []byte{0xDE, 0xAD, 0xBE, 0xEF}
-	codec := proxy.Codec()
+	d := []byte{0xDE, 0xAD, 0xBE, 0xEF}
 
-	buffer := mem.Copy(data, mem.DefaultBufferPool())
-	defer buffer.Free()
+	for key, val := range map[string][]byte{
+		"short message": d,
+		"long message":  bytes.Repeat(d, 3072),
+	} {
+		t.Run(key, func(t *testing.T) {
+			framePtr := proxy.NewFrame(nil)
+			codec := proxy.Codec()
 
-	require.NoError(t, codec.Unmarshal(mem.BufferSlice{buffer}, framePtr), "unmarshalling must go ok")
-	out, err := codec.Marshal(framePtr)
-	require.NoError(t, err, "no marshal error")
-	require.Equal(t, data, out.Materialize(), "output and data must be the same")
+			buffer := mem.Copy(val, mem.DefaultBufferPool())
+			defer func() { buffer.Free() }()
 
-	out.Free()
-	buffer.Free()
-	buffer = mem.Copy([]byte{0x55}, mem.DefaultBufferPool())
+			require.NoError(t, codec.Unmarshal(mem.BufferSlice{buffer}, framePtr), "unmarshalling must go ok")
+			out, err := codec.Marshal(framePtr)
+			require.NoError(t, err, "no marshal error")
+			require.Equal(t, val, out.Materialize(), "output and data must be the same")
 
-	// reuse
-	require.NoError(t, codec.Unmarshal(mem.BufferSlice{buffer}, framePtr), "unmarshalling must go ok")
-	out, err = codec.Marshal(framePtr)
-	require.NoError(t, err, "no marshal error")
-	require.Equal(t, []byte{0x55}, out.Materialize(), "output and data must be the same")
+			out.Free()
+			buffer.Free()
+			buffer = mem.Copy([]byte{0x55}, mem.DefaultBufferPool())
 
-	out.Free()
+			// reuse
+			require.NoError(t, codec.Unmarshal(mem.BufferSlice{buffer}, framePtr), "unmarshalling must go ok")
+			out, err = codec.Marshal(framePtr)
+			require.NoError(t, err, "no marshal error")
+			require.Equal(t, []byte{0x55}, out.Materialize(), "output and data must be the same")
+
+			out.Free()
+		})
+	}
+}
+
+func TestCodecUsualMessage(t *testing.T) {
+	const msg = "short message"
+
+	for key, val := range map[string]string{
+		"short message": "edbca",
+		"long message":  strings.Repeat(msg, 3072),
+	} {
+		t.Run(key, func(t *testing.T) {
+			msg := &talos_testproto.PingRequest{Value: val}
+
+			codec := proxy.Codec()
+
+			out, err := codec.Marshal(msg)
+			require.NoError(t, err, "no marshal error")
+
+			defer out.Free()
+
+			var dst talos_testproto.PingRequest
+
+			require.NoError(t, codec.Unmarshal(out, &dst), "unmarshalling must go ok")
+			require.NotZero(t, dst.Value, "output must not be zero")
+			require.Equal(t, msg.Value, dst.Value, "output and data must be the same")
+		})
+	}
 }