Added ArgsTransform to ChainWriter

smartcontractkit · Jan 15, 2025 · aa1ee7a · aa1ee7a
1 parent 74b7a72
commit aa1ee7a
Show file tree

Hide file tree

Showing 3 changed files with 87 additions and 12 deletions.
diff --git a/pkg/solana/chainwriter/ccip_example_config.go b/pkg/solana/chainwriter/ccip_example_config.go
@@ -1,9 +1,15 @@
 package chainwriter
 
 import (
+	"context"
+	"fmt"
+	"reflect"
+
 	"github.com/gagliardetto/solana-go"
 )
 
+const registryAddress = "4Nn9dsYBcSTzRbK9hg9kzCUdrCSkMZq1UR6Vw1Tkaf6A"
+
 func TestConfig() {
 	// Fake constant addresses for the purpose of this example.
 	registryAddress := "4Nn9dsYBcSTzRbK9hg9kzCUdrCSkMZq1UR6Vw1Tkaf6A"
@@ -331,3 +337,53 @@ func TestConfig() {
 	}
 	_ = chainWriterConfig
 }
+
+func CCIPArgsTransform(ctx context.Context, cw *SolanaChainWriterService, args any, accounts solana.AccountMetaSlice) (any, error) {
+	TokenPoolLookupTable := LookupTables{
+		DerivedLookupTables: []DerivedLookupTable{
+			{
+				Name: "RegistryTokenState",
+				Accounts: PDALookups{
+					Name: "RegistryTokenState",
+					PublicKey: AccountConstant{
+						Address:    registryAddress,
+						IsSigner:   false,
+						IsWritable: false,
+					},
+					Seeds: []Seed{
+						{Dynamic: AccountLookup{Location: "Message.TokenAmounts.DestTokenAddress"}},
+					},
+					IsSigner:   false,
+					IsWritable: false,
+					InternalField: InternalField{
+						Type:     reflect.TypeOf(DataAccount{}),
+						Location: "LookupTable",
+					},
+				},
+			},
+		}}
+	tableMap, _, err := cw.ResolveLookupTables(ctx, args, TokenPoolLookupTable)
+	if err != nil {
+		return nil, err
+	}
+	registryTables := tableMap["RegistryTokenState"]
+	tokenPoolAddresses := []solana.PublicKey{}
+	for _, table := range registryTables {
+		tokenPoolAddresses = append(tokenPoolAddresses, table[0].PublicKey)
+	}
+
+	tokenIndexes := []uint8{}
+	for i, account := range accounts {
+		for _, address := range tokenPoolAddresses {
+			if account.PublicKey == address {
+				tokenIndexes = append(tokenIndexes, uint8(i))
+			}
+		}
+	}
+
+	if len(tokenIndexes) != len(tokenPoolAddresses) {
+		return nil, fmt.Errorf("missing token pools in accounts")
+	}
+
+	args["token_indexes"] = tokenIndexes
+}
diff --git a/pkg/solana/chainwriter/chain_writer.go b/pkg/solana/chainwriter/chain_writer.go
@@ -59,6 +59,7 @@ type MethodConfig struct {
 	Accounts           []Lookup
 	// Location in the args where the debug ID is stored
 	DebugIDLocation string
+	ArgsTransform   func(ctx context.Context, cw *SolanaChainWriterService, args any, accounts solana.AccountMetaSlice) (any, error)
 }
 
 func NewSolanaChainWriterService(logger logger.Logger, reader client.Reader, txm txm.TxManager, ge fees.Estimator, config ChainWriterConfig) (*SolanaChainWriterService, error) {
@@ -256,15 +257,6 @@ func (s *SolanaChainWriterService) SubmitTransaction(ctx context.Context, contra
 		}
 	}
 
-	encodedPayload, err := s.encoder.Encode(ctx, args, codec.WrapItemType(true, contractName, method, ""))
-
-	if err != nil {
-		return errorWithDebugID(fmt.Errorf("error encoding transaction payload: %w", err), debugID)
-	}
-
-	discriminator := GetDiscriminator(methodConfig.ChainSpecificName)
-	encodedPayload = append(discriminator[:], encodedPayload...)
-
 	// Fetch derived and static table maps
 	derivedTableMap, staticTableMap, err := s.ResolveLookupTables(ctx, args, methodConfig.LookupTables)
 	if err != nil {
@@ -282,12 +274,17 @@ func (s *SolanaChainWriterService) SubmitTransaction(ctx context.Context, contra
 		return errorWithDebugID(fmt.Errorf("error parsing fee payer address: %w", err), debugID)
 	}
 
-	accounts = append([]*solana.AccountMeta{solana.Meta(feePayer).SIGNER().WRITE()}, accounts...)
-	accounts = append(accounts, solana.Meta(solana.SystemProgramID))
-
 	// Filter the lookup table addresses based on which accounts are actually used
 	filteredLookupTableMap := s.FilterLookupTableAddresses(accounts, derivedTableMap, staticTableMap)
 
+	// Transform args if necessary
+	if methodConfig.ArgsTransform != nil {
+		args, err = methodConfig.ArgsTransform(ctx, s, args, accounts)
+		if err != nil {
+			return errorWithDebugID(fmt.Errorf("error transforming args: %w", err), debugID)
+		}
+	}
+
 	// Fetch latest blockhash
 	blockhash, err := s.reader.LatestBlockhash(ctx)
 	if err != nil {
@@ -300,6 +297,15 @@ func (s *SolanaChainWriterService) SubmitTransaction(ctx context.Context, contra
 		return errorWithDebugID(fmt.Errorf("error parsing program ID: %w", err), debugID)
 	}
 
+	encodedPayload, err := s.encoder.Encode(ctx, args, codec.WrapItemType(true, contractName, method, ""))
+
+	if err != nil {
+		return errorWithDebugID(fmt.Errorf("error encoding transaction payload: %w", err), debugID)
+	}
+
+	discriminator := GetDiscriminator(methodConfig.ChainSpecificName)
+	encodedPayload = append(discriminator[:], encodedPayload...)
+
 	tx, err := solana.NewTransaction(
 		[]solana.Instruction{
 			solana.NewInstruction(programID, accounts, encodedPayload),

diff --git a/pkg/solana/chainwriter/chain_writer_test.go b/pkg/solana/chainwriter/chain_writer_test.go
@@ -454,6 +454,12 @@ func TestChainWriter_SubmitTransaction(t *testing.T) {
 							StaticLookupTables: []solana.PublicKey{staticLookupTablePubkey},
 						},
 						Accounts: []chainwriter.Lookup{
+							chainwriter.AccountConstant{
+								Name:       "feepayer",
+								Address:    admin.String(),
+								IsSigner:   false,
+								IsWritable: false,
+							},
 							chainwriter.AccountConstant{
 								Name:       "Constant",
 								Address:    account1.String(),
@@ -482,7 +488,14 @@ func TestChainWriter_SubmitTransaction(t *testing.T) {
 								LookupTableName: "DerivedTable",
 								IncludeIndexes:  []int{0},
 							},
+							chainwriter.AccountConstant{
+								Name:       "systemprogram",
+								Address:    solana.SystemProgramID.String(),
+								IsSigner:   false,
+								IsWritable: false,
+							},
 						},
+						ArgsTransform: nil,
 					},
 				},
 				IDL: testContractIDLJson,