• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 
17 static HINSTANCE hModule;
18 static CLIPFORMAT cfDropDescription;
19 static CLIPFORMAT cfDragWindow;
20 
21 #define CLASS_GUID "{BEA218D2-6950-497B-9434-61683EC065FE}"
22 static const LPCWSTR CLASS_SUBKEY = L"Software\\Classes\\CLSID\\" CLASS_GUID;
23 static const LPCWSTR DRAG_MESSAGE = L"Open with %1";
24 
25 using namespace Microsoft::WRL;
26 
FilenameListCchLengthA(LPCSTR pszSource,size_t cchMax,size_t * pcchLength,size_t * pcchCount)27 HRESULT 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 
FilenameListCchLengthW(LPCWSTR pszSource,size_t cchMax,size_t * pcchLength,size_t * pcchCount)48 HRESULT 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 
FilenameListCchCopyA(STRSAFE_LPSTR pszDest,size_t cchDest,LPCSTR pszSource,LPCSTR pszSeparator)69 HRESULT 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 
FilenameListCchCopyW(STRSAFE_LPWSTR pszDest,size_t cchDest,LPCWSTR pszSource,LPCWSTR pszSeparator)96 HRESULT 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 
123 class 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 
134 public:
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 
151 private:
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 
328 public:
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 
489 CoCreatableClass(PyShellExt);
490 
491 STDAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, _COM_Outptr_ void** ppv) {
492     return Module<InProc>::GetModule().GetClassObject(rclsid, riid, ppv);
493 }
494 
495 STDAPI DllCanUnloadNow() {
496     return Module<InProc>::GetModule().Terminate() ? S_OK : S_FALSE;
497 }
498 
499 STDAPI 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 
566 STDAPI 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 
592 STDAPI_(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 }