feat: Automatic consumer registration (#34)

contracts/babylon/benches/main.rs

@@ -44,6 +44,8 @@ pub fn setup_instance() -> Instance<MockApi, MockStorage, MockQuerier> {
     let msg = InstantiateMsg {
         network: babylon_bitcoin::chain_params::Network::Regtest,
         babylon_tag: "01020304".to_string(),
+        consumer_name: None,
+        consumer_description: None,
         btc_confirmation_depth: 10,
         checkpoint_finalization_timeout: 1,
         notify_cosmos_zone: false,
@@ -172,4 +174,4 @@ criterion_group!(
     config = make_config();
     targets = bench_btc_light_client
 );
-criterion_main!(btc_light_client);
+criterion_main!(btc_light_client);

contracts/babylon/src/contract.rs

@@ -26,20 +26,25 @@ pub fn instantiate(
 ) -> Result<Response<BabylonMsg>, ContractError> {
     msg.validate()?;
 
-    // initialise config
-    let cfg = Config {
+    // Initialize config with None values for consumer fields
+    let mut cfg = Config {
         network: msg.network.clone(),
         babylon_tag: msg.babylon_tag_to_bytes()?,
         btc_confirmation_depth: msg.btc_confirmation_depth,
         checkpoint_finalization_timeout: msg.checkpoint_finalization_timeout,
         notify_cosmos_zone: msg.notify_cosmos_zone,
         btc_staking: None, // Will be set in `reply` if `btc_staking_code_id` is provided
+        consumer_name: None,
+        consumer_description: None,
     };
-    CONFIG.save(deps.storage, &cfg)?;
 
     let mut res = Response::new().add_attribute("action", "instantiate");
 
     if let Some(btc_staking_code_id) = msg.btc_staking_code_id {
+        // Update config with consumer information
+        cfg.consumer_name = msg.consumer_name;
+        cfg.consumer_description = msg.consumer_description;
+
         // Instantiate BTC staking contract
         let init_msg = WasmMsg::Instantiate {
             admin: msg.admin,
@@ -52,6 +57,10 @@ pub fn instantiate(
 
         res = res.add_submessage(init_msg);
     }
+
+    // Save the config after potentially updating it
+    CONFIG.save(deps.storage, &cfg)?;
+
     Ok(res)
 }
 
@@ -214,9 +223,11 @@ mod tests {
             btc_staking_code_id: None,
             btc_staking_msg: None,
             admin: None,
+            consumer_name: None,
+            consumer_description: None,
         };
         let info = message_info(&deps.api.addr_make(CREATOR), &[]);
         let res = instantiate(deps.as_mut(), mock_env(), info, msg).unwrap();
         assert_eq!(0, res.messages.len());
     }
-}
+}

contracts/babylon/src/error.rs

@@ -34,6 +34,8 @@ pub enum ContractError {
     Unauthorized {},
     #[error("The BTC staking contract is not set")]
     BtcStakingNotSet {},
+    #[error("Invalid configuration: {msg}")]
+    InvalidConfig { msg: String },
 }
 
 #[derive(Error, Debug, PartialEq)]
@@ -127,4 +129,4 @@ pub enum BabylonEpochChainError {
     EmptyTxKey {},
     #[error("The BTC header cannot be decoded")]
     BTCHeaderDecodeError {},
-}
+}

contracts/babylon/src/ibc.rs

@@ -3,14 +3,14 @@ use babylon_bindings::BabylonMsg;
 use babylon_proto::babylon::zoneconcierge::v1::{
     zoneconcierge_packet_data::Packet, BtcTimestamp, ZoneconciergePacketData,
 };
+use babylon_proto::babylon::btcstkconsumer::v1::ConsumerRegisterIbcPacket;
+
 use cosmwasm_std::{
-    DepsMut, Env, Event, Ibc3ChannelOpenResponse, IbcBasicResponse, IbcChannel, IbcChannelCloseMsg,
-    IbcChannelConnectMsg, IbcChannelOpenMsg, IbcChannelOpenResponse, IbcOrder, IbcPacketAckMsg,
-    IbcPacketReceiveMsg, IbcPacketTimeoutMsg, IbcReceiveResponse, IbcTimeout, Never, StdAck,
-    StdError, StdResult,
+    Binary, DepsMut, Env, Event, Ibc3ChannelOpenResponse, IbcBasicResponse, IbcChannel, IbcChannelCloseMsg, IbcChannelConnectMsg, IbcChannelOpenMsg, IbcChannelOpenResponse, IbcMsg, IbcOrder, IbcPacketAckMsg, IbcPacketReceiveMsg, IbcPacketTimeoutMsg, IbcReceiveResponse, IbcTimeout, Never, StdAck, StdError, StdResult
 };
 use cw_storage_plus::Item;
 use prost::Message;
+use crate::state::config::CONFIG;
 
 pub const IBC_VERSION: &str = "zoneconcierge-1";
 pub const IBC_ORDERING: IbcOrder = IbcOrder::Ordered;
@@ -57,7 +57,7 @@ pub fn ibc_channel_open(
 /// Second part of the 4-step handshake, i.e. ChannelOpenAck and ChannelOpenConfirm.
 pub fn ibc_channel_connect(
     deps: DepsMut,
-    _env: Env,
+    env: Env,
     msg: IbcChannelConnectMsg,
 ) -> Result<IbcBasicResponse, ContractError> {
     // Ensure we have no channel yet
@@ -69,11 +69,39 @@ pub fn ibc_channel_connect(
     // Store the channel
     IBC_CHANNEL.save(deps.storage, channel)?;
 
+    // Load the config
+    let cfg = CONFIG.load(deps.storage)?;
+
     let chan_id = &channel.endpoint.channel_id;
-    Ok(IbcBasicResponse::new()
+    let mut response = IbcBasicResponse::new()
         .add_attribute("action", "ibc_connect")
         .add_attribute("channel_id", chan_id)
-        .add_event(Event::new("ibc").add_attribute("channel", "connect")))
+        .add_event(Event::new("ibc").add_attribute("channel", "connect"));
+
+    // If the consumer name and description are set, create and send a ConsumerRegister packet
+    if let (Some(name), Some(description)) = (&cfg.consumer_name, &cfg.consumer_description) {
+        let consumer_register_packet = ConsumerRegisterIbcPacket {
+            consumer_name: name.clone(),
+            consumer_description: description.clone(),
+        };
+
+        let packet_data = ZoneconciergePacketData {
+            packet: Some(Packet::ConsumerRegister(consumer_register_packet)),
+        };
+
+        let ibc_msg = IbcMsg::SendPacket {
+            channel_id: channel.endpoint.channel_id.clone(),
+            data: Binary::new(packet_data.encode_to_vec()),
+            timeout: packet_timeout(&env),
+        };
+
+        response = response
+            .add_message(ibc_msg)
+            .add_attribute("consumer_name", name)
+            .add_attribute("consumer_description", description);
+    }
+
+    Ok(response)
 }
 
 /// This is invoked on the IBC Channel Close message
@@ -123,7 +151,10 @@ pub fn ibc_packet_receive(
             Packet::BtcTimestamp(btc_ts) => ibc_packet::handle_btc_timestamp(deps, caller, &btc_ts),
             Packet::BtcStaking(btc_staking) => {
                 ibc_packet::handle_btc_staking(deps, caller, &btc_staking)
-            }
+            },
+            Packet::ConsumerRegister(_) => {
+                return Err(StdError::generic_err("ConsumerRegister packet should not be received").into())
+            },
         }
     })()
     .or_else(|e| {
@@ -395,6 +426,8 @@ mod tests {
             btc_staking_code_id: None,
             btc_staking_msg: None,
             admin: None,
+            consumer_name: None,
+            consumer_description: None,
         };
         let info = message_info(&deps.api.addr_make(CREATOR), &[]);
         let res = instantiate(deps.as_mut(), mock_env(), info, msg).unwrap();