From bbf72e6f048f03ffb26b5443d811c523998629b3 Mon Sep 17 00:00:00 2001
From: blockiosaurus <blockiosaurus@gmail.com>
Date: Fri, 1 Mar 2024 14:48:03 -0500
Subject: [PATCH] Making fetch_plugin generic for program.

---
 programs/mpl-core/src/plugins/utils.rs        | 31 ++++++++++++++-----
 .../mpl-core/src/state/update_authority.rs    |  8 +++--
 2 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/programs/mpl-core/src/plugins/utils.rs b/programs/mpl-core/src/plugins/utils.rs
index 07c38e50..a411f590 100644
--- a/programs/mpl-core/src/plugins/utils.rs
+++ b/programs/mpl-core/src/plugins/utils.rs
@@ -79,11 +79,11 @@ pub fn assert_plugins_initialized(account: &AccountInfo) -> ProgramResult {
 }
 
 /// Fetch the plugin from the registry.
-pub fn fetch_plugin(
+pub fn fetch_plugin<T: DataBlob + SolanaAccount, U: BorshDeserialize>(
     account: &AccountInfo,
     plugin_type: PluginType,
-) -> Result<(Vec<Authority>, Plugin, usize), ProgramError> {
-    let asset = Asset::load(account, 0)?;
+) -> Result<(Vec<Authority>, U, usize), ProgramError> {
+    let asset = T::load(account, 0)?;
 
     let header = PluginHeader::load(account, asset.get_size())?;
     let PluginRegistry { registry, .. } =
@@ -98,19 +98,30 @@ pub fn fetch_plugin(
     // Deserialize the plugin.
     let plugin = Plugin::deserialize(&mut &(*account.data).borrow()[registry_record.offset..])?;
 
+    if PluginType::from(&plugin) != plugin_type {
+        return Err(MplCoreError::PluginNotFound.into());
+    }
+
+    let inner = U::deserialize(
+        &mut &(*account.data).borrow()[registry_record
+            .offset
+            .checked_add(1)
+            .ok_or(MplCoreError::NumericalOverflow)?..],
+    )?;
+
     // Return the plugin and its authorities.
     Ok((
         registry_record.authorities.clone(),
-        plugin,
+        inner,
         registry_record.offset,
     ))
 }
 
 /// Fetch the collection plugin from the registry.
-pub fn fetch_collection_plugin(
+pub fn fetch_collection_plugin<T: BorshDeserialize>(
     account: &AccountInfo,
     plugin_type: PluginType,
-) -> Result<(Vec<Authority>, Plugin, usize), ProgramError> {
+) -> Result<(Vec<Authority>, T, usize), ProgramError> {
     let collection = CollectionData::load(account, 0)?;
 
     let header = PluginHeader::load(account, collection.get_size())?;
@@ -126,10 +137,16 @@ pub fn fetch_collection_plugin(
     // Deserialize the plugin.
     let plugin = Plugin::deserialize(&mut &(*account.data).borrow()[registry_record.offset..])?;
 
+    if PluginType::from(&plugin) != plugin_type {
+        return Err(MplCoreError::PluginNotFound.into());
+    }
+
+    let inner = T::deserialize(&mut &(*account.data).borrow()[registry_record.offset..])?;
+
     // Return the plugin and its authorities.
     Ok((
         registry_record.authorities.clone(),
-        plugin,
+        inner,
         registry_record.offset,
     ))
 }
diff --git a/programs/mpl-core/src/state/update_authority.rs b/programs/mpl-core/src/state/update_authority.rs
index 8c2aefbc..8d02a315 100644
--- a/programs/mpl-core/src/state/update_authority.rs
+++ b/programs/mpl-core/src/state/update_authority.rs
@@ -8,7 +8,7 @@ use crate::{
         BurnAccounts, CompressAccounts, CreateAccounts, DecompressAccounts, TransferAccounts,
         UpdateAccounts,
     },
-    plugins::{fetch_collection_plugin, CheckResult, PluginType, ValidationResult},
+    plugins::{fetch_plugin, CheckResult, PluginType, UpdateDelegate, ValidationResult},
     processor::CreateArgs,
     state::{Authority, CollectionData, SolanaAccount},
     utils::assert_collection_authority,
@@ -68,8 +68,10 @@ impl UpdateAuthority {
                     None => ctx.payer,
                 };
 
-                let maybe_update_delegate =
-                    fetch_collection_plugin(collection_info, PluginType::UpdateDelegate);
+                let maybe_update_delegate = fetch_plugin::<CollectionData, UpdateDelegate>(
+                    collection_info,
+                    PluginType::UpdateDelegate,
+                );
 
                 if let Ok((mut authorities, _, _)) = maybe_update_delegate {
                     authorities.push(Authority::UpdateAuthority);