From d50b871a9403309b3911a913a9d1fccb12aefd44 Mon Sep 17 00:00:00 2001
From: Samuel Vanderwaal <samuel.vanderwaal@gmail.com>
Date: Sat, 5 Nov 2022 08:45:03 +0000
Subject: [PATCH] check edition is derived from mint for initialize

---
 .../src/instructions/set_collection.rs        |  6 ++++-
 candy-machine-core/program/src/utils.rs       | 24 +++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/candy-machine-core/program/src/instructions/set_collection.rs b/candy-machine-core/program/src/instructions/set_collection.rs
index 34aba849cc..627d86ba44 100644
--- a/candy-machine-core/program/src/instructions/set_collection.rs
+++ b/candy-machine-core/program/src/instructions/set_collection.rs
@@ -6,7 +6,10 @@ use mpl_token_metadata::{
 };
 use solana_program::program::{invoke, invoke_signed};
 
-use crate::{cmp_pubkeys, constants::AUTHORITY_SEED, CandyError, CandyMachine};
+use crate::{
+    cmp_pubkeys, constants::AUTHORITY_SEED, utils::assert_edition_from_mint, CandyError,
+    CandyMachine,
+};
 
 pub fn set_collection(ctx: Context<SetCollection>) -> Result<()> {
     let accounts = ctx.accounts;
@@ -93,6 +96,7 @@ pub fn approve_collection_authority_helper(
     }
 
     assert_master_edition(&collection_data, &collection_master_edition)?;
+    assert_edition_from_mint(&collection_master_edition, &collection_mint)?;
 
     let approve_collection_authority_ix = approve_collection_authority(
         token_metadata_program.key(),
diff --git a/candy-machine-core/program/src/utils.rs b/candy-machine-core/program/src/utils.rs
index aed8773379..1808febce8 100644
--- a/candy-machine-core/program/src/utils.rs
+++ b/candy-machine-core/program/src/utils.rs
@@ -1,5 +1,10 @@
 use anchor_lang::prelude::*;
 use arrayref::array_ref;
+use mpl_token_metadata::{
+    error::MetadataError,
+    state::{EDITION, PREFIX},
+    utils::assert_derivation,
+};
 use solana_program::{
     account_info::AccountInfo,
     program::invoke,
@@ -8,6 +13,7 @@ use solana_program::{
     pubkey::{Pubkey, PUBKEY_BYTES},
     system_instruction,
 };
+use std::result::Result as StdResult;
 
 use crate::{
     constants::{HIDDEN_SECTION, NULL_STRING, REPLACEMENT_INDEX, REPLACEMENT_INDEX_INCREMENT},
@@ -80,6 +86,24 @@ pub fn replace_patterns(value: String, index: usize) -> String {
     mutable
 }
 
+pub fn assert_edition_from_mint(
+    edition_account: &AccountInfo,
+    mint_account: &AccountInfo,
+) -> StdResult<(), ProgramError> {
+    assert_derivation(
+        &mpl_token_metadata::id(),
+        edition_account,
+        &[
+            PREFIX.as_bytes(),
+            mpl_token_metadata::id().as_ref(),
+            mint_account.key().as_ref(),
+            EDITION.as_bytes(),
+        ],
+    )
+    .map_err(|_| MetadataError::CollectionMasterEditionAccountInvalid)?;
+    Ok(())
+}
+
 #[cfg(test)]
 pub mod tests {
     use super::*;