From 2a55e07b1ffe1b4e4e39786e9942e0fed1c3f4ee Mon Sep 17 00:00:00 2001
From: Shaopeng Wang <spxwang@gmail.com>
Date: Thu, 17 Sep 2020 20:24:58 +1200
Subject: [PATCH] Add VestedTransferOrigin. (#287)

---
 vesting/src/lib.rs   |  7 +++++--
 vesting/src/mock.rs  | 20 ++++++++++++++++++++
 vesting/src/tests.rs | 18 +++++++++++++++++-
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/vesting/src/lib.rs b/vesting/src/lib.rs
index ddc07e992..74a347a4b 100644
--- a/vesting/src/lib.rs
+++ b/vesting/src/lib.rs
@@ -29,7 +29,7 @@
 use codec::{Decode, Encode, HasCompact};
 use frame_support::{
 	decl_error, decl_event, decl_module, decl_storage, ensure,
-	traits::{Currency, ExistenceRequirement, Get, LockIdentifier, LockableCurrency, WithdrawReasons},
+	traits::{Currency, EnsureOrigin, ExistenceRequirement, Get, LockIdentifier, LockableCurrency, WithdrawReasons},
 	weights::Weight,
 };
 use frame_system::{ensure_root, ensure_signed};
@@ -122,6 +122,9 @@ pub trait Trait: frame_system::Trait {
 	/// The minimum amount transferred to call `vested_transfer`.
 	type MinVestedTransfer: Get<BalanceOf<Self>>;
 
+	/// Required origin for vested transfer.
+	type VestedTransferOrigin: EnsureOrigin<Self::Origin, Success = Self::AccountId>;
+
 	/// Weight information for extrinsics in this module.
 	type WeightInfo: WeightInfo;
 }
@@ -220,7 +223,7 @@ decl_module! {
 			dest: <T::Lookup as StaticLookup>::Source,
 			schedule: VestingScheduleOf<T>,
 		) {
-			let from = ensure_signed(origin)?;
+			let from = T::VestedTransferOrigin::ensure_origin(origin)?;
 			let to = T::Lookup::lookup(dest)?;
 			Self::do_vested_transfer(&from, &to, schedule.clone())?;
 
diff --git a/vesting/src/mock.rs b/vesting/src/mock.rs
index 868ca7874..fef068308 100644
--- a/vesting/src/mock.rs
+++ b/vesting/src/mock.rs
@@ -3,6 +3,7 @@
 #![cfg(test)]
 
 use frame_support::{impl_outer_event, impl_outer_origin, parameter_types};
+use frame_system::RawOrigin;
 use pallet_balances;
 use sp_core::H256;
 use sp_runtime::{testing::Header, traits::IdentityLookup, Perbill};
@@ -81,10 +82,29 @@ impl pallet_balances::Trait for Runtime {
 }
 pub type PalletBalances = pallet_balances::Module<Runtime>;
 
+pub struct EnsureAliceOrBob;
+impl EnsureOrigin<Origin> for EnsureAliceOrBob {
+	type Success = AccountId;
+
+	fn try_origin(o: Origin) -> Result<Self::Success, Origin> {
+		Into::<Result<RawOrigin<AccountId>, Origin>>::into(o).and_then(|o| match o {
+			RawOrigin::Signed(ALICE) => Ok(ALICE),
+			RawOrigin::Signed(BOB) => Ok(BOB),
+			r => Err(Origin::from(r)),
+		})
+	}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn successful_origin() -> Origin {
+		Origin::from(RawOrigin::Signed(Default::default()))
+	}
+}
+
 impl Trait for Runtime {
 	type Event = TestEvent;
 	type Currency = PalletBalances;
 	type MinVestedTransfer = MinVestedTransfer;
+	type VestedTransferOrigin = EnsureAliceOrBob;
 	type WeightInfo = ();
 }
 pub type Vesting = Module<Runtime>;
diff --git a/vesting/src/tests.rs b/vesting/src/tests.rs
index 97d56fc1c..71faec09e 100644
--- a/vesting/src/tests.rs
+++ b/vesting/src/tests.rs
@@ -3,7 +3,7 @@
 #![cfg(test)]
 
 use super::*;
-use frame_support::{assert_err, assert_ok, traits::WithdrawReason};
+use frame_support::{assert_err, assert_noop, assert_ok, error::BadOrigin, traits::WithdrawReason};
 use mock::{ExtBuilder, Origin, PalletBalances, Runtime, System, TestEvent, Vesting, ALICE, BOB, CHARLIE};
 use pallet_balances::{BalanceLock, Reasons};
 
@@ -188,6 +188,22 @@ fn vested_transfer_fails_if_overflow() {
 	});
 }
 
+#[test]
+fn vested_transfer_fails_if_bad_origin() {
+	ExtBuilder::build().execute_with(|| {
+		let schedule = VestingSchedule {
+			start: 0u64,
+			period: 10u64,
+			period_count: 1u32,
+			per_period: 100u64,
+		};
+		assert_noop!(
+			Vesting::vested_transfer(Origin::signed(CHARLIE), BOB, schedule.clone()),
+			BadOrigin
+		);
+	});
+}
+
 #[test]
 fn claim_works() {
 	ExtBuilder::build().execute_with(|| {