1  #include "pch.h"
  2  #include "StlThumbnailProvider.h"
  3  
  4  #include <filesystem>
  5  #include <fstream>
  6  #include <shellapi.h>
  7  #include <Shlwapi.h>
  8  #include <string>
  9  
 10  #include <wil/com.h>
 11  
 12  #include <common/interop/shared_constants.h>
 13  #include <common/logger/logger.h>
 14  #include <common/SettingsAPI/settings_helpers.h>
 15  #include <common/utils/process_path.h>
 16  
 17  extern HINSTANCE g_hInst;
 18  extern long g_cDllRef;
 19  
 20  StlThumbnailProvider::StlThumbnailProvider() :
 21      m_cRef(1), m_pStream(NULL), m_process(NULL)
 22  {
 23      std::filesystem::path logFilePath(PTSettingsHelper::get_local_low_folder_location());
 24      logFilePath.append(LogSettings::stlThumbLogPath);
 25      Logger::init(LogSettings::stlThumbLoggerName, logFilePath.wstring(), PTSettingsHelper::get_log_settings_file_location());
 26  
 27      InterlockedIncrement(&g_cDllRef);
 28  }
 29  
 30  StlThumbnailProvider::~StlThumbnailProvider()
 31  {
 32      InterlockedDecrement(&g_cDllRef);
 33  }
 34  
 35  #pragma region IUnknown
 36  
 37  IFACEMETHODIMP StlThumbnailProvider::QueryInterface(REFIID riid, void** ppv)
 38  {
 39      static const QITAB qit[] = {
 40          QITABENT(StlThumbnailProvider, IThumbnailProvider),
 41          QITABENT(StlThumbnailProvider, IInitializeWithStream),
 42          { 0 },
 43      };
 44      return QISearch(this, qit, riid, ppv);
 45  }
 46  
 47  IFACEMETHODIMP_(ULONG)
 48  StlThumbnailProvider::AddRef()
 49  {
 50      return InterlockedIncrement(&m_cRef);
 51  }
 52  
 53  IFACEMETHODIMP_(ULONG)
 54  StlThumbnailProvider::Release()
 55  {
 56      ULONG cRef = InterlockedDecrement(&m_cRef);
 57      if (0 == cRef)
 58      {
 59          delete this;
 60      }
 61      return cRef;
 62  }
 63  
 64  #pragma endregion
 65  
 66  #pragma region IInitializationWithStream
 67  
 68  IFACEMETHODIMP StlThumbnailProvider::Initialize(IStream* pStream, DWORD grfMode)
 69  {
 70      HRESULT hr = E_INVALIDARG;
 71      if (pStream)
 72      {
 73          // Initialize can be called more than once, so release existing valid
 74          // m_pStream.
 75          if (m_pStream)
 76          {
 77              m_pStream->Release();
 78              m_pStream = NULL;
 79          }
 80  
 81          m_pStream = pStream;
 82          m_pStream->AddRef();
 83          hr = S_OK;
 84      }
 85      return hr;
 86  }
 87  
 88  #pragma endregion
 89  
 90  #pragma region IThumbnailProvider
 91  
 92  IFACEMETHODIMP StlThumbnailProvider::GetThumbnail(UINT cx, HBITMAP* phbmp, WTS_ALPHATYPE* pdwAlpha)
 93  {
 94      // Read stream into the buffer
 95      char buffer[4096];
 96      ULONG cbRead;
 97  
 98      Logger::trace(L"Begin");
 99  
100      GUID guid;
101      if (CoCreateGuid(&guid) == S_OK)
102      {
103          wil::unique_cotaskmem_string guidString;
104          if (SUCCEEDED(StringFromCLSID(guid, &guidString)))
105          {
106              Logger::info(L"Read stream and save to tmp file.");
107              
108              // {CLSID} -> CLSID
109              std::wstring guid = std::wstring(guidString.get()).substr(1, std::wstring(guidString.get()).size() - 2);
110              std::wstring filePath = PTSettingsHelper::get_local_low_folder_location() + L"\\StlThumbnail-Temp\\";
111              if (!std::filesystem::exists(filePath))
112              {
113                  std::filesystem::create_directories(filePath);
114              }
115  
116              std::wstring fileName = filePath + guid + L".stl";
117  
118              // Write data to tmp file
119              std::fstream file;
120              file.open(fileName, std::ios_base::out | std::ios_base::binary);
121  
122              if (!file.is_open())
123              {
124                  return 0;
125              }
126  
127              while (true)
128              {
129                  auto result = m_pStream->Read(buffer, 4096, &cbRead);
130  
131                  file.write(buffer, cbRead);
132                  if (result == S_FALSE)
133                  {
134                      break;
135                  }
136              }
137              file.close();
138  
139              m_pStream->Release();
140              m_pStream = NULL;
141  
142              try
143              {
144                  Logger::info(L"Start StlThumbnailProvider.exe");
145                  
146                  STARTUPINFO info = { sizeof(info) };
147                  std::wstring cmdLine{ L"\"" + fileName + L"\"" };
148                  cmdLine += L" ";
149                  cmdLine += std::to_wstring(cx);
150  
151                  std::wstring appPath = get_module_folderpath(g_hInst) + L"\\PowerToys.StlThumbnailProvider.exe";
152  
153                  SHELLEXECUTEINFO sei{ sizeof(sei) };
154                  sei.fMask = { SEE_MASK_NOCLOSEPROCESS | SEE_MASK_FLAG_NO_UI };
155                  sei.lpFile = appPath.c_str();
156                  sei.lpParameters = cmdLine.c_str();
157                  sei.nShow = SW_SHOWDEFAULT;
158                  ShellExecuteEx(&sei);
159                  m_process = sei.hProcess;
160                  WaitForSingleObject(m_process, INFINITE);
161                  std::filesystem::remove(fileName);
162  
163                  std::wstring fileNameBmp = filePath + guid + L".bmp";
164                  if (std::filesystem::exists(fileNameBmp))
165                  {
166                      *phbmp = static_cast<HBITMAP>(LoadImage(NULL, fileNameBmp.c_str(), IMAGE_BITMAP, 0, 0, LR_LOADFROMFILE));
167                      *pdwAlpha = WTS_ALPHATYPE::WTSAT_ARGB;
168                      std::filesystem::remove(fileNameBmp);
169                  }
170                  else
171                  {
172                      Logger::info(L"Bmp file not generated.");
173                      return E_FAIL;
174                  }
175              }
176              catch (std::exception& e)
177              {
178                  std::wstring errorMessage = std::wstring{ winrt::to_hstring(e.what()) };
179                  Logger::error(L"Failed to start StlThumbnailProvider.exe. Error: {}", errorMessage);
180              }
181          }
182      }
183  
184      // ensure releasing the stream (not all if branches contain it)
185      if (m_pStream)
186      {
187          m_pStream->Release();
188          m_pStream = NULL;
189      }
190  
191      return S_OK;
192  }
193  
194  #pragma endregion
195  
196  #pragma region Helper Functions
197  
198  #pragma endregion