Add Codama IDL type (#85)

* Add Codama IDL type

* Refactor

include-idl-cli/src/main.rs

@@ -1,7 +1,7 @@
 use std::path::PathBuf;
 
 // use goblin::error::Result;
-use include_idl::parse::{parse_idl_from_program_binary, IdlType};
+use include_idl::parse::parse_idl_from_program_binary;
 
 use clap::{Error, Parser, Subcommand};
 
@@ -17,24 +17,21 @@ enum Commands {
     Parse {
         /// Read IDL from a solana program binary
         path: PathBuf,
-        idl_type: IdlType,
     },
 }
 
-// This example uses ArgEnum, so this might not be necessary.
-
 pub fn main() -> Result<(), Error> {
     let cli = Cli::parse();
 
     match &cli.command {
-        Some(Commands::Parse { path, idl_type }) => {
+        Some(Commands::Parse { path }) => {
             let buffer = std::fs::read(path).expect("Could not read file.");
-            if let Ok(idl) = parse_idl_from_program_binary(&buffer, idl_type.clone()) {
-                println!("        Program IDL");
+            if let Ok((idl_type, idl_data)) = parse_idl_from_program_binary(&buffer) {
+                println!("Program IDL ({})", idl_type);
                 println!("============================");
-                println!("{}", idl);
+                println!("{}", idl_data);
             } else {
-                println!("Could not find {:?} IDL in program binary", idl_type);
+                println!("Could not find IDL in program binary");
             }
         }
         None => {}

include-idl/src/lib.rs

@@ -1,4 +1,3 @@
-#[cfg(feature = "parse")]
 pub mod parse;
 
 #[cfg(feature = "shrink")]
@@ -9,26 +8,21 @@ pub use shrink::compress_idl;
 
 #[macro_export]
 macro_rules! include_idl {
-    ($s:expr) => {
+    ($type:path, $file:expr) => {
         #[cfg_attr(
             any(target_arch = "sbf", target_arch = "bpf"),
-            link_section = ".solana.idl"
+            link_section = ".idl.type"
         )]
         #[allow(dead_code)]
         #[no_mangle]
-        pub static IDL_BYTES: &[u8] = include_bytes!($s);
-    };
-}
+        pub static IDL_TYPE: &[u8] = $type.as_str().as_bytes();
 
-#[macro_export]
-macro_rules! include_kinobi_idl {
-    ($s:expr) => {
         #[cfg_attr(
             any(target_arch = "sbf", target_arch = "bpf"),
-            link_section = ".kinobi.idl"
+            link_section = ".idl.data"
         )]
         #[allow(dead_code)]
         #[no_mangle]
-        pub static KINOBI_IDL_BYTES: &[u8] = include_bytes!($s);
+        pub static IDL_BYTES: &[u8] = include_bytes!($file);
     };
 }

include-idl/src/parse.rs

@@ -1,75 +1,110 @@
 #[cfg(feature = "parse")]
 use {
-    flate2::bufread::ZlibDecoder, goblin::elf::Elf, serde_json::Value, std::fmt, std::io::Read,
-    std::str::FromStr,
+    flate2::bufread::ZlibDecoder,
+    goblin::elf::Elf,
+    serde_json::Value,
+    std::io::Read,
+    std::str::{from_utf8, FromStr},
 };
 
+/// Name of the section containing the IDL type value.
+pub const IDL_TYPE_SECTION: &str = ".idl.type";
+
+/// Name of the section containing the IDL data.
+pub const IDL_DATA_SECTION: &str = ".idl.data";
+
+const ANCHOR_IDL_TYPE: &str = "anchor";
+const CODAMA_IDL_TYPE: &str = "codama";
+
+/// Defines the IDL type.
 #[derive(Clone, Debug)]
 pub enum IdlType {
     Anchor,
-    Kinobi,
+    Codama,
+}
+
+impl IdlType {
+    pub const fn as_str(&self) -> &'static str {
+        match self {
+            IdlType::Anchor => ANCHOR_IDL_TYPE,
+            IdlType::Codama => CODAMA_IDL_TYPE,
+        }
+    }
 }
 
-impl fmt::Display for IdlType {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+impl std::fmt::Display for IdlType {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
-            IdlType::Anchor => write!(f, "Anchor"),
-            IdlType::Kinobi => write!(f, "Kinobi"),
+            IdlType::Anchor => write!(f, "{ANCHOR_IDL_TYPE}"),
+            IdlType::Codama => write!(f, "{CODAMA_IDL_TYPE}"),
         }
     }
 }
 
-impl FromStr for IdlType {
+impl std::str::FromStr for IdlType {
     type Err = &'static str;
 
     fn from_str(s: &str) -> Result<Self, &'static str> {
         match s.to_lowercase().as_str() {
-            "anchor" => Ok(IdlType::Anchor),
-            "kinobi" => Ok(IdlType::Kinobi),
+            ANCHOR_IDL_TYPE => Ok(IdlType::Anchor),
+            CODAMA_IDL_TYPE => Ok(IdlType::Codama),
             _ => Err("Invalid IDL type"),
         }
     }
 }
 
-fn get_section_name(idl_type: IdlType) -> String {
-    match idl_type {
-        IdlType::Anchor => ".solana.idl".to_string(),
-        IdlType::Kinobi => ".kinobi.idl".to_string(),
-    }
-}
-
-pub fn parse_idl_from_program_binary(
-    buffer: &[u8],
-    idl_type: IdlType,
-) -> goblin::error::Result<Value> {
+#[cfg(feature = "parse")]
+pub fn parse_idl_from_program_binary(buffer: &[u8]) -> goblin::error::Result<(IdlType, Value)> {
     let elf = Elf::parse(buffer)?;
 
-    let section_name = get_section_name(idl_type);
+    let mut idl_type: Option<IdlType> = None;
+    let mut idl_json: Option<Value> = None;
 
-    // Iterate over section headers and print information
+    // Iterate over section headers and retrieve the IDL data.
     for sh in &elf.section_headers {
         let name = elf.shdr_strtab.get_at(sh.sh_name).unwrap_or("<invalid>");
-        if name == section_name {
-            // Get offset of .solana.idl section data
-            let offset = sh.sh_offset as usize;
-
-            // Get offset & size of the compressed IDL bytes
-            let _data_loc = &buffer[offset + 4..offset + 8];
-            let data_loc = u32::from_le_bytes(_data_loc.try_into().unwrap());
-            let _data_size = &buffer[offset + 8..offset + 16];
-            let data_size = u64::from_le_bytes(_data_size.try_into().unwrap());
-
-            let compressed_data =
-                &buffer[data_loc as usize..(data_loc as u64 + data_size) as usize];
-            let mut d = ZlibDecoder::new(compressed_data);
-            let mut decompressed_data = Vec::new();
-            d.read_to_end(&mut decompressed_data).unwrap();
-
-            let json: Value = serde_json::from_slice(&decompressed_data).unwrap();
-            return Ok(json);
+
+        match name {
+            IDL_DATA_SECTION => {
+                let (location, size) = get_section_data_offset(buffer, sh.sh_offset as usize);
+
+                let slice = &buffer[location..location + size];
+                let mut compressed_data = ZlibDecoder::new(slice);
+                let mut data = Vec::new();
+                compressed_data.read_to_end(&mut data).unwrap();
+
+                idl_json = Some(serde_json::from_slice(&data).unwrap());
+            }
+            IDL_TYPE_SECTION => {
+                let (location, size) = get_section_data_offset(buffer, sh.sh_offset as usize);
+                let slice = &buffer[location..location + size];
+
+                idl_type = Some(IdlType::from_str(from_utf8(slice).unwrap()).unwrap());
+            }
+            // Ignore other sections.
+            _ => (),
         }
     }
-    Err(goblin::error::Error::Malformed(
-        "Could not find .solana.idl section".to_string(),
-    ))
+
+    if idl_type.is_some() && idl_json.is_some() {
+        #[allow(clippy::unnecessary_unwrap)]
+        Ok((idl_type.unwrap(), idl_json.unwrap()))
+    } else {
+        // Returns an error if we could not find the IDL information.
+        Err(goblin::error::Error::Malformed(
+            "Could not find .idl.* sections".to_string(),
+        ))
+    }
+}
+
+#[cfg(feature = "parse")]
+#[inline(always)]
+fn get_section_data_offset(buffer: &[u8], offset: usize) -> (usize, usize) {
+    let slice = &buffer[offset + 4..offset + 8];
+    let location = u32::from_le_bytes(slice.try_into().unwrap());
+
+    let slice = &buffer[offset + 8..offset + 16];
+    let size = u64::from_le_bytes(slice.try_into().unwrap());
+
+    (location as usize, size as usize)
 }

programs/asset/program/src/lib.rs

@@ -12,18 +12,12 @@ pub mod instruction;
 pub mod processor;
 pub mod utils;
 
+use include_idl::{include_idl, parse::IdlType};
 pub use solana_program;
-
-#[cfg(not(feature = "no-entrypoint"))]
-use include_idl::include_kinobi_idl;
-
-#[cfg(not(feature = "no-entrypoint"))]
 use solana_security_txt::security_txt;
 
-#[cfg(not(feature = "no-entrypoint"))]
-include_kinobi_idl!(concat!(env!("OUT_DIR"), "/kinobi.idl.zip"));
+include_idl!(IdlType::Codama, concat!(env!("OUT_DIR"), "/kinobi.idl.zip"));
 
-#[cfg(not(feature = "no-entrypoint"))]
 security_txt! {
     // Required fields
     name: "Nifty Asset",