chore: optimize checkTx (#3954)

## Overview

this PR makes a simple optimization for checkTx

---------

Co-authored-by: Rootul P <[email protected]>
Co-authored-by: CHAMI Rachid <[email protected]>

app/process_proposal.go

@@ -67,6 +67,9 @@ func (app *App) ProcessProposal(req abci.RequestProcessProposal) (resp abci.Resp
 			tx = blobTx.Tx
 		}
 
+		// todo: uncomment once we're sure this isn't consensus breaking
+		// sdkCtx = sdkCtx.WithTxBytes(tx)
+
 		sdkTx, err := app.txConfig.TxDecoder()(tx)
 		if err != nil {
 			if req.Header.Version.App == v1 {

app/test/prepare_proposal_test.go

@@ -137,6 +137,22 @@ func TestPrepareProposalFiltering(t *testing.T) {
 	require.NoError(t, err)
 	noAccountTx := []byte(testutil.SendTxWithManualSequence(t, encConf.TxConfig, kr, nilAccount, accounts[0], 1000, "", 0, 6))
 
+	// create a tx that can't be included in a 64 x 64 when accounting for the
+	// pfb along with the shares
+	tooManyShareBtx := blobfactory.ManyMultiBlobTx(
+		t,
+		encConf.TxConfig,
+		kr,
+		testutil.ChainID,
+		accounts[3:4],
+		infos[3:4],
+		blobfactory.NestedBlobs(
+			t,
+			testfactory.RandomBlobNamespaces(tmrand.NewRand(), 4000),
+			[][]int{repeat(4000, 1)},
+		),
+	)[0]
+
 	type test struct {
 		name      string
 		txs       func() [][]byte
@@ -181,6 +197,13 @@ func TestPrepareProposalFiltering(t *testing.T) {
 			},
 			prunedTxs: [][]byte{noAccountTx},
 		},
+		{
+			name: "blob tx with too many shares",
+			txs: func() [][]byte {
+				return [][]byte{tooManyShareBtx}
+			},
+			prunedTxs: [][]byte{tooManyShareBtx},
+		},
 	}
 
 	for _, tt := range tests {
@@ -216,3 +239,12 @@ func queryAccountInfo(capp *app.App, accs []string, kr keyring.Keyring) []blobfa
 	}
 	return infos
 }
+
+// repeat returns a slice of length n with each element set to val.
+func repeat[T any](n int, val T) []T {
+	result := make([]T, n)
+	for i := range result {
+		result[i] = val
+	}
+	return result
+}

app/test/process_proposal_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/celestiaorg/celestia-app/v3/pkg/appconsts"
 	v1 "github.com/celestiaorg/celestia-app/v3/pkg/appconsts/v1"
 	v2 "github.com/celestiaorg/celestia-app/v3/pkg/appconsts/v2"
+	v3 "github.com/celestiaorg/celestia-app/v3/pkg/appconsts/v3"
 	"github.com/celestiaorg/celestia-app/v3/pkg/da"
 	"github.com/celestiaorg/celestia-app/v3/pkg/user"
 	testutil "github.com/celestiaorg/celestia-app/v3/test/util"
@@ -80,6 +81,20 @@ func TestProcessProposal(t *testing.T) {
 	ns1 := share.MustNewV0Namespace(bytes.Repeat([]byte{1}, share.NamespaceVersionZeroIDSize))
 	data := bytes.Repeat([]byte{1}, 13)
 
+	tooManyShareBtx := blobfactory.ManyMultiBlobTx(
+		t,
+		enc,
+		kr,
+		testutil.ChainID,
+		accounts[3:4],
+		infos[3:4],
+		blobfactory.NestedBlobs(
+			t,
+			testfactory.RandomBlobNamespaces(tmrand.NewRand(), 4000),
+			[][]int{repeat(4000, 1)},
+		),
+	)[0]
+
 	type test struct {
 		name           string
 		input          *tmproto.Data
@@ -299,6 +314,19 @@ func TestProcessProposal(t *testing.T) {
 			appVersion:     appconsts.LatestVersion,
 			expectedResult: abci.ResponseProcessProposal_REJECT,
 		},
+		{
+			name: "blob tx that takes up too many shares",
+			input: &tmproto.Data{
+				Txs: [][]byte{},
+			},
+			mutator: func(d *tmproto.Data) {
+				// this tx will get filtered out by prepare proposal before this
+				// so we add it here
+				d.Txs = append(d.Txs, tooManyShareBtx)
+			},
+			appVersion:     v3.Version,
+			expectedResult: abci.ResponseProcessProposal_REJECT,
+		},
 	}
 
 	for _, tt := range tests {

app/validate_txs.go

@@ -45,6 +45,7 @@ func FilterTxs(logger log.Logger, ctx sdk.Context, handler sdk.AnteHandler, txCo
 func filterStdTxs(logger log.Logger, dec sdk.TxDecoder, ctx sdk.Context, handler sdk.AnteHandler, txs [][]byte) ([][]byte, sdk.Context) {
 	n := 0
 	for _, tx := range txs {
+		ctx = ctx.WithTxBytes(tx)
 		sdkTx, err := dec(tx)
 		if err != nil {
 			logger.Error("decoding already checked transaction", "tx", tmbytes.HexBytes(coretypes.Tx(tx).Hash()), "error", err)
@@ -78,6 +79,7 @@ func filterStdTxs(logger log.Logger, dec sdk.TxDecoder, ctx sdk.Context, handler
 func filterBlobTxs(logger log.Logger, dec sdk.TxDecoder, ctx sdk.Context, handler sdk.AnteHandler, txs []*tx.BlobTx) ([]*tx.BlobTx, sdk.Context) {
 	n := 0
 	for _, tx := range txs {
+		ctx = ctx.WithTxBytes(tx.Tx)
 		sdkTx, err := dec(tx.Tx)
 		if err != nil {
 			logger.Error("decoding already checked blob transaction", "tx", tmbytes.HexBytes(coretypes.Tx(tx.Tx).Hash()), "error", err)

x/blob/ante/blob_share_decorator.go

@@ -36,7 +36,7 @@ func (d BlobShareDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool
 	maxBlobShares := d.getMaxBlobShares(ctx)
 	for _, m := range tx.GetMsgs() {
 		if pfb, ok := m.(*blobtypes.MsgPayForBlobs); ok {
-			if sharesNeeded := getSharesNeeded(pfb.BlobSizes); sharesNeeded > maxBlobShares {
+			if sharesNeeded := getSharesNeeded(uint32(len(ctx.TxBytes())), pfb.BlobSizes); sharesNeeded > maxBlobShares {
 				return ctx, errors.Wrapf(blobtypes.ErrBlobsTooLarge, "the number of shares occupied by blobs in this MsgPayForBlobs %d exceeds the max number of shares available for blob data %d", sharesNeeded, maxBlobShares)
 			}
 		}
@@ -49,10 +49,8 @@ func (d BlobShareDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool
 func (d BlobShareDecorator) getMaxBlobShares(ctx sdk.Context) int {
 	squareSize := d.getMaxSquareSize(ctx)
 	totalShares := squareSize * squareSize
-	// The PFB tx share must occupy at least one share so the number of blob shares
-	// is at most one less than totalShares.
-	blobShares := totalShares - 1
-	return blobShares
+	// the shares used up by the tx are calculated in `getSharesNeeded`
+	return totalShares
 }
 
 // getMaxSquareSize returns the maximum square size based on the current values
@@ -74,8 +72,9 @@ func (d BlobShareDecorator) getMaxSquareSize(ctx sdk.Context) int {
 }
 
 // getSharesNeeded returns the total number of shares needed to represent all of
-// the blobs described by blobSizes.
-func getSharesNeeded(blobSizes []uint32) (sum int) {
+// the blobs described by blobSizes along with the shares used by the tx
+func getSharesNeeded(txSize uint32, blobSizes []uint32) (sum int) {
+	sum = share.CompactSharesNeeded(txSize)
 	for _, blobSize := range blobSizes {
 		sum += share.SparseSharesNeeded(blobSize)
 	}