From b7c3d4c58077a6a9dc2cabbb805491cf8395b303 Mon Sep 17 00:00:00 2001
From: blockiosaurus <blockiosaurus@gmail.com>
Date: Sat, 9 Mar 2024 16:29:18 -0500
Subject: [PATCH] Add validation checks to remove_plugin.

---
 programs/mpl-core/src/plugins/utils.rs        | 24 ++++++
 .../mpl-core/src/processor/remove_plugin.rs   | 73 +++++++++++++++----
 2 files changed, 81 insertions(+), 16 deletions(-)

diff --git a/programs/mpl-core/src/plugins/utils.rs b/programs/mpl-core/src/plugins/utils.rs
index fa0826e2..1ade83b3 100644
--- a/programs/mpl-core/src/plugins/utils.rs
+++ b/programs/mpl-core/src/plugins/utils.rs
@@ -138,6 +138,30 @@ pub fn fetch_plugin<T: DataBlob + SolanaAccount, U: BorshDeserialize>(
     ))
 }
 
+/// Fetch the plugin from the registry.
+pub fn fetch_wrapped_plugin<T: DataBlob + SolanaAccount>(
+    account: &AccountInfo,
+    plugin_type: PluginType,
+) -> Result<(Authority, Plugin), ProgramError> {
+    let asset = T::load(account, 0)?;
+
+    let header = PluginHeader::load(account, asset.get_size())?;
+    let PluginRegistry { registry, .. } =
+        PluginRegistry::load(account, header.plugin_registry_offset)?;
+
+    // Find the plugin in the registry.
+    let registry_record = registry
+        .iter()
+        .find(|record| record.plugin_type == plugin_type)
+        .ok_or(MplCoreError::PluginNotFound)?;
+
+    // Deserialize the plugin.
+    let plugin = Plugin::deserialize(&mut &(*account.data).borrow()[registry_record.offset..])?;
+
+    // Return the plugin and its authorities.
+    Ok((registry_record.authority.clone(), plugin))
+}
+
 /// Fetch the plugin registry.
 pub fn fetch_plugins(account: &AccountInfo) -> Result<Vec<RegistryRecord>, ProgramError> {
     let asset = Asset::load(account, 0)?;
diff --git a/programs/mpl-core/src/processor/remove_plugin.rs b/programs/mpl-core/src/processor/remove_plugin.rs
index 96bb2ff6..3fd0682d 100644
--- a/programs/mpl-core/src/processor/remove_plugin.rs
+++ b/programs/mpl-core/src/processor/remove_plugin.rs
@@ -5,9 +5,12 @@ use solana_program::{account_info::AccountInfo, entrypoint::ProgramResult, msg};
 use crate::{
     error::MplCoreError,
     instruction::accounts::{RemoveCollectionPluginAccounts, RemovePluginAccounts},
-    plugins::{delete_plugin, PluginType},
-    state::{Asset, Authority, Collection, Key},
-    utils::{fetch_core_data, load_key, resolve_payer, resolve_to_authority},
+    plugins::{delete_plugin, fetch_wrapped_plugin, Plugin, PluginType},
+    state::{Asset, Authority, Collection, DataBlob, Key},
+    utils::{
+        fetch_core_data, load_key, resolve_payer, resolve_to_authority, validate_asset_permissions,
+        validate_collection_permissions,
+    },
 };
 
 #[repr(C)]
@@ -42,16 +45,32 @@ pub(crate) fn remove_plugin<'a>(
     let authority_type =
         resolve_to_authority(ctx.accounts.authority, ctx.accounts.collection, &asset)?;
 
-    delete_plugin(
+    let (_, plugin_to_remove) =
+        fetch_wrapped_plugin::<Asset>(ctx.accounts.asset, args.plugin_type)?;
+
+    // Validate asset permissions.
+    let _ = validate_asset_permissions(
+        ctx.accounts.authority,
+        ctx.accounts.asset,
+        ctx.accounts.collection,
+        None,
+        Some(&plugin_to_remove),
+        Asset::check_add_plugin,
+        Collection::check_add_plugin,
+        PluginType::check_add_plugin,
+        Asset::validate_add_plugin,
+        Collection::validate_add_plugin,
+        Plugin::validate_add_plugin,
+    )?;
+
+    process_remove_plugin(
         &args.plugin_type,
         &asset,
         &authority_type,
         ctx.accounts.asset,
         payer,
         ctx.accounts.system_program,
-    )?;
-
-    process_remove_plugin()
+    )
 }
 
 #[repr(C)]
@@ -78,23 +97,45 @@ pub(crate) fn remove_collection_plugin<'a>(
         return Err(MplCoreError::PluginNotFound.into());
     }
 
-    if ctx.accounts.authority.key != &collection.update_authority {
-        return Err(MplCoreError::InvalidAuthority.into());
-    }
+    let (_, plugin_to_remove) =
+        fetch_wrapped_plugin::<Collection>(ctx.accounts.collection, args.plugin_type)?;
 
-    delete_plugin(
+    // Validate collection permissions.
+    let _ = validate_collection_permissions(
+        ctx.accounts.authority,
+        ctx.accounts.collection,
+        Some(&plugin_to_remove),
+        Collection::check_add_plugin,
+        PluginType::check_add_plugin,
+        Collection::validate_add_plugin,
+        Plugin::validate_add_plugin,
+    )?;
+
+    process_remove_plugin(
         &args.plugin_type,
         &collection,
         &Authority::UpdateAuthority,
         ctx.accounts.collection,
         payer,
         ctx.accounts.system_program,
-    )?;
-
-    process_remove_plugin()
+    )
 }
 
 //TODO
-fn process_remove_plugin() -> ProgramResult {
-    Ok(())
+fn process_remove_plugin<'a, T: DataBlob>(
+    plugin_type: &PluginType,
+    core: &T,
+    authority_type: &Authority,
+    account: &AccountInfo<'a>,
+    payer: &AccountInfo<'a>,
+    system_program: &AccountInfo<'a>,
+) -> ProgramResult {
+    delete_plugin(
+        plugin_type,
+        core,
+        authority_type,
+        account,
+        payer,
+        system_program,
+    )
 }