From 4531ba61d7315fee2d09a16a5be576ea3c045ee9 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Wed, 10 Jan 2024 10:56:14 +0000
Subject: [PATCH] Serialize storage keys in metadata as SCALE encoded (#2048)

* Use scale encoding in metadata for keys

* Fix tests
---
 Cargo.lock                          |  1 +
 crates/metadata/Cargo.toml          |  1 +
 crates/metadata/src/layout/mod.rs   | 16 +++++++++++++---
 crates/metadata/src/layout/tests.rs | 26 +++++++++++++-------------
 4 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 89fac014845..b7d3574c6cc 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2767,6 +2767,7 @@ dependencies = [
  "ink_prelude 5.0.0-rc",
  "ink_primitives 5.0.0-rc",
  "linkme",
+ "parity-scale-codec",
  "pretty_assertions",
  "scale-info",
  "schemars",
diff --git a/crates/metadata/Cargo.toml b/crates/metadata/Cargo.toml
index 97db8b3206b..9036f079758 100644
--- a/crates/metadata/Cargo.toml
+++ b/crates/metadata/Cargo.toml
@@ -22,6 +22,7 @@ serde = { workspace = true, features = ["derive", "alloc"] }
 impl-serde = { workspace = true, default-features = true }
 derive_more = { workspace = true, features = ["from"] }
 linkme = { workspace = true }
+scale = { workspace = true }
 scale-info = { workspace = true, features = ["derive", "serde", "decode", "schema"] }
 schemars = { workspace = true }
 
diff --git a/crates/metadata/src/layout/mod.rs b/crates/metadata/src/layout/mod.rs
index c667daa423f..fea4d129899 100644
--- a/crates/metadata/src/layout/mod.rs
+++ b/crates/metadata/src/layout/mod.rs
@@ -29,6 +29,10 @@ use crate::{
 use derive_more::From;
 use ink_prelude::collections::btree_map::BTreeMap;
 use ink_primitives::Key;
+use scale::{
+    Decode,
+    Encode,
+};
 use scale_info::{
     form::{
         Form,
@@ -42,7 +46,10 @@ use scale_info::{
 };
 use schemars::JsonSchema;
 use serde::{
-    de::DeserializeOwned,
+    de::{
+        DeserializeOwned,
+        Error,
+    },
     Deserialize,
     Serialize,
 };
@@ -87,7 +94,7 @@ impl serde::Serialize for LayoutKey {
     where
         S: serde::Serializer,
     {
-        serde_hex::serialize(&self.key.to_be_bytes(), serializer)
+        serde_hex::serialize(&self.key.encode(), serializer)
     }
 }
 
@@ -98,7 +105,10 @@ impl<'de> serde::Deserialize<'de> for LayoutKey {
     {
         let mut arr = [0; 4];
         serde_hex::deserialize_check_len(d, serde_hex::ExpectedLen::Exact(&mut arr[..]))?;
-        Ok(Key::from_be_bytes(arr).into())
+        let key = Key::decode(&mut &arr[..]).map_err(|err| {
+            Error::custom(format!("Error decoding layout key: {}", err))
+        })?;
+        Ok(key.into())
     }
 }
 
diff --git a/crates/metadata/src/layout/tests.rs b/crates/metadata/src/layout/tests.rs
index 72993dfec10..0518068b761 100644
--- a/crates/metadata/src/layout/tests.rs
+++ b/crates/metadata/src/layout/tests.rs
@@ -20,7 +20,7 @@ use scale_info::Path;
 fn layout_key_works() {
     let layout_key = LayoutKey::from(&1);
     let json = serde_json::to_string(&layout_key).unwrap();
-    assert_eq!(json, "\"0x00000001\"",);
+    assert_eq!(json, "\"0x01000000\"",);
 }
 
 fn named_fields_struct_layout(key: &Key) -> Layout {
@@ -47,7 +47,7 @@ fn named_fields_work() {
                     {
                         "layout": {
                             "leaf": {
-                                "key": "0x00000159",
+                                "key": "0x59010000",
                                 "ty": 0,
                             }
                         },
@@ -56,7 +56,7 @@ fn named_fields_work() {
                     {
                         "layout": {
                             "leaf": {
-                                "key": "0x00000159",
+                                "key": "0x59010000",
                                 "ty": 1,
                             }
                         },
@@ -94,7 +94,7 @@ fn tuple_struct_work() {
                     {
                         "layout": {
                             "leaf": {
-                                "key": "0x000000ea",
+                                "key": "0xea000000",
                                 "ty": 0,
                             }
                         },
@@ -103,7 +103,7 @@ fn tuple_struct_work() {
                     {
                         "layout": {
                             "leaf": {
-                                "key": "0x000000ea",
+                                "key": "0xea000000",
                                 "ty": 1,
                             }
                         },
@@ -139,7 +139,7 @@ fn clike_enum_work() {
     let expected = serde_json::json! {
         {
             "enum": {
-                "dispatchKey": "0x0000007b",
+                "dispatchKey": "0x7b000000",
                 "name": "Enum",
                 "variants": {
                     "0": {
@@ -219,7 +219,7 @@ fn mixed_enum_work() {
     let expected = serde_json::json! {
         {
             "enum": {
-                "dispatchKey": "0x000001c8",
+                "dispatchKey": "0xc8010000",
                 "name": "Enum",
                 "variants": {
                     "0": {
@@ -231,7 +231,7 @@ fn mixed_enum_work() {
                             {
                                 "layout": {
                                     "leaf": {
-                                        "key": "0x000001c8",
+                                        "key": "0xc8010000",
                                         "ty": 0,
                                     }
                                 },
@@ -240,7 +240,7 @@ fn mixed_enum_work() {
                             {
                                 "layout": {
                                     "leaf": {
-                                        "key": "0x000001c8",
+                                        "key": "0xc8010000",
                                         "ty": 1,
                                     }
                                 },
@@ -254,7 +254,7 @@ fn mixed_enum_work() {
                             {
                                 "layout": {
                                     "leaf": {
-                                        "key": "0x000001c8",
+                                        "key": "0xc8010000",
                                         "ty": 0,
                                     }
                                 },
@@ -263,7 +263,7 @@ fn mixed_enum_work() {
                             {
                                 "layout": {
                                     "leaf": {
-                                        "key": "0x000001c8",
+                                        "key": "0xc8010000",
                                         "ty": 1,
                                     }
                                 },
@@ -304,11 +304,11 @@ fn unbounded_layout_works() {
             "hash": {
                 "layout": {
                     "leaf": {
-                        "key": "0x00000237",
+                        "key": "0x37020000",
                         "ty": 0
                     }
                 },
-                "offset": "0x00000237",
+                "offset": "0x37020000",
                 "strategy": {
                         "hasher": "Blake2x256",
                         "prefix": "0x696e6b2073746f7261676520686173686d6170",