From 68e011efcef9640fbca4b71b4a4fa1ce3dcebf18 Mon Sep 17 00:00:00 2001
From: Alexander Chuprov <a.chuprov@yadro.com>
Date: Wed, 30 Oct 2024 18:02:08 +0300
Subject: [PATCH] [#118] frostfsid: Restrict keys to a single subject

Signed-off-by: Alexander Chuprov <a.chuprov@yadro.com>
---
 frostfsid/frostfsid_contract.go | 20 +++++++++++++++++
 tests/frostfsid_test.go         | 38 +++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/frostfsid/frostfsid_contract.go b/frostfsid/frostfsid_contract.go
index a9cca0e..1a79476 100644
--- a/frostfsid/frostfsid_contract.go
+++ b/frostfsid/frostfsid_contract.go
@@ -96,6 +96,8 @@ const (
 	groupSubjectsKeysPrefix      = 'G'
 	groupCounterKey              = 'c'
 	namespaceGroupsNamesPrefix   = 'm'
+	// allAddressPrefix contains all keys and signifies that the key has been used.
+	allAddressPrefix = 'A'
 )
 
 func _deploy(data any, isUpdate bool) {
@@ -182,6 +184,11 @@ func CreateSubject(ns string, key interop.PublicKey) {
 		panic("key is occupied")
 	}
 
+	allAddressKey := allAddress(addr)
+	if storage.Get(ctx, allAddressKey) != nil {
+		panic("key is occupied")
+	}
+
 	nsKey := namespaceKey(ns)
 	data = storage.Get(ctx, nsKey).([]byte)
 	if data == nil {
@@ -197,6 +204,7 @@ func CreateSubject(ns string, key interop.PublicKey) {
 
 	nsSubjKey := namespaceSubjectKey(ns, addr)
 	storage.Put(ctx, nsSubjKey, []byte{1})
+	storage.Put(ctx, allAddressKey, true)
 
 	runtime.Notify("CreateSubject", interop.Hash160(addr))
 }
@@ -213,6 +221,11 @@ func AddSubjectKey(addr interop.Hash160, key interop.PublicKey) {
 		panic("incorrect public key length")
 	}
 
+	allAddressKey := allAddress(contract.CreateStandardAccount(key))
+	if storage.Get(ctx, allAddressKey) != nil {
+		panic("key is occupied")
+	}
+
 	saKey := subjectAdditionalKey(key, addr)
 	data := storage.Get(ctx, saKey).([]byte)
 	if data != nil {
@@ -230,6 +243,7 @@ func AddSubjectKey(addr interop.Hash160, key interop.PublicKey) {
 	subject.AdditionalKeys = append(subject.AdditionalKeys, key)
 
 	storage.Put(ctx, sKey, std.Serialize(subject))
+	storage.Put(ctx, allAddressKey, true)
 	runtime.Notify("AddSubjectKey", addr, key)
 }
 
@@ -269,6 +283,7 @@ func RemoveSubjectKey(addr interop.Hash160, key interop.PublicKey) {
 	subject.AdditionalKeys = additionalKeys
 
 	storage.Put(ctx, sKey, std.Serialize(subject))
+	storage.Delete(ctx, allAddress(contract.CreateStandardAccount(key)))
 	runtime.Notify("RemoveSubjectKey", addr, key)
 }
 
@@ -362,6 +377,7 @@ func DeleteSubject(addr interop.Hash160) {
 
 	for i := 0; i < len(subj.AdditionalKeys); i++ {
 		storage.Delete(ctx, subjectAdditionalKey(subj.AdditionalKeys[i], addr))
+		storage.Delete(ctx, allAddress(contract.CreateStandardAccount(subj.AdditionalKeys[i])))
 	}
 	storage.Delete(ctx, sKey)
 
@@ -1029,3 +1045,7 @@ func idToBytes(itemID int) []byte {
 	zeros := make([]byte, 8-ln)
 	return append(b, zeros...)
 }
+
+func allAddress(address []byte) []byte {
+	return append([]byte{allAddressPrefix}, address...)
+}
diff --git a/tests/frostfsid_test.go b/tests/frostfsid_test.go
index 0e5a38e..419fb9f 100644
--- a/tests/frostfsid_test.go
+++ b/tests/frostfsid_test.go
@@ -604,6 +604,44 @@ func TestFrostFSID_GroupManagement(t *testing.T) {
 	})
 }
 
+func TestAdditionalKeyFromPrimarySubject(t *testing.T) {
+	f := newFrostFSIDInvoker(t)
+	invoker := f.OwnerInvoker()
+
+	subjAPrimaryKey, err := keys.NewPrivateKey()
+	require.NoError(t, err)
+	subjAKeyAddr := subjAPrimaryKey.PublicKey().GetScriptHash()
+
+	subjBPrimaryKey, err := keys.NewPrivateKey()
+	require.NoError(t, err)
+	subjBKeyAddr := subjBPrimaryKey.PublicKey().GetScriptHash()
+
+	subjCPrimaryKey, err := keys.NewPrivateKey()
+	require.NoError(t, err)
+
+	subjDPrimaryKey, err := keys.NewPrivateKey()
+	require.NoError(t, err)
+
+	invoker.Invoke(t, stackitem.Null{}, createSubjectMethod, defaultNamespace, subjAPrimaryKey.PublicKey().Bytes())
+
+	invoker.Invoke(t, stackitem.Null{}, createSubjectMethod, defaultNamespace, subjBPrimaryKey.PublicKey().Bytes())
+
+	invoker.InvokeFail(t, "key is occupied", addSubjectKeyMethod, subjBKeyAddr, subjAPrimaryKey.PublicKey().Bytes())
+
+	invoker.Invoke(t, stackitem.Null{}, addSubjectKeyMethod, subjAKeyAddr, subjCPrimaryKey.PublicKey().Bytes())
+	invoker.InvokeFail(t, "key is occupied", addSubjectKeyMethod, subjBKeyAddr, subjCPrimaryKey.PublicKey().Bytes())
+
+	invoker.Invoke(t, stackitem.Null{}, addSubjectKeyMethod, subjAKeyAddr, subjDPrimaryKey.PublicKey().Bytes())
+
+	invoker.InvokeFail(t, "key is occupied", addSubjectKeyMethod, subjBKeyAddr, subjDPrimaryKey.PublicKey().Bytes())
+	invoker.Invoke(t, stackitem.Null{}, removeSubjectKeyMethod, subjAKeyAddr, subjDPrimaryKey.PublicKey().Bytes())
+	invoker.Invoke(t, stackitem.Null{}, addSubjectKeyMethod, subjBKeyAddr, subjDPrimaryKey.PublicKey().Bytes())
+
+	invoker.InvokeFail(t, "key is occupied", addSubjectKeyMethod, subjAKeyAddr, subjDPrimaryKey.PublicKey().Bytes())
+	invoker.Invoke(t, stackitem.Null{}, deleteSubjectMethod, subjBKeyAddr)
+	invoker.Invoke(t, stackitem.Null{}, addSubjectKeyMethod, subjAKeyAddr, subjDPrimaryKey.PublicKey().Bytes())
+}
+
 func checkPublicKeyResult(t *testing.T, s *vm.Stack, err error, key *keys.PrivateKey) {
 	if key == nil {
 		require.ErrorContains(t, err, "not found")