Use Extensions to register offchain worker custom extensions

substrate/primitives/externalities/src/extensions.rs

@@ -35,17 +35,24 @@ use sp_std::{
 ///
 /// As extensions are stored as `Box<Any>`, this trait should give more confidence that the correct
 /// type is registered and requested.
-pub trait Extension: Send + Any {
+pub trait Extension: Send + 'static {
 	/// Return the extension as `&mut dyn Any`.
 	///
 	/// This is a trick to make the trait type castable into an `Any`.
 	fn as_mut_any(&mut self) -> &mut dyn Any;
+
+	/// Get the [`TypeId`] of this `Extension`.
+	fn type_id(&self) -> TypeId;
 }
 
 impl Extension for Box<dyn Extension> {
 	fn as_mut_any(&mut self) -> &mut dyn Any {
 		(**self).as_mut_any()
 	}
+
+	fn type_id(&self) -> TypeId {
+		(**self).type_id()
+	}
 }
 
 /// Macro for declaring an extension that usable with [`Extensions`].
@@ -74,6 +81,10 @@ macro_rules! decl_extension {
 			fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
 				self
 			}
+
+			fn type_id(&self) -> std::any::TypeId {
+				std::any::Any::type_id(self)
+			}
 		}
 
 		impl std::ops::Deref for $ext_name {
@@ -107,6 +118,10 @@ macro_rules! decl_extension {
 			fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
 				self
 			}
+
+			fn type_id(&self) -> std::any::TypeId {
+				std::any::Any::type_id(self)
+			}
 		}
 	}
 }
@@ -235,4 +250,25 @@ mod tests {
 
 		assert_eq!(ext_ty.0, 1);
 	}
+
+	#[test]
+	fn register_box_extension() {
+		let mut exts = Extensions::new();
+		let box1: Box<dyn Extension> = Box::new(DummyExt(1));
+		let box2: Box<dyn Extension> = Box::new(DummyExt2(2));
+		exts.register(box1);
+		exts.register(box2);
+
+		{
+			let ext = exts.get_mut(TypeId::of::<DummyExt>()).expect("Extension 1 is registered");
+			let ext_ty = ext.downcast_mut::<DummyExt>().expect("Downcasting works for Extension 1");
+			assert_eq!(ext_ty.0, 1);
+		}
+		{
+			let ext2 = exts.get_mut(TypeId::of::<DummyExt2>()).expect("Extension 2 is registered");
+			let ext_ty2 =
+				ext2.downcast_mut::<DummyExt2>().expect("Downcasting works for Extension 2");
+			assert_eq!(ext_ty2.0, 2);
+		}
+	}
 }