From 734fc4a5dd10207a5855477219e473277b313d8c Mon Sep 17 00:00:00 2001
From: Evan Batsell <ebatsell@gmail.com>
Date: Wed, 13 Nov 2024 16:03:04 -0500
Subject: [PATCH] pr

---
 core/src/tracked_mints.rs    |  4 +++-
 program/src/register_mint.rs | 12 ++++++++----
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/core/src/tracked_mints.rs b/core/src/tracked_mints.rs
index e1ea35c..16d9402 100644
--- a/core/src/tracked_mints.rs
+++ b/core/src/tracked_mints.rs
@@ -12,6 +12,7 @@ use crate::{discriminators::Discriminators, error::TipRouterError};
 pub struct MintEntry {
     st_mint: Pubkey,
     vault_index: PodU64,
+    reserved: [u8; 32],
 }
 
 impl MintEntry {
@@ -19,6 +20,7 @@ impl MintEntry {
         Self {
             st_mint: mint,
             vault_index: PodU64::from(vault_index),
+            reserved: [0; 32],
         }
     }
 
@@ -38,7 +40,7 @@ impl Default for MintEntry {
 pub struct TrackedMints {
     pub ncn: Pubkey,
     pub bump: u8,
-    pub reserved: [u8; 7],
+    pub reserved: [u8; 7], // TODO extend to 127; figure out serde issue
     pub st_mint_list: [MintEntry; 16], // TODO extend to 64; figure out serde issue
 }
 
diff --git a/program/src/register_mint.rs b/program/src/register_mint.rs
index 8fe2b0e..c3b3f70 100644
--- a/program/src/register_mint.rs
+++ b/program/src/register_mint.rs
@@ -48,13 +48,17 @@ pub fn process_register_mint(program_id: &Pubkey, accounts: &[AccountInfo]) -> P
     // Once tracked_mints.mint_count() == ncn.vault_count, the weight table can be initialized
     // Once the weight table is initialized, you can't add any more mints
     if weight_table.owner.eq(&system_program::ID) {
+        let expected_pubkey = WeightTable::find_program_address(program_id, ncn.key, ncn_epoch).0;
+        if weight_table.key.ne(&expected_pubkey) {
+            msg!("Weight table incorrect PDA");
+            return Err(ProgramError::InvalidAccountData);
+        }
         load_system_account(weight_table, false)?;
-    } else if weight_table.owner.eq(program_id) {
+    }
+
+    if weight_table.owner.eq(program_id) {
         WeightTable::load(program_id, weight_table, ncn, ncn_epoch, false)?;
         return Err(TipRouterError::TrackedMintsLocked.into());
-    } else {
-        msg!("Weight table account is not owned by this program or the system program");
-        return Err(ProgramError::InvalidAccountOwner);
     }
 
     // Verify tickets are active