spl: Improve Token 2022 support (#503)

michaelhly · Feb 6, 2025 · 39f185a · 39f185a
1 parent ab9623b
commit 39f185a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/src/spl/token/core.py b/src/spl/token/core.py
@@ -357,7 +357,9 @@ def _create_mint_info(self, info: GetAccountInfoResp) -> MintInfo:
             raise AttributeError(f"Invalid mint owner: {owner}")
 
         bytes_data = value.data
-        if len(bytes_data) != MINT_LAYOUT.sizeof():
+
+        # TODO: This condition is necessary but not sufficient
+        if len(bytes_data) < MINT_LAYOUT.sizeof():
             raise ValueError("Invalid mint size")
 
         decoded_data = MINT_LAYOUT.parse(bytes_data)
@@ -380,7 +382,9 @@ def _create_account_info(self, info: GetAccountInfoResp) -> AccountInfo:
             raise AttributeError("Invalid account owner")
 
         bytes_data = value.data
-        if len(bytes_data) != ACCOUNT_LAYOUT.sizeof():
+
+        # TODO: This condition is necessary but not sufficient
+        if len(bytes_data) < ACCOUNT_LAYOUT.sizeof():
             raise ValueError("Invalid account size")
 
         decoded_data = ACCOUNT_LAYOUT.parse(bytes_data)

diff --git a/src/spl/token/instructions.py b/src/spl/token/instructions.py
@@ -1271,21 +1271,25 @@ def create_associated_token_account(
     )
 
 
-def create_idempotent_associated_token_account(payer: Pubkey, owner: Pubkey, mint: Pubkey) -> Instruction:
+def create_idempotent_associated_token_account(
+    payer: Pubkey, owner: Pubkey, mint: Pubkey, token_program_id: Pubkey = TOKEN_PROGRAM_ID
+) -> Instruction:
     """Creates an associated token account for the given address/token mint if it not exists.
 
     Returns:
         The instruction to create the associated token account.
     """
-    associated_token_address = get_associated_token_address(owner, mint)
+    if token_program_id not in [TOKEN_PROGRAM_ID, TOKEN_2022_PROGRAM_ID]:
+        raise ValueError("token_program_id must be one of TOKEN_PROGRAM_ID or TOKEN_2022_PROGRAM_ID.")
+    associated_token_address = get_associated_token_address(owner, mint, token_program_id)
     return Instruction(
         accounts=[
             AccountMeta(pubkey=payer, is_signer=True, is_writable=True),
             AccountMeta(pubkey=associated_token_address, is_signer=False, is_writable=True),
             AccountMeta(pubkey=owner, is_signer=False, is_writable=False),
             AccountMeta(pubkey=mint, is_signer=False, is_writable=False),
             AccountMeta(pubkey=SYS_PROGRAM_ID, is_signer=False, is_writable=False),
-            AccountMeta(pubkey=TOKEN_PROGRAM_ID, is_signer=False, is_writable=False),
+            AccountMeta(pubkey=token_program_id, is_signer=False, is_writable=False),
         ],
         program_id=ASSOCIATED_TOKEN_PROGRAM_ID,
         data=bytes([1]),