network: proposal payload compression

go.mod

@@ -39,6 +39,7 @@ require (
 )
 
 require (
+	github.com/DataDog/zstd v1.5.2 // indirect
 	github.com/chrismcguire/gobberish v0.0.0-20150821175641-1d8adb509a0e // indirect
 	github.com/cpuguy83/go-md2man v1.0.8 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect

go.sum

@@ -1,3 +1,5 @@
+github.com/DataDog/zstd v1.5.2 h1:vUG4lAyuPCXO0TLbXvPv7EB7cNK1QV/luu55UHLrrn8=
+github.com/DataDog/zstd v1.5.2/go.mod h1:g4AWEaM3yOg3HYfnJ3YIawPnVdXJh9QME85blwSAmyw=
 github.com/algorand/avm-abi v0.1.0 h1:znZFQXpSUVYz37vXbaH5OZG2VK4snTyXwnc/tV9CVr4=
 github.com/algorand/avm-abi v0.1.0/go.mod h1:+CgwM46dithy850bpTeHh9MC99zpn2Snirb3QTl2O/g=
 github.com/algorand/falcon v0.0.0-20220727072124-02a2a64c4414 h1:nwYN+GQ7Z5OOfZwqBO1ma7DSlP7S1YrKWICOyjkwqrc=

network/msgCompressor.go

@@ -0,0 +1,157 @@
+// Copyright (C) 2019-2022 Algorand, Inc.
+// This file is part of go-algorand
+//
+// go-algorand is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// go-algorand is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with go-algorand.  If not, see <https://www.gnu.org/licenses/>.
+
+package network
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+
+	"github.com/DataDog/zstd"
+
+	"github.com/algorand/go-algorand/logging"
+	"github.com/algorand/go-algorand/protocol"
+)
+
+var zstdCompressionMagic = [4]byte{0x28, 0xb5, 0x2f, 0xfd}
+
+const zstdCompressionLevel = zstd.BestSpeed
+
+// checkCanCompress checks if there is an proposal payload message and peers supporting compression
+func checkCanCompress(request broadcastRequest, peers []*wsPeer) bool {
+	canCompress := false
+	hasPP := false
+	for _, tag := range request.tags {
+		if tag == protocol.ProposalPayloadTag {
+			hasPP = true
+			break
+		}
+	}
+	// if have proposal payload check if there are any peers supporting compression
+	if hasPP {
+		for _, peer := range peers {
+			if peer.pfProposalCompressionSupported() {
+				canCompress = true
+				break
+			}
+		}
+	}
+	return canCompress
+}
+
+// zstdCompressMsg returns a concatenation of a tag and compressed data
+func zstdCompressMsg(tbytes []byte, d []byte) ([]byte, string) {
+	bound := zstd.CompressBound(len(d))
+	if bound < len(d) {
+		// although CompressBound allocated more than the src size, this is an implementation detail.
+		// increase the buffer size to always have enough space for the raw data if compression fails.
+		bound = len(d)
+	}
+	mbytesComp := make([]byte, len(tbytes)+bound)
+	copy(mbytesComp, tbytes)
+	comp, err := zstd.CompressLevel(mbytesComp[len(tbytes):], d, zstdCompressionLevel)
+	if err != nil {
+		// fallback and reuse non-compressed original data
+		logMsg := fmt.Sprintf("failed to compress into buffer of len %d: %v", len(d), err)
+		copied := copy(mbytesComp[len(tbytes):], d)
+		return mbytesComp[:len(tbytes)+copied], logMsg
+	}
+	mbytesComp = mbytesComp[:len(tbytes)+len(comp)]
+	return mbytesComp, ""
+}
+
+// MaxDecompressedMessageSize defines a maximum decompressed data size
+// to prevent zip bombs
+const MaxDecompressedMessageSize = 20 * 1024 * 1024 // some large enough value
+
+// wsPeerMsgDataConverter performs optional incoming messages conversion.
+// At the moment it only supports zstd decompression for payload proposal
+type wsPeerMsgDataConverter struct {
+	log    logging.Logger
+	origin string
+
+	// actual converter(s)
+	ppdec zstdProposalDecompressor
+}
+
+type zstdProposalDecompressor struct {
+	active bool
+}
+
+func (dec zstdProposalDecompressor) enabled() bool {
+	return dec.active
+}
+
+func (dec zstdProposalDecompressor) accept(data []byte) bool {
+	return len(data) > 4 && bytes.Equal(data[:4], zstdCompressionMagic[:])
+}
+
+func (dec zstdProposalDecompressor) convert(data []byte) ([]byte, error) {
+	r := zstd.NewReader(bytes.NewReader(data))
+	defer r.Close()
+	b := make([]byte, 0, 3*len(data))
+	for {
+		if len(b) == cap(b) {
+			// grow capacity, retain length
+			b = append(b, 0)[:len(b)]
+		}
+		n, err := r.Read(b[len(b):cap(b)])
+		b = b[:len(b)+n]
+		if err != nil {
+			if err == io.EOF {
+				return b, nil
+			}
+			return nil, err
+		}
+		if len(b) > MaxDecompressedMessageSize {
+			return nil, fmt.Errorf("proposal data is too large: %d", len(b))
+		}
+	}
+}
+
+func (c *wsPeerMsgDataConverter) convert(tag protocol.Tag, data []byte) ([]byte, error) {
+	if tag == protocol.ProposalPayloadTag {
+		if c.ppdec.enabled() {
+			// sender might support compressed payload but fail to compress for whatever reason,
+			// in this case it sends non-compressed payload - the receiver decompress only if it is compressed.
+			if c.ppdec.accept(data) {
+				res, err := c.ppdec.convert(data)
+				if err != nil {
+					return nil, fmt.Errorf("peer %s: %w", c.origin, err)
+				}
+				return res, nil
+			}
+			c.log.Warnf("peer %s supported zstd but sent non-compressed data", c.origin)
+		}
+	}
+	return data, nil
+}
+
+func makeWsPeerMsgDataConverter(wp *wsPeer) *wsPeerMsgDataConverter {
+	c := wsPeerMsgDataConverter{
+		log:    wp.net.log,
+		origin: wp.originAddress,
+	}
+
+	if wp.pfProposalCompressionSupported() {
+		c.ppdec = zstdProposalDecompressor{
+			active: true,
+		}
+	}
+
+	return &c
+}

network/msgCompressor_test.go

@@ -0,0 +1,142 @@
+// Copyright (C) 2019-2022 Algorand, Inc.
+// This file is part of go-algorand
+//
+// go-algorand is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// go-algorand is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with go-algorand.  If not, see <https://www.gnu.org/licenses/>.
+
+package network
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/DataDog/zstd"
+	"github.com/algorand/go-algorand/logging"
+	"github.com/algorand/go-algorand/protocol"
+	"github.com/algorand/go-algorand/test/partitiontest"
+	"github.com/stretchr/testify/require"
+)
+
+func TestZstdDecompress(t *testing.T) {
+	partitiontest.PartitionTest(t)
+
+	// happy case - small message
+	msg := []byte(strings.Repeat("1", 2048))
+	compressed, err := zstd.Compress(nil, msg)
+	require.NoError(t, err)
+	d := zstdProposalDecompressor{}
+	decompressed, err := d.convert(compressed)
+	require.NoError(t, err)
+	require.Equal(t, msg, decompressed)
+
+	// error case - large message
+	msg = []byte(strings.Repeat("1", MaxDecompressedMessageSize+10))
+	compressed, err = zstd.Compress(nil, msg)
+	require.NoError(t, err)
+	decompressed, err = d.convert(compressed)
+	require.Error(t, err)
+	require.Nil(t, decompressed)
+}
+
+func TestCheckCanCompress(t *testing.T) {
+	partitiontest.PartitionTest(t)
+
+	req := broadcastRequest{}
+	peers := []*wsPeer{}
+	r := checkCanCompress(req, peers)
+	require.False(t, r)
+
+	req.tags = []protocol.Tag{protocol.AgreementVoteTag}
+	r = checkCanCompress(req, peers)
+	require.False(t, r)
+
+	req.tags = []protocol.Tag{protocol.AgreementVoteTag, protocol.ProposalPayloadTag}
+	r = checkCanCompress(req, peers)
+	require.False(t, r)
+
+	peer1 := wsPeer{
+		features: 0,
+	}
+	peers = []*wsPeer{&peer1}
+	r = checkCanCompress(req, peers)
+	require.False(t, r)
+
+	peer2 := wsPeer{
+		features: pfCompressedProposal,
+	}
+	peers = []*wsPeer{&peer1, &peer2}
+	r = checkCanCompress(req, peers)
+	require.True(t, r)
+}
+
+func TestZstdCompressMsg(t *testing.T) {
+	partitiontest.PartitionTest(t)
+
+	ppt := len(protocol.ProposalPayloadTag)
+	data := []byte("data")
+	comp, msg := zstdCompressMsg([]byte(protocol.ProposalPayloadTag), data)
+	require.Empty(t, msg)
+	require.Equal(t, []byte(protocol.ProposalPayloadTag), comp[:ppt])
+	require.Equal(t, zstdCompressionMagic[:], comp[ppt:ppt+len(zstdCompressionMagic)])
+	d := zstdProposalDecompressor{}
+	decompressed, err := d.convert(comp[ppt:])
+	require.NoError(t, err)
+	require.Equal(t, data, decompressed)
+}
+
+type converterTestLogger struct {
+	logging.Logger
+	WarnfCallback func(string, ...interface{})
+	warnMsgCount  int
+}
+
+func (cl *converterTestLogger) Warnf(s string, args ...interface{}) {
+	cl.warnMsgCount++
+}
+
+func TestWsPeerMsgDataConverterConvert(t *testing.T) {
+	partitiontest.PartitionTest(t)
+
+	c := wsPeerMsgDataConverter{}
+	c.ppdec = zstdProposalDecompressor{active: false}
+	tag := protocol.AgreementVoteTag
+	data := []byte("data")
+
+	r, err := c.convert(tag, data)
+	require.NoError(t, err)
+	require.Equal(t, data, r)
+
+	tag = protocol.ProposalPayloadTag
+	r, err = c.convert(tag, data)
+	require.NoError(t, err)
+	require.Equal(t, data, r)
+
+	l := converterTestLogger{}
+	c.log = &l
+	c.ppdec = zstdProposalDecompressor{active: true}
+	r, err = c.convert(tag, data)
+	require.NoError(t, err)
+	require.Equal(t, data, r)
+	require.Equal(t, 1, l.warnMsgCount)
+
+	l = converterTestLogger{}
+	c.log = &l
+
+	comp, err := zstd.Compress(nil, data)
+	require.NoError(t, err)
+
+	r, err = c.convert(tag, comp)
+	require.NoError(t, err)
+	require.Equal(t, data, r)
+	require.Equal(t, 0, l.warnMsgCount)
+}