From 2153a8b536208ec72dc69298e8bace7cff8b53b9 Mon Sep 17 00:00:00 2001
From: "Scott E. Graves" <scott.e.graves@gmail.com>
Date: Thu, 13 Apr 2017 11:39:32 -0500
Subject: [PATCH] Fixes

---
 src/siadrive_dokan_api/siadokandrive.cpp | 85 ++++++++++++++++++------
 1 file changed, 65 insertions(+), 20 deletions(-)

diff --git a/src/siadrive_dokan_api/siadokandrive.cpp b/src/siadrive_dokan_api/siadokandrive.cpp
index 550ecb5..606e909 100644
--- a/src/siadrive_dokan_api/siadokandrive.cpp
+++ b/src/siadrive_dokan_api/siadokandrive.cpp
@@ -1217,6 +1217,8 @@ private:
 	static NTSTATUS _mountStatus;
 	static SString _mountPoint;
   static std::vector<OpenFileInfo*> _openFiles;
+  static std::mutex _downloadMutex;
+  static std::vector<SString> _activeDownloads;
 
 private:
 	inline static const FilePath& GetCacheLocation()
@@ -1246,9 +1248,31 @@ private:
 
   static bool AddFileToCache(OpenFileInfo& openFileInfo, PDOKAN_FILE_INFO dokanFileInfo)
   {    
-    // TODO taking too long
+    bool found = true;
+    bool wasFound = false;
+    // Wait for active download to complete
+    while (found)
+    {
+      {
+        std::lock_guard<std::mutex> l(_downloadMutex);
+        found = std::find(_activeDownloads.begin(), _activeDownloads.end(), openFileInfo.SiaPath) != _activeDownloads.end();
+        wasFound = wasFound || found;
+        if (!found && !wasFound)
+        {
+          // Add to active download tracking
+          _activeDownloads.push_back(openFileInfo.SiaPath);
+        }
+      }
+
+      if (found)
+      {
+        ::Sleep(10);
+      }
+    }
+
     bool ret = true;
-    if (openFileInfo.Dummy)
+    // Only download if necessary
+    if (openFileInfo.Dummy && !wasFound)
     {
       FilePath tempFilePath = FilePath::GetTempDirectory();
       tempFilePath.Append(GenerateSha256(openFileInfo.SiaPath) + ".siatmp");
@@ -1270,13 +1294,9 @@ private:
               matched.push_back(ofi);
             }
           });
-        }
 
-        // Close all to allow move to complete
-        for (auto& ofi : matched)
-        {
-          std::lock_guard<std::mutex> l(_openFileMutex);
-          if (std::find(_openFiles.begin(), _openFiles.end(), ofi) != _openFiles.end())
+          // Close all to allow move to complete
+          for (auto& ofi : matched)
           {
             ::CloseHandle(ofi->FileHandle);
           }
@@ -1303,6 +1323,10 @@ private:
             }
           }
         }
+
+        // Remove from active download tracking
+        std::lock_guard<std::mutex> l(_downloadMutex);
+        _activeDownloads.erase(std::remove(_activeDownloads.begin(), _activeDownloads.end(), openFileInfo.SiaPath), _activeDownloads.end());
       }
     }
 
@@ -1318,6 +1342,10 @@ private:
 
   static void HandleSiaFileClose(const OpenFileInfo& openFileInfo, const std::uint64_t& fileSize, const bool& deleteOnClose)
   {
+    std::function<void()> notifyCritical;
+
+    // Lock while closing
+    std::lock_guard<std::mutex> l(_openFileMutex);
     CEventSystem::EventSystem.NotifyEvent(CreateSystemEvent(DokanCloseFile(openFileInfo.CacheFilePath)));
     if (deleteOnClose)
     {
@@ -1325,7 +1353,11 @@ private:
       auto result = _uploadManager->Remove(openFileInfo.SiaPath);
       if (!ApiSuccess(result))
       {
-        CEventSystem::EventSystem.NotifyEvent(CreateSystemCriticalEvent(result.GetReason()));
+        SString reason = result.GetReason();
+        notifyCritical = [=]()
+        {
+          CEventSystem::EventSystem.NotifyEvent(CreateSystemCriticalEvent(reason));
+        };
       }
     }
     else if (openFileInfo.Changed)
@@ -1339,7 +1371,11 @@ private:
         auto result = _uploadManager->AddOrUpdate(openFileInfo.SiaPath, openFileInfo.CacheFilePath, *reinterpret_cast<std::uint64_t*>(&fileTimes[2]));
         if (!ApiSuccess(result))
         {
-          CEventSystem::EventSystem.NotifyEvent(CreateSystemCriticalEvent(result.GetReason()));
+          SString reason = result.GetReason();
+          notifyCritical = [=]()
+          {
+            CEventSystem::EventSystem.NotifyEvent(CreateSystemCriticalEvent(reason));
+          };
         }
       }
       else
@@ -1350,7 +1386,11 @@ private:
         auto result = _uploadManager->Remove(openFileInfo.SiaPath);
         if (!ApiSuccess(result))
         {
-          CEventSystem::EventSystem.NotifyEvent(CreateSystemCriticalEvent(result.GetReason()));
+          SString reason = result.GetReason();
+          notifyCritical = [=]()
+          {
+            CEventSystem::EventSystem.NotifyEvent(CreateSystemCriticalEvent(reason));
+          };
         }
       }
     }
@@ -1359,8 +1399,13 @@ private:
       ::CloseHandle(openFileInfo.FileHandle);
     }
 
-    std::lock_guard<std::mutex> l(_openFileMutex);
     _openFiles.erase(std::remove(_openFiles.begin(), _openFiles.end(), &openFileInfo), _openFiles.end());
+
+    // Make sure cleanup is complete before notifying
+    if (notifyCritical)
+    {
+      notifyCritical();
+    }
   }
 
 	static void StartFileListThread()
@@ -1954,15 +1999,13 @@ private:
 		PULONGLONG TotalNumberOfFreeBytes, PDOKAN_FILE_INFO dokanFileInfo)
 	{
 		UNREFERENCED_PARAMETER(dokanFileInfo);
-    SiaCurrency allocatedFunds = _siaApi->GetRenter()->GetFunds();
-    SiaCurrency unspentFunds = _siaApi->GetRenter()->GetUnspent();
-    SiaCurrency totalUsedGb = _siaApi->GetRenter()->GetTotalUsedBytes() ? _siaApi->GetRenter()->GetTotalUsedBytes() : 0.0;
-    auto totalAvailable = (totalUsedGb / (allocatedFunds - unspentFunds)) * allocatedFunds;
-    auto totalRemainGb = totalAvailable - totalUsedGb;
+    
+    SiaCurrency totalBytes;
+    _siaApi->GetRenter()->CalculateEstimatedStorage(_siaApi->GetRenter()->GetFunds(), totalBytes);
+    SiaCurrency totalUsed = _siaApi->GetRenter()->GetTotalUsedBytes();
 
-		*FreeBytesAvailable = totalRemainGb.ToUInt();
-		*TotalNumberOfBytes = totalAvailable.ToUInt();
-		*TotalNumberOfFreeBytes = totalAvailable.ToUInt();
+    *FreeBytesAvailable = (totalBytes == 0) ? 0 : (totalBytes - totalUsed).ToUInt();
+		*TotalNumberOfFreeBytes = *TotalNumberOfBytes = totalBytes.ToUInt();
 
 		return STATUS_SUCCESS;
 	}
@@ -2693,6 +2736,8 @@ std::unique_ptr<std::thread> DokanImpl::_mountThread;
 NTSTATUS DokanImpl::_mountStatus = STATUS_SUCCESS;
 SString DokanImpl::_mountPoint;
 std::vector<DokanImpl::OpenFileInfo*> DokanImpl::_openFiles;
+std::mutex DokanImpl::_downloadMutex;
+std::vector<SString> DokanImpl::_activeDownloads;
 
 CSiaDokanDrive::CSiaDokanDrive(CSiaApi& siaApi, CSiaDriveConfig* siaDriveConfig) :
 	_siaApi(siaApi),