7 files changed, 136 insertions, 106 deletions

diff --git a/src/core/function/gpg/GpgAbstractKeyGetter.cpp b/src/core/function/gpg/GpgAbstractKeyGetter.cpp
index 0b544469..193fec0e 100644
--- a/src/core/function/gpg/GpgAbstractKeyGetter.cpp
+++ b/src/core/function/gpg/GpgAbstractKeyGetter.cpp
@@ -66,7 +66,7 @@ auto GpgAbstractKeyGetter::GetKey(const QString& key_id) -> GpgAbstractKeyPtr {
   if (IsKeyGroupID(key_id)) {
     return kg_.KeyGroup(key_id);
   }
-  return key_.GetKeyPtr(key_id);
+  return key_.GetKeyORSubkeyPtr(key_id);
 }
 
 auto GpgAbstractKeyGetter::GetKeys(const QStringList& key_ids)
diff --git a/src/core/function/gpg/GpgAssuanHelper.cpp b/src/core/function/gpg/GpgAssuanHelper.cpp
index 8fa311dc..2fa836e5 100644
--- a/src/core/function/gpg/GpgAssuanHelper.cpp
+++ b/src/core/function/gpg/GpgAssuanHelper.cpp
@@ -28,8 +28,6 @@
 
 #include "GpgAssuanHelper.h"
 
-#include <utility>
-
 #include "core/module/ModuleManager.h"
 #include "core/utils/GpgUtils.h"
 
diff --git a/src/core/function/gpg/GpgAutomatonHandler.cpp b/src/core/function/gpg/GpgAutomatonHandler.cpp
index b2d9d562..33c94436 100644
--- a/src/core/function/gpg/GpgAutomatonHandler.cpp
+++ b/src/core/function/gpg/GpgAutomatonHandler.cpp
@@ -28,8 +28,6 @@
 
 #include "GpgAutomatonHandler.h"
 
-#include <utility>
-
 #include "core/model/GpgData.h"
 #include "core/model/GpgKey.h"
 #include "core/utils/GpgUtils.h"
diff --git a/src/core/function/gpg/GpgKeyGetter.cpp b/src/core/function/gpg/GpgKeyGetter.cpp
index 86218a94..3419c18d 100644
--- a/src/core/function/gpg/GpgKeyGetter.cpp
+++ b/src/core/function/gpg/GpgKeyGetter.cpp
@@ -32,9 +32,7 @@
 
 #include <mutex>
 
-#include "core/GpgModel.h"
 #include "core/function/gpg/GpgContext.h"
-#include "core/function/gpg/GpgKeyGroupGetter.h"
 #include "core/utils/GpgUtils.h"
 
 namespace GpgFrontend {
@@ -44,48 +42,65 @@ class GpgKeyGetter::Impl : public SingletonFunctionObject<GpgKeyGetter::Impl> {
   explicit Impl(int channel)
       : SingletonFunctionObject<GpgKeyGetter::Impl>(channel) {}
 
-  auto GetKey(const QString& fpr, bool use_cache) -> GpgKey {
+  auto GetKeyPtr(const QString& key_id, bool cache) -> GpgKeyPtr {
     // find in cache first
-    if (use_cache) {
-      auto key = get_key_in_cache(fpr);
-      if (key.IsGood()) return key;
+    if (cache) {
+      auto key = get_key_in_cache(key_id);
+      if (key != nullptr) return qSharedPointerDynamicCast<GpgKey>(key);
+
+      LOG_W() << "get gpg key" << key_id
+              << "from cache failed, channel: " << GetChannel();
     }
 
     gpgme_key_t p_key = nullptr;
-    gpgme_get_key(ctx_.DefaultContext(), fpr.toUtf8(), &p_key, 1);
+    gpgme_get_key(ctx_.DefaultContext(), key_id.toUtf8(), &p_key, 1);
     if (p_key == nullptr) {
-      LOG_W() << "GpgKeyGetter GetKey Private _p_key Null, fpr: " << fpr;
-      return GetPubkey(fpr, true);
+      LOG_W() << "GpgKeyGetter GetKey p_key is null, fpr: " << key_id;
+      return GetPubkeyPtr(key_id, true);
     }
-    return GpgKey(std::move(p_key));
+
+    return QSharedPointer<GpgKey>::create(p_key);
+  }
+
+  auto GetKey(const QString& key_id, bool cache) -> GpgKey {
+    auto key = GetKeyPtr(key_id, cache);
+
+    if (key != nullptr) return *key;
+    return {};
+  }
+
+  auto GetPubkey(const QString& key_id, bool cache) -> GpgKey {
+    auto key = GetKeyPtr(key_id, cache);
+
+    if (key != nullptr) return *key;
+    return {};
   }
 
-  auto GetPubkey(const QString& fpr, bool use_cache) -> GpgKey {
+  auto GetPubkeyPtr(const QString& key_id, bool cache) -> GpgKeyPtr {
     // find in cache first
-    if (use_cache) {
-      auto key = get_key_in_cache(fpr);
-      if (key.IsGood()) return key;
+    if (cache) {
+      auto key = get_key_in_cache(key_id);
+      if (key != nullptr) return qSharedPointerDynamicCast<GpgKey>(key);
+
+      LOG_W() << "get public gpg key" << key_id
+              << "from cache failed, channel: " << GetChannel();
     }
 
     gpgme_key_t p_key = nullptr;
-    gpgme_get_key(ctx_.DefaultContext(), fpr.toUtf8(), &p_key, 0);
-    if (p_key == nullptr)
-      LOG_W() << "GpgKeyGetter GetKey _p_key Null, fpr: " << fpr;
-    return GpgKey(p_key);
-  }
-
-  auto GetPubkeyPtr(const QString& fpr, bool use_cache) -> GpgKeyPtr {
-    auto key = GetPubkey(fpr, use_cache);
-    if (!key.IsGood()) return nullptr;
-    return QSharedPointer<GpgKey>::create(key);
+    gpgme_get_key(ctx_.DefaultContext(), key_id.toUtf8(), &p_key, 0);
+    if (p_key == nullptr) {
+      LOG_W() << "GpgKeyGetter GetKey p_key is null, key id: " << key_id;
+      return nullptr;
+    }
+    return QSharedPointer<GpgKey>::create(p_key);
   }
 
-  auto FetchKey() -> GpgKeyList {
-    if (keys_search_cache_.empty()) {
+  auto FetchKey() -> GpgKeyPtrList {
+    if (keys_cache_.empty() || keys_search_cache_.empty()) {
       FlushKeyCache();
     }
 
-    auto keys_list = GpgKeyList{};
+    auto keys_list = GpgKeyPtrList{};
     {
       // get the lock
       std::lock_guard<std::mutex> lock(keys_cache_mutex_);
@@ -106,7 +121,7 @@ class GpgKeyGetter::Impl : public SingletonFunctionObject<GpgKeyGetter::Impl> {
       // get the lock
       std::lock_guard<std::mutex> lock(keys_cache_mutex_);
       for (const auto& key : keys_cache_) {
-        keys_list.push_back(QSharedPointer<GpgKey>::create(key));
+        keys_list.push_back(key);
       }
     }
 
@@ -133,18 +148,25 @@ class GpgKeyGetter::Impl : public SingletonFunctionObject<GpgKeyGetter::Impl> {
       gpgme_key_t key;
       while ((err = gpgme_op_keylist_next(ctx_.DefaultContext(), &key)) ==
              GPG_ERR_NO_ERROR) {
-        auto gpg_key = GpgKey(std::move(key));
+        auto g_key = QSharedPointer<GpgKey>::create(key);
 
         // detect if the key is in a smartcard
         // if so, try to get full information using gpgme_get_key()
         // this maybe a bug in gpgme
-        if (gpg_key.IsHasCardKey()) {
-          gpg_key = GetKey(gpg_key.ID(), false);
+        if (g_key->IsHasCardKey()) {
+          g_key = GetKeyPtr(g_key->ID(), false);
         }
 
-        keys_cache_.push_back(gpg_key);
-        keys_search_cache_.insert(gpg_key.ID(), gpg_key);
-        keys_search_cache_.insert(gpg_key.Fingerprint(), gpg_key);
+        keys_cache_.push_back(g_key);
+        keys_search_cache_.insert(g_key->ID(), g_key);
+        keys_search_cache_.insert(g_key->Fingerprint(), g_key);
+
+        for (const auto& s_key : g_key->SubKeys()) {
+          if (s_key.ID() == g_key->ID()) continue;
+          auto p_s_key = QSharedPointer<GpgSubKey>::create(s_key);
+          keys_search_cache_.insert(s_key.ID(), p_s_key);
+          keys_search_cache_.insert(s_key.Fingerprint(), p_s_key);
+        }
       }
     }
 
@@ -163,29 +185,25 @@ class GpgKeyGetter::Impl : public SingletonFunctionObject<GpgKeyGetter::Impl> {
     return keys;
   }
 
-  auto GetKeysCopy(const GpgKeyList& keys) -> GpgKeyList {
-    // get the lock
-    std::lock_guard<std::mutex> lock(ctx_mutex_);
-    auto keys_copy = GpgKeyList{};
-    for (const auto& key : keys) keys_copy.push_back(key);
-    return keys_copy;
-  }
-
   auto Fetch() -> QContainer<QSharedPointer<GpgKey>> {
     auto keys = FetchKey();
 
     auto ret = QContainer<QSharedPointer<GpgKey>>();
     for (const auto& key : keys) {
-      ret.append(QSharedPointer<GpgKey>::create(key));
+      ret.append(key);
     }
     return ret;
   }
 
-  auto GetKeyPtr(const QString& key_id,
-                 bool use_cache) -> QSharedPointer<GpgKey> {
-    auto key = GetKey(key_id, use_cache);
-    if (!key.IsGood()) return nullptr;
-    return QSharedPointer<GpgKey>::create(key);
+  auto GetKeyORSubkeyPtr(const QString& key_id) -> GpgAbstractKeyPtr {
+    auto key = get_key_in_cache(key_id);
+
+    if (key != nullptr) return key;
+
+    LOG_W() << "get key" << key_id
+            << "from cache failed, channel: " << GetChannel();
+
+    return nullptr;
   }
 
  private:
@@ -206,13 +224,13 @@ class GpgKeyGetter::Impl : public SingletonFunctionObject<GpgKeyGetter::Impl> {
    * @brief cache the keys with key id
    *
    */
-  QMap<QString, GpgKey> keys_search_cache_;
+  QMap<QString, GpgAbstractKeyPtr> keys_search_cache_;
 
   /**
    * @brief
    *
    */
-  QContainer<GpgKey> keys_cache_;
+  QContainer<GpgKeyPtr> keys_cache_;
 
   /**
    * @brief shared mutex for the keys cache
@@ -226,8 +244,9 @@ class GpgKeyGetter::Impl : public SingletonFunctionObject<GpgKeyGetter::Impl> {
    * @param id
    * @return GpgKey
    */
-  auto get_key_in_cache(const QString& key_id) -> GpgKey {
+  auto get_key_in_cache(const QString& key_id) -> GpgAbstractKeyPtr {
     std::lock_guard<std::mutex> lock(keys_cache_mutex_);
+
     if (keys_search_cache_.find(key_id) != keys_search_cache_.end()) {
       std::lock_guard<std::mutex> lock(ctx_mutex_);
       // return a copy of the key in cache
@@ -269,13 +288,10 @@ auto GpgKeyGetter::GetKeys(const KeyIdArgsList& ids) -> GpgKeyList {
   return p_->GetKeys(ids);
 }
 
-auto GpgKeyGetter::GetKeysCopy(const GpgKeyList& keys) -> GpgKeyList {
-  return p_->GetKeysCopy(keys);
-}
-
-auto GpgKeyGetter::FetchKey() -> GpgKeyList { return p_->FetchKey(); }
+auto GpgKeyGetter::Fetch() -> GpgKeyPtrList { return p_->Fetch(); }
 
-auto GpgKeyGetter::Fetch() -> QContainer<QSharedPointer<GpgKey>> {
-  return p_->Fetch();
+auto GpgKeyGetter::GetKeyORSubkeyPtr(const QString& key_id)
+    -> GpgAbstractKeyPtr {
+  return p_->GetKeyORSubkeyPtr(key_id);
 }
 }  // namespace GpgFrontend
diff --git a/src/core/function/gpg/GpgKeyGetter.h b/src/core/function/gpg/GpgKeyGetter.h
index b266bf17..1fe4aeb7 100644
--- a/src/core/function/gpg/GpgKeyGetter.h
+++ b/src/core/function/gpg/GpgKeyGetter.h
@@ -90,11 +90,13 @@ class GPGFRONTEND_CORE_EXPORT GpgKeyGetter
   auto GetPubkeyPtr(const QString& key_id, bool use_cache = true) -> GpgKeyPtr;
 
   /**
-   * @brief Get all the keys by receiving a linked list
+   * @brief Get the Pubkey Ptr object
    *
-   * @return KeyLinkListPtr
+   * @param key_id
+   * @param use_cache
+   * @return GpgKeyPtr
    */
-  auto FetchKey() -> GpgKeyList;
+  auto GetKeyORSubkeyPtr(const QString& key_id) -> GpgAbstractKeyPtr;
 
   /**
    * @brief
@@ -117,14 +119,6 @@ class GPGFRONTEND_CORE_EXPORT GpgKeyGetter
    */
   auto GetKeys(const KeyIdArgsList& key_ids) -> GpgKeyList;
 
-  /**
-   * @brief Get the Keys Copy object
-   *
-   * @param keys
-   * @return KeyListPtr
-   */
-  auto GetKeysCopy(const GpgKeyList& keys) -> GpgKeyList;
-
  private:
   class Impl;
   SecureUniquePtr<Impl> p_;
diff --git a/src/core/function/gpg/GpgKeyGroupGetter.cpp b/src/core/function/gpg/GpgKeyGroupGetter.cpp
index 47878414..4d4a18cc 100644
--- a/src/core/function/gpg/GpgKeyGroupGetter.cpp
+++ b/src/core/function/gpg/GpgKeyGroupGetter.cpp
@@ -76,9 +76,10 @@ void GpgKeyGroupGetter::fetch_key_groups() {
     LOG_D() << "load raw key group:" << key_group.id
             << "key ids: " << key_group.key_ids;
 
-    key_groups_forest_.insert(
-        key_group.id,
-        QSharedPointer<GpgKeyGroupTreeNode>::create(GpgKeyGroup{key_group}));
+    auto node =
+        QSharedPointer<GpgKeyGroupTreeNode>::create(GpgKeyGroup{key_group});
+    node->key_group->SetKeyGroupGetter(this);
+    key_groups_forest_.insert(key_group.id, node);
   }
 
   build_gpg_key_group_tree();
@@ -108,30 +109,30 @@ auto GpgKeyGroupGetter::KeyGroup(const QString& id)
 }
 
 void GpgKeyGroupGetter::check_key_group(
-    const QSharedPointer<GpgKeyGroup>& key_group) {
-  if (key_group == nullptr || key_group->IsDisabled()) return;
+    const QSharedPointer<GpgKeyGroupTreeNode>& node) {
+  if (node == nullptr || node->disabled) return;
 
-  for (const auto& key_id : key_group->KeyIds()) {
-    LOG_D() << "check" << key_id << "of" << key_group->UID();
+  for (const auto& key_id : node->KeyIds()) {
+    assert(key_id != node->key_group->ID());
+    if (key_id == node->key_group->ID()) continue;
 
-    if (IsKeyGroupID(key_id) || key_id == key_group->ID()) {
+    if (IsKeyGroupID(key_id)) {
       if (!key_groups_forest_.contains(key_id)) {
-        key_group->SetDisabled(true);
         return;
       }
 
       auto s_node = key_groups_forest_.value(key_id, nullptr);
-      check_key_group(s_node->key_group);
+      check_key_group(s_node);
 
       if (s_node->key_group->IsDisabled()) {
-        key_group->SetDisabled(true);
+        node->disabled = true;
         return;
       }
 
     } else {
       auto key = GpgKeyGetter::GetInstance(GetChannel()).GetKeyPtr(key_id);
       if (key == nullptr || !key->IsHasEncrCap()) {
-        key_group->SetDisabled(true);
+        node->disabled = true;
         return;
       }
     }
@@ -139,20 +140,23 @@ void GpgKeyGroupGetter::check_key_group(
 }
 
 void GpgKeyGroupGetter::check_all_key_groups() {
-  for (const auto& node : key_groups_forest_) {
-    node->key_group->SetDisabled(false);
-  }
+  for (const auto& node : key_groups_forest_) node->disabled = false;
 
   for (const auto& node : key_groups_forest_) {
-    auto key_group = node->key_group;
-    check_key_group(key_group);
+    check_key_group(node);
 
-    LOG_D() << "key group" << key_group->ID() << "ids: " << key_group->KeyIds()
-            << "status: " << key_group->IsDisabled();
+    assert(node->key_group != nullptr);
+    LOG_D() << "key group" << node->key_group->ID()
+            << "ids: " << node->key_group->KeyIds()
+            << "status: " << node->disabled;
   }
 }
 
 auto GpgKeyGroupGetter::FlushCache() -> bool {
+  for (const auto& node : key_groups_forest_.values()) {
+    node->Apply();
+  }
+
   check_all_key_groups();
   persist_key_groups();
   return true;
@@ -180,6 +184,7 @@ void GpgKeyGroupGetter::build_gpg_key_group_tree() {
 
 void GpgKeyGroupGetter::AddKeyGroup(const GpgKeyGroup& key_group) {
   auto node = QSharedPointer<GpgKeyGroupTreeNode>::create(key_group);
+  node->key_group->SetKeyGroupGetter(this);
 
   key_groups_forest_.insert(node->key_group->ID(), node);
   LOG_D() << "add new key group" << key_group.ID()
@@ -249,18 +254,7 @@ GpgKeyGroupTreeNode::GpgKeyGroupTreeNode(GpgKeyGroup kg)
   }
 }
 
-void GpgKeyGroupTreeNode::Apply() {
-  QStringList key_ids;
-  for (const auto& child : children) {
-    key_ids.push_back(child->key_group->ID());
-  }
-
-  for (const auto& key_id : non_key_group_ids) {
-    key_ids.push_back(key_id);
-  }
-
-  key_group->SetKeyIds(key_ids);
-}
+void GpgKeyGroupTreeNode::Apply() { key_group->SetKeyIds(KeyIds()); }
 
 auto GpgKeyGroupTreeNode::AddNonKeyGroupKey(const GpgAbstractKeyPtr& key)
     -> bool {
@@ -308,4 +302,19 @@ auto GpgKeyGroupTreeNode::RemoveNonKeyGroupKey(const QString& key) -> bool {
   Apply();
   return true;
 }
+
+auto GpgKeyGroupGetter::IsKeyGroupDisabled(const QString& id) -> bool {
+  if (!key_groups_forest_.contains(id)) return false;
+  auto node = key_groups_forest_.value(id);
+  return node->disabled;
+}
+
+auto GpgKeyGroupTreeNode::KeyIds() const -> QStringList {
+  QStringList ret;
+  for (const auto& child : children) {
+    ret.push_back(child->key_group->ID());
+  }
+  ret.append(non_key_group_ids);
+  return ret;
+}
 }  // namespace GpgFrontend
\ No newline at end of file
diff --git a/src/core/function/gpg/GpgKeyGroupGetter.h b/src/core/function/gpg/GpgKeyGroupGetter.h
index 9eabccde..823fe338 100644
--- a/src/core/function/gpg/GpgKeyGroupGetter.h
+++ b/src/core/function/gpg/GpgKeyGroupGetter.h
@@ -42,6 +42,7 @@ struct GpgKeyGroupTreeNode {
 
   // over take
   QStringList non_key_group_ids;
+  bool disabled;
 
   /**
    * @brief Construct a new Gpg Key Group Tree Node object
@@ -106,6 +107,13 @@ struct GpgKeyGroupTreeNode {
    * @return false
    */
   auto RemoveNonKeyGroupKey(const QString& key) -> bool;
+
+  /**
+   * @brief
+   *
+   * @return QStringList
+   */
+  [[nodiscard]] auto KeyIds() const -> QStringList;
 };
 
 class GPGFRONTEND_CORE_EXPORT GpgKeyGroupGetter
@@ -173,6 +181,13 @@ class GPGFRONTEND_CORE_EXPORT GpgKeyGroupGetter
    */
   auto KeyGroup(const QString& id) -> QSharedPointer<GpgKeyGroup>;
 
+  /**
+   * @brief
+   *
+   * @param id
+   */
+  auto IsKeyGroupDisabled(const QString& id) -> bool;
+
  private:
   GpgContext& ctx_ =
       GpgContext::GetInstance(SingletonFunctionObject::GetChannel());
@@ -203,7 +218,7 @@ class GPGFRONTEND_CORE_EXPORT GpgKeyGroupGetter
    * @brief
    *
    */
-  void check_key_group(const QSharedPointer<GpgKeyGroup>&);
+  void check_key_group(const QSharedPointer<GpgKeyGroupTreeNode>&);
 
   /**
    * @brief