feat: add share index to malleated (wrapped) txs (celestiaorg#819)

* feat: add share index to malleated (wrapped) txs

pass the share index when wrapping

* chore: refactor unrwap function to return the entire malleated tx

* fix compile issue

Author: evan-forbes

mempool/v0/clist_mempool.go

@@ -616,9 +616,9 @@ func (mem *CListMempool) Update(
 			mem.removeTx(tx, e.(*clist.CElement), false)
 			// see if the transaction is a malleated transaction of a some parent
 			// transaction that exists in the mempool
-		} else if parentHash, _, isMalleated := types.UnwrapMalleatedTx(tx); isMalleated {
+		} else if malleatedTx, isMalleated := types.UnwrapMalleatedTx(tx); isMalleated {
 			var parentKey [types.TxKeySize]byte
-			copy(parentKey[:], parentHash)
+			copy(parentKey[:], malleatedTx.OriginalTxHash)
 			err := mem.RemoveTxByKey(parentKey)
 			if err != nil {
 				return err

mempool/v1/mempool.go

@@ -406,8 +406,8 @@ func (txmp *TxMempool) Update(
 
 		// Regardless of outcome, remove the transaction from the mempool.
 		if err := txmp.removeTxByKey(originalKey); err != nil {
-			if originalHash, _, isMalleated := types.UnwrapMalleatedTx(tx); isMalleated {
-				copy(originalKey[:], originalHash)
+			if malleatedTx, isMalleated := types.UnwrapMalleatedTx(tx); isMalleated {
+				copy(originalKey[:], malleatedTx.OriginalTxHash)
 				_ = txmp.removeTxByKey(originalKey)
 			}
 		}

mempool/v1/mempool_test.go

@@ -661,7 +661,7 @@ func TestMalleatedTxRemoval(t *testing.T) {
 	originalHash := sha256.Sum256(originalTx)
 
 	// create the wrapped child transaction
-	wTx, err := types.WrapMalleatedTx(originalHash[:], malleatedTx)
+	wTx, err := types.WrapMalleatedTx(originalHash[:], 0, malleatedTx)
 	require.NoError(t, err)
 
 	// add the parent transaction to the mempool

proto/tendermint/types/types.pb.go

@@ -224,6 +224,7 @@ message TxProof {
 message MalleatedTx {
   bytes original_tx_hash = 1;
   bytes tx               = 2;
+  uint32 share_index     = 3;
 }
 
 // Proof represents proof of a namespace.ID in an NMT.

proto/tendermint/types/types.proto

@@ -584,9 +584,9 @@ func fireEvents(
 	for i, tx := range block.Data.Txs {
 		var txHash []byte
 		var rawTx []byte
-		if originalHash, malleatedTx, ismalleated := types.UnwrapMalleatedTx(tx); ismalleated {
-			txHash = originalHash
-			rawTx = malleatedTx
+		if malleatedTx, ismalleated := types.UnwrapMalleatedTx(tx); ismalleated {
+			txHash = malleatedTx.OriginalTxHash
+			rawTx = malleatedTx.Tx
 		} else {
 			txHash = tx.Hash()
 			rawTx = tx

state/execution.go

@@ -220,30 +220,30 @@ func ComputeProtoSizeForTxs(txs []Tx) int64 {
 // passed is not a tmproto.MalleatedTx, since the schema for PayForMessage is kept
 // in the app, we cannot perform further checks without creating an import
 // cycle.
-func UnwrapMalleatedTx(tx Tx) (originalHash []byte, unwrapped Tx, isMalleated bool) {
+func UnwrapMalleatedTx(tx Tx) (malleatedTx tmproto.MalleatedTx, isMalleated bool) {
 	// attempt to unmarshal into a a malleated transaction
-	var malleatedTx tmproto.MalleatedTx
 	err := proto.Unmarshal(tx, &malleatedTx)
 	if err != nil {
-		return nil, nil, false
+		return malleatedTx, false
 	}
 	// this check will fail to catch unwanted types should those unmarshalled
 	// types happen to have a hash sized slice of bytes in the same field number
 	// as originalTxHash. TODO(evan): either fix this, or better yet use a different
 	// mechanism
 	if len(malleatedTx.OriginalTxHash) != tmhash.Size {
-		return nil, nil, false
+		return malleatedTx, false
 	}
-	return malleatedTx.OriginalTxHash, malleatedTx.Tx, true
+	return malleatedTx, true
 }
 
 // WrapMalleatedTx creates a wrapped Tx that includes the original transaction's hash
 // so that it can be easily removed from the mempool. note: must be unwrapped to
 // be a viable sdk.Tx
-func WrapMalleatedTx(originalHash []byte, malleated Tx) (Tx, error) {
+func WrapMalleatedTx(originalHash []byte, shareIndex uint32, malleated Tx) (Tx, error) {
 	wTx := tmproto.MalleatedTx{
 		OriginalTxHash: originalHash,
 		Tx:             malleated,
+		ShareIndex:     shareIndex,
 	}
 	return proto.Marshal(&wTx)
 }

types/tx.go

@@ -48,7 +48,7 @@ func TestUnwrapMalleatedTx(t *testing.T) {
 	// perform a simple test for being unable to decode a non
 	// malleated transaction
 	tx := Tx{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0}
-	_, _, ok := UnwrapMalleatedTx(tx)
+	_, ok := UnwrapMalleatedTx(tx)
 	require.False(t, ok)
 
 	data := Data{
@@ -81,16 +81,16 @@ func TestUnwrapMalleatedTx(t *testing.T) {
 
 	// due to protobuf not actually requiring type compatibility
 	// we need to make sure that there is some check
-	_, _, ok = UnwrapMalleatedTx(rawBlock)
+	_, ok = UnwrapMalleatedTx(rawBlock)
 	require.False(t, ok)
 
 	pHash := sha256.Sum256(rawBlock)
-	MalleatedTx, err := WrapMalleatedTx(pHash[:], rawBlock)
+	MalleatedTx, err := WrapMalleatedTx(pHash[:], 0, rawBlock)
 	require.NoError(t, err)
 
 	// finally, ensure that the unwrapped bytes are identical to the input
-	unwrappedHash, unwrapped, ok := UnwrapMalleatedTx(MalleatedTx)
+	malleated, ok := UnwrapMalleatedTx(MalleatedTx)
 	require.True(t, ok)
-	assert.Equal(t, 32, len(unwrappedHash))
-	require.Equal(t, rawBlock, []byte(unwrapped))
+	assert.Equal(t, 32, len(malleated.OriginalTxHash))
+	require.Equal(t, rawBlock, malleated.Tx)
 }