feat: reject tx if total blob size too large (#2202)

Closes #2156 by
implementing **Option C**

(cherry picked from commit b6f3968)

# Conflicts:
#	app/ante.go
#	app/test/check_tx_test.go

app/ante.go

@@ -36,6 +36,11 @@ func NewAnteHandler(
 		ante.NewSigGasConsumeDecorator(accountKeeper, sigGasConsumer),
 		ante.NewSigVerificationDecorator(accountKeeper, signModeHandler),
 		blobante.NewMinGasPFBDecorator(blobKeeper),
+<<<<<<< HEAD:app/ante.go
+=======
+		blobante.NewMaxBlobSizeDecorator(blobKeeper),
+		NewGovProposalDecorator(),
+>>>>>>> b6f3968 (feat: reject tx if total blob size too large (#2202)):app/ante/ante.go
 		ante.NewIncrementSequenceDecorator(accountKeeper),
 		ibcante.NewRedundantRelayDecorator(channelKeeper),
 	)

app/test/check_tx_test.go

@@ -107,6 +107,54 @@ func TestCheckTx(t *testing.T) {
 			},
 			expectedABCICode: abci.CodeTypeOK,
 		},
+<<<<<<< HEAD
+=======
+		{
+			name:      "1,000 byte blob",
+			checkType: abci.CheckTxType_New,
+			getTx: func() []byte {
+				tx := blobfactory.RandBlobTxsWithAccounts(encCfg.TxConfig.TxEncoder(), tmrand.NewRand(), kr, nil, 1_000, 1, false, testutil.ChainID, accs[4:5])[0]
+				return tx
+			},
+			expectedABCICode: abci.CodeTypeOK,
+		},
+		{
+			name:      "10,000 byte blob",
+			checkType: abci.CheckTxType_New,
+			getTx: func() []byte {
+				tx := blobfactory.RandBlobTxsWithAccounts(encCfg.TxConfig.TxEncoder(), tmrand.NewRand(), kr, nil, 10_000, 1, false, testutil.ChainID, accs[5:6])[0]
+				return tx
+			},
+			expectedABCICode: abci.CodeTypeOK,
+		},
+		{
+			name:      "100,000 byte blob",
+			checkType: abci.CheckTxType_New,
+			getTx: func() []byte {
+				tx := blobfactory.RandBlobTxsWithAccounts(encCfg.TxConfig.TxEncoder(), tmrand.NewRand(), kr, nil, 100_000, 1, false, testutil.ChainID, accs[6:7])[0]
+				return tx
+			},
+			expectedABCICode: abci.CodeTypeOK,
+		},
+		{
+			name:      "1,000,000 byte blob",
+			checkType: abci.CheckTxType_New,
+			getTx: func() []byte {
+				tx := blobfactory.RandBlobTxsWithAccounts(encCfg.TxConfig.TxEncoder(), tmrand.NewRand(), kr, nil, 1_000_000, 1, false, testutil.ChainID, accs[7:8])[0]
+				return tx
+			},
+			expectedABCICode: abci.CodeTypeOK,
+		},
+		{
+			name:      "10,000,000 byte blob",
+			checkType: abci.CheckTxType_New,
+			getTx: func() []byte {
+				tx := blobfactory.RandBlobTxsWithAccounts(encCfg.TxConfig.TxEncoder(), tmrand.NewRand(), kr, nil, 10_000_000, 1, false, testutil.ChainID, accs[8:9])[0]
+				return tx
+			},
+			expectedABCICode: blobtypes.ErrTotalBlobSizeTooLarge.ABCICode(),
+		},
+>>>>>>> b6f3968 (feat: reject tx if total blob size too large (#2202))
 	}
 
 	for _, tt := range tests {

app/test/max_total_blob_size_test.go

@@ -0,0 +1,141 @@
+package app_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/celestiaorg/celestia-app/app"
+	"github.com/celestiaorg/celestia-app/app/encoding"
+	"github.com/celestiaorg/celestia-app/pkg/appconsts"
+	"github.com/celestiaorg/celestia-app/pkg/square"
+	"github.com/celestiaorg/celestia-app/test/util/testfactory"
+	"github.com/celestiaorg/celestia-app/test/util/testnode"
+	blobtypes "github.com/celestiaorg/celestia-app/x/blob/types"
+	sdk_tx "github.com/cosmos/cosmos-sdk/types/tx"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/suite"
+	abci "github.com/tendermint/tendermint/abci/types"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
+	coretypes "github.com/tendermint/tendermint/types"
+	"google.golang.org/grpc"
+)
+
+const (
+	mebibyte   = 1_048_576 // one mebibyte in bytes
+	squareSize = 64
+)
+
+func TestMaxTotalBlobSizeSuite(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping max total blob size suite in short mode.")
+	}
+	suite.Run(t, &MaxTotalBlobSizeSuite{})
+}
+
+type MaxTotalBlobSizeSuite struct {
+	suite.Suite
+
+	ecfg     encoding.Config
+	accounts []string
+	cctx     testnode.Context
+}
+
+func (s *MaxTotalBlobSizeSuite) SetupSuite() {
+	t := s.T()
+
+	s.accounts = testfactory.GenerateAccounts(1)
+
+	tmConfig := testnode.DefaultTendermintConfig()
+	tmConfig.Mempool.MaxTxBytes = 10 * mebibyte
+
+	cParams := testnode.DefaultParams()
+	cParams.Block.MaxBytes = 10 * mebibyte
+
+	cfg := testnode.DefaultConfig().
+		WithAccounts(s.accounts).
+		WithTendermintConfig(tmConfig).
+		WithConsensusParams(cParams)
+
+	cctx, _, _ := testnode.NewNetwork(t, cfg)
+	s.cctx = cctx
+	s.ecfg = encoding.MakeConfig(app.ModuleEncodingRegisters...)
+
+	require.NoError(t, cctx.WaitForNextBlock())
+}
+
+// TestSubmitPayForBlob_blobSizes verifies the tx response ABCI code when
+// SubmitPayForBlob is invoked with different blob sizes.
+func (s *MaxTotalBlobSizeSuite) TestSubmitPayForBlob_blobSizes() {
+	t := s.T()
+
+	type testCase struct {
+		name string
+		blob *tmproto.Blob
+		// want is the expected tx response ABCI code.
+		want uint32
+	}
+	testCases := []testCase{
+		{
+			name: "1 byte blob",
+			blob: mustNewBlob(t, 1),
+			want: abci.CodeTypeOK,
+		},
+		{
+			name: "1 mebibyte blob",
+			blob: mustNewBlob(t, mebibyte),
+			want: abci.CodeTypeOK,
+		},
+		{
+			name: "2 mebibyte blob",
+			blob: mustNewBlob(t, 2*mebibyte),
+			want: blobtypes.ErrTotalBlobSizeTooLarge.ABCICode(),
+		},
+	}
+
+	signer := blobtypes.NewKeyringSigner(s.cctx.Keyring, s.accounts[0], s.cctx.ChainID)
+
+	for _, tc := range testCases {
+		s.Run(tc.name, func() {
+			blobTx := newBlobTx(t, signer, s.cctx.GRPCClient, tc.blob)
+			res, err := blobtypes.BroadcastTx(context.TODO(), s.cctx.GRPCClient, sdk_tx.BroadcastMode_BROADCAST_MODE_BLOCK, blobTx)
+			require.NoError(t, err)
+			require.NotNil(t, res)
+			require.Equal(t, tc.want, res.TxResponse.Code, res.TxResponse.Logs)
+
+			sq, err := square.Construct([][]byte{blobTx}, appconsts.LatestVersion, squareSize)
+			if tc.want == abci.CodeTypeOK {
+				// verify that if the tx was accepted, the blob can fit in a square
+				assert.NoError(t, err)
+				assert.False(t, sq.IsEmpty())
+			} else {
+				// verify that if the tx was rejected, the blob can not fit in a square
+				assert.Error(t, err)
+			}
+		})
+	}
+}
+
+func newBlobTx(t *testing.T, signer *blobtypes.KeyringSigner, conn *grpc.ClientConn, blob *tmproto.Blob) coretypes.Tx {
+	addr, err := signer.GetSignerInfo().GetAddress()
+	require.NoError(t, err)
+
+	msg, err := blobtypes.NewMsgPayForBlobs(addr.String(), blob)
+	require.NoError(t, err)
+
+	err = signer.QueryAccountNumber(context.TODO(), conn)
+	require.NoError(t, err)
+
+	options := []blobtypes.TxBuilderOption{blobtypes.SetGasLimit(1e9)} // set gas limit to 1 billion to avoid gas exhaustion
+	builder := signer.NewTxBuilder(options...)
+	stx, err := signer.BuildSignedTx(builder, msg)
+	require.NoError(t, err)
+
+	rawTx, err := signer.EncodeTx(stx)
+	require.NoError(t, err)
+
+	blobTx, err := coretypes.MarshalBlobTx(rawTx, blob)
+	require.NoError(t, err)
+
+	return blobTx
+}

x/blob/ante/ante.go

@@ -48,4 +48,5 @@ func (d MinGasPFBDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool
 
 type BlobKeeper interface {
 	GasPerBlobByte(ctx sdk.Context) uint32
+	GovMaxSquareSize(ctx sdk.Context) uint64
 }

x/blob/ante/ante_test.go

@@ -13,7 +13,10 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-const testGasPerBlobByte = 10
+const (
+	testGasPerBlobByte   = 10
+	testGovMaxSquareSize = 64
+)
 
 func TestPFBAnteHandler(t *testing.T) {
 	txConfig := encoding.MakeConfig(app.ModuleEncodingRegisters...).TxConfig
@@ -106,3 +109,7 @@ type mockBlobKeeper struct{}
 func (mockBlobKeeper) GasPerBlobByte(_ sdk.Context) uint32 {
 	return testGasPerBlobByte
 }
+
+func (mockBlobKeeper) GovMaxSquareSize(_ sdk.Context) uint64 {
+	return testGovMaxSquareSize
+}

x/blob/ante/max_total_blob_size_ante.go

@@ -0,0 +1,88 @@
+package ante
+
+import (
+	"github.com/celestiaorg/celestia-app/pkg/appconsts"
+	"github.com/celestiaorg/celestia-app/pkg/shares"
+	blobtypes "github.com/celestiaorg/celestia-app/x/blob/types"
+
+	"cosmossdk.io/errors"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+// MaxTotalBlobSizeDecorator helps to prevent a PFB from being included in a block
+// but not fitting in a data square.
+type MaxTotalBlobSizeDecorator struct {
+	k BlobKeeper
+}
+
+func NewMaxBlobSizeDecorator(k BlobKeeper) MaxTotalBlobSizeDecorator {
+	return MaxTotalBlobSizeDecorator{k}
+}
+
+// AnteHandle implements the Cosmos SDK AnteHandler function signature. It
+// returns an error if tx contains a MsgPayForBlobs where the total blob size is
+// greater than the max total blob size.
+func (d MaxTotalBlobSizeDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool, next sdk.AnteHandler) (sdk.Context, error) {
+	if !ctx.IsCheckTx() {
+		return next(ctx, tx, simulate)
+	}
+
+	max := d.maxTotalBlobSize(ctx)
+	for _, m := range tx.GetMsgs() {
+		if pfb, ok := m.(*blobtypes.MsgPayForBlobs); ok {
+			if total := getTotal(pfb.BlobSizes); total > max {
+				return ctx, errors.Wrapf(blobtypes.ErrTotalBlobSizeTooLarge, "total blob size %d exceeds max %d", total, max)
+			}
+		}
+	}
+
+	return next(ctx, tx, simulate)
+}
+
+// maxTotalBlobSize returns the max the number of bytes available for blobs in a
+// data square based on the max square size. Note it is possible that txs with a
+// total blob size less than this max still fail to be included in a block due
+// to overhead from the PFB tx and/or padding shares.
+func (d MaxTotalBlobSizeDecorator) maxTotalBlobSize(ctx sdk.Context) int {
+	squareSize := d.getMaxSquareSize(ctx)
+	totalShares := squareSize * squareSize
+	// The PFB tx share must occupy at least one share so the # of blob shares
+	// is at least one less than totalShares.
+	blobShares := totalShares - 1
+	return shares.AvailableBytesFromSparseShares(blobShares)
+}
+
+// getMaxSquareSize returns the maximum square size based on the current values
+// for the relevant governance parameter and the versioned constant.
+func (d MaxTotalBlobSizeDecorator) getMaxSquareSize(ctx sdk.Context) int {
+	// TODO: fix hack that forces the max square size for the first height to
+	// 64. This is due to our fork of the sdk not initializing state before
+	// BeginBlock of the first block. This is remedied in versions of the sdk
+	// and comet that have full support of PreparePropsoal, although
+	// celestia-app does not currently use those. see this PR for more details
+	// https://github.com/cosmos/cosmos-sdk/pull/14505
+	if ctx.BlockHeader().Height <= 1 {
+		return int(appconsts.DefaultGovMaxSquareSize)
+	}
+
+	upperBound := appconsts.SquareSizeUpperBound(ctx.ConsensusParams().Version.AppVersion)
+	govParam := d.k.GovMaxSquareSize(ctx)
+	return min(upperBound, int(govParam))
+}
+
+// getTotal returns the sum of the given sizes.
+func getTotal(sizes []uint32) (sum int) {
+	for _, size := range sizes {
+		sum += int(size)
+	}
+	return sum
+}
+
+// min returns the minimum of two ints. This function can be removed once we
+// upgrade to Go 1.21.
+func min(a, b int) int {
+	if a < b {
+		return a
+	}
+	return b
+}