Add a context menu icon to the Windows shell extension

5 files changed, 119 insertions, 15 deletions

diff --git a/shell_integration/windows/OCContextMenu/CMakeLists.txt b/shell_integration/windows/OCContextMenu/CMakeLists.txt
index 74c12621b..f84ed4721 100644
--- a/shell_integration/windows/OCContextMenu/CMakeLists.txt
+++ b/shell_integration/windows/OCContextMenu/CMakeLists.txt
@@ -9,7 +9,10 @@ add_library(OCContextMenu MODULE
 )
 
 target_link_libraries(OCContextMenu
-    OCUtil)
+                      OCUtil
+                      crypt32
+                      gdiplus)
+target_compile_definitions(OCContextMenu PRIVATE JSON_NOEXCEPTION)
 
 install(TARGETS OCContextMenu
     RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
diff --git a/shell_integration/windows/OCContextMenu/OCClientInterface.cpp b/shell_integration/windows/OCContextMenu/OCClientInterface.cpp
index ad7a35564..98f4230cc 100644
--- a/shell_integration/windows/OCContextMenu/OCClientInterface.cpp
+++ b/shell_integration/windows/OCContextMenu/OCClientInterface.cpp
@@ -24,13 +24,85 @@
 #include <algorithm>
 #include <iostream>
 #include <sstream>
+#include <string>
 #include <iterator>
 #include <unordered_set>
 
+// use std::min in gdiplus
 using namespace std;
 
+#include <comdef.h>
+#include <gdiplus.h>
+#include <wincrypt.h>
+#include <shlwapi.h>
+#include <wrl/client.h>
+
+#include "../3rdparty/nlohmann-json/json.hpp"
+
+using Microsoft::WRL::ComPtr;
+
 #define PIPE_TIMEOUT  5*1000 //ms
 
+namespace {
+
+template <typename T = wstring>
+void log(const wstring &msg, const T &error = {})
+{
+    wstringstream tmp;
+    tmp << L"ownCloud: " << msg;
+    if (!error.empty()) {
+        tmp << L" " << error.data();
+    }
+    OutputDebugStringW(tmp.str().data());
+}
+void logWinError(const wstring &msg, const DWORD &error = GetLastError())
+{
+    log(msg, wstring(_com_error(error).ErrorMessage()));
+}
+
+void sendV2(const CommunicationSocket &socket, const wstring &command, const nlohmann::json &j)
+{
+    static int messageId = 0;
+    const nlohmann::json json { { "id", to_string(messageId++) }, { "arguments", j } };
+    const auto data = json.dump();
+    wstringstream tmp;
+    tmp << command << L":" << StringUtil::toUtf16(data.data(), data.size()) << L"\n";
+    socket.SendMsg(tmp.str().data());
+}
+
+pair<wstring, nlohmann::json> parseV2(const wstring &data)
+{
+    const auto index = data.find(L":");
+    const auto argStart = data.cbegin() + index + 1;
+    const auto cData = StringUtil::toUtf8(&*argStart, distance(argStart, data.cend()));
+    return { data.substr(0, index), nlohmann::json::parse(cData) };
+}
+
+std::shared_ptr<HBITMAP> saveImage(const string &data)
+{
+    DWORD size = 2 * 1024;
+    std::vector<BYTE> buf(size, 0);
+    DWORD skipped;
+    if (!CryptStringToBinaryA(data.data(), 0, CRYPT_STRING_BASE64, buf.data(), &size, &skipped, nullptr)) {
+        logWinError(L"Failed to decode icon");
+        return {};
+    }
+    ComPtr<IStream> stream = SHCreateMemStream(buf.data(), size);
+    if (!stream) {
+        log(L"Failed to create stream");
+        return {};
+    };
+    HBITMAP result;
+    Gdiplus::Bitmap bitmap(stream.Get(), true);
+    const auto status = bitmap.GetHBITMAP(0, &result);
+    if (status != Gdiplus::Ok) {
+        log(L"Failed to get HBITMAP", to_wstring(status));
+        return {};
+    }
+    return std::shared_ptr<HBITMAP> { new HBITMAP(result), &DeleteObject };
+}
+}
+
 OCClientInterface::ContextMenuInfo OCClientInterface::FetchInfo(const std::wstring &files)
 {
     auto pipename = CommunicationSocket::DefaultPipePath();
@@ -42,6 +114,7 @@ OCClientInterface::ContextMenuInfo OCClientInterface::FetchInfo(const std::wstri
     if (!socket.Connect(pipename)) {
         return {};
     }
+    sendV2(socket, L"V2/GET_CLIENT_ICON", { { "size", 16 } });
     socket.SendMsg(L"GET_STRINGS:CONTEXT_MENU_TITLE\n");
     socket.SendMsg((L"GET_MENU_ITEMS:" + files + L"\n").data());
 
@@ -50,11 +123,21 @@ OCClientInterface::ContextMenuInfo OCClientInterface::FetchInfo(const std::wstri
     int sleptCount = 0;
     while (sleptCount < 5) {
         if (socket.ReadLine(&response)) {
-            if (StringUtil::begins_with(response, wstring(L"REGISTER_PATH:"))) {
+            if (StringUtil::begins_with(response, wstring(L"V2/"))) {
+                const auto msg = parseV2(response);
+                const auto &arguments = msg.second["arguments"];
+                if (msg.first == L"V2/GET_CLIENT_ICON_RESULT") {
+                    if (arguments.contains("error")) {
+                        log(L"V2/GET_CLIENT_ICON failed", arguments["error"].get<string>());
+                    } else {
+                        info.icon = saveImage(arguments["png"].get<string>());
+                    }
+                }
+
+            } else if (StringUtil::begins_with(response, wstring(L"REGISTER_PATH:"))) {
                 wstring responsePath = response.substr(14); // length of REGISTER_PATH
                 info.watchedDirectories.push_back(responsePath);
-            }
-            else if (StringUtil::begins_with(response, wstring(L"STRING:"))) {
+            } else if (StringUtil::begins_with(response, wstring(L"STRING:"))) {
                 wstring stringName, stringValue;
                 if (!StringUtil::extractChunks(response, stringName, stringValue))
                     continue;
diff --git a/shell_integration/windows/OCContextMenu/OCClientInterface.h b/shell_integration/windows/OCContextMenu/OCClientInterface.h
index 586a03f2f..b656a2dad 100644
--- a/shell_integration/windows/OCContextMenu/OCClientInterface.h
+++ b/shell_integration/windows/OCContextMenu/OCClientInterface.h
@@ -38,6 +38,8 @@
 #include <atomic>
 #include <condition_variable>
 
+#include <windows.h>
+
 class CommunicationSocket;
 
 class OCClientInterface
@@ -46,6 +48,7 @@ public:
     struct ContextMenuInfo {
         std::vector<std::wstring> watchedDirectories;
         std::wstring contextMenuTitle;
+        std::shared_ptr<HBITMAP> icon;
         struct MenuItem
         {
             std::wstring command, flags, title;
diff --git a/shell_integration/windows/OCContextMenu/OCContextMenu.cpp b/shell_integration/windows/OCContextMenu/OCContextMenu.cpp
index a17df3248..d950ebe92 100644
--- a/shell_integration/windows/OCContextMenu/OCContextMenu.cpp
+++ b/shell_integration/windows/OCContextMenu/OCContextMenu.cpp
@@ -149,6 +149,10 @@ IFACEMETHODIMP OCContextMenu::QueryContextMenu(HMENU hMenu, UINT indexMenu, UINT
         mii.hSubMenu = hSubmenu;
         mii.fType = MFT_STRING;
         mii.dwTypeData = &m_info.contextMenuTitle[0];
+        if (m_info.icon) {
+            mii.fMask |= MIIM_BITMAP;
+            mii.hbmpItem = *m_info.icon;
+        }
 
         if (!InsertMenuItem(hMenu, indexMenu++, TRUE, &mii))
             return HRESULT_FROM_WIN32(GetLastError());
diff --git a/shell_integration/windows/OCContextMenu/dllmain.cpp b/shell_integration/windows/OCContextMenu/dllmain.cpp
index 7d6d77dd0..26de855af 100644
--- a/shell_integration/windows/OCContextMenu/dllmain.cpp
+++ b/shell_integration/windows/OCContextMenu/dllmain.cpp
@@ -17,6 +17,11 @@
 #include "OCContextMenuRegHandler.h"
 #include "OCContextMenuFactory.h"
 
+// gdiplus min/max
+using namespace std;
+#include <algorithm>
+#include <gdiplus.h>
+
 // {841A0AAD-AA11-4B50-84D9-7F8E727D77D7}
 static const GUID CLSID_FileContextMenuExt = { 0x841a0aad, 0xaa11, 0x4b50, { 0x84, 0xd9, 0x7f, 0x8e, 0x72, 0x7d, 0x77, 0xd7 } };
 
@@ -25,18 +30,24 @@ long        g_cDllRef = 0;
 
 BOOL APIENTRY DllMain(HMODULE hModule, DWORD dwReason, LPVOID lpReserved)
 {
-	switch (dwReason)
-	{
-	case DLL_PROCESS_ATTACH:
-		// Hold the instance of this DLL module, we will use it to get the 
-		// path of the DLL to register the component.
-		g_hInst = hModule;
-		DisableThreadLibraryCalls(hModule);
-		break;
-	case DLL_THREAD_ATTACH:
+    static ULONG_PTR gdiplusToken = 0;
+    switch (dwReason) {
+    case DLL_PROCESS_ATTACH: {
+        // Hold the instance of this DLL module, we will use it to get the
+        // path of the DLL to register the component.
+        g_hInst = hModule;
+        DisableThreadLibraryCalls(hModule);
+
+        Gdiplus::GdiplusStartupInput gdiplusStartupInput;
+        Gdiplus::GdiplusStartup(&gdiplusToken, &gdiplusStartupInput, NULL);
+        break;
+    }
+        case DLL_THREAD_ATTACH:
 	case DLL_THREAD_DETACH:
-	case DLL_PROCESS_DETACH:
-		break;
+            break;
+        case DLL_PROCESS_DETACH:
+            Gdiplus::GdiplusShutdown(gdiplusToken);
+            break;
 	}
 	return TRUE;
 }