Merge pull request #358 from skalenetwork/bug/SKALE-4262-register-existing-key-owner

Bug/skale 4262 register existing key owner

sgxwallet_common.h

@@ -185,6 +185,8 @@ extern bool autoconfirm;
 #define CORRUPT_DATABASE -112
 #define INVALID_SEK -113
 #define INVALID_DECRYPTION_VALUE_FORMAT -114
+#define INVALID_KEY_FORMAT -115
+#define KEY_ALREADY_REGISTERED -116
 
 #define SGX_ENCLAVE_ERROR -666
 

zmq_src/ReqMessage.cpp

@@ -32,9 +32,15 @@ Json::Value ECDSASignReqMessage::process() {
     auto base = getInt64Rapid("base");
     auto keyName = getStringRapid("keyName");
     auto hash = getStringRapid("messageHash");
-    if (checkKeyOwnership && !isKeyByOwner(keyName, getStringRapid("cert"))) {
-        spdlog::error("Cert {} try to access key {} which does not belong to it", getStringRapid("cert"), keyName);
-        throw std::invalid_argument("Only owner of the key can access it");
+    if (checkKeyOwnership) {
+        if (!isKeyRegistered(keyName)) {
+            addKeyByOwner(keyName, getStringRapid("cert"));
+        } else {
+            if (!isKeyByOwner(keyName, getStringRapid("cert"))) {
+                spdlog::error("Cert {} try to access key {} which does not belong to it", getStringRapid("cert"), keyName);
+                throw std::invalid_argument("Only owner of the key can access it");
+            }
+        }
     }
     auto result = SGXWalletServer::ecdsaSignMessageHashImpl(base, keyName, hash);
     result["type"] = ZMQMessage::ECDSA_SIGN_RSP;
@@ -46,9 +52,15 @@ Json::Value BLSSignReqMessage::process() {
     auto hash = getStringRapid("messageHash");
     auto t = getInt64Rapid("t");
     auto n = getInt64Rapid("n");
-    if (checkKeyOwnership && !isKeyByOwner(keyName, getStringRapid("cert"))) {
-        spdlog::error("Cert {} try to access key {} which does not belong to it", getStringRapid("cert"), keyName);
-        throw std::invalid_argument("Only owner of the key can access it");
+    if (checkKeyOwnership) {
+        if (!isKeyRegistered(keyName)) {
+            addKeyByOwner(keyName, getStringRapid("cert"));
+        } else {
+            if (!isKeyByOwner(keyName, getStringRapid("cert"))) {
+                spdlog::error("Cert {} try to access key {} which does not belong to it", getStringRapid("cert"), keyName);
+                throw std::invalid_argument("Only owner of the key can access it");
+            }
+        }
     }
     auto result = SGXWalletServer::blsSignMessageHashImpl(keyName, hash, t, n);
     result["type"] = ZMQMessage::BLS_SIGN_RSP;

zmq_src/RspMessage.h

@@ -259,5 +259,4 @@ public:
     }
 };
 
-
 #endif //SGXWALLET_RSPMESSAGE_H

zmq_src/ZMQMessage.cpp

@@ -334,6 +334,10 @@ void ZMQMessage::addKeyByOwner(const string& keyName, const string& cert) {
     SGXWalletServer::writeDataToDB(keyName + ":OWNER", cert);
 }
 
+bool ZMQMessage::isKeyRegistered(const string& keyName) {
+    return LevelDB::getLevelDb()->readString(keyName  + ":OWNER") != nullptr;
+}
+
 cache::lru_cache<string, pair < EVP_PKEY * , X509 *>> ZMQMessage::verifiedCerts(256);
 
 const std::map<string, int> ZMQMessage::requests{

zmq_src/ZMQMessage.h

@@ -57,6 +57,8 @@ protected:
 
     static void addKeyByOwner(const string& keyName, const string& cert);
 
+    static bool isKeyRegistered(const std::string& keyName);
+
 public:
 
     static constexpr const char *BLS_SIGN_REQ = "BLSSignReq";