"OR gate" for EnsureOrigin (#6237)

* 'OR gate' for EnsureOrigin.

* Formatting.

* More formatting.

* Add docstring; Update 'Success' type.

* Bump runtime impl_version.

* Fix successful_origin.

* Add either into std feature list.

* Update docs.

substrate/frame/identity/src/lib.rs

@@ -78,7 +78,7 @@ use frame_support::{
 	traits::{Currency, ReservableCurrency, OnUnbalanced, Get, BalanceStatus, EnsureOrigin},
 	weights::Weight,
 };
-use frame_system::{self as system, ensure_signed, ensure_root};
+use frame_system::{self as system, ensure_signed};
 
 mod benchmarking;
 
@@ -635,9 +635,7 @@ decl_module! {
 		/// # </weight>
 		#[weight = weight_for::add_registrar::<T>(T::MaxRegistrars::get().into()) ]
 		fn add_registrar(origin, account: T::AccountId) -> DispatchResultWithPostInfo {
-			T::RegistrarOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::RegistrarOrigin::ensure_origin(origin)?;
 
 			let (i, registrar_count) = <Registrars<T>>::try_mutate(
 				|registrars| -> Result<(RegistrarIndex, usize), DispatchError> {
@@ -1108,9 +1106,7 @@ decl_module! {
 			T::MaxAdditionalFields::get().into(), // X
 		)]
 		fn kill_identity(origin, target: <T::Lookup as StaticLookup>::Source) -> DispatchResultWithPostInfo {
-			T::ForceOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::ForceOrigin::ensure_origin(origin)?;
 
 			// Figure out who we're meant to be clearing.
 			let target = T::Lookup::lookup(target)?;
@@ -1435,7 +1431,7 @@ mod tests {
 		new_test_ext().execute_with(|| {
 			assert_ok!(Identity::set_identity(Origin::signed(10), ten()));
 			assert_ok!(Identity::set_subs(Origin::signed(10), vec![(20, Data::Raw(vec![40; 1]))]));
-			assert_ok!(Identity::kill_identity(Origin::ROOT, 10));
+			assert_ok!(Identity::kill_identity(Origin::signed(2), 10));
 			assert_eq!(Balances::free_balance(10), 80);
 			assert!(Identity::super_of(20).is_none());
 		});

substrate/frame/membership/src/lib.rs

@@ -28,7 +28,7 @@ use frame_support::{
 	decl_module, decl_storage, decl_event, decl_error,
 	traits::{ChangeMembers, InitializeMembers, EnsureOrigin},
 };
-use frame_system::{self as system, ensure_root, ensure_signed};
+use frame_system::{self as system, ensure_signed};
 
 pub trait Trait<I=DefaultInstance>: frame_system::Trait {
 	/// The overarching event type.
@@ -120,9 +120,7 @@ decl_module! {
 		/// May only be called from `AddOrigin` or root.
 		#[weight = 50_000_000]
 		pub fn add_member(origin, who: T::AccountId) {
-			T::AddOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::AddOrigin::ensure_origin(origin)?;
 
 			let mut members = <Members<T, I>>::get();
 			let location = members.binary_search(&who).err().ok_or(Error::<T, I>::AlreadyMember)?;
@@ -139,9 +137,7 @@ decl_module! {
 		/// May only be called from `RemoveOrigin` or root.
 		#[weight = 50_000_000]
 		pub fn remove_member(origin, who: T::AccountId) {
-			T::RemoveOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::RemoveOrigin::ensure_origin(origin)?;
 
 			let mut members = <Members<T, I>>::get();
 			let location = members.binary_search(&who).ok().ok_or(Error::<T, I>::NotMember)?;
@@ -161,9 +157,7 @@ decl_module! {
 		/// Prime membership is *not* passed from `remove` to `add`, if extant.
 		#[weight = 50_000_000]
 		pub fn swap_member(origin, remove: T::AccountId, add: T::AccountId) {
-			T::SwapOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::SwapOrigin::ensure_origin(origin)?;
 
 			if remove == add { return Ok(()) }
 
@@ -190,9 +184,7 @@ decl_module! {
 		/// May only be called from `ResetOrigin` or root.
 		#[weight = 50_000_000]
 		pub fn reset_members(origin, members: Vec<T::AccountId>) {
-			T::ResetOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::ResetOrigin::ensure_origin(origin)?;
 
 			let mut members = members;
 			members.sort();
@@ -241,9 +233,7 @@ decl_module! {
 		/// Set the prime member. Must be a current member.
 		#[weight = 50_000_000]
 		pub fn set_prime(origin, who: T::AccountId) {
-			T::PrimeOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::PrimeOrigin::ensure_origin(origin)?;
 			Self::members().binary_search(&who).ok().ok_or(Error::<T, I>::NotMember)?;
 			Prime::<T, I>::put(&who);
 			T::MembershipChanged::set_prime(Some(who));
@@ -252,9 +242,7 @@ decl_module! {
 		/// Remove the prime member if it exists.
 		#[weight = 50_000_000]
 		pub fn clear_prime(origin) {
-			T::PrimeOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::PrimeOrigin::ensure_origin(origin)?;
 			Prime::<T, I>::kill();
 			T::MembershipChanged::set_prime(None);
 		}

substrate/frame/nicks/src/lib.rs

@@ -47,7 +47,7 @@ use frame_support::{
 	decl_module, decl_event, decl_storage, ensure, decl_error,
 	traits::{Currency, EnsureOrigin, ReservableCurrency, OnUnbalanced, Get},
 };
-use frame_system::{self as system, ensure_signed, ensure_root};
+use frame_system::{self as system, ensure_signed};
 
 type BalanceOf<T> = <<T as Trait>::Currency as Currency<<T as frame_system::Trait>::AccountId>>::Balance;
 type NegativeImbalanceOf<T> = <<T as Trait>::Currency as Currency<<T as frame_system::Trait>::AccountId>>::NegativeImbalance;
@@ -197,9 +197,7 @@ decl_module! {
 		/// # </weight>
 		#[weight = 70_000_000]
 		fn kill_name(origin, target: <T::Lookup as StaticLookup>::Source) {
-			T::ForceOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::ForceOrigin::ensure_origin(origin)?;
 
 			// Figure out who we're meant to be clearing.
 			let target = T::Lookup::lookup(target)?;
@@ -225,9 +223,7 @@ decl_module! {
 		/// # </weight>
 		#[weight = 70_000_000]
 		fn force_name(origin, target: <T::Lookup as StaticLookup>::Source, name: Vec<u8>) {
-			T::ForceOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::ForceOrigin::ensure_origin(origin)?;
 
 			let target = T::Lookup::lookup(target)?;
 			let deposit = <NameOf<T>>::get(&target).map(|x| x.1).unwrap_or_else(Zero::zero);

substrate/frame/scored-pool/src/lib.rs

@@ -318,9 +318,7 @@ decl_module! {
 			dest: <T::Lookup as StaticLookup>::Source,
 			index: u32
 		) {
-			T::KickOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::KickOrigin::ensure_origin(origin)?;
 
 			let who = T::Lookup::lookup(dest)?;
 
@@ -344,9 +342,7 @@ decl_module! {
 			index: u32,
 			score: T::Score
 		) {
-			T::ScoreOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::ScoreOrigin::ensure_origin(origin)?;
 
 			let who = T::Lookup::lookup(dest)?;
 

substrate/frame/staking/src/lib.rs

@@ -1906,9 +1906,7 @@ decl_module! {
 			.saturating_add((35 * WEIGHT_PER_MICROS).saturating_mul(slash_indices.len() as Weight))
 		]
 		fn cancel_deferred_slash(origin, era: EraIndex, slash_indices: Vec<u32>) {
-			T::SlashCancelOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::SlashCancelOrigin::ensure_origin(origin)?;
 
 			ensure!(!slash_indices.is_empty(), Error::<T>::EmptyTargets);
 			ensure!(is_sorted_and_unique(&slash_indices), Error::<T>::NotSortedAndUnique);

substrate/frame/system/src/lib.rs

@@ -102,7 +102,7 @@ use sp_std::marker::PhantomData;
 use sp_std::fmt::Debug;
 use sp_version::RuntimeVersion;
 use sp_runtime::{
-	RuntimeDebug, Perbill, DispatchError, DispatchResult,
+	RuntimeDebug, Perbill, DispatchError, DispatchResult, Either,
 	generic::{self, Era},
 	transaction_validity::{
 		ValidTransaction, TransactionPriority, TransactionLongevity, TransactionValidityError,
@@ -847,6 +847,30 @@ impl<O, T> EnsureOrigin<O> for EnsureNever<T> {
 	}
 }
 
+/// The "OR gate" implementation of `EnsureOrigin`.
+///
+/// Origin check will pass if `L` or `R` origin check passes. `L` is tested first.
+pub struct EnsureOneOf<AccountId, L, R>(sp_std::marker::PhantomData<(AccountId, L, R)>);
+impl<
+	AccountId,
+	O: Into<Result<RawOrigin<AccountId>, O>> + From<RawOrigin<AccountId>>,
+	L: EnsureOrigin<O>,
+	R: EnsureOrigin<O>,
+> EnsureOrigin<O> for EnsureOneOf<AccountId, L, R> {
+	type Success = Either<L::Success, R::Success>;
+	fn try_origin(o: O) -> Result<Self::Success, O> {
+		L::try_origin(o).map_or_else(
+			|o| R::try_origin(o).map(|o| Either::Right(o)),
+			|o| Ok(Either::Left(o)),
+		)
+	}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn successful_origin() -> O {
+		L::successful_origin()
+	}
+}
+
 /// Ensure that the origin `o` represents a signed extrinsic (i.e. transaction).
 /// Returns `Ok` with the account that signed the extrinsic or an `Err` otherwise.
 pub fn ensure_signed<OuterOrigin, AccountId>(o: OuterOrigin) -> Result<AccountId, BadOrigin>
@@ -1879,7 +1903,7 @@ pub(crate) mod tests {
 	use sp_core::H256;
 	use sp_runtime::{traits::{BlakeTwo256, IdentityLookup, SignedExtension}, testing::Header, DispatchError};
 	use frame_support::{
-		impl_outer_origin, parameter_types, assert_ok, assert_noop,
+		impl_outer_origin, parameter_types, assert_ok, assert_noop, assert_err,
 		weights::WithPostDispatchInfo,
 	};
 
@@ -2701,4 +2725,15 @@ pub(crate) mod tests {
 			assert!(System::events().len() == 1);
 		});
 	}
+
+	#[test]
+	fn ensure_one_of_works() {
+		fn ensure_root_or_signed(o: RawOrigin<u64>) -> Result<Either<(), u64>, Origin> {
+			EnsureOneOf::<u64, EnsureRoot<u64>, EnsureSigned<u64>>::try_origin(o.into())
+		}
+
+		assert_ok!(ensure_root_or_signed(RawOrigin::Root), Either::Left(()));
+		assert_ok!(ensure_root_or_signed(RawOrigin::Signed(0)), Either::Right(0));
+		assert_err!(ensure_root_or_signed(RawOrigin::None), Origin::from(RawOrigin::None));
+	}
 }

substrate/frame/treasury/src/lib.rs

@@ -102,7 +102,7 @@ use sp_runtime::{Permill, ModuleId, Percent, RuntimeDebug, traits::{
 use frame_support::weights::{Weight, DispatchClass};
 use frame_support::traits::{Contains, ContainsLengthBound, EnsureOrigin};
 use codec::{Encode, Decode};
-use frame_system::{self as system, ensure_signed, ensure_root};
+use frame_system::{self as system, ensure_signed};
 
 mod tests;
 mod benchmarking;
@@ -362,9 +362,7 @@ decl_module! {
 		/// # </weight>
 		#[weight = (130_000_000 + T::DbWeight::get().reads_writes(2, 2), DispatchClass::Operational)]
 		fn reject_proposal(origin, #[compact] proposal_id: ProposalIndex) {
-			T::RejectOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::RejectOrigin::ensure_origin(origin)?;
 
 			let proposal = <Proposals<T>>::take(&proposal_id).ok_or(Error::<T>::InvalidProposalIndex)?;
 			let value = proposal.bond;
@@ -384,9 +382,7 @@ decl_module! {
 		/// # </weight>
 		#[weight = (34_000_000 + T::DbWeight::get().reads_writes(2, 1), DispatchClass::Operational)]
 		fn approve_proposal(origin, #[compact] proposal_id: ProposalIndex) {
-			T::ApproveOrigin::try_origin(origin)
-				.map(|_| ())
-				.or_else(ensure_root)?;
+			T::ApproveOrigin::ensure_origin(origin)?;
 
 			ensure!(<Proposals<T>>::contains_key(proposal_id), Error::<T>::InvalidProposalIndex);
 			Approvals::mutate(|v| v.push(proposal_id));