1
0

deviceprovider.cpp 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343
  1. #include "api.h"
  2. #include "./deviceprovider.h"
  3. #include "../devices/api_devicemanager.h"
  4. extern PMPDevicePlugin plugin;
  5. void connectDrive(wchar_t drive, bool checkSize, bool checkBlacklist);
  6. static size_t tlsIndex = (size_t)-1;
  7. static BOOL
  8. DiscoveryProvider_RegisterCancelSwitch(BOOL *cancelSwitch)
  9. {
  10. if ((size_t)-1 != tlsIndex &&
  11. NULL != WASABI_API_APP)
  12. {
  13. WASABI_API_APP->SetThreadStorage(tlsIndex, cancelSwitch);
  14. return TRUE;
  15. }
  16. return FALSE;
  17. }
  18. static BOOL
  19. DiscoveryProvider_GetCancelSwitchOn()
  20. {
  21. if ((size_t)-1 != tlsIndex &&
  22. NULL != WASABI_API_APP)
  23. {
  24. BOOL *cancelSwitch = (BOOL*)WASABI_API_APP->GetThreadStorage(tlsIndex);
  25. if (NULL != cancelSwitch &&
  26. FALSE != *cancelSwitch)
  27. {
  28. return TRUE;
  29. }
  30. }
  31. return FALSE;
  32. }
  33. static void
  34. DeviceProvider_DriverEnumCb(wchar_t drive, unsigned int type)
  35. {
  36. if (DRIVE_REMOVABLE == type &&
  37. FALSE == DiscoveryProvider_GetCancelSwitchOn())
  38. {
  39. connectDrive(drive,true,true);
  40. }
  41. }
  42. DeviceProvider::DeviceProvider()
  43. : ref(1), activity(0), manager(NULL), readyEvent(NULL), cancelDiscovery(FALSE)
  44. {
  45. InitializeCriticalSection(&lock);
  46. enumerator = (ENUMDRIVES)SendMessageW(plugin.hwndPortablesParent,
  47. WM_PMP_IPC, 0, PMP_IPC_ENUM_ACTIVE_DRIVES);
  48. }
  49. DeviceProvider::~DeviceProvider()
  50. {
  51. CancelDiscovery();
  52. if (NULL != readyEvent)
  53. CloseHandle(readyEvent);
  54. DeleteCriticalSection(&lock);
  55. }
  56. HRESULT DeviceProvider::CreateInstance(DeviceProvider **instance)
  57. {
  58. if (NULL == instance)
  59. return E_POINTER;
  60. *instance = new DeviceProvider();
  61. if (NULL == *instance)
  62. return E_OUTOFMEMORY;
  63. return S_OK;
  64. }
  65. size_t DeviceProvider::AddRef()
  66. {
  67. return InterlockedIncrement((LONG*)&ref);
  68. }
  69. size_t DeviceProvider::Release()
  70. {
  71. if (0 == ref)
  72. return ref;
  73. LONG r = InterlockedDecrement((LONG*)&ref);
  74. if (0 == r)
  75. delete(this);
  76. return r;
  77. }
  78. int DeviceProvider::QueryInterface(GUID interface_guid, void **object)
  79. {
  80. if (NULL == object)
  81. return E_POINTER;
  82. if (IsEqualIID(interface_guid, IFC_DeviceProvider))
  83. *object = static_cast<ifc_deviceprovider*>(this);
  84. else
  85. {
  86. *object = NULL;
  87. return E_NOINTERFACE;
  88. }
  89. if (NULL == *object)
  90. return E_UNEXPECTED;
  91. AddRef();
  92. return S_OK;
  93. }
  94. void DeviceProvider::Lock()
  95. {
  96. EnterCriticalSection(&lock);
  97. }
  98. void DeviceProvider::Unlock()
  99. {
  100. LeaveCriticalSection(&lock);
  101. }
  102. DWORD DeviceProvider::DiscoveryThread()
  103. {
  104. IncrementActivity();
  105. if (NULL != enumerator &&
  106. FALSE == cancelDiscovery)
  107. {
  108. DiscoveryProvider_RegisterCancelSwitch(&cancelDiscovery);
  109. enumerator(DeviceProvider_DriverEnumCb);
  110. DiscoveryProvider_RegisterCancelSwitch(NULL);
  111. }
  112. DecrementActivity();
  113. Lock();
  114. if (NULL != readyEvent)
  115. SetEvent(readyEvent);
  116. Unlock();
  117. return 0;
  118. }
  119. static int DeviceProvider_DiscoveryThreadStarter(HANDLE handle, void *user, intptr_t id)
  120. {
  121. DeviceProvider *self;
  122. DWORD result;
  123. self = (DeviceProvider*)user;
  124. if (NULL != self)
  125. result = self->DiscoveryThread();
  126. else
  127. result = -2;
  128. return result;
  129. }
  130. HRESULT DeviceProvider::BeginDiscovery(api_devicemanager *manager)
  131. {
  132. HRESULT hr;
  133. if (NULL == enumerator)
  134. return E_UNEXPECTED;
  135. Lock();
  136. if (NULL != readyEvent &&
  137. WAIT_TIMEOUT == WaitForSingleObject(readyEvent, 0))
  138. {
  139. hr = E_PENDING;
  140. }
  141. else
  142. {
  143. hr = S_OK;
  144. if (NULL == readyEvent)
  145. {
  146. readyEvent = CreateEvent(NULL, TRUE, TRUE, NULL);
  147. if (NULL == readyEvent)
  148. hr = E_FAIL;
  149. }
  150. if ((size_t)-1 == tlsIndex &&
  151. NULL != WASABI_API_APP)
  152. {
  153. tlsIndex = WASABI_API_APP->AllocateThreadStorage();
  154. }
  155. if (SUCCEEDED(hr))
  156. {
  157. cancelDiscovery = FALSE;
  158. ResetEvent(readyEvent);
  159. if (0 != WASABI_API_THREADPOOL->RunFunction(0, DeviceProvider_DiscoveryThreadStarter,
  160. this, 0, api_threadpool::FLAG_LONG_EXECUTION))
  161. {
  162. SetEvent(readyEvent);
  163. hr = E_FAIL;
  164. }
  165. }
  166. }
  167. Unlock();
  168. return hr;
  169. }
  170. HRESULT DeviceProvider::CancelDiscovery()
  171. {
  172. HRESULT hr;
  173. hr = S_FALSE;
  174. Lock();
  175. if (NULL != readyEvent)
  176. {
  177. cancelDiscovery = TRUE;
  178. if (WAIT_OBJECT_0 == WaitForSingleObject(readyEvent, 0))
  179. hr = S_OK;
  180. cancelDiscovery = FALSE;
  181. }
  182. Unlock();
  183. return hr;
  184. }
  185. HRESULT DeviceProvider::GetActive()
  186. {
  187. HRESULT hr;
  188. Lock();
  189. if (0 != activity)
  190. hr = S_OK;
  191. else
  192. hr = S_FALSE;
  193. Unlock();
  194. return hr;
  195. }
  196. HRESULT DeviceProvider::Register(api_devicemanager *manager)
  197. {
  198. HRESULT hr;
  199. if (NULL != this->manager)
  200. return E_UNEXPECTED;
  201. if (NULL == manager)
  202. return E_POINTER;
  203. hr = manager->RegisterProvider(this);
  204. if (SUCCEEDED(hr))
  205. {
  206. this->manager = manager;
  207. manager->AddRef();
  208. }
  209. return hr;
  210. }
  211. HRESULT DeviceProvider::Unregister()
  212. {
  213. HRESULT hr;
  214. if (NULL == manager)
  215. return E_UNEXPECTED;
  216. hr = manager->UnregisterProvider(this);
  217. manager->Release();
  218. manager = NULL;
  219. return hr;
  220. }
  221. size_t DeviceProvider::IncrementActivity()
  222. {
  223. size_t a;
  224. Lock();
  225. activity++;
  226. if (1 == activity &&
  227. NULL != manager)
  228. {
  229. manager->SetProviderActive(this, TRUE);
  230. }
  231. a = activity;
  232. Unlock();
  233. return a;
  234. }
  235. size_t DeviceProvider::DecrementActivity()
  236. {
  237. size_t a;
  238. Lock();
  239. if (0 != activity)
  240. {
  241. activity--;
  242. if (0 == activity &&
  243. NULL != manager)
  244. {
  245. manager->SetProviderActive(this, FALSE);
  246. }
  247. }
  248. a = activity;
  249. Unlock();
  250. return a;
  251. }
  252. #define CBCLASS DeviceProvider
  253. START_DISPATCH;
  254. CB(ADDREF, AddRef)
  255. CB(RELEASE, Release)
  256. CB(QUERYINTERFACE, QueryInterface)
  257. CB(API_BEGINDISCOVERY, BeginDiscovery)
  258. CB(API_CANCELDISCOVERY, CancelDiscovery)
  259. CB(API_GETACTIVE, GetActive)
  260. END_DISPATCH;
  261. #undef CBCLASS