fix(router): correct JSON schema to generate the contract schema from (…

…#375)
axelarnetwork · May 1, 2024 · 01f9705 · 01f9705
1 parent 2cbdfe1
commit 01f9705
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 3 deletions.
diff --git a/packages/axelar-wasm-std/src/flagset.rs b/packages/axelar-wasm-std/src/flagset.rs
@@ -69,14 +69,14 @@ where
 
 impl<T> JsonSchema for FlagSet<T>
 where
-    T: flagset::Flags + Serialize,
+    T: flagset::Flags + Serialize + JsonSchema,
     <T as flagset::Flags>::Type: Serialize,
 {
     fn schema_name() -> String {
         "FlagSet".to_string()
     }
 
     fn json_schema(gen: &mut SchemaGenerator) -> Schema {
-        gen.root_schema_for::<FlagSet<T>>().schema.into()
+        gen.subschema_for::<T>()
     }
 }
diff --git a/packages/router-api/src/primitives.rs b/packages/router-api/src/primitives.rs
@@ -10,6 +10,8 @@ use cw_storage_plus::{Key, KeyDeserialize, Prefixer, PrimaryKey};
 use error_stack::Report;
 use error_stack::ResultExt;
 use flagset::flags;
+use schemars::gen::SchemaGenerator;
+use schemars::schema::Schema;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use sha3::{Digest, Keccak256};
@@ -245,7 +247,7 @@ impl AsRef<str> for ChainName {
 
 flags! {
     #[repr(u8)]
-    #[derive(Deserialize, Serialize, Hash, JsonSchema)]
+    #[derive(Deserialize, Serialize, Hash)]
     pub enum GatewayDirection: u8 {
         None = 0,
         Incoming = 1,
@@ -254,6 +256,16 @@ flags! {
     }
 }
 
+impl JsonSchema for GatewayDirection {
+    fn schema_name() -> String {
+        "GatewayDirection".to_string()
+    }
+
+    fn json_schema(gen: &mut SchemaGenerator) -> Schema {
+        gen.subschema_for::<u8>()
+    }
+}
+
 #[cw_serde]
 pub struct Gateway {
     pub address: Addr,
@@ -378,6 +390,16 @@ mod tests {
         assert!(!chain_name.eq(&"Ethereum-1".to_string()));
     }
 
+    #[test]
+    fn json_schema_for_gateway_direction_flag_set_does_not_panic() {
+        let gen = &mut SchemaGenerator::default();
+        // check it doesn't panic
+        let _ = FlagSet::<GatewayDirection>::json_schema(gen);
+
+        // make sure it's the same as the underlying type
+        assert_eq!(GatewayDirection::json_schema(gen), u8::json_schema(gen));
+    }
+
     fn dummy_message() -> Message {
         Message {
             cc_id: CrossChainId {