From 20b63f63e5548ee517ea89a26e777e6d5ec9e6f0 Mon Sep 17 00:00:00 2001
From: glihm <dev@glihm.net>
Date: Thu, 7 Dec 2023 16:14:39 -0600
Subject: [PATCH] feat: add new rpc endpoint to set addresses

---
 crates/katana/core/src/hooker.rs    | 13 +++++++++++++
 crates/katana/core/src/sequencer.rs |  6 +++++-
 crates/katana/rpc/src/api/katana.rs |  4 ++++
 crates/katana/rpc/src/katana.rs     |  6 ++++++
 4 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/crates/katana/core/src/hooker.rs b/crates/katana/core/src/hooker.rs
index 4c6b33a344..1c4f5d37df 100644
--- a/crates/katana/core/src/hooker.rs
+++ b/crates/katana/core/src/hooker.rs
@@ -8,6 +8,12 @@ use starknet::core::types::{BroadcastedInvokeTransaction, FieldElement};
 
 use crate::sequencer::KatanaSequencer;
 
+#[derive(Debug, Clone, serde::Deserialize, serde::Serialize, Copy, PartialEq, Eq)]
+pub struct HookerAddresses {
+    orderbook_arkchain: FieldElement,
+    executor_starknet: FieldElement,
+}
+
 #[async_trait]
 pub trait KatanaHooker {
     /// Sets a reference to the underlying sequencer.
@@ -59,4 +65,11 @@ pub trait KatanaHooker {
     /// * `call` - The `Call` of the transaction that has failed. Usually the same as
     ///   `verify_message_to_starknet_before_tx`.
     async fn on_starknet_tx_failed(&self, call: Call);
+
+    /// Sets important addresses.
+    ///
+    /// # Arguments
+    ///
+    /// * `addresses` - Important addresses related to solis.
+    fn set_addresses(&self, addresses: &HookerAddresses);
 }
diff --git a/crates/katana/core/src/sequencer.rs b/crates/katana/core/src/sequencer.rs
index 7514cff379..1c2d5b599c 100644
--- a/crates/katana/core/src/sequencer.rs
+++ b/crates/katana/core/src/sequencer.rs
@@ -25,7 +25,7 @@ use crate::backend::storage::transaction::{
 use crate::backend::{Backend, ExternalFunctionCall};
 use crate::db::{AsStateRefDb, StateExtRef, StateRefDb};
 use crate::execution::{MaybeInvalidExecutedTransaction, PendingState};
-use crate::hooker::KatanaHooker;
+use crate::hooker::{HookerAddresses, KatanaHooker};
 use crate::pool::TransactionPool;
 use crate::sequencer_error::SequencerError;
 use crate::service::block_producer::{BlockProducer, BlockProducerMode};
@@ -96,6 +96,10 @@ impl KatanaSequencer {
         Self { pool, config, backend, block_producer, hooker }
     }
 
+    pub async fn set_addresses(&self, addresses: &HookerAddresses) {
+        self.hooker.read().await.set_addresses(addresses);
+    }
+
     /// Returns the pending state if the sequencer is running in _interval_ mode. Otherwise `None`.
     pub fn pending_state(&self) -> Option<Arc<PendingState>> {
         match &*self.block_producer.inner.read() {
diff --git a/crates/katana/rpc/src/api/katana.rs b/crates/katana/rpc/src/api/katana.rs
index 4414494e92..a091adb761 100644
--- a/crates/katana/rpc/src/api/katana.rs
+++ b/crates/katana/rpc/src/api/katana.rs
@@ -3,6 +3,7 @@ use jsonrpsee::proc_macros::rpc;
 use jsonrpsee::types::error::CallError;
 use jsonrpsee::types::ErrorObject;
 use katana_core::accounts::Account;
+use katana_core::hooker::HookerAddresses;
 use starknet::core::types::FieldElement;
 
 #[derive(thiserror::Error, Clone, Copy, Debug)]
@@ -24,6 +25,9 @@ impl From<KatanaApiError> for Error {
 
 #[rpc(server, namespace = "katana")]
 pub trait KatanaApi {
+    #[method(name = "setSolisAddresses")]
+    async fn set_addresses(&self, addresses: HookerAddresses) -> Result<(), Error>;
+
     #[method(name = "generateBlock")]
     async fn generate_block(&self) -> Result<(), Error>;
 
diff --git a/crates/katana/rpc/src/katana.rs b/crates/katana/rpc/src/katana.rs
index 966d9f7384..681b3bfa93 100644
--- a/crates/katana/rpc/src/katana.rs
+++ b/crates/katana/rpc/src/katana.rs
@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use jsonrpsee::core::{async_trait, Error};
 use katana_core::accounts::Account;
+use katana_core::hooker::HookerAddresses;
 use katana_core::sequencer::KatanaSequencer;
 use starknet::core::types::FieldElement;
 use starknet_api::core::{ContractAddress, PatriciaKey};
@@ -23,6 +24,11 @@ impl KatanaApi {
 
 #[async_trait]
 impl KatanaApiServer for KatanaApi {
+    async fn set_addresses(&self, addresses: HookerAddresses) -> Result<(), Error> {
+        self.sequencer.set_addresses(&addresses).await;
+        Ok(())
+    }
+
     async fn generate_block(&self) -> Result<(), Error> {
         self.sequencer.block_producer().force_mine();
         Ok(())