Merge pull request #19 from drift-labs/ahzam/add-oracle-override

Add pyth price overrides to margin check

Author: 0xahzam

src/margin.rs

@@ -85,7 +85,19 @@ pub fn calculate_simplified_margin_requirement(
         }
 
         let spot_market = market_state.get_spot_market(spot_position.market_index);
-        let oracle_price = market_state.get_spot_oracle_price(spot_position.market_index);
+        let oracle_price = match market_state.get_spot_pyth_price(spot_position.market_index) {
+            Some(pyth_price) => {
+                let oracle = market_state.get_spot_oracle_price(spot_position.market_index);
+                let diff_bps = (pyth_price.price.abs_diff(oracle.price) * 10_000)
+                    / oracle.price.unsigned_abs();
+                if diff_bps > market_state.pyth_oracle_diff_threshold_bps {
+                    pyth_price
+                } else {
+                    *oracle
+                }
+            }
+            None => *market_state.get_spot_oracle_price(spot_position.market_index),
+        };
 
         let signed_token_amount = spot_position.get_signed_token_amount(spot_market).unwrap();
 
@@ -186,7 +198,19 @@ pub fn calculate_simplified_margin_requirement(
         }
 
         let perp_market = market_state.get_perp_market(perp_position.market_index);
-        let oracle_price = market_state.get_perp_oracle_price(perp_position.market_index);
+        let oracle_price = match market_state.get_perp_pyth_price(perp_position.market_index) {
+            Some(pyth_price) => {
+                let oracle = market_state.get_perp_oracle_price(perp_position.market_index);
+                let diff_bps = (pyth_price.price.abs_diff(oracle.price) * 10_000)
+                    / oracle.price.unsigned_abs();
+                if diff_bps > market_state.pyth_oracle_diff_threshold_bps {
+                    pyth_price
+                } else {
+                    *oracle
+                }
+            }
+            None => *market_state.get_perp_oracle_price(perp_position.market_index),
+        };
 
         let strict_quote_price = {
             let quote_price_data =
@@ -213,7 +237,7 @@ pub fn calculate_simplified_margin_requirement(
         ) = calculate_perp_position_value_and_pnl(
             perp_position,
             perp_market,
-            oracle_price,
+            &oracle_price,
             &strict_quote_price,
             margin_type,
             user_custom_margin_ratio.max(perp_position_custom_margin_ratio),

src/types.rs

@@ -227,6 +227,9 @@ pub struct MarketState {
     pub perp_markets: HashMap<u16, PerpMarket, FxBuildHasher>,
     pub spot_oracle_prices: HashMap<u16, OraclePriceData, FxBuildHasher>,
     pub perp_oracle_prices: HashMap<u16, OraclePriceData, FxBuildHasher>,
+    pub spot_pyth_prices: HashMap<u16, i64, FxBuildHasher>, // Override spot with pyth price
+    pub perp_pyth_prices: HashMap<u16, i64, FxBuildHasher>, // Override perp with pyth price
+    pub pyth_oracle_diff_threshold_bps: u64, // Min bps diff to prefer pyth price over oracle. Defaults to 0 (always use pyth when set).
 }
 
 impl MarketState {
@@ -246,6 +249,30 @@ impl MarketState {
         self.perp_oracle_prices.get(&market_index).unwrap()
     }
 
+    pub fn get_spot_pyth_price(&self, market_index: u16) -> Option<OraclePriceData> {
+        self.spot_pyth_prices
+            .get(&market_index)
+            .map(|&price| OraclePriceData {
+                price,
+                confidence: 0,
+                delay: 0,
+                has_sufficient_number_of_data_points: true,
+                sequence_id: None,
+            })
+    }
+
+    pub fn get_perp_pyth_price(&self, market_index: u16) -> Option<OraclePriceData> {
+        self.perp_pyth_prices
+            .get(&market_index)
+            .map(|&price| OraclePriceData {
+                price,
+                confidence: 0,
+                delay: 0,
+                has_sufficient_number_of_data_points: true,
+                sequence_id: None,
+            })
+    }
+
     pub fn set_spot_market(&mut self, market: SpotMarket) {
         self.spot_markets.insert(market.market_index, market);
     }
@@ -261,4 +288,12 @@ impl MarketState {
     pub fn set_perp_oracle_price(&mut self, market_index: u16, price_data: OraclePriceData) {
         self.perp_oracle_prices.insert(market_index, price_data);
     }
+
+    pub fn set_spot_pyth_price(&mut self, market_index: u16, price_data: i64) {
+        self.spot_pyth_prices.insert(market_index, price_data);
+    }
+
+    pub fn set_perp_pyth_price(&mut self, market_index: u16, price_data: i64) {
+        self.perp_pyth_prices.insert(market_index, price_data);
+    }
 }