From a2c7f9eb1b40ca75139c74409da35ce1b18d5125 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20K=C3=B6cher?= <git@kchr.de>
Date: Tue, 6 Dec 2022 15:58:52 +0100
Subject: [PATCH] frame-support: Introduce `EnsureOriginOrHigherPrivilege`
 (#12844)

* frame-support: Introduce `EnsureOriginOrHigherPrivilege`

This adds a new `EnsureOrigin` implementation that checks if a given origin matches or if the origin
is has a higher or equal origin matches or if the origin is has a higher or equal privilege.

* FMT
---
 frame/support/src/traits.rs          |  4 +-
 frame/support/src/traits/dispatch.rs | 92 +++++++++++++++++++++++++++-
 2 files changed, 91 insertions(+), 5 deletions(-)

diff --git a/frame/support/src/traits.rs b/frame/support/src/traits.rs
index 1b1d5d68cb52c..e5ba98fe0c5bb 100644
--- a/frame/support/src/traits.rs
+++ b/frame/support/src/traits.rs
@@ -99,8 +99,8 @@ mod dispatch;
 pub use dispatch::EnsureOneOf;
 pub use dispatch::{
 	AsEnsureOriginWithArg, CallerTrait, EitherOf, EitherOfDiverse, EnsureOrigin,
-	EnsureOriginWithArg, MapSuccess, NeverEnsureOrigin, OriginTrait, TryMapSuccess,
-	UnfilteredDispatchable,
+	EnsureOriginEqualOrHigherPrivilege, EnsureOriginWithArg, MapSuccess, NeverEnsureOrigin,
+	OriginTrait, TryMapSuccess, UnfilteredDispatchable,
 };
 
 mod voting;
diff --git a/frame/support/src/traits/dispatch.rs b/frame/support/src/traits/dispatch.rs
index b96cfae4500e2..36ddf5b507c0c 100644
--- a/frame/support/src/traits/dispatch.rs
+++ b/frame/support/src/traits/dispatch.rs
@@ -20,10 +20,12 @@
 use crate::dispatch::{DispatchResultWithPostInfo, Parameter, RawOrigin};
 use codec::MaxEncodedLen;
 use sp_runtime::{
-	traits::{BadOrigin, Member, Morph, TryMorph},
+	traits::{BadOrigin, Get, Member, Morph, TryMorph},
 	Either,
 };
-use sp_std::marker::PhantomData;
+use sp_std::{cmp::Ordering, marker::PhantomData};
+
+use super::misc;
 
 /// Some sort of check on the origin is performed by this object.
 pub trait EnsureOrigin<OuterOrigin> {
@@ -59,7 +61,7 @@ pub trait EnsureOrigin<OuterOrigin> {
 	}
 }
 
-/// `EnsureOrigin` implementation that always fails.
+/// [`EnsureOrigin`] implementation that always fails.
 pub struct NeverEnsureOrigin<Success>(sp_std::marker::PhantomData<Success>);
 impl<OO, Success> EnsureOrigin<OO> for NeverEnsureOrigin<Success> {
 	type Success = Success;
@@ -72,6 +74,90 @@ impl<OO, Success> EnsureOrigin<OO> for NeverEnsureOrigin<Success> {
 	}
 }
 
+/// [`EnsureOrigin`] implementation that checks that an origin has equal or higher privilege
+/// compared to the expected `Origin`.
+///
+/// It will take the shortcut of comparing the incoming origin with the expected `Origin` and if
+/// both are the same the origin is accepted.
+///
+/// # Example
+///
+/// ```rust
+/// # use frame_support::traits::{EnsureOriginEqualOrHigherPrivilege, PrivilegeCmp, EnsureOrigin as _};
+/// # use sp_runtime::traits::{parameter_types, Get};
+/// # use sp_std::cmp::Ordering;
+///
+/// #[derive(Eq, PartialEq, Debug)]
+/// pub enum Origin {
+///     Root,
+///     SomethingBelowRoot,
+///     NormalUser,
+/// }
+///
+/// struct OriginPrivilegeCmp;
+///
+/// impl PrivilegeCmp<Origin> for OriginPrivilegeCmp {
+///     fn cmp_privilege(left: &Origin, right: &Origin) -> Option<Ordering> {
+///         match (left, right) {
+///             (Origin::Root, Origin::Root) => Some(Ordering::Equal),
+///             (Origin::Root, _) => Some(Ordering::Greater),
+///             (Origin::SomethingBelowRoot, Origin::SomethingBelowRoot) => Some(Ordering::Equal),
+///             (Origin::SomethingBelowRoot, Origin::Root) => Some(Ordering::Less),
+///             (Origin::SomethingBelowRoot, Origin::NormalUser) => Some(Ordering::Greater),
+///             (Origin::NormalUser, Origin::NormalUser) => Some(Ordering::Equal),
+///             (Origin::NormalUser, _) => Some(Ordering::Less),
+///         }
+///     }
+/// }
+///
+/// parameter_types! {
+///     pub const ExpectedOrigin: Origin = Origin::SomethingBelowRoot;
+/// }
+///
+/// type EnsureOrigin = EnsureOriginEqualOrHigherPrivilege<ExpectedOrigin, OriginPrivilegeCmp>;
+///
+/// // `Root` has an higher privilege as our expected origin.
+/// assert!(EnsureOrigin::ensure_origin(Origin::Root).is_ok());
+/// // `SomethingBelowRoot` is exactly the expected origin.
+/// assert!(EnsureOrigin::ensure_origin(Origin::SomethingBelowRoot).is_ok());
+/// // The `NormalUser` origin is not allowed.
+/// assert!(EnsureOrigin::ensure_origin(Origin::NormalUser).is_err());
+/// ```
+pub struct EnsureOriginEqualOrHigherPrivilege<Origin, PrivilegeCmp>(
+	sp_std::marker::PhantomData<(Origin, PrivilegeCmp)>,
+);
+
+impl<OuterOrigin, Origin, PrivilegeCmp> EnsureOrigin<OuterOrigin>
+	for EnsureOriginEqualOrHigherPrivilege<Origin, PrivilegeCmp>
+where
+	Origin: Get<OuterOrigin>,
+	OuterOrigin: Eq,
+	PrivilegeCmp: misc::PrivilegeCmp<OuterOrigin>,
+{
+	type Success = ();
+
+	fn try_origin(o: OuterOrigin) -> Result<Self::Success, OuterOrigin> {
+		let expected_origin = Origin::get();
+
+		// If this is the expected origin, it has the same privilege.
+		if o == expected_origin {
+			return Ok(())
+		}
+
+		let cmp = PrivilegeCmp::cmp_privilege(&o, &expected_origin);
+
+		match cmp {
+			Some(Ordering::Equal) | Some(Ordering::Greater) => Ok(()),
+			None | Some(Ordering::Less) => Err(o),
+		}
+	}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn try_successful_origin() -> Result<OuterOrigin, ()> {
+		Ok(Origin::get())
+	}
+}
+
 /// Some sort of check on the origin is performed by this object.
 pub trait EnsureOriginWithArg<OuterOrigin, Argument> {
 	/// A return type.