trimming binary, remove PropertyStoreWrapper

Author: Bush2021

src/appid.cc

@@ -3,10 +3,8 @@
 #include <windows.h>
 
 #include <propkey.h>
-#include <propvarutil.h>
 #include <shobjidl.h>
 
-#include <atomic>
 #include <functional>
 #include <string>
 
@@ -16,100 +14,54 @@
 
 namespace {
 
-const std::wstring& GetCustomAppUserModelID() {
+std::wstring_view GetCustomAppUserModelID() {
   static const std::wstring custom_appid = [] {
     auto hash = std::hash<std::wstring>{}(GetAppDir());
-    wchar_t buf[48];
-    swprintf_s(buf, L"ChromePlusNext.%016llX",
-               static_cast<unsigned long long>(hash));
-    return std::wstring(buf);
+    constexpr wchar_t hex[] = L"0123456789ABCDEF";
+    std::wstring result = L"CPN.";
+    for (int i = 60; i >= 0; i -= 4) {
+      result += hex[(hash >> i) & 0xF];
+    }
+    return result;
   }();
   return custom_appid;
 }
 
-class PropertyStoreWrapper final : public IPropertyStore {
- public:
-  explicit PropertyStoreWrapper(IPropertyStore* real_store)
-      : real_store_(real_store) {}
-
-  ~PropertyStoreWrapper() {
-    if (real_store_) {
-      real_store_->Release();
-    }
-  }
-
-  PropertyStoreWrapper(const PropertyStoreWrapper&) = delete;
-  PropertyStoreWrapper& operator=(const PropertyStoreWrapper&) = delete;
-
-  HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void** ppv) override {
-    if (!ppv) {
-      return E_POINTER;
-    }
-    if (riid == IID_IUnknown || riid == IID_IPropertyStore) {
-      *ppv = static_cast<IPropertyStore*>(this);
-      AddRef();
-      return S_OK;
-    }
-    return real_store_->QueryInterface(riid, ppv);
-  }
-
-  ULONG STDMETHODCALLTYPE AddRef() override {
-    return ref_count_.fetch_add(1, std::memory_order_relaxed) + 1;
-  }
-
-  ULONG STDMETHODCALLTYPE Release() override {
-    auto count = ref_count_.fetch_sub(1, std::memory_order_acq_rel) - 1;
-    if (count == 0) {
-      delete this;
-    }
-    return count;
-  }
-
-  HRESULT STDMETHODCALLTYPE GetCount(DWORD* cProps) override {
-    return real_store_->GetCount(cProps);
+PROPVARIANT MakeAppIdVariant() {
+  const auto& src = GetCustomAppUserModelID();
+  auto len = src.size() + 1;
+  PROPVARIANT pv{};
+  pv.vt = VT_LPWSTR;
+  pv.pwszVal = static_cast<LPWSTR>(CoTaskMemAlloc(len * sizeof(wchar_t)));
+  if (pv.pwszVal) {
+    memcpy(pv.pwszVal, src.data(), len * sizeof(wchar_t));
   }
-
-  HRESULT STDMETHODCALLTYPE GetAt(DWORD iProp, PROPERTYKEY* pkey) override {
-    return real_store_->GetAt(iProp, pkey);
-  }
-
-  HRESULT STDMETHODCALLTYPE GetValue(REFPROPERTYKEY key,
-                                     PROPVARIANT* pv) override {
-    if (IsEqualPropertyKey(key, PKEY_AppUserModel_ID)) {
-      return InitPropVariantFromString(GetCustomAppUserModelID().c_str(), pv);
-    }
-    return real_store_->GetValue(key, pv);
-  }
-
-  HRESULT STDMETHODCALLTYPE SetValue(REFPROPERTYKEY key,
-                                     REFPROPVARIANT propvar) override {
-    if (IsEqualPropertyKey(key, PKEY_AppUserModel_ID)) {
-      PROPVARIANT pv{};
-      if (auto hr =
-              InitPropVariantFromString(GetCustomAppUserModelID().c_str(), &pv);
-          FAILED(hr)) {
-        return hr;
-      }
-      auto hr = real_store_->SetValue(key, pv);
-      PropVariantClear(&pv);
-      return hr;
-    }
-    return real_store_->SetValue(key, propvar);
-  }
-
-  HRESULT STDMETHODCALLTYPE Commit() override { return real_store_->Commit(); }
-
- private:
-  IPropertyStore* real_store_;
-  std::atomic<ULONG> ref_count_{1};
-};
+  return pv;
+}
 
 static auto RawSetCurrentProcessExplicitAppUserModelID =
     SetCurrentProcessExplicitAppUserModelID;
 
 HRESULT WINAPI MySetCurrentProcessExplicitAppUserModelID(PCWSTR appid) {
   return RawSetCurrentProcessExplicitAppUserModelID(
-      GetCustomAppUserModelID().c_str());
+      GetCustomAppUserModelID().data());
+}
+
+using SetValueFn = HRESULT(STDMETHODCALLTYPE*)(IPropertyStore*,
+                                               REFPROPERTYKEY,
+                                               REFPROPVARIANT);
+static SetValueFn RawPropertyStoreSetValue = nullptr;
+
+HRESULT STDMETHODCALLTYPE MyPropertyStoreSetValue(IPropertyStore* self,
+                                                  REFPROPERTYKEY key,
+                                                  REFPROPVARIANT propvar) {
+  if (IsEqualPropertyKey(key, PKEY_AppUserModel_ID)) {
+    auto pv = MakeAppIdVariant();
+    auto hr = RawPropertyStoreSetValue(self, key, pv);
+    PropVariantClear(&pv);
+    return hr;
+  }
+  return RawPropertyStoreSetValue(self, key, propvar);
 }
 
 static auto RawSHGetPropertyStoreForWindow = SHGetPropertyStoreForWindow;
@@ -118,11 +70,17 @@ HRESULT WINAPI MySHGetPropertyStoreForWindow(HWND hwnd,
                                              REFIID riid,
                                              void** ppv) {
   HRESULT hr = RawSHGetPropertyStoreForWindow(hwnd, riid, ppv);
-  if (SUCCEEDED(hr) && ppv && *ppv && riid == IID_IPropertyStore) {
-    auto* wrapper = new (std::nothrow)
-        PropertyStoreWrapper(static_cast<IPropertyStore*>(*ppv));
-    if (wrapper) {
-      *ppv = static_cast<IPropertyStore*>(wrapper);
+  if (SUCCEEDED(hr) && ppv && *ppv && riid == IID_IPropertyStore &&
+      !RawPropertyStoreSetValue) {
+    auto** vtable = *reinterpret_cast<void***>(*ppv);
+    RawPropertyStoreSetValue = reinterpret_cast<SetValueFn>(vtable[6]);
+    DetourTransactionBegin();
+    DetourUpdateThread(GetCurrentThread());
+    DetourAttach(reinterpret_cast<LPVOID*>(&RawPropertyStoreSetValue),
+                 reinterpret_cast<void*>(MyPropertyStoreSetValue));
+    auto status = DetourTransactionCommit();
+    if (status != NO_ERROR) {
+      DebugLog(L"Hook PropertyStore::SetValue failed {}", status);
     }
   }
   return hr;