Anton
392 lines
15 KiB
Python
392 lines
15 KiB
Python
"""
|
|
IDL Parser module for Solana programs.
|
|
Provides functionality to load and parse Anchor IDL files and decode instruction data.
|
|
"""
|
|
|
|
import json
|
|
import struct
|
|
from typing import Any
|
|
|
|
import base58
|
|
|
|
# Constants for Anchor data layout
|
|
DISCRIMINATOR_SIZE = 8
|
|
PUBLIC_KEY_SIZE = 32
|
|
STRING_LENGTH_PREFIX_SIZE = 4
|
|
ENUM_DISCRIMINATOR_SIZE = 1
|
|
|
|
|
|
class IDLParser:
|
|
"""Parser for automatically decoding instructions using IDL definitions."""
|
|
|
|
# A single source of truth for primitive type information, mapping the type name
|
|
# to its struct format character and size in bytes.
|
|
_PRIMITIVE_TYPE_INFO = {
|
|
# type_name: (format_char, size_in_bytes)
|
|
"u8": ("<B", 1),
|
|
"u16": ("<H", 2),
|
|
"u32": ("<I", 4),
|
|
"u64": ("<Q", 8),
|
|
"i8": ("<b", 1),
|
|
"i16": ("<h", 2),
|
|
"i32": ("<i", 4),
|
|
"i64": ("<q", 8),
|
|
"bool": ("<?", 1),
|
|
"pubkey": (None, PUBLIC_KEY_SIZE),
|
|
"string": (
|
|
None,
|
|
STRING_LENGTH_PREFIX_SIZE,
|
|
), # Min size is for the length prefix
|
|
}
|
|
|
|
def __init__(self, idl_path: str, verbose: bool = False):
|
|
"""
|
|
Initialize the IDL parser.
|
|
|
|
Args:
|
|
idl_path: Path to the IDL JSON file
|
|
verbose: Whether to print debug information during initialization
|
|
"""
|
|
self.verbose = verbose
|
|
with open(idl_path) as f:
|
|
self.idl = json.load(f)
|
|
self.instructions: dict[bytes, dict[str, Any]] = {}
|
|
self.types: dict[str, dict[str, Any]] = {}
|
|
self.instruction_min_sizes: dict[bytes, int] = {}
|
|
self._build_instruction_map()
|
|
self._build_type_map()
|
|
self._calculate_instruction_sizes()
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Public Methods (External API)
|
|
# --------------------------------------------------------------------------
|
|
|
|
def get_instruction_discriminators(self) -> dict[str, bytes]:
|
|
"""Get a mapping of instruction names to their discriminators."""
|
|
return {instr["name"]: disc for disc, instr in self.instructions.items()}
|
|
|
|
def get_instruction_names(self) -> list[str]:
|
|
"""Get a list of all available instruction names."""
|
|
return [instr["name"] for instr in self.instructions.values()]
|
|
|
|
def validate_instruction_data_length(
|
|
self, ix_data: bytes, discriminator: bytes
|
|
) -> bool:
|
|
"""Validate that instruction data meets minimum length requirements."""
|
|
if discriminator not in self.instruction_min_sizes:
|
|
return True # Allow if we don't know the expected size
|
|
|
|
expected_min_size = self.instruction_min_sizes[discriminator]
|
|
actual_size = len(ix_data)
|
|
|
|
if actual_size < expected_min_size:
|
|
instruction_name = self.instructions[discriminator]["name"]
|
|
if self.verbose:
|
|
print(
|
|
f"⚠️ Instruction data for '{instruction_name}' is shorter than the expected minimum "
|
|
f"({actual_size}/{expected_min_size} bytes)."
|
|
)
|
|
return False
|
|
|
|
return True
|
|
|
|
def decode_instruction(
|
|
self, ix_data: bytes, keys: list[bytes], accounts: list[int]
|
|
) -> dict[str, Any] | None:
|
|
"""Decode instruction data using IDL definitions."""
|
|
if len(ix_data) < DISCRIMINATOR_SIZE:
|
|
return None
|
|
|
|
discriminator = ix_data[:DISCRIMINATOR_SIZE]
|
|
if discriminator not in self.instructions:
|
|
return None
|
|
|
|
if not self.validate_instruction_data_length(ix_data, discriminator):
|
|
return None
|
|
|
|
instruction = self.instructions[discriminator]
|
|
data_args = ix_data[DISCRIMINATOR_SIZE:]
|
|
|
|
# Decode instruction arguments
|
|
args = {}
|
|
decode_offset = 0
|
|
for arg in instruction.get("args", []):
|
|
try:
|
|
value, decode_offset = self._decode_type(
|
|
data_args, decode_offset, arg["type"]
|
|
)
|
|
args[arg["name"]] = value
|
|
except Exception as e:
|
|
if self.verbose:
|
|
print(f"❌ Decode error in argument '{arg['name']}': {e}")
|
|
return None
|
|
|
|
# Helper to safely retrieve account public keys
|
|
def get_account_key(index: int) -> str | None:
|
|
if index < len(accounts):
|
|
account_index = accounts[index]
|
|
if account_index < len(keys):
|
|
return base58.b58encode(keys[account_index]).decode("utf-8")
|
|
return None # Return None for invalid indices
|
|
|
|
# Build account info based on instruction definition
|
|
account_info = {}
|
|
instruction_accounts = instruction.get("accounts", [])
|
|
for i, account_def in enumerate(instruction_accounts):
|
|
account_info[account_def["name"]] = get_account_key(i)
|
|
|
|
return {
|
|
"instruction_name": instruction["name"],
|
|
"args": args,
|
|
"accounts": account_info,
|
|
}
|
|
|
|
def decode_account_data(
|
|
self,
|
|
account_data: bytes,
|
|
account_type_name: str,
|
|
skip_discriminator: bool = True,
|
|
) -> dict[str, Any] | None:
|
|
"""
|
|
Decode account data using a specific account type from the IDL.
|
|
|
|
Args:
|
|
account_data: Raw account data bytes.
|
|
account_type_name: Name of the account type in the IDL (e.g., "MyAccount").
|
|
skip_discriminator: Whether to skip the first 8 bytes, which Anchor uses as a
|
|
type discriminator for account data. Set to False if your
|
|
data does not have this prefix.
|
|
|
|
Returns:
|
|
Decoded account data as a dictionary, or None if decoding fails.
|
|
"""
|
|
try:
|
|
if account_type_name not in self.types:
|
|
if self.verbose:
|
|
print(f"Account type '{account_type_name}' not found in IDL")
|
|
return None
|
|
|
|
data = account_data
|
|
if skip_discriminator:
|
|
if len(account_data) < DISCRIMINATOR_SIZE:
|
|
if self.verbose:
|
|
print(
|
|
f"Account data too short to contain a discriminator: {len(account_data)} bytes"
|
|
)
|
|
return None
|
|
data = account_data[DISCRIMINATOR_SIZE:]
|
|
|
|
decoded_data, _ = self._decode_defined_type(data, 0, account_type_name)
|
|
return decoded_data
|
|
|
|
except Exception as e:
|
|
if self.verbose:
|
|
print(f"Error decoding account data for {account_type_name}: {e}")
|
|
return None
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Internal Helper Methods
|
|
# --------------------------------------------------------------------------
|
|
|
|
def _build_instruction_map(self):
|
|
"""Build a map of discriminators to instruction definitions."""
|
|
for instruction in self.idl.get("instructions", []):
|
|
# The discriminator from the JSON IDL is a list of u8 integers.
|
|
discriminator = bytes(instruction["discriminator"])
|
|
self.instructions[discriminator] = instruction
|
|
|
|
def _build_type_map(self):
|
|
"""Build a map of type names to their definitions."""
|
|
for type_def in self.idl.get("types", []):
|
|
self.types[type_def["name"]] = type_def
|
|
|
|
def _calculate_instruction_sizes(self):
|
|
"""Calculate minimum data sizes for each instruction."""
|
|
for discriminator, instruction in self.instructions.items():
|
|
try:
|
|
min_size = DISCRIMINATOR_SIZE
|
|
for arg in instruction.get("args", []):
|
|
min_size += self._calculate_type_min_size(arg["type"])
|
|
self.instruction_min_sizes[discriminator] = min_size
|
|
if self.verbose and instruction["name"] == "initialize":
|
|
print(f"📏 Initialize instruction min size: {min_size} bytes")
|
|
except Exception as e:
|
|
if self.verbose:
|
|
print(f"⚠️ Could not calculate size for {instruction['name']}: {e}")
|
|
self.instruction_min_sizes[discriminator] = DISCRIMINATOR_SIZE
|
|
|
|
def _calculate_type_min_size(self, type_def: str | dict) -> int:
|
|
"""Calculate minimum size in bytes for a type definition."""
|
|
if isinstance(type_def, str):
|
|
return self._get_primitive_size(type_def)
|
|
|
|
if isinstance(type_def, dict):
|
|
if "defined" in type_def:
|
|
type_name = self._get_defined_type_name(type_def)
|
|
return self._calculate_defined_type_min_size(type_name)
|
|
if "array" in type_def:
|
|
element_type, array_length = type_def["array"]
|
|
element_size = self._calculate_type_min_size(element_type)
|
|
return element_size * array_length
|
|
|
|
raise ValueError(
|
|
f"Invalid or unknown type definition for size calculation: {type_def}"
|
|
)
|
|
|
|
def _get_primitive_size(self, type_name: str) -> int:
|
|
"""Get size in bytes for primitive types from the central map."""
|
|
info = self._PRIMITIVE_TYPE_INFO.get(type_name)
|
|
return info[1] if info else 0
|
|
|
|
def _get_defined_type_name(self, type_def: dict[str, Any]) -> str:
|
|
"""Extracts the type name from a 'defined' type, handling old and new IDL formats."""
|
|
defined_value = type_def["defined"]
|
|
# New format: {'defined': {'name': 'MyType'}}
|
|
# Old format: {'defined': 'MyType'}
|
|
return (
|
|
defined_value["name"] if isinstance(defined_value, dict) else defined_value
|
|
)
|
|
|
|
def _calculate_defined_type_min_size(self, type_name: str) -> int:
|
|
"""Calculate minimum size for user-defined types (structs and enums)."""
|
|
if type_name not in self.types:
|
|
raise ValueError(f"Unknown defined type: {type_name}")
|
|
|
|
type_def = self.types[type_name]["type"]
|
|
|
|
if type_def["kind"] == "struct":
|
|
return sum(
|
|
self._calculate_type_min_size(field["type"])
|
|
for field in type_def["fields"]
|
|
)
|
|
|
|
if type_def["kind"] == "enum":
|
|
# The size of an enum is its discriminator plus the size of its LARGEST variant,
|
|
# as the data layout must accommodate any possible variant.
|
|
max_variant_size = 0
|
|
for variant in type_def["variants"]:
|
|
variant_size = 0
|
|
for field in variant.get("fields", []):
|
|
# A field can be a type string/dict (tuple variant) or a dict with a 'type' key (struct variant)
|
|
field_type = field["type"] if isinstance(field, dict) else field
|
|
variant_size += self._calculate_type_min_size(field_type)
|
|
max_variant_size = max(max_variant_size, variant_size)
|
|
return ENUM_DISCRIMINATOR_SIZE + max_variant_size
|
|
|
|
raise ValueError(
|
|
f"Unsupported type kind for size calculation: {type_def['kind']}"
|
|
)
|
|
|
|
def _decode_type(
|
|
self, data: bytes, offset: int, type_def: str | dict
|
|
) -> tuple[Any, int]:
|
|
"""Decode a value based on its type definition."""
|
|
if isinstance(type_def, str):
|
|
return self._decode_primitive(data, offset, type_def)
|
|
|
|
if isinstance(type_def, dict):
|
|
if "defined" in type_def:
|
|
type_name = self._get_defined_type_name(type_def)
|
|
return self._decode_defined_type(data, offset, type_name)
|
|
if "array" in type_def:
|
|
return self._decode_array(data, offset, type_def["array"])
|
|
|
|
raise ValueError(f"Invalid or unknown type definition for decoding: {type_def}")
|
|
|
|
def _decode_array(
|
|
self, data: bytes, offset: int, array_def: list
|
|
) -> tuple[list[Any], int]:
|
|
"""Decode fixed-size array types."""
|
|
element_type, array_length = array_def
|
|
array_data = []
|
|
for _ in range(array_length):
|
|
value, offset = self._decode_type(data, offset, element_type)
|
|
array_data.append(value)
|
|
return array_data, offset
|
|
|
|
def _decode_primitive(
|
|
self, data: bytes, offset: int, type_name: str
|
|
) -> tuple[Any, int]:
|
|
"""Decode primitive types."""
|
|
if type_name not in self._PRIMITIVE_TYPE_INFO:
|
|
raise ValueError(f"Unknown primitive type: {type_name}")
|
|
|
|
if type_name == "string":
|
|
length = struct.unpack_from("<I", data, offset)[0]
|
|
offset += STRING_LENGTH_PREFIX_SIZE
|
|
value = data[offset : offset + length].decode("utf-8")
|
|
return value, offset + length
|
|
|
|
if type_name == "pubkey":
|
|
end = offset + PUBLIC_KEY_SIZE
|
|
value = base58.b58encode(data[offset:end]).decode("utf-8")
|
|
return value, end
|
|
|
|
# Handle all numeric and bool types from the map
|
|
fmt, size = self._PRIMITIVE_TYPE_INFO[type_name]
|
|
value = struct.unpack_from(fmt, data, offset)[0]
|
|
return value, offset + size
|
|
|
|
def _decode_defined_type(
|
|
self, data: bytes, offset: int, type_name: str
|
|
) -> tuple[dict[str, Any], int]:
|
|
"""Decode user-defined types (structs and enums)."""
|
|
if type_name not in self.types:
|
|
raise ValueError(f"Unknown defined type: {type_name}")
|
|
|
|
type_def = self.types[type_name]["type"]
|
|
|
|
if type_def["kind"] == "struct":
|
|
struct_data = {}
|
|
for field in type_def["fields"]:
|
|
value, offset = self._decode_type(data, offset, field["type"])
|
|
struct_data[field["name"]] = value
|
|
return struct_data, offset
|
|
|
|
if type_def["kind"] == "enum":
|
|
variant_index = struct.unpack_from("<B", data, offset)[0]
|
|
offset += ENUM_DISCRIMINATOR_SIZE
|
|
|
|
variants = type_def["variants"]
|
|
if variant_index >= len(variants):
|
|
raise ValueError(
|
|
f"Invalid enum variant index {variant_index} for type {type_name}"
|
|
)
|
|
|
|
variant = variants[variant_index]
|
|
result = {"variant": variant["name"]}
|
|
variant_fields = variant.get("fields", [])
|
|
|
|
if variant_fields:
|
|
# Check if it's a struct variant (fields are dicts) or tuple variant (fields are strings/dicts)
|
|
if isinstance(variant_fields[0], dict):
|
|
struct_data = {}
|
|
for field in variant_fields:
|
|
value, offset = self._decode_type(data, offset, field["type"])
|
|
struct_data[field["name"]] = value
|
|
result["data"] = struct_data
|
|
else: # Tuple variant
|
|
tuple_data = []
|
|
for field_type in variant_fields:
|
|
value, offset = self._decode_type(data, offset, field_type)
|
|
tuple_data.append(value)
|
|
result["data"] = tuple_data
|
|
|
|
return result, offset
|
|
|
|
raise ValueError(f"Unsupported type kind for decoding: {type_def['kind']}")
|
|
|
|
|
|
def load_idl_parser(idl_path: str, verbose: bool = False) -> IDLParser:
|
|
"""
|
|
Convenience function to load an IDL parser.
|
|
|
|
Args:
|
|
idl_path: Path to the IDL JSON file
|
|
verbose: Whether to print debug information
|
|
|
|
Returns:
|
|
Initialized IDLParser instance
|
|
"""
|
|
return IDLParser(idl_path, verbose)
|