Check available disk space before extracting

clowd · Apr 7, 2022 · daddc02 · daddc02
1 parent 5c5dbdc
commit daddc02
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 6 deletions.
diff --git a/src/Setup/Setup.cpp b/src/Setup/Setup.cpp
@@ -50,6 +50,32 @@ void unzipSingleFile(BYTE* zipBuf, DWORD cZipBuf, wstring fileLocation, std::fun
     if (!unzipSuccess) throw wstring(L"Unable to extract embedded package (predicate not found).");
 }
 
+// Prints to the provided buffer a nice number of bytes (KB, MB, GB, etc)
+wstring pretty_bytes(uint64_t bytes)
+{
+    wchar_t buf[128];
+    const wchar_t* suffixes[7];
+    suffixes[0] = L"B";
+    suffixes[1] = L"KB";
+    suffixes[2] = L"MB";
+    suffixes[3] = L"GB";
+    suffixes[4] = L"TB";
+    suffixes[5] = L"PB";
+    suffixes[6] = L"EB";
+    uint64_t s = 0; // which suffix to use
+    double count = bytes;
+    while (count >= 1000 && s < 7) {
+        s++;
+        count /= 1000;
+    }
+    if (count - floor(count) == 0.0)
+        swprintf(buf, 128, L"%d %s", (int)count, suffixes[s]);
+    else
+        swprintf(buf, 128, L"%.1f %s", count, suffixes[s]);
+
+    return wstring(buf);
+}
+
 int WINAPI wWinMain(_In_ HINSTANCE hInstance, _In_opt_ HINSTANCE hPrevInstance, _In_ PWSTR pCmdLine, _In_ int nCmdShow)
 {
     if (!IsWindows7SP1OrGreater()) {
@@ -65,7 +91,7 @@ int WINAPI wWinMain(_In_ HINSTANCE hInstance, _In_opt_ HINSTANCE hPrevInstance,
         // locate bundled package and map to memory
         memAddr = util::mmap_read(util::get_current_process_path(), 0);
         if (!memAddr) {
-            throw new wstring(L"Unable to map executable to memory");
+            throw wstring(L"Unable to map executable to memory");
         }
 
         int64_t packageOffset, packageLength;
@@ -75,6 +101,16 @@ int WINAPI wWinMain(_In_ HINSTANCE hInstance, _In_opt_ HINSTANCE hPrevInstance,
             throw wstring(L"The embedded package was not found");
         }
 
+        // rough check for sufficient disk space before extracting anything
+        // required space is size of compressed nupkg, size of extracted app, 
+        // and squirrel overheads (incl temp files). the constant 0.38 is a
+        // aggressive estimate on what the compression ratio might be.
+        int64_t squirrelOverhead = 50 * 1000 * 1000; 
+        int64_t requiredSpace = squirrelOverhead + packageLength + (packageLength / (double)0.38);
+        if (!util::check_diskspace(requiredSpace)) {
+            throw wstring(L"Insufficient disk space. This application requires at least " + pretty_bytes(requiredSpace) + L" free space to be installed");
+        }
+
         // extract Squirrel installer from package
         std::function<bool(ZIPENTRY&)> endsWithSquirrel([](ZIPENTRY& z) {
             return hasEnding(std::wstring(z.name), L"Squirrel.exe");

diff --git a/src/Setup/platform_util.cpp b/src/Setup/platform_util.cpp
@@ -1,5 +1,6 @@
 #include "platform_util.h"
 #include <windows.h>
+#include <shlobj_core.h>
 #include <tchar.h>
 #include <string>
 
@@ -16,20 +17,25 @@ wstring get_filename_from_path(wstring& path)
     return path.substr(idx + 1);
 }
 
-// https://stackoverflow.com/a/17387176/184746
-void throwLastWin32Error(wstring addedInfo)
+void throwWin32Error(HRESULT hr, wstring addedInfo)
 {
-    DWORD errorMessageID = ::GetLastError();
-    if (errorMessageID == 0) {
+    if (hr == 0) {
         return;
     }
 
+    // https://stackoverflow.com/a/17387176/184746
+    // https://stackoverflow.com/a/455533/184746
     LPWSTR messageBuffer = nullptr;
     size_t size = FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
-                                NULL, errorMessageID, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (LPWSTR)&messageBuffer, 0, NULL);
+                                NULL, hr, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (LPWSTR)&messageBuffer, 0, NULL);
 
     wstring message(messageBuffer, size);
 
+    if (messageBuffer) {
+        LocalFree(messageBuffer);
+        messageBuffer = nullptr;
+    }
+
     if (addedInfo.empty()) {
         throw message;
     }
@@ -38,6 +44,11 @@ void throwLastWin32Error(wstring addedInfo)
     }
 }
 
+void throwLastWin32Error(wstring addedInfo)
+{
+    throwWin32Error(::GetLastError(), addedInfo);
+}
+
 std::wstring util::get_temp_file_path(wstring extension)
 {
     wchar_t tempFolderBuf[MAX_PATH];
@@ -53,6 +64,20 @@ std::wstring util::get_temp_file_path(wstring extension)
     return tempFile;
 }
 
+bool util::check_diskspace(uint64_t requiredSpace)
+{
+    TCHAR szPath[MAX_PATH];
+    auto hr = SHGetFolderPath(NULL, CSIDL_LOCAL_APPDATA, NULL, 0, szPath);
+    if (FAILED(hr))
+        throwWin32Error(hr, L"Unable to locate %localappdata%.");
+
+    ULARGE_INTEGER freeSpace;
+    if (!GetDiskFreeSpaceEx(szPath, 0, 0, &freeSpace))
+        throwLastWin32Error(L"Unable to verify sufficient available free space on disk.");
+
+    return freeSpace.QuadPart > requiredSpace;
+}
+
 std::wstring util::get_current_process_path()
 {
     wchar_t ourFile[MAX_PATH];

diff --git a/src/Setup/platform_util.h b/src/Setup/platform_util.h
@@ -7,6 +7,7 @@ namespace util
 {
     std::wstring get_temp_file_path(std::wstring extension);
     std::wstring get_current_process_path();
+    bool check_diskspace(uint64_t requiredSpace);
     void wexec(const wchar_t* cmd);
     void show_error_dialog(std::wstring msg);
     uint8_t* mmap_read(const std::wstring& filePath, size_t* length);