SubAccounts

//! # SubAccounts Pallet
//! <!-- Original author of paragraph: @BMateo
//!
//! ## Overview
//!
//! Pallet allows users to add and remove sub-accounts to send extrinsics for their profiles, as
//! well as add and remove verified addresses associated with asset IDs.
//!
//! ### Goals
//!
//! The pallet is designed to make the following possible:
//!
//! * Allow the main account or any subaccount of a profile to add/remove sub accounts.
//! * Allow the user to add or remove verified addresses associated with an asset ID to their
//!   account.
//!
//! ## Interface
//!
//! ### Privileged Functions
//!
//! - `add_sub_account`: The sender can add a sub account for its profile
//! - `remove_sub_account`: The sender can remove any sub account of its profile.
//! - `add_verified_address`: The sender has the ability to add a verified address associated with
//!   an asset ID to an account.
//! - `remove_verified_address`:The sender has the ability to remove a verified address associated
//!   with an asset ID
//! to an account.
//!
//!//! Please refer to the [`Call`] enum and its associated variants for documentation on each
//! function.

#![cfg_attr(not(feature = "std"), no_std)]

pub use pallet::*;
mod types;
pub use types::*;

use base64ct::{ Base64, Encoding };
use bitcoin::{ self, address::{ NetworkChecked, NetworkUnchecked }, Address, Network };
use frame_support::{
	pallet_prelude::*,
	traits::{ fungibles, tokens::fungibles::Inspect, UnixTime },
};
use frame_system::{ ensure_signed, pallet_prelude::OriginFor };
use sp_io::{ crypto::secp256k1_ecdsa_recover, hashing::keccak_256 };
use sp_runtime::traits::Hash;
use sp_std::vec::Vec;
use traits::{
	asset::{ AssetInterface, TokenType },
	profile::{ ProfileInspect, ProfileInterface },
	subaccounts::{ AccountOrigin, ChargeFees, SubAccounts, SubAccountsAddress },
};


#[cfg(test)]
mod mock;

#[cfg(test)]
mod tests;

#[cfg(feature = "runtime-benchmarks")]
mod benchmarking;

pub mod weights;
pub use weights::WeightInfo;

#[frame_support::pallet]
pub mod pallet {
	use super::*;

	pub type BalanceOf<T> =
		<<T as Config>::Fungibles as Inspect<<T as frame_system::Config>::AccountId>>::Balance;

	pub type AssetIdOf<T> =
		<<T as Config>::Fungibles as Inspect<<T as frame_system::Config>::AccountId>>::AssetId;

	/// Type used to represent the token decimals.
	pub type Decimals = u8;

	#[pallet::pallet]
	pub struct Pallet<T>(_);

	/// Configure the pallet by specifying the parameters and types on which it depends.
	#[pallet::config]
	pub trait Config: frame_system::Config {
		/// Because this pallet emits events, it depends on the runtime's definition of an event.
		type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;

		#[pallet::constant]
		type StringLimit: Get<u32>;

		/// Type to access the profile pallet
		type Profile: ProfileInspect<Self::AccountId, Self::StringLimit> +
			ProfileInterface<Self::AccountId, Self::StringLimit>;

		/// Type to access the Assets Pallet.
		type Fungibles: fungibles::Inspect<Self::AccountId> +
			fungibles::Mutate<Self::AccountId> +
			fungibles::metadata::Inspect<Self::AccountId> +
			AssetInterface<
				<Self::Fungibles as Inspect<Self::AccountId>>::AssetId,
				BalanceOf<Self>,
				Self::AccountId,
				Decimals,
				TokenType
			>;

		/// Provides actual time
		type TimeProvider: UnixTime;

		#[pallet::constant]
		type BitcoinAssetId: Get<AssetIdOf<Self>>;

		#[pallet::constant]
		type EthereumAssetId: Get<AssetIdOf<Self>>;

		#[pallet::constant]
		type NearAssetId: Get<AssetIdOf<Self>>;

		#[pallet::constant]
		type PolygonAssetId: Get<AssetIdOf<Self>>;

		/// Weight information for extrinsics in this pallet.
		type WeightInfo: WeightInfo;
	}

	/// Store the main account of a given address
	#[pallet::storage]
	#[pallet::getter(fn sub_account)]
	pub type SubAccount<T: Config> = StorageMap<
		_,
		Blake2_128Concat,
		T::AccountId,
		(T::AccountId, AccountOrigin)
	>;

	/// Store the verified address of a account for a specific asset ID and the timestamp it was
	/// added.
	#[pallet::storage]
	#[pallet::getter(fn get_account_asset_address)]
	pub type AccountAssetVerifiedAddress<T: Config> = StorageNMap<
		_,
		(
			NMapKey<Blake2_128Concat, T::AccountId>,
			NMapKey<Blake2_128Concat, AssetIdOf<T>>,
			NMapKey<Blake2_128Concat, T::Hash>,
		),
		(BoundedVec<u8, T::StringLimit>, u64),
		OptionQuery
	>;

	/// Store the addresses used.
	#[pallet::storage]
	#[pallet::getter(fn get_addresses)]
	pub type AddressesUsed<T: Config> = StorageDoubleMap<
		_,
		Blake2_128Concat,
		AssetIdOf<T>,
		Blake2_128Concat,
		T::Hash,
		bool,
		OptionQuery
	>;

	#[pallet::event]
	#[pallet::generate_deposit(pub(super) fn deposit_event)]
	pub enum Event<T: Config> {
		/// A sub account has been added
		SubAccountAdded {
			main: T::AccountId,
			sub: T::AccountId,
		},
		/// A sub account has been removed
		SubAccountRemoved {
			main: T::AccountId,
			sub: T::AccountId,
		},
		/// A verified address has been associated with an account and a specific asset ID.
		VerifiedAddressAdded {
			who: T::AccountId,
			asset_id: AssetIdOf<T>,
			address: BoundedVec<u8, T::StringLimit>,
		},
		/// A verified address has been dissociated from an account and a specific asset ID.
		VerifiedAddressRemoved {
			who: T::AccountId,
			asset_id: AssetIdOf<T>,
			address: BoundedVec<u8, T::StringLimit>,
		},
	}

	#[pallet::error]
	pub enum Error<T> {
		/// Sender is not a sub account
		NoSubAccount,
		/// Sender is not a sub account of the given address
		NotAllowed,
		/// Cannot remove all sub-accounts
		NoAccountsLeft,
		/// Cannot add a sub account twice
		AlreadySubAccount,
		/// The signature verification failed.
		VerificationFailed,
		/// The signed message contains an error.
		InvalidSignedMessage,
		/// The public key contains an error.
		InvalidPublicKey,
		/// The asset ID used to register an address is currently not functioning to verify the
		/// signature..
		UnprotectedAsset,
		/// The address being attempted to be added is already in use.
		AddressAlreadyBeingUsed,
		/// The address being attempted to be removed was not found.
		AddressNotFound,
		/// The account associated with the asset ID already has a verified address.
		AccountAssetAddressAlreadyRegistered,
		/// The account associated with the asset ID does not have a specific verified address.
		AccountAssetAddressNotFound,
		/// The account address linked to the specified asset ID could not be located.
		AccountAddressNotFound,
		/// The account is not the address owner.
		AccountNotAddressOwner,
	}

	#[pallet::call(weight(<T as Config>::WeightInfo))]
	impl<T: Config> Pallet<T> {
		/// The origin can add a sub account for the given main account.
		///
		/// The origin must be Signed and the sender should have access to 'main'
		///
		/// Parameters:
		/// - `main`: The address that has a profile associated
		/// - `new_sub_account`: The address that will be added as a connected account of 'main'
		/// - `account_origin`: The method used to create this account.
		/// Emits `SubAccountAdded` event when successful.
		#[pallet::call_index(0)]
		pub fn add_sub_account(
			origin: OriginFor<T>,
			main: T::AccountId,
			new_sub_account: T::AccountId,
			account_origin: AccountOrigin
		) -> DispatchResultWithPostInfo {
			let sender = ensure_signed(origin)?;

			// Check if the sender has permission to add a sub account
			Self::is_sub_account(sender, main.clone())?;

			// Check that the sub account wasn't added before.
			ensure!(
				!SubAccount::<T>::contains_key(new_sub_account.clone()),
				Error::<T>::AlreadySubAccount
			);

			SubAccount::<T>::insert(new_sub_account.clone(), (main.clone(), account_origin));

			// Store the sub account in the profile info
			T::Profile::add_sub_account(main.clone(), new_sub_account.clone())?;

			// Emit an event
			Self::deposit_event(Event::SubAccountAdded { main, sub: new_sub_account });

			Ok(().into())
		}

		/// The origin can remove a sub account for the given main account.
		///
		/// The origin must be Signed and the sender should have access to 'main'
		///
		///	Can't remove all the connected accounts for a profile
		///
		/// Parameters:
		/// - `main`: The address that has a profile associated
		/// - `sub_account_to_remove`: The address that will be removed as a connected account of
		///   'main'
		///
		/// Emits `SubAccountRemoved` event when successful.
		#[pallet::call_index(1)]
		pub fn remove_sub_account(
			origin: OriginFor<T>,
			main: T::AccountId,
			sub_account_to_remove: T::AccountId
		) -> DispatchResultWithPostInfo {
			let sender = ensure_signed(origin)?;

			// Verify that the sender and the account to be removed are sub accounts
			Self::is_sub_account(sender, main.clone())?;
			Self::is_sub_account(sub_account_to_remove.clone(), main.clone())?;

			// verify that the profile will have a connected account after removing this one.
			// Otherwise the user will lose access to the main account
			let sub_accounts = T::Profile::get_sub_accounts(main.clone())?;
			ensure!(sub_accounts.len() > 1, Error::<T>::NoAccountsLeft);

			SubAccount::<T>::remove(sub_account_to_remove.clone());

			// Remove sub account from the profile info
			T::Profile::remove_sub_account(main.clone(), sub_account_to_remove.clone())?;

			// Emit an event
			Self::deposit_event(Event::SubAccountRemoved { main, sub: sub_account_to_remove });

			Ok(().into())
		}

		/// The origin can add a verified address associated with an asset ID.
		///
		/// Parameters:
		/// - `asset_id`: The asset ID that corresponds to the chain of the address.
		/// - `address`: The address to verified and associated to an account.
		/// - `signature`: The signature that serves as proof of the user's ownership of the
		///   address.
		///
		/// Emits `VerifiedAddressAdded` event when successful.
		#[pallet::call_index(2)]
		pub fn add_verified_address(
			origin: OriginFor<T>,
			asset_id: AssetIdOf<T>,
			address: BoundedVec<u8, T::StringLimit>,
			signature: BoundedVec<u8, T::StringLimit>
		) -> DispatchResult {
			let mut who = ensure_signed(origin)?;
			// Mutate the origin to use the main account
			who = <pallet::Pallet<T> as traits::subaccounts::SubAccounts<
				T::AccountId,
				AccountOrigin
			>>::get_main_account(who)?;

			<T::Fungibles as traits::asset::AssetInterface<
				AssetIdOf<T>,
				BalanceOf<T>,
				T::AccountId,
				Decimals,
				TokenType
			>>::asset_exists(asset_id.clone())?;

			Self::verify_signature(asset_id.clone(), address.clone(), signature)?;

			let address_hash = T::Hashing::hash(&address);

			let timestamp_now = T::TimeProvider::now().as_secs();

			ensure!(
				!AddressesUsed::<T>::contains_key(&asset_id, address_hash),
				Error::<T>::AddressAlreadyBeingUsed
			);
			ensure!(
				!AccountAssetVerifiedAddress::<T>::contains_key((&who, &asset_id, &address_hash)),
				Error::<T>::AccountAssetAddressAlreadyRegistered
			);

			AddressesUsed::<T>::insert(&asset_id, address_hash, true);
			AccountAssetVerifiedAddress::<T>::insert(
				(&who, &asset_id, &address_hash),
				(address.clone(), timestamp_now)
			);

			Self::deposit_event(Event::VerifiedAddressAdded { who, asset_id, address });

			Ok(())
		}

		/// The origin can remove a verified address associated with an asset ID.
		///
		/// Parameters:
		/// - `asset_id`: The asset ID that corresponds to the chain of the address.
		/// - `address`: The address to verified and associated to an account.
		///
		/// Emits `VerifiedAddressRemoved` event when successful.
		#[pallet::call_index(3)]
		pub fn remove_verified_address(
			origin: OriginFor<T>,
			asset_id: AssetIdOf<T>,
			address: BoundedVec<u8, T::StringLimit>
		) -> DispatchResult {
			let mut who = ensure_signed(origin)?;
			// Mutate the origin to use the main account
			who = <pallet::Pallet<T> as traits::subaccounts::SubAccounts<
				T::AccountId,
				AccountOrigin
			>>::get_main_account(who)?;

			<T::Fungibles as traits::asset::AssetInterface<
				AssetIdOf<T>,
				BalanceOf<T>,
				T::AccountId,
				Decimals,
				TokenType
			>>::asset_exists(asset_id.clone())?;

			let address_hash = T::Hashing::hash(&address);

			ensure!(
				AddressesUsed::<T>::contains_key(&asset_id, address_hash),
				Error::<T>::AddressNotFound
			);
			ensure!(
				AccountAssetVerifiedAddress::<T>::contains_key((&who, &asset_id, &address_hash)),
				Error::<T>::AccountAssetAddressNotFound
			);

			AddressesUsed::<T>::remove(&asset_id, address_hash);
			AccountAssetVerifiedAddress::<T>::remove((&who, &asset_id, &address_hash));

			Self::deposit_event(Event::VerifiedAddressRemoved { who, asset_id, address });

			Ok(())
		}

		#[pallet::weight({ 0 })]
		#[pallet::call_index(4)]
		pub fn add_no_verified_address(
			origin: OriginFor<T>,
			asset_id: AssetIdOf<T>,
			address: BoundedVec<u8, T::StringLimit>
		) -> DispatchResult {
			let mut who = ensure_signed(origin)?;
			// Mutate the origin to use the main account
			who = <pallet::Pallet<T> as traits::subaccounts::SubAccounts<
				T::AccountId,
				AccountOrigin
			>>::get_main_account(who)?;

			<T::Fungibles as traits::asset::AssetInterface<
				AssetIdOf<T>,
				BalanceOf<T>,
				T::AccountId,
				Decimals,
				TokenType
			>>::asset_exists(asset_id.clone())?;

			let address_hash = T::Hashing::hash(&address);

			let timestamp_now = T::TimeProvider::now().as_secs();

			ensure!(
				!AddressesUsed::<T>::contains_key(&asset_id, address_hash),
				Error::<T>::AddressAlreadyBeingUsed
			);
			ensure!(
				!AccountAssetVerifiedAddress::<T>::contains_key((&who, &asset_id, &address_hash)),
				Error::<T>::AccountAssetAddressAlreadyRegistered
			);

			AddressesUsed::<T>::insert(&asset_id, address_hash, true);
			AccountAssetVerifiedAddress::<T>::insert(
				(&who, &asset_id, &address_hash),
				(address.clone(), timestamp_now)
			);

			Self::deposit_event(Event::VerifiedAddressAdded { who, asset_id, address });

			Ok(())
		}
	}

	impl<T: Config> Pallet<T> {
		fn verify_signature(
			asset_id: AssetIdOf<T>,
			address: BoundedVec<u8, T::StringLimit>,
			signed_message: BoundedVec<u8, T::StringLimit>
		) -> Result<(), DispatchError> {
			let btc = T::BitcoinAssetId::get();
			let eth = T::EthereumAssetId::get();
			let polygon = T::PolygonAssetId::get();

			if asset_id == btc {
				Self::verify_bitcoin(signed_message, address)?;
				Ok(())
			} else if asset_id == eth || asset_id == polygon {
				Self::verify_eth(signed_message, address)?;
				Ok(())
			} else {
				Err(Error::<T>::UnprotectedAsset.into())
			}
		}

		// Constructs the message that Ethereum RPC's `personal_sign` and `eth_sign` would sign.
		pub fn ethereum_signable_message() -> Vec<u8> {
			let what = b"Hello Unit".to_vec();
			let mut l = what.len();
			let mut rev = Vec::new();
			// convert decimal into bytes
			while l > 0 {
				rev.push(b'0' + ((l % 10) as u8));
				l /= 10;
			}

			let mut v = b"\x19Ethereum Signed Message:\n".to_vec();
			v.extend(rev.into_iter().rev());
			v.extend_from_slice(&what);
			v
		}

		// Attempts to recover the Ethereum address from a message signature signed by using
		// the Ethereum RPC's `personal_sign` and `eth_sign`.
		pub fn eth_recover(s: &[u8; 65]) -> Option<[u8; 20]> {
			let msg = keccak_256(&Self::ethereum_signable_message());
			let mut res = [0u8; 20];
			res[0..20].copy_from_slice(
				&keccak_256(&secp256k1_ecdsa_recover(s, &msg).ok()?[..])[12..]
			);
			Some(res)
		}

		pub fn verify_eth(
			signed_message: BoundedVec<u8, T::StringLimit>,
			address: BoundedVec<u8, T::StringLimit>
		) -> Result<(), DispatchError> {
			let address_array: [u8; 20] = address
				.as_slice()
				.try_into()
				.map_err(|_| Error::<T>::InvalidPublicKey)?;

			let signed_message_array: [u8; 65] = signed_message
				.as_slice()
				.try_into()
				.map_err(|_| Error::<T>::InvalidSignedMessage)?;

			let address = Self::eth_recover(&signed_message_array);

			ensure!(address == Some(address_array), Error::<T>::VerificationFailed);
			Ok(())
		}

		pub fn verify_bitcoin(
			signed_message: BoundedVec<u8, T::StringLimit>,
			address: BoundedVec<u8, T::StringLimit>
		) -> Result<(), DispatchError> {
			// Construct signature
			let mut dec_buf = [0u8; 65];
			let signature_message = Base64::decode(
				sp_std::str
					::from_utf8(&signed_message[..])
					.map_err(|_| Error::<T>::InvalidSignedMessage)?,
				&mut dec_buf
			).map_err(|_| Error::<T>::InvalidSignedMessage)?;

			let signature = bitcoin::sign_message::MessageSignature
				::from_slice(signature_message)
				.map_err(|_| Error::<T>::InvalidSignedMessage)?;

			// Construct message
			let message = "Hello Unit";
			let msg_hash = bitcoin::sign_message::signed_msg_hash(message);

			// Construct address
			let address_str = sp_std::str
				::from_utf8(&address[..])
				.map_err(|_| Error::<T>::InvalidPublicKey)?;
			let address: Address<NetworkUnchecked> = address_str
				.parse()
				.map_err(|_| Error::<T>::InvalidPublicKey)?;
			let address: Address<NetworkChecked> = address
				.require_network(Network::Bitcoin)
				.map_err(|_| Error::<T>::InvalidPublicKey)?;

			// Verify signature is signed by address
			let secp = bitcoin::secp256k1::Secp256k1::new();
			// Returns Ok(false) if the proccess was successful but the message was not signed by
			// the address
			match signature.is_signed_by_address(&secp, &address, msg_hash) {
				Ok(true) => Ok(()),
				_ => Err(Error::<T>::VerificationFailed.into()),
			}
		}
	}

	impl<T: Config> SubAccounts<T::AccountId, AccountOrigin> for Pallet<T> {
		fn get_main_account(who: T::AccountId) -> Result<T::AccountId, DispatchError> {
			let address_main = SubAccount::<T>::get(who).ok_or(Error::<T>::NoSubAccount)?;

			let profile_main = T::Profile::get_profile_address(address_main.0)?;

			Ok(profile_main)
		}

		fn add_sub_account(
			main: T::AccountId,
			sub: T::AccountId,
			account_origin: AccountOrigin
		) -> Result<(), DispatchError> {
			SubAccount::<T>::insert(sub, (main, account_origin));
			Ok(())
		}

		fn is_sub_account(sender: T::AccountId, main: T::AccountId) -> Result<(), DispatchError> {
			let main_account_of_sender = <Self as SubAccounts<
				T::AccountId,
				AccountOrigin
			>>::get_main_account(sender)?;

			ensure!(main_account_of_sender == main, Error::<T>::NotAllowed);

			Ok(())
		}

		fn already_sub_account(who: T::AccountId) -> Result<(), DispatchError> {
			ensure!(!SubAccount::<T>::contains_key(who), Error::<T>::AlreadySubAccount);
			Ok(())
		}
	}

	impl<T: Config> SubAccountsAddress<AssetIdOf<T>, T::AccountId, T::Hash> for Pallet<T> {
		fn is_account_address_owner(
			who: T::AccountId,
			asset_id: AssetIdOf<T>,
			address_hash: T::Hash
		) -> Result<(), DispatchError> {
			AccountAssetVerifiedAddress::<T>
				::get((&who, &asset_id, &address_hash))
				.ok_or(Error::<T>::AccountAddressNotFound)?;

			Ok(())
		}
	}

	impl<T: Config> ChargeFees<T::AccountId> for Pallet<T> {
		fn get_main_account(who: &T::AccountId) -> Option<T::AccountId> {
			if let Some(x) = SubAccount::<T>::get(who) { Some(x.0) } else { None }
		}
	}
}