vialabs-stellar-common 0.1.3

use crate::{
  errors::Error,
  message_gateway_v4::{MessageGatewayV4Client, SendRequest},
  storage::{DataKey, GATEWAY_CONTRACT_ADDRESS},
  utils::is_zero_address,
};
use soroban_sdk::{
  contracttype, panic_with_error, symbol_short, vec, xdr::ToXdr, Address, Bytes, Env, IntoVal,
  Symbol, Vec,
};

const MESSAGE_OWNER_KEY: Symbol = symbol_short!("owner");

#[contracttype]
#[derive(Clone, Debug, PartialEq)]
pub struct ProcessFromGatewayRequest {
  pub tx_id: u128,
  pub source_chain_id: u64,
  pub sender: Bytes,
  pub recipient: Address,
  pub on_chain_data: Bytes,
  pub off_chain_data: Bytes,
  pub gas_fee: u64,
}

pub struct Base;

impl Base {
  /// Retrieves the current message owner address.
  ///
  /// # Panics
  ///
  /// Panics if the message owner has not been set.
  pub fn get_message_owner(env: &Env) -> Address {
    env.storage().instance().get(&MESSAGE_OWNER_KEY).unwrap()
  }

  /// Sets the message owner address.
  pub fn set_message_owner(env: &Env, owner: Address) {
    env.storage().instance().set(&MESSAGE_OWNER_KEY, &owner);
  }

  /// Verifies that the caller is the message owner.
  ///
  /// # Panics
  ///
  /// Panics if the caller is not the message owner.
  pub fn only_message_owner(env: &Env) {
    let owner = Self::get_message_owner(env);

    owner.require_auth();
  }

  /// Sets the message gateway contract address.
  ///
  /// # Panics
  ///
  /// Panics if the caller is not the message owner.
  pub fn set_message_gateway(env: &Env, contract_address: Address) {
    Self::only_message_owner(env);

    env
      .storage()
      .instance()
      .set(&GATEWAY_CONTRACT_ADDRESS, &contract_address);
  }

  /// Retrieves the message gateway contract address.
  ///
  /// # Panics
  ///
  /// Panics if the message gateway has not been set.
  pub fn get_message_gateway(env: &Env) -> Address {
    env
      .storage()
      .instance()
      .get(&GATEWAY_CONTRACT_ADDRESS)
      .unwrap_or_else(|| panic_with_error!(env, Error::MissingMessageGateway))
  }

  /// Sets message endpoints for different chains.
  ///
  /// # Arguments
  ///
  /// * `chains` - Vector of chain IDs
  /// * `endpoints` - Vector of endpoint addresses for each chain
  ///
  /// # Panics
  ///
  /// Panics if the caller is not the message owner or if the chain and endpoint vectors have different lengths.
  pub fn set_message_endpoints(env: &Env, chains: Vec<u64>, endpoints: Vec<Bytes>) {
    Self::only_message_owner(env);

    let chains_length = chains.len();

    if chains_length != endpoints.len() {
      panic_with_error!(env, Error::ChainsEndpointsMislength);
    }

    for id in 0..chains_length {
      env.storage().instance().set(
        &DataKey::ChainsEndpoints(chains.get(id).unwrap()),
        &endpoints.get(id),
      );
    }
  }

  /// Retrieves the endpoint address for a specific chain ID.
  ///
  /// Returns an empty `Bytes` if no endpoint is set for the given chain ID.
  pub fn get_endpoint_by_chain_id(env: &Env, chain_id: u64) -> Bytes {
    env
      .storage()
      .instance()
      .get(&DataKey::ChainsEndpoints(chain_id))
      .unwrap_or(Bytes::new(env))
  }

  /// Sends a message to a destination chain through the message gateway.
  ///
  /// # Arguments
  ///
  /// * `destination_chain` - Chain ID of the destination chain
  /// * `chain_data` - Data to send to the destination chain
  /// * `confirmations` - Number of confirmations required
  ///
  /// # Returns
  ///
  /// Returns the transaction ID of the sent message.
  ///
  /// # Panics
  ///
  /// Panics if the message gateway is not set or if no endpoint is configured for the destination chain.
  pub fn message_send(
    env: &Env,
    destination_chain: u64,
    chain_data: Bytes,
    confirmations: u32,
  ) -> u128 {
    let gateway_address = Self::get_message_gateway(env);

    if is_zero_address(env, gateway_address.clone().to_xdr(env)) {
      panic_with_error!(env, Error::MissingMessageGateway);
    }

    let chain_endpoint = Self::get_endpoint_by_chain_id(env, destination_chain);

    if chain_endpoint.is_empty() {
      panic_with_error!(env, Error::MissingChainEndpoint);
    }

    let request = SendRequest {
      sender: env.current_contract_address().to_xdr(env),
      recipient: chain_endpoint,
      destination_chain,
      chain_data,
      confirmations,
    };

    let tx_id = MessageGatewayV4Client::new(env, &gateway_address).send(&request);

    tx_id
  }

  /// Validates that a sender is authorized for a given source chain.
  ///
  /// # Arguments
  ///
  /// * `source_chain_id` - Chain ID where the message originated
  /// * `sender` - Sender address in bytes format
  ///
  /// # Panics
  ///
  /// Panics if the sender length doesn't match the endpoint length or if the sender hash doesn't match the endpoint hash.
  pub fn validate_chain_sender(env: &Env, source_chain_id: u64, sender: Bytes) {
    let chain_endpoint = Self::get_endpoint_by_chain_id(env, source_chain_id);

    if chain_endpoint.len() != sender.len() {
      panic_with_error!(env, Error::SenderLengthMismatch);
    }

    let endpoint_hash = env.crypto().keccak256(&chain_endpoint);
    let sender_hash = env.crypto().keccak256(&sender);

    if sender_hash.to_array() != endpoint_hash.to_array() {
      panic_with_error!(env, Error::InvalidChainSender);
    }
  }

  /// Processes a message received from the gateway.
  ///
  /// This function invokes the `message_process` method on the recipient contract.
  ///
  /// # Arguments
  ///
  /// * `request` - The process request containing message data
  pub fn message_process(env: &Env, request: ProcessFromGatewayRequest) {
    let method = Symbol::new(env, "message_process");
    let args = vec![env, request.clone().into_val(env)];

    env.invoke_contract::<()>(&request.recipient, &method, args);
  }

  /// Processes a message received from the gateway.
  ///
  /// This function invokes the `message_process` method on the recipient contract.
  ///
  /// # Arguments
  ///
  /// * `request` - The process request containing message data
  ///
  /// # Panics
  ///
  /// Panics if the message process fails.
  pub fn message_process_from_gateway(env: &Env, request: ProcessFromGatewayRequest) {
    Self::message_process(env, request);
  }
}

pub trait MessageClientV4Interface {
  fn only_message_owner(env: &Env);

  fn set_message_owner(env: &Env, owner: Address);
  fn set_message_gateway(env: &Env, contract_address: Address);
  fn set_message_endpoints(env: &Env, chains: Vec<u64>, endpoints: Vec<Bytes>);

  fn get_endpoint_by_chain_id(env: &Env, chain_id: u64) -> Bytes;
  fn get_message_gateway(env: &Env) -> Address;
  fn get_message_owner(env: &Env) -> Address;

  fn validate_chain_sender(env: &Env, source_chain_id: u64, sender: Bytes);

  fn message_send(
    env: &Env,
    destination_chain: u64,
    chain_data: Bytes,
    confirmations: u32,
  ) -> u128;
  fn message_process(env: &Env, message: ProcessFromGatewayRequest);
  fn message_process_from_gateway(env: &Env, message: ProcessFromGatewayRequest);
}