From 217dede988b29eee0e3363e727913bbb12866b84 Mon Sep 17 00:00:00 2001
From: Foivos <foivos@umich.edu>
Date: Tue, 26 Nov 2024 14:08:25 +0100
Subject: [PATCH] add allow and disallow function to relayer_discovery

---
 .changeset/proud-pumas-breathe.md             |  5 ++
 move/relayer_discovery/sources/discovery.move | 60 ++++++++++++++++++-
 .../sources/types/owner_cap.move              | 23 +++++++
 .../versioned/relayer_discovery_v0.move       | 17 ++++++
 4 files changed, 104 insertions(+), 1 deletion(-)
 create mode 100644 .changeset/proud-pumas-breathe.md
 create mode 100644 move/relayer_discovery/sources/types/owner_cap.move

diff --git a/.changeset/proud-pumas-breathe.md b/.changeset/proud-pumas-breathe.md
new file mode 100644
index 00000000..bf7b5dc3
--- /dev/null
+++ b/.changeset/proud-pumas-breathe.md
@@ -0,0 +1,5 @@
+---
+'@axelar-network/axelar-cgp-sui': minor
+---
+
+Add allow_function and disallow_function to relayer_discovery.
diff --git a/move/relayer_discovery/sources/discovery.move b/move/relayer_discovery/sources/discovery.move
index cb7e57ac..c86789a9 100644
--- a/move/relayer_discovery/sources/discovery.move
+++ b/move/relayer_discovery/sources/discovery.move
@@ -7,9 +7,10 @@
 module relayer_discovery::discovery;
 
 use axelar_gateway::channel::Channel;
+use relayer_discovery::owner_cap::{Self, OwnerCap};
 use relayer_discovery::relayer_discovery_v0::{Self, RelayerDiscovery_v0};
 use relayer_discovery::transaction::Transaction;
-use std::ascii;
+use std::ascii::{Self, String};
 use sui::versioned::{Self, Versioned};
 use version_control::version_control::{Self, VersionControl};
 
@@ -43,6 +44,7 @@ fun init(ctx: &mut TxContext) {
         id: object::new(ctx),
         inner,
     });
+    transfer::public_transfer(owner_cap::create(ctx), ctx.sender());
 }
 
 /// ------
@@ -68,6 +70,32 @@ macro fun value_mut(
     value
 }
 
+// ---------------
+// Entry Functions
+// ---------------
+entry fun allow_function(
+    self: &mut RelayerDiscovery,
+    _: &OwnerCap,
+    version: u64,
+    function_name: String,
+) {
+    self.value_mut!(b"allow_function").allow_function(version, function_name);
+}
+
+entry fun disallow_function(
+    self: &mut RelayerDiscovery,
+    _: &OwnerCap,
+    version: u64,
+    function_name: String,
+) {
+    self
+        .value_mut!(b"disallow_function")
+        .disallow_function(version, function_name);
+}
+
+// ----------------
+// Public Functions
+// ----------------
 /// During the creation of the object, the UID should be passed here to
 /// receive the Channel and emit an event which will be handled by the
 /// Relayer.
@@ -118,6 +146,8 @@ fun version_control(): VersionControl {
             b"register_transaction",
             b"remove_transaction",
             b"get_transaction",
+            b"allow_function",
+            b"disallow_function",
         ].map!(|function_name| function_name.to_ascii_string()),
     ])
 }
@@ -165,3 +195,31 @@ fun test_register_and_get() {
     sui::test_utils::destroy(self);
     sui::test_utils::destroy(channel);
 }
+
+#[test]
+fun test_allow_function() {
+    let ctx = &mut sui::tx_context::dummy();
+    let mut self = new(ctx);
+    let owner_cap = owner_cap::create(ctx);
+    let version = 0;
+    let function_name = b"function_name".to_ascii_string();
+
+    self.allow_function(&owner_cap, version, function_name);
+
+    sui::test_utils::destroy(self);
+    owner_cap.destroy_for_testing();
+}
+
+#[test]
+fun test_disallow_function() {
+    let ctx = &mut sui::tx_context::dummy();
+    let mut self = new(ctx);
+    let owner_cap = owner_cap::create(ctx);
+    let version = 0;
+    let function_name = b"register_transaction".to_ascii_string();
+
+    self.disallow_function(&owner_cap, version, function_name);
+
+    sui::test_utils::destroy(self);
+    owner_cap.destroy_for_testing();
+}
diff --git a/move/relayer_discovery/sources/types/owner_cap.move b/move/relayer_discovery/sources/types/owner_cap.move
new file mode 100644
index 00000000..c6a53c9f
--- /dev/null
+++ b/move/relayer_discovery/sources/types/owner_cap.move
@@ -0,0 +1,23 @@
+module relayer_discovery::owner_cap;
+
+// -----
+// Types
+// -----
+public struct OwnerCap has key, store {
+    id: UID,
+}
+
+public(package) fun create(ctx: &mut TxContext): OwnerCap {
+    OwnerCap {
+        id: object::new(ctx),
+    }
+}
+
+/// ---------
+/// Test Only
+/// ---------
+#[test_only]
+public(package) fun destroy_for_testing(self: OwnerCap) {
+    let OwnerCap { id } = self;
+    id.delete();
+}
diff --git a/move/relayer_discovery/sources/versioned/relayer_discovery_v0.move b/move/relayer_discovery/sources/versioned/relayer_discovery_v0.move
index 8149049b..2bde49b7 100644
--- a/move/relayer_discovery/sources/versioned/relayer_discovery_v0.move
+++ b/move/relayer_discovery/sources/versioned/relayer_discovery_v0.move
@@ -2,6 +2,7 @@ module relayer_discovery::relayer_discovery_v0;
 
 use relayer_discovery::events;
 use relayer_discovery::transaction::Transaction;
+use std::ascii::String;
 use sui::table::{Self, Table};
 use version_control::version_control::VersionControl;
 
@@ -74,6 +75,22 @@ public(package) fun version_control(
     &self.version_control
 }
 
+public(package) fun allow_function(
+    self: &mut RelayerDiscovery_v0,
+    version: u64,
+    function_name: String,
+) {
+    self.version_control.allow_function(version, function_name);
+}
+
+public(package) fun disallow_function(
+    self: &mut RelayerDiscovery_v0,
+    version: u64,
+    function_name: String,
+) {
+    self.version_control.disallow_function(version, function_name);
+}
+
 // ---------
 // Test Only
 // ---------