xref: /third_party/python/PC/pyshellext.cpp (revision 7db96d56) 
1// Support back to Vista
2#define _WIN32_WINNT _WIN32_WINNT_VISTA
3#include <sdkddkver.h>
4
5// Use WRL to define a classic COM class
6#define __WRL_CLASSIC_COM__
7#include <wrl.h>
8
9#include <windows.h>
10#include <shlobj.h>
11#include <shlwapi.h>
12#include <olectl.h>
13#include <strsafe.h>
14
15#define DDWM_UPDATEWINDOW (WM_USER+3)
16
17static HINSTANCE hModule;
18static CLIPFORMAT cfDropDescription;
19static CLIPFORMAT cfDragWindow;
20
21#define CLASS_GUID "{BEA218D2-6950-497B-9434-61683EC065FE}"
22static const LPCWSTR CLASS_SUBKEY = L"Software\\Classes\\CLSID\\" CLASS_GUID;
23static const LPCWSTR DRAG_MESSAGE = L"Open with %1";
24
25using namespace Microsoft::WRL;
26
27HRESULT FilenameListCchLengthA(LPCSTR pszSource, size_t cchMax, size_t *pcchLength, size_t *pcchCount) {
28    HRESULT hr = S_OK;
29    size_t count = 0;
30    size_t length = 0;
31
32    while (pszSource && pszSource[0]) {
33        size_t oneLength;
34        hr = StringCchLengthA(pszSource, cchMax - length, &oneLength);
35        if (FAILED(hr)) {
36            return hr;
37        }
38        count += 1;
39        length += oneLength + (strchr(pszSource, ' ') ? 3 : 1);
40        pszSource = &pszSource[oneLength + 1];
41    }
42
43    *pcchCount = count;
44    *pcchLength = length;
45    return hr;
46}
47
48HRESULT FilenameListCchLengthW(LPCWSTR pszSource, size_t cchMax, size_t *pcchLength, size_t *pcchCount) {
49    HRESULT hr = S_OK;
50    size_t count = 0;
51    size_t length = 0;
52
53    while (pszSource && pszSource[0]) {
54        size_t oneLength;
55        hr = StringCchLengthW(pszSource, cchMax - length, &oneLength);
56        if (FAILED(hr)) {
57            return hr;
58        }
59        count += 1;
60        length += oneLength + (wcschr(pszSource, ' ') ? 3 : 1);
61        pszSource = &pszSource[oneLength + 1];
62    }
63
64    *pcchCount = count;
65    *pcchLength = length;
66    return hr;
67}
68
69HRESULT FilenameListCchCopyA(STRSAFE_LPSTR pszDest, size_t cchDest, LPCSTR pszSource, LPCSTR pszSeparator) {
70    HRESULT hr = S_OK;
71    size_t count = 0;
72    size_t length = 0;
73
74    while (pszSource[0]) {
75        STRSAFE_LPSTR newDest;
76
77        hr = StringCchCopyExA(pszDest, cchDest, pszSource, &newDest, &cchDest, 0);
78        if (FAILED(hr)) {
79            return hr;
80        }
81        pszSource += (newDest - pszDest) + 1;
82        pszDest = PathQuoteSpacesA(pszDest) ? newDest + 2 : newDest;
83
84        if (pszSource[0]) {
85            hr = StringCchCopyExA(pszDest, cchDest, pszSeparator, &newDest, &cchDest, 0);
86            if (FAILED(hr)) {
87                return hr;
88            }
89            pszDest = newDest;
90        }
91    }
92
93    return hr;
94}
95
96HRESULT FilenameListCchCopyW(STRSAFE_LPWSTR pszDest, size_t cchDest, LPCWSTR pszSource, LPCWSTR pszSeparator) {
97    HRESULT hr = S_OK;
98    size_t count = 0;
99    size_t length = 0;
100
101    while (pszSource[0]) {
102        STRSAFE_LPWSTR newDest;
103
104        hr = StringCchCopyExW(pszDest, cchDest, pszSource, &newDest, &cchDest, 0);
105        if (FAILED(hr)) {
106            return hr;
107        }
108        pszSource += (newDest - pszDest) + 1;
109        pszDest = PathQuoteSpacesW(pszDest) ? newDest + 2 : newDest;
110
111        if (pszSource[0]) {
112            hr = StringCchCopyExW(pszDest, cchDest, pszSeparator, &newDest, &cchDest, 0);
113            if (FAILED(hr)) {
114                return hr;
115            }
116            pszDest = newDest;
117        }
118    }
119
120    return hr;
121}
122
123class DECLSPEC_UUID(CLASS_GUID) PyShellExt : public RuntimeClass<
124    RuntimeClassFlags<ClassicCom>,
125    IDropTarget,
126    IPersistFile
127>
128{
129    LPOLESTR target, target_dir;
130    DWORD target_mode;
131
132    IDataObject *data_obj;
133
134public:
135    PyShellExt() : target(NULL), target_dir(NULL), target_mode(0), data_obj(NULL) {
136        OutputDebugString(L"PyShellExt::PyShellExt");
137    }
138
139    ~PyShellExt() {
140        if (target) {
141            CoTaskMemFree(target);
142        }
143        if (target_dir) {
144            CoTaskMemFree(target_dir);
145        }
146        if (data_obj) {
147            data_obj->Release();
148        }
149    }
150
151private:
152    HRESULT UpdateDropDescription(IDataObject *pDataObj) {
153        STGMEDIUM medium;
154        FORMATETC fmt = {
155            cfDropDescription,
156            NULL,
157            DVASPECT_CONTENT,
158            -1,
159            TYMED_HGLOBAL
160        };
161
162        auto hr = pDataObj->GetData(&fmt, &medium);
163        if (FAILED(hr)) {
164            OutputDebugString(L"PyShellExt::UpdateDropDescription - failed to get DROPDESCRIPTION format");
165            return hr;
166        }
167        if (!medium.hGlobal) {
168            OutputDebugString(L"PyShellExt::UpdateDropDescription - DROPDESCRIPTION format had NULL hGlobal");
169            ReleaseStgMedium(&medium);
170            return E_FAIL;
171        }
172        auto dd = (DROPDESCRIPTION*)GlobalLock(medium.hGlobal);
173        if (!dd) {
174            OutputDebugString(L"PyShellExt::UpdateDropDescription - failed to lock DROPDESCRIPTION hGlobal");
175            ReleaseStgMedium(&medium);
176            return E_FAIL;
177        }
178        StringCchCopy(dd->szMessage, sizeof(dd->szMessage) / sizeof(dd->szMessage[0]), DRAG_MESSAGE);
179        StringCchCopy(dd->szInsert, sizeof(dd->szInsert) / sizeof(dd->szInsert[0]), PathFindFileNameW(target));
180        dd->type = DROPIMAGE_MOVE;
181
182        GlobalUnlock(medium.hGlobal);
183        ReleaseStgMedium(&medium);
184
185        return S_OK;
186    }
187
188    HRESULT GetDragWindow(IDataObject *pDataObj, HWND *phWnd) {
189        HRESULT hr;
190        HWND *pMem;
191        STGMEDIUM medium;
192        FORMATETC fmt = {
193            cfDragWindow,
194            NULL,
195            DVASPECT_CONTENT,
196            -1,
197            TYMED_HGLOBAL
198        };
199
200        hr = pDataObj->GetData(&fmt, &medium);
201        if (FAILED(hr)) {
202            OutputDebugString(L"PyShellExt::GetDragWindow - failed to get DragWindow format");
203            return hr;
204        }
205        if (!medium.hGlobal) {
206            OutputDebugString(L"PyShellExt::GetDragWindow - DragWindow format had NULL hGlobal");
207            ReleaseStgMedium(&medium);
208            return E_FAIL;
209        }
210
211        pMem = (HWND*)GlobalLock(medium.hGlobal);
212        if (!pMem) {
213            OutputDebugString(L"PyShellExt::GetDragWindow - failed to lock DragWindow hGlobal");
214            ReleaseStgMedium(&medium);
215            return E_FAIL;
216        }
217
218        *phWnd = *pMem;
219
220        GlobalUnlock(medium.hGlobal);
221        ReleaseStgMedium(&medium);
222
223        return S_OK;
224    }
225
226    HRESULT GetArguments(IDataObject *pDataObj, LPCWSTR *pArguments) {
227        HRESULT hr;
228        DROPFILES *pdropfiles;
229
230        STGMEDIUM medium;
231        FORMATETC fmt = {
232            CF_HDROP,
233            NULL,
234            DVASPECT_CONTENT,
235            -1,
236            TYMED_HGLOBAL
237        };
238
239        hr = pDataObj->GetData(&fmt, &medium);
240        if (FAILED(hr)) {
241            OutputDebugString(L"PyShellExt::GetArguments - failed to get CF_HDROP format");
242            return hr;
243        }
244        if (!medium.hGlobal) {
245            OutputDebugString(L"PyShellExt::GetArguments - CF_HDROP format had NULL hGlobal");
246            ReleaseStgMedium(&medium);
247            return E_FAIL;
248        }
249
250        pdropfiles = (DROPFILES*)GlobalLock(medium.hGlobal);
251        if (!pdropfiles) {
252            OutputDebugString(L"PyShellExt::GetArguments - failed to lock CF_HDROP hGlobal");
253            ReleaseStgMedium(&medium);
254            return E_FAIL;
255        }
256
257        if (pdropfiles->fWide) {
258            LPCWSTR files = (LPCWSTR)((char*)pdropfiles + pdropfiles->pFiles);
259            size_t len, count;
260            hr = FilenameListCchLengthW(files, 32767, &len, &count);
261            if (SUCCEEDED(hr)) {
262                LPWSTR args = (LPWSTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1));
263                if (args) {
264                    hr = FilenameListCchCopyW(args, 32767, files, L" ");
265                    if (SUCCEEDED(hr)) {
266                        *pArguments = args;
267                    } else {
268                        CoTaskMemFree(args);
269                    }
270                } else {
271                    hr = E_OUTOFMEMORY;
272                }
273            }
274        } else {
275            LPCSTR files = (LPCSTR)((char*)pdropfiles + pdropfiles->pFiles);
276            size_t len, count;
277            hr = FilenameListCchLengthA(files, 32767, &len, &count);
278            if (SUCCEEDED(hr)) {
279                LPSTR temp = (LPSTR)CoTaskMemAlloc(sizeof(CHAR) * (len + 1));
280                if (temp) {
281                    hr = FilenameListCchCopyA(temp, 32767, files, " ");
282                    if (SUCCEEDED(hr)) {
283                        int wlen = MultiByteToWideChar(CP_ACP, 0, temp, (int)len, NULL, 0);
284                        if (wlen) {
285                            LPWSTR args = (LPWSTR)CoTaskMemAlloc(sizeof(WCHAR) * (wlen + 1));
286                            if (MultiByteToWideChar(CP_ACP, 0, temp, (int)len, args, wlen + 1)) {
287                                *pArguments = args;
288                            } else {
289                                OutputDebugString(L"PyShellExt::GetArguments - failed to convert multi-byte to wide-char path");
290                                CoTaskMemFree(args);
291                                hr = E_FAIL;
292                            }
293                        } else {
294                            OutputDebugString(L"PyShellExt::GetArguments - failed to get length of wide-char path");
295                            hr = E_FAIL;
296                        }
297                    }
298                    CoTaskMemFree(temp);
299                } else {
300                    hr = E_OUTOFMEMORY;
301                }
302            }
303        }
304
305        GlobalUnlock(medium.hGlobal);
306        ReleaseStgMedium(&medium);
307
308        return hr;
309    }
310
311    HRESULT NotifyDragWindow(HWND hwnd) {
312        LRESULT res;
313
314        if (!hwnd) {
315            return S_FALSE;
316        }
317
318        res = SendMessage(hwnd, DDWM_UPDATEWINDOW, 0, NULL);
319
320        if (res) {
321            OutputDebugString(L"PyShellExt::NotifyDragWindow - failed to post DDWM_UPDATEWINDOW");
322            return E_FAIL;
323        }
324
325        return S_OK;
326    }
327
328public:
329    // IDropTarget implementation
330
331    STDMETHODIMP DragEnter(IDataObject *pDataObj, DWORD grfKeyState, POINTL pt, DWORD *pdwEffect) {
332        HWND hwnd;
333
334        OutputDebugString(L"PyShellExt::DragEnter");
335
336        pDataObj->AddRef();
337        data_obj = pDataObj;
338
339        *pdwEffect = DROPEFFECT_MOVE;
340
341        if (FAILED(UpdateDropDescription(data_obj))) {
342            OutputDebugString(L"PyShellExt::DragEnter - failed to update drop description");
343        }
344        if (FAILED(GetDragWindow(data_obj, &hwnd))) {
345            OutputDebugString(L"PyShellExt::DragEnter - failed to get drag window");
346        }
347        if (FAILED(NotifyDragWindow(hwnd))) {
348            OutputDebugString(L"PyShellExt::DragEnter - failed to notify drag window");
349        }
350
351        return S_OK;
352    }
353
354    STDMETHODIMP DragLeave() {
355        return S_OK;
356    }
357
358    STDMETHODIMP DragOver(DWORD grfKeyState, POINTL pt, DWORD *pdwEffect) {
359        return S_OK;
360    }
361
362    STDMETHODIMP Drop(IDataObject *pDataObj, DWORD grfKeyState, POINTL pt, DWORD *pdwEffect) {
363        LPCWSTR args;
364
365        OutputDebugString(L"PyShellExt::Drop");
366        *pdwEffect = DROPEFFECT_NONE;
367
368        if (pDataObj != data_obj) {
369            OutputDebugString(L"PyShellExt::Drop - unexpected data object");
370            return E_FAIL;
371        }
372
373        data_obj->Release();
374        data_obj = NULL;
375
376        if (SUCCEEDED(GetArguments(pDataObj, &args))) {
377            OutputDebugString(args);
378            ShellExecute(NULL, NULL, target, args, target_dir, SW_NORMAL);
379
380            CoTaskMemFree((LPVOID)args);
381        } else {
382            OutputDebugString(L"PyShellExt::Drop - failed to get launch arguments");
383        }
384
385        return S_OK;
386    }
387
388    // IPersistFile implementation
389
390    STDMETHODIMP GetCurFile(LPOLESTR *ppszFileName) {
391        HRESULT hr;
392        size_t len;
393
394        if (!ppszFileName) {
395            return E_POINTER;
396        }
397
398        hr = StringCchLength(target, STRSAFE_MAX_CCH - 1, &len);
399        if (FAILED(hr)) {
400            return E_FAIL;
401        }
402
403        *ppszFileName = (LPOLESTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1));
404        if (!*ppszFileName) {
405            return E_OUTOFMEMORY;
406        }
407
408        hr = StringCchCopy(*ppszFileName, len + 1, target);
409        if (FAILED(hr)) {
410            CoTaskMemFree(*ppszFileName);
411            *ppszFileName = NULL;
412            return E_FAIL;
413        }
414
415        return S_OK;
416    }
417
418    STDMETHODIMP IsDirty() {
419        return S_FALSE;
420    }
421
422    STDMETHODIMP Load(LPCOLESTR pszFileName, DWORD dwMode) {
423        HRESULT hr;
424        size_t len;
425
426        OutputDebugString(L"PyShellExt::Load");
427        OutputDebugString(pszFileName);
428
429        hr = StringCchLength(pszFileName, STRSAFE_MAX_CCH - 1, &len);
430        if (FAILED(hr)) {
431            OutputDebugString(L"PyShellExt::Load - failed to get string length");
432            return hr;
433        }
434
435        if (target) {
436            CoTaskMemFree(target);
437        }
438        if (target_dir) {
439            CoTaskMemFree(target_dir);
440        }
441
442        target = (LPOLESTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1));
443        if (!target) {
444            OutputDebugString(L"PyShellExt::Load - E_OUTOFMEMORY");
445            return E_OUTOFMEMORY;
446        }
447        target_dir = (LPOLESTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1));
448        if (!target_dir) {
449            OutputDebugString(L"PyShellExt::Load - E_OUTOFMEMORY");
450            return E_OUTOFMEMORY;
451        }
452
453        hr = StringCchCopy(target, len + 1, pszFileName);
454        if (FAILED(hr)) {
455            OutputDebugString(L"PyShellExt::Load - failed to copy string");
456            return hr;
457        }
458
459        hr = StringCchCopy(target_dir, len + 1, pszFileName);
460        if (FAILED(hr)) {
461            OutputDebugString(L"PyShellExt::Load - failed to copy string");
462            return hr;
463        }
464        if (!PathRemoveFileSpecW(target_dir)) {
465            OutputDebugStringW(L"PyShellExt::Load - failed to remove filespec from target");
466            return E_FAIL;
467        }
468
469        OutputDebugString(target);
470        target_mode = dwMode;
471        OutputDebugString(L"PyShellExt::Load - S_OK");
472        return S_OK;
473    }
474
475    STDMETHODIMP Save(LPCOLESTR pszFileName, BOOL fRemember) {
476        return E_NOTIMPL;
477    }
478
479    STDMETHODIMP SaveCompleted(LPCOLESTR pszFileName) {
480        return E_NOTIMPL;
481    }
482
483    STDMETHODIMP GetClassID(CLSID *pClassID) {
484        *pClassID = __uuidof(PyShellExt);
485        return S_OK;
486    }
487};
488
489CoCreatableClass(PyShellExt);
490
491STDAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, _COM_Outptr_ void** ppv) {
492    return Module<InProc>::GetModule().GetClassObject(rclsid, riid, ppv);
493}
494
495STDAPI DllCanUnloadNow() {
496    return Module<InProc>::GetModule().Terminate() ? S_OK : S_FALSE;
497}
498
499STDAPI DllRegisterServer() {
500    LONG res;
501    SECURITY_ATTRIBUTES secattr = { sizeof(SECURITY_ATTRIBUTES), NULL, FALSE };
502    LPSECURITY_ATTRIBUTES psecattr = NULL;
503    HKEY key, ipsKey;
504    WCHAR modname[MAX_PATH];
505    DWORD modname_len;
506
507    OutputDebugString(L"PyShellExt::DllRegisterServer");
508    if (!hModule) {
509        OutputDebugString(L"PyShellExt::DllRegisterServer - module handle was not set");
510        return SELFREG_E_CLASS;
511    }
512    modname_len = GetModuleFileName(hModule, modname, MAX_PATH);
513    if (modname_len == 0 ||
514        (modname_len == MAX_PATH && GetLastError() == ERROR_INSUFFICIENT_BUFFER)) {
515        OutputDebugString(L"PyShellExt::DllRegisterServer - failed to get module file name");
516        return SELFREG_E_CLASS;
517    }
518
519    DWORD disp;
520    res = RegCreateKeyEx(HKEY_LOCAL_MACHINE, CLASS_SUBKEY, 0, NULL, 0,
521        KEY_ALL_ACCESS, psecattr, &key, &disp);
522    if (res == ERROR_ACCESS_DENIED) {
523        OutputDebugString(L"PyShellExt::DllRegisterServer - failed to write per-machine registration. Attempting per-user instead.");
524        res = RegCreateKeyEx(HKEY_CURRENT_USER, CLASS_SUBKEY, 0, NULL, 0,
525            KEY_ALL_ACCESS, psecattr, &key, &disp);
526    }
527    if (res != ERROR_SUCCESS) {
528        OutputDebugString(L"PyShellExt::DllRegisterServer - failed to create class key");
529        return SELFREG_E_CLASS;
530    }
531
532    res = RegCreateKeyEx(key, L"InProcServer32", 0, NULL, 0,
533        KEY_ALL_ACCESS, psecattr, &ipsKey, NULL);
534    if (res != ERROR_SUCCESS) {
535        RegCloseKey(key);
536        OutputDebugString(L"PyShellExt::DllRegisterServer - failed to create InProcServer32 key");
537        return SELFREG_E_CLASS;
538    }
539
540    res = RegSetValueEx(ipsKey, NULL, 0,
541        REG_SZ, (LPBYTE)modname, modname_len * sizeof(modname[0]));
542
543    if (res != ERROR_SUCCESS) {
544        RegCloseKey(ipsKey);
545        RegCloseKey(key);
546        OutputDebugString(L"PyShellExt::DllRegisterServer - failed to set server path");
547        return SELFREG_E_CLASS;
548    }
549
550    res = RegSetValueEx(ipsKey, L"ThreadingModel", 0,
551        REG_SZ, (LPBYTE)(L"Apartment"), sizeof(L"Apartment"));
552
553    RegCloseKey(ipsKey);
554    RegCloseKey(key);
555    if (res != ERROR_SUCCESS) {
556        OutputDebugString(L"PyShellExt::DllRegisterServer - failed to set threading model");
557        return SELFREG_E_CLASS;
558    }
559
560    SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, NULL, NULL);
561
562    OutputDebugString(L"PyShellExt::DllRegisterServer - S_OK");
563    return S_OK;
564}
565
566STDAPI DllUnregisterServer() {
567    LONG res_lm, res_cu;
568
569    res_lm = RegDeleteTree(HKEY_LOCAL_MACHINE, CLASS_SUBKEY);
570    if (res_lm != ERROR_SUCCESS && res_lm != ERROR_FILE_NOT_FOUND) {
571        OutputDebugString(L"PyShellExt::DllUnregisterServer - failed to delete per-machine registration");
572        return SELFREG_E_CLASS;
573    }
574
575    res_cu = RegDeleteTree(HKEY_CURRENT_USER, CLASS_SUBKEY);
576    if (res_cu != ERROR_SUCCESS && res_cu != ERROR_FILE_NOT_FOUND) {
577        OutputDebugString(L"PyShellExt::DllUnregisterServer - failed to delete per-user registration");
578        return SELFREG_E_CLASS;
579    }
580
581    if (res_lm == ERROR_FILE_NOT_FOUND && res_cu == ERROR_FILE_NOT_FOUND) {
582        OutputDebugString(L"PyShellExt::DllUnregisterServer - extension was not registered");
583        return SELFREG_E_CLASS;
584    }
585
586    SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, NULL, NULL);
587
588    OutputDebugString(L"PyShellExt::DllUnregisterServer - S_OK");
589    return S_OK;
590}
591
592STDAPI_(BOOL) DllMain(_In_opt_ HINSTANCE hinst, DWORD reason, _In_opt_ void*) {
593    if (reason == DLL_PROCESS_ATTACH) {
594        hModule = hinst;
595
596        cfDropDescription = RegisterClipboardFormat(CFSTR_DROPDESCRIPTION);
597        if (!cfDropDescription) {
598            OutputDebugString(L"PyShellExt::DllMain - failed to get CFSTR_DROPDESCRIPTION format");
599        }
600        cfDragWindow = RegisterClipboardFormat(L"DragWindow");
601        if (!cfDragWindow) {
602            OutputDebugString(L"PyShellExt::DllMain - failed to get DragWindow format");
603        }
604
605        DisableThreadLibraryCalls(hinst);
606    }
607    return TRUE;
608}