From dc45cc29167ee6358b527f3a0bcc0617dc9e4c99 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20K=C3=B6cher?= <bkchr@users.noreply.github.com>
Date: Thu, 11 Aug 2022 16:53:40 +0200
Subject: [PATCH] Use correct length for Pallet and Storage prefix when
 calculating `KeyLenOf` (#11994)

* Use correct length for Pallet and Storage prefix when calculating `KeyLenOf`

* Add tests

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
---
 frame/support/src/storage/types/double_map.rs | 19 ++++++++++---
 frame/support/src/storage/types/map.rs        | 27 +++++++++++++++++--
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/frame/support/src/storage/types/double_map.rs b/frame/support/src/storage/types/double_map.rs
index 2e090d30119aa..9ba4cf052e82a 100644
--- a/frame/support/src/storage/types/double_map.rs
+++ b/frame/support/src/storage/types/double_map.rs
@@ -25,6 +25,7 @@ use crate::{
 		KeyLenOf, StorageAppend, StorageDecodeLength, StoragePrefixedMap, StorageTryAppend,
 	},
 	traits::{Get, GetDefault, StorageInfo, StorageInstance},
+	StorageHasher, Twox128,
 };
 use codec::{Decode, Encode, EncodeLike, FullCodec, MaxEncodedLen};
 use sp_arithmetic::traits::SaturatedConversion;
@@ -91,10 +92,10 @@ impl<Prefix, Hasher1, Key1, Hasher2, Key2, Value, QueryKind, OnEmpty, MaxValues>
 	Key2: MaxEncodedLen,
 {
 	fn get() -> u32 {
-		let z = Hasher1::max_len::<Key1>() +
-			Hasher2::max_len::<Key2>() +
-			Prefix::pallet_prefix().len() +
-			Prefix::STORAGE_PREFIX.len();
+		// The `max_len` of both key hashes plus the pallet prefix and storage prefix (which both
+		// are hashed with `Twox128`).
+		let z =
+			Hasher1::max_len::<Key1>() + Hasher2::max_len::<Key2>() + Twox128::max_len::<()>() * 2;
 		z as u32
 	}
 }
@@ -755,6 +756,16 @@ mod test {
 		}
 	}
 
+	#[test]
+	fn keylenof_works() {
+		// Works with Blake2_128Concat and Twox64Concat.
+		type A = StorageDoubleMap<Prefix, Blake2_128Concat, u64, Twox64Concat, u32, u32>;
+		let size = 16 * 2 // Two Twox128
+			+ 16 + 8 // Blake2_128Concat = hash + key
+			+ 8 + 4; // Twox64Concat = hash + key
+		assert_eq!(KeyLenOf::<A>::get(), size);
+	}
+
 	#[test]
 	fn test() {
 		type A =
diff --git a/frame/support/src/storage/types/map.rs b/frame/support/src/storage/types/map.rs
index f4ac83c22663b..0f89e2378a55d 100644
--- a/frame/support/src/storage/types/map.rs
+++ b/frame/support/src/storage/types/map.rs
@@ -25,6 +25,7 @@ use crate::{
 		KeyLenOf, StorageAppend, StorageDecodeLength, StoragePrefixedMap, StorageTryAppend,
 	},
 	traits::{Get, GetDefault, StorageInfo, StorageInstance},
+	StorageHasher, Twox128,
 };
 use codec::{Decode, Encode, EncodeLike, FullCodec, MaxEncodedLen};
 use sp_arithmetic::traits::SaturatedConversion;
@@ -61,8 +62,9 @@ where
 	Key: FullCodec + MaxEncodedLen,
 {
 	fn get() -> u32 {
-		let z =
-			Hasher::max_len::<Key>() + Prefix::pallet_prefix().len() + Prefix::STORAGE_PREFIX.len();
+		// The `max_len` of the key hash plus the pallet prefix and storage prefix (which both are
+		// hashed with `Twox128`).
+		let z = Hasher::max_len::<Key>() + Twox128::max_len::<()>() * 2;
 		z as u32
 	}
 }
@@ -501,6 +503,27 @@ mod test {
 		}
 	}
 
+	#[test]
+	fn keylenof_works() {
+		// Works with Blake2_128Concat.
+		type A = StorageMap<Prefix, Blake2_128Concat, u32, u32>;
+		let size = 16 * 2 // Two Twox128
+			+ 16 + 4; // Blake2_128Concat = hash + key
+		assert_eq!(KeyLenOf::<A>::get(), size);
+
+		// Works with Blake2_256.
+		type B = StorageMap<Prefix, Blake2_256, u32, u32>;
+		let size = 16 * 2 // Two Twox128
+			+ 32; // Blake2_256
+		assert_eq!(KeyLenOf::<B>::get(), size);
+
+		// Works with Twox64Concat.
+		type C = StorageMap<Prefix, Twox64Concat, u32, u32>;
+		let size = 16 * 2 // Two Twox128
+			+ 8 + 4; // Twox64Concat = hash + key
+		assert_eq!(KeyLenOf::<C>::get(), size);
+	}
+
 	#[test]
 	fn test() {
 		type A = StorageMap<Prefix, Blake2_128Concat, u16, u32, OptionQuery>;