Bug 2061: Drive that was ever opened in local file panel cannot be safely removed until WinSCP is closed

https://winscp.net/tracker/2061

Source commit: 49bf274d2e7f928a92cfa8170e807d4e6164dd99

source/packages/filemng/DriveView.pas

@@ -68,6 +68,8 @@ type
     DiscMonitor: TDiscMonitor; {Monitor thread}
     ChangeTimer: TTimer;       {Change timer for the monitor thread}
     DefaultDir: string;        {Current directory}
+    DriveHandle: THandle;
+    NotificationHandle: HDEVNOTIFY;
   end;
 
   TDriveStatusPair = TPair<string, TDriveStatus>;
@@ -185,6 +187,9 @@ type
     function WatchThreadActive: Boolean; overload;
     function WatchThreadActive(Drive: string): Boolean; overload;
     procedure InternalWndProc(var Msg: TMessage);
+    procedure UpdateDriveNotifications(Drive: string);
+    procedure DriveRemoved(Drive: string);
+    procedure DriveRemoving(Drive: string);
 
     function DirAttrMask: Integer;
     function CreateDriveStatus: TDriveStatus;
@@ -506,19 +511,20 @@ end; {Create}
 
 destructor TDriveView.Destroy;
 var
-  DriveStatus: TDriveStatus;
+  DriveStatusPair: TDriveStatusPair;
 begin
   Classes.DeallocateHWnd(FInternalWindowHandle);
 
-  for DriveStatus in FDriveStatus.Values do
+  for DriveStatusPair in FDriveStatus do
   begin
-    with DriveStatus do
+    with DriveStatusPair.Value do
     begin
       if Assigned(DiscMonitor) then
-        DiscMonitor.Free;
+        FreeAndNil(DiscMonitor);
       if Assigned(ChangeTimer) then
-        ChangeTimer.Free;
+        FreeAndNil(ChangeTimer);
     end;
+    UpdateDriveNotifications(DriveStatusPair.Key);
   end;
   FDriveStatus.Free;
 
@@ -544,9 +550,17 @@ begin
     ChangeTimer.Interval := 0;
     ChangeTimer.Enabled := False;
     ChangeTimer.OnTimer := ChangeTimerOnTimer;
+    DriveHandle := INVALID_HANDLE_VALUE;
+    NotificationHandle := nil;
   end;
 end;
 
+procedure TDriveView.DriveRemoving(Drive: string);
+begin
+  DriveRemoved(Drive);
+  TerminateWatchThread(Drive);
+end;
+
 type
   PDevBroadcastHdr = ^TDevBroadcastHdr;
   TDevBroadcastHdr = record
@@ -564,17 +578,34 @@ type
     dbcv_flags: WORD;
   end;
 
+  PDEV_BROADCAST_HANDLE = ^DEV_BROADCAST_HANDLE;
+  DEV_BROADCAST_HANDLE = record
+    dbch_size       : DWORD;
+    dbch_devicetype : DWORD;
+    dbch_reserved   : DWORD;
+    dbch_handle     : THandle;
+    dbch_hdevnotify : HDEVNOTIFY  ;
+    dbch_eventguid  : TGUID;
+    dbch_nameoffset : LongInt;
+    dbch_data       : Byte;
+  end;
+
 const
+  DBT_DEVTYP_HANDLE = $00000006;
   DBT_CONFIGCHANGED = $0018;
   DBT_DEVICEARRIVAL = $8000;
+  DBT_DEVICEQUERYREMOVE = $8001;
   DBT_DEVICEREMOVEPENDING = $8003;
   DBT_DEVICEREMOVECOMPLETE = $8004;
   DBT_DEVTYP_VOLUME = $00000002;
 
 procedure TDriveView.InternalWndProc(var Msg: TMessage);
 var
+  DeviceType: DWORD;
   UnitMask: DWORD;
+  DeviceHandle: THandle;
   Drive: Char;
+  DriveStatusPair: TDriveStatusPair;
 begin
   with Msg do
   begin
@@ -592,9 +623,14 @@ begin
         SetTimer(FInternalWindowHandle, 1, MSecsPerSec, nil);
       end
         else
-      if wParam = DBT_DEVICEREMOVEPENDING then
+      if (wParam = DBT_DEVICEQUERYREMOVE) or
+         (wParam = DBT_DEVICEREMOVEPENDING) then
       begin
-        if PDevBroadcastHdr(lParam)^.dbch_devicetype = DBT_DEVTYP_VOLUME then
+        DeviceType := PDevBroadcastHdr(lParam)^.dbch_devicetype;
+        // This is specifically for VeraCrypt.
+        // For normal drives, see DBT_DEVTYP_HANDLE below
+        // (and maybe now that we have generic implementation, this specific code for VeraCrypt might not be needed anymore)
+        if DeviceType = DBT_DEVTYP_VOLUME then
         begin
           UnitMask := PDevBroadcastVolume(lParam)^.dbcv_unitmask;
           Drive := FirstDrive;
@@ -602,23 +638,21 @@ begin
           begin
             if UnitMask and $01 <> 0 then
             begin
-              // Disable disk monitor to release the handle to the drive.
-              // It may happen that the dirve is not removed in the end. In this case we do not currently resume the
-              // monitoring. We can watch for DBT_DEVICEQUERYREMOVEFAILED to resume the monitoring.
-              // But currently we implement this for VeraCrypt, which does not send this notification.
-              with GetDriveStatus(Drive) do
-              begin
-                if Assigned(DiscMonitor) then
-                begin
-                  DiscMonitor.Enabled := False;
-                  DiscMonitor.Free;
-                  DiscMonitor := nil;
-                end;
-              end;
+              DriveRemoving(Drive);
             end;
             UnitMask := UnitMask shr 1;
             Drive := Chr(Ord(Drive) + 1);
           end;
+        end
+          else
+        if DeviceType = DBT_DEVTYP_HANDLE then
+        begin
+          DeviceHandle := PDEV_BROADCAST_HANDLE(lParam)^.dbch_handle;
+          for DriveStatusPair in FDriveStatus do
+            if DriveStatusPair.Value.DriveHandle = DeviceHandle then
+            begin
+              DriveRemoving(DriveStatusPair.Key);
+            end;
         end;
       end;
     end
@@ -1292,6 +1326,41 @@ begin
   Result := Drives;
 end;
 
+procedure TDriveView.DriveRemoved(Drive: string);
+var
+  NewDrive: Char;
+begin
+  if (Directory <> '') and (Directory[1] = Drive) then
+  begin
+    if DriveInfo.IsRealDrive(Drive) then NewDrive := Drive[1]
+      else NewDrive := SystemDrive;
+
+    repeat
+      if NewDrive < SystemDrive then NewDrive := SystemDrive
+        else
+      if NewDrive = SystemDrive then NewDrive := LastDrive
+        else Dec(NewDrive);
+      DriveInfo.ReadDriveStatus(NewDrive, dsSize or dsImageIndex);
+
+      if NewDrive = Drive then
+      begin
+        Break;
+      end;
+
+      if DriveInfo.Get(NewDrive).Valid and DriveInfo.Get(NewDrive).DriveReady and Assigned(GetDriveStatus(NewDrive).RootNode) then
+      begin
+        Directory := NodePathName(GetDriveStatus(NewDrive).RootNode);
+        break;
+      end;
+    until False;
+
+    if not Assigned(Selected) then
+    begin
+      Directory := NodePathName(GetDriveStatus(SystemDrive).RootNode);
+    end;
+  end;
+end;
+
 procedure TDriveView.RefreshRootNodes(dsFlags: Integer);
 var
   Drives: TStrings;
@@ -1299,7 +1368,6 @@ var
   SaveCursor: TCursor;
   WasValid: Boolean;
   NodeData: TNodeData;
-  NewDrive: Char;
   DriveStatus: TDriveStatus;
   NextDriveNode: TTreeNode;
   Index: Integer;
@@ -1376,35 +1444,7 @@ begin
           if WasValid then
           {Drive has been removed => delete rootnode:}
           begin
-            if (Directory <> '') and (Directory[1] = Drive) then
-            begin
-              if DriveInfo.IsRealDrive(Drive) then NewDrive := Drive[1]
-                else NewDrive := SystemDrive;
-
-              repeat
-                if NewDrive < SystemDrive then NewDrive := SystemDrive
-                  else
-                if NewDrive = SystemDrive then NewDrive := LastDrive
-                  else Dec(NewDrive);
-                DriveInfo.ReadDriveStatus(NewDrive, dsSize or dsImageIndex);
-
-                if NewDrive = Drive then
-                begin
-                  Break;
-                end;
-
-                if DriveInfo.Get(NewDrive).Valid and DriveInfo.Get(NewDrive).DriveReady and Assigned(GetDriveStatus(NewDrive).RootNode) then
-                begin
-                  Directory := NodePathName(GetDriveStatus(NewDrive).RootNode);
-                  break;
-                end;
-              until False;
-
-              if not Assigned(Selected) then
-              begin
-                Directory := NodePathName(GetDriveStatus(SystemDrive).RootNode);
-              end;
-            end;
+            DriveRemoved(Drive);
             Scanned := False;
             Verified := False;
             RootNode.Delete;
@@ -1965,6 +2005,7 @@ begin
       DiscMonitor.SetDirectory(DriveInfo.GetDriveRoot(Drive));
       DiscMonitor.Open;
     end;
+    UpdateDriveNotifications(Drive); // probably noop, as the monitor is not enabled yet
   end;
 end; {CreateWatchThread}
 
@@ -2002,13 +2043,14 @@ end; {NodeWatched}
 procedure TDriveView.ChangeInvalid(Sender: TObject; const Directory: string;
   const ErrorStr: string);
 var
-  Dir: string;
+  Drive: string;
 begin
-  Dir := (Sender as TDiscMonitor).Directories[0];
-  with GetDriveStatus(DriveInfo.GetDriveKey(Dir)) do
+  Drive := DriveInfo.GetDriveKey((Sender as TDiscMonitor).Directories[0]);
+  with GetDriveStatus(Drive) do
   begin
     DiscMonitor.Close;
   end;
+  UpdateDriveNotifications(Drive);
 end; {DirWatchChangeInvalid}
 
 procedure TDriveView.ChangeDetected(Sender: TObject; const Directory: string;
@@ -2065,6 +2107,60 @@ begin
   end;
 end; {ChangeTimerOnTimer}
 
+procedure TDriveView.UpdateDriveNotifications(Drive: string);
+var
+  NeedNotifications: Boolean;
+  Path: string;
+  DevBroadcastHandle: DEV_BROADCAST_HANDLE;
+  Size: Integer;
+begin
+  if DriveInfo.IsFixedDrive(Drive) then
+  begin
+    with GetDriveStatus(Drive) do
+    begin
+      NeedNotifications :=
+        WatchThreadActive(Drive) and
+        (DriveInfo.Get(Drive).DriveType <> DRIVE_REMOTE) and
+        DriveInfo.Get(Drive).DriveReady;
+
+      if NeedNotifications <> (DriveHandle <> INVALID_HANDLE_VALUE) then
+      begin
+        if NeedNotifications then
+        begin
+          Path := DriveInfo.GetDriveRoot(Drive);
+          DriveHandle :=
+            CreateFile(PChar(Path), GENERIC_READ, FILE_SHARE_READ or FILE_SHARE_WRITE, nil,
+            OPEN_EXISTING, FILE_FLAG_BACKUP_SEMANTICS or FILE_ATTRIBUTE_NORMAL, 0);
+          if DriveHandle <> INVALID_HANDLE_VALUE then
+          begin
+            Size := SizeOf(DevBroadcastHandle);
+            ZeroMemory(@DevBroadcastHandle, Size);
+            DevBroadcastHandle.dbch_size := Size;
+            DevBroadcastHandle.dbch_devicetype := DBT_DEVTYP_HANDLE;
+            DevBroadcastHandle.dbch_handle := DriveHandle;
+
+            NotificationHandle :=
+              RegisterDeviceNotification(FInternalWindowHandle, @DevBroadcastHandle, DEVICE_NOTIFY_WINDOW_HANDLE);
+            if NotificationHandle = nil then
+            begin
+              CloseHandle(DriveHandle);
+              DriveHandle := INVALID_HANDLE_VALUE;
+            end;
+          end;
+        end
+          else
+        begin
+          UnregisterDeviceNotification(NotificationHandle);
+          NotificationHandle := nil;
+
+          CloseHandle(DriveHandle);
+          DriveHandle := INVALID_HANDLE_VALUE;
+        end;
+      end;
+    end;
+  end;
+end;
+
 procedure TDriveView.StartWatchThread;
 var
   Drive: string;
@@ -2082,14 +2178,22 @@ begin
     if Assigned(DiscMonitor) and not DiscMonitor.Enabled then
       DiscMonitor.Enabled := True;
   end;
+  UpdateDriveNotifications(Drive);
 end; {StartWatchThread}
 
 procedure TDriveView.StopWatchThread;
+var
+  Drive: string;
 begin
   if Assigned(Selected) then
-    with GetDriveStatus(GetDriveToNode(Selected)) do
+  begin
+    Drive := GetDriveToNode(Selected);
+    with GetDriveStatus(Drive) do
       if Assigned(DiscMonitor) then
         DiscMonitor.Enabled := False;
+
+    UpdateDriveNotifications(Drive);
+  end;
 end; {StopWatchThread}
 
 procedure TDriveView.SuspendChangeTimer;
@@ -2111,6 +2215,8 @@ begin
       DiscMonitor.Free;
       DiscMonitor := nil;
     end;
+
+  UpdateDriveNotifications(Drive);
 end; {StopWatchThread}
 
 procedure TDriveView.StartAllWatchThreads;
@@ -2128,7 +2234,10 @@ begin
         if not Assigned(DiscMonitor) then
           CreateWatchThread(DriveStatusPair.Key);
         if Assigned(DiscMonitor) and (not DiscMonitor.Active) then
+        begin
           DiscMonitor.Open;
+          UpdateDriveNotifications(Drive);
+        end;
       end;
 
   if Assigned(Selected) then
@@ -2152,7 +2261,10 @@ begin
     with DriveStatusPair.Value do
     begin
       if Assigned(DiscMonitor) then
+      begin
         DiscMonitor.Close;
+        UpdateDriveNotifications(DriveStatusPair.Key);
+      end;
     end;
 end; {StopAllWatchThreads}