From 1eea24d54ce60fb1c4938b924942efde9c49111f Mon Sep 17 00:00:00 2001
From: Christian Krueger <christian@jito.wtf>
Date: Tue, 17 Dec 2024 16:41:14 -0700
Subject: [PATCH] added set weight

---
 program/src/set_weight.rs | 69 +++++++++++++++++++++++++++------------
 1 file changed, 48 insertions(+), 21 deletions(-)

diff --git a/program/src/set_weight.rs b/program/src/set_weight.rs
index b51c214..9bb4494 100644
--- a/program/src/set_weight.rs
+++ b/program/src/set_weight.rs
@@ -1,50 +1,77 @@
 use jito_bytemuck::AccountDeserialize;
-use jito_jsm_core::loader::{load_signer, load_token_mint};
 use jito_restaking_core::ncn::Ncn;
-use jito_tip_router_core::{error::TipRouterError, weight_table::WeightTable};
+use jito_tip_router_core::{
+    constants::{JTO_MINT, JTO_USD_FEED, MAX_STALE_SLOTS, WEIGHT_PRECISION},
+    error::TipRouterError,
+    weight_table::WeightTable,
+};
 use solana_program::{
     account_info::AccountInfo, clock::Clock, entrypoint::ProgramResult, msg,
     program_error::ProgramError, pubkey::Pubkey, sysvar::Sysvar,
 };
+use switchboard_on_demand::{
+    prelude::rust_decimal::{prelude::ToPrimitive, Decimal},
+    PullFeedAccountData,
+};
 
 /// Updates weight table
-pub fn process_set_weight(
+pub fn process_set_jto_weight(
     program_id: &Pubkey,
     accounts: &[AccountInfo],
     epoch: u64,
-    weight: u128,
 ) -> ProgramResult {
-    let [ncn, weight_table, weight_table_admin, mint, restaking_program] = accounts else {
+    let [ncn, weight_table, jto_usd_feed] = accounts else {
         return Err(ProgramError::NotEnoughAccountKeys);
     };
 
-    Ncn::load(restaking_program.key, ncn, false)?;
-    let ncn_weight_table_admin = {
-        let ncn_data = ncn.data.borrow();
-        let ncn = Ncn::try_from_slice_unchecked(&ncn_data)?;
-        ncn.weight_table_admin
-    };
+    Ncn::load(&jito_restaking_program::id(), ncn, false)?;
 
-    load_signer(weight_table_admin, true)?;
-    load_token_mint(mint)?;
-    WeightTable::load(program_id, weight_table, ncn, ncn_epoch, true)?;
+    WeightTable::load(program_id, weight_table, ncn, epoch, true)?;
 
-    if restaking_program.key.ne(&jito_restaking_program::id()) {
-        msg!("Incorrect restaking program ID");
+    if jto_usd_feed.key.ne(&JTO_USD_FEED) {
+        msg!("Incorrect jto usd feed");
         return Err(ProgramError::InvalidAccountData);
     }
 
-    if ncn_weight_table_admin.ne(weight_table_admin.key) {
-        msg!("Vault update delegations ticket is not at the correct PDA");
-        return Err(TipRouterError::IncorrectWeightTableAdmin.into());
-    }
+    let weight: u128 = {
+        let feed = PullFeedAccountData::parse(jto_usd_feed.data.borrow())
+            .map_err(|_| TipRouterError::BadSwitchboardFeed)?;
+        let price: Decimal = feed.value().ok_or(TipRouterError::BadSwitchboardValue)?;
+
+        let current_slot = Clock::get()?.slot;
+        let stale_slot = {
+            feed.result
+                .slot
+                .checked_add(MAX_STALE_SLOTS)
+                .ok_or(TipRouterError::ArithmeticOverflow)?
+        };
+
+        if current_slot > stale_slot {
+            msg!("Stale feed");
+            return Err(ProgramError::InvalidAccountData);
+        }
+
+        msg!("Oracle Price: {}", price);
+        let weight = price
+            .checked_mul(WEIGHT_PRECISION.into())
+            .ok_or(TipRouterError::ArithmeticOverflow)?
+            .round();
+
+        msg!("Weight: {}", weight);
+
+        weight.to_u128().ok_or(TipRouterError::CastToU128Error)?
+    };
 
     let mut weight_table_data = weight_table.try_borrow_mut_data()?;
     let weight_table_account = WeightTable::try_from_slice_unchecked_mut(&mut weight_table_data)?;
 
     weight_table_account.check_initialized()?;
+    if weight_table_account.finalized() {
+        msg!("Weight table is finalized");
+        return Err(ProgramError::InvalidAccountData);
+    }
 
-    weight_table_account.set_weight(mint.key, weight, Clock::get()?.slot)?;
+    weight_table_account.set_weight(&JTO_MINT, weight, Clock::get()?.slot)?;
 
     Ok(())
 }