From c7cb1f25d1f8bb1a922d466e39ee935f5f027266 Mon Sep 17 00:00:00 2001
From: joe petrowski <25483142+joepetrowski@users.noreply.github.com>
Date: Wed, 22 May 2024 09:21:12 +0200
Subject: [PATCH] Add Extra Check in Primary Username Setter (#4534)

---
 prdoc/pr_4534.prdoc                   | 13 +++++
 substrate/frame/identity/src/lib.rs   |  4 +-
 substrate/frame/identity/src/tests.rs | 72 ++++++++++++++++++++++++++-
 3 files changed, 87 insertions(+), 2 deletions(-)
 create mode 100644 prdoc/pr_4534.prdoc

diff --git a/prdoc/pr_4534.prdoc b/prdoc/pr_4534.prdoc
new file mode 100644
index 000000000000..417e4d3dace0
--- /dev/null
+++ b/prdoc/pr_4534.prdoc
@@ -0,0 +1,13 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: Add Extra Check in Primary Username Setter
+
+doc:
+  - audience: Runtime User
+    description: |
+      Setting primary usernames requires an additional verification.
+
+crates:
+  - name: pallet-identity
+    bump: patch
diff --git a/substrate/frame/identity/src/lib.rs b/substrate/frame/identity/src/lib.rs
index 4a977880b315..5a36101cc2f7 100644
--- a/substrate/frame/identity/src/lib.rs
+++ b/substrate/frame/identity/src/lib.rs
@@ -1169,7 +1169,9 @@ pub mod pallet {
 		pub fn set_primary_username(origin: OriginFor<T>, username: Username<T>) -> DispatchResult {
 			// ensure `username` maps to `origin` (i.e. has already been set by an authority).
 			let who = ensure_signed(origin)?;
-			ensure!(AccountOfUsername::<T>::contains_key(&username), Error::<T>::NoUsername);
+			let account_of_username =
+				AccountOfUsername::<T>::get(&username).ok_or(Error::<T>::NoUsername)?;
+			ensure!(who == account_of_username, Error::<T>::InvalidUsername);
 			let (registration, _maybe_username) =
 				IdentityOf::<T>::get(&who).ok_or(Error::<T>::NoIdentity)?;
 			IdentityOf::<T>::insert(&who, (registration, Some(username.clone())));
diff --git a/substrate/frame/identity/src/tests.rs b/substrate/frame/identity/src/tests.rs
index 0a9464256ce3..60579a23b91b 100644
--- a/substrate/frame/identity/src/tests.rs
+++ b/substrate/frame/identity/src/tests.rs
@@ -25,7 +25,7 @@ use crate::{
 
 use codec::{Decode, Encode};
 use frame_support::{
-	assert_noop, assert_ok, derive_impl, parameter_types,
+	assert_err, assert_noop, assert_ok, derive_impl, parameter_types,
 	traits::{ConstU32, ConstU64, Get, OnFinalize, OnInitialize},
 	BoundedVec,
 };
@@ -1491,6 +1491,76 @@ fn setting_primary_should_work() {
 	});
 }
 
+#[test]
+fn must_own_primary() {
+	new_test_ext().execute_with(|| {
+		// set up authority
+		let [authority, _] = unfunded_accounts();
+		let suffix: Vec<u8> = b"test".to_vec();
+		let allocation: u32 = 10;
+		assert_ok!(Identity::add_username_authority(
+			RuntimeOrigin::root(),
+			authority.clone(),
+			suffix.clone(),
+			allocation
+		));
+
+		// Set up first user ("pi") and a username.
+		let pi_public = sr25519_generate(0.into(), None);
+		let pi_account: AccountIdOf<Test> = MultiSigner::Sr25519(pi_public).into_account().into();
+		let (pi_username, pi_to_sign) =
+			test_username_of(b"username314159".to_vec(), suffix.clone());
+		let encoded_pi_username = Encode::encode(&pi_to_sign.to_vec());
+		let pi_signature = MultiSignature::Sr25519(
+			sr25519_sign(0.into(), &pi_public, &encoded_pi_username).unwrap(),
+		);
+		assert_ok!(Identity::set_username_for(
+			RuntimeOrigin::signed(authority.clone()),
+			pi_account.clone(),
+			pi_username.clone(),
+			Some(pi_signature)
+		));
+
+		// Set up second user ("e") and a username.
+		let e_public = sr25519_generate(1.into(), None);
+		let e_account: AccountIdOf<Test> = MultiSigner::Sr25519(e_public).into_account().into();
+		let (e_username, e_to_sign) = test_username_of(b"username271828".to_vec(), suffix.clone());
+		let encoded_e_username = Encode::encode(&e_to_sign.to_vec());
+		let e_signature = MultiSignature::Sr25519(
+			sr25519_sign(1.into(), &e_public, &encoded_e_username).unwrap(),
+		);
+		assert_ok!(Identity::set_username_for(
+			RuntimeOrigin::signed(authority.clone()),
+			e_account.clone(),
+			e_username.clone(),
+			Some(e_signature)
+		));
+
+		// Ensure that both users have their usernames.
+		assert_eq!(
+			AccountOfUsername::<Test>::get::<&Username<Test>>(&pi_to_sign),
+			Some(pi_account.clone())
+		);
+		assert_eq!(
+			AccountOfUsername::<Test>::get::<&Username<Test>>(&e_to_sign),
+			Some(e_account.clone())
+		);
+
+		// Cannot set primary to a username that does not exist.
+		let (_, c_username) = test_username_of(b"speedoflight".to_vec(), suffix.clone());
+		assert_err!(
+			Identity::set_primary_username(RuntimeOrigin::signed(pi_account.clone()), c_username,),
+			Error::<Test>::NoUsername
+		);
+
+		// Cannot take someone else's username as your primary.
+		assert_err!(
+			Identity::set_primary_username(RuntimeOrigin::signed(pi_account.clone()), e_to_sign,),
+			Error::<Test>::InvalidUsername
+		);
+	});
+}
+
 #[test]
 fn unaccepted_usernames_should_expire() {
 	new_test_ext().execute_with(|| {