From 21c07b7bf3f63542e1218d28c9fc2d7f6df7588e Mon Sep 17 00:00:00 2001
From: Wojtek Figat <wojtek@figat.pl>
Date: Tue, 11 Mar 2025 16:41:27 +0100
Subject: [PATCH] Fix regression from f256b0670fd91bd519665343084d1a32e7a86acc

#3277
---
 .../Gizmo/ViewportRubberBandSelector.cs       | 17 ++++++++++-
 .../Viewport/MainEditorGizmoViewport.cs       | 29 ++++++++++++-------
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/Source/Editor/Gizmo/ViewportRubberBandSelector.cs b/Source/Editor/Gizmo/ViewportRubberBandSelector.cs
index 8b0fb1e97..0381c6535 100644
--- a/Source/Editor/Gizmo/ViewportRubberBandSelector.cs
+++ b/Source/Editor/Gizmo/ViewportRubberBandSelector.cs
@@ -13,6 +13,7 @@ namespace FlaxEngine.Gizmo;
 /// </summary>
 public class ViewportRubberBandSelector
 {
+    private bool _isMosueCaptured;
     private bool _isRubberBandSpanning;
     private bool _tryStartRubberBand;
     private Float2 _cachedStartingMousePosition;
@@ -51,11 +52,15 @@ public class ViewportRubberBandSelector
     /// <returns>Returns true if rubber band is currently spanning</returns>
     public bool ReleaseRubberBandSelection()
     {
+        if (_isMosueCaptured)
+        {
+            _isMosueCaptured = false;
+            _owner.Viewport.EndMouseCapture();
+        }
         if (_tryStartRubberBand)
         {
             _tryStartRubberBand = false;
         }
-
         if (_isRubberBandSpanning)
         {
             _isRubberBandSpanning = false;
@@ -91,6 +96,11 @@ public class ViewportRubberBandSelector
             _rubberBandRect.Height = mousePosition.Y - _cachedStartingMousePosition.Y;
             if (_lastRubberBandRect != _rubberBandRect)
             {
+                if (!_isMosueCaptured)
+                {
+                    _isMosueCaptured = true;
+                    _owner.Viewport.StartMouseCapture();
+                }
                 UpdateRubberBand(ref viewFrustum);
             }
         }
@@ -254,6 +264,11 @@ public class ViewportRubberBandSelector
     /// <returns>True if rubber band was active before stopping.</returns>
     public bool StopRubberBand()
     {
+        if (_isMosueCaptured)
+        {
+            _isMosueCaptured = false;
+            _owner.Viewport.EndMouseCapture();
+        }
         var result = _tryStartRubberBand;
         _isRubberBandSpanning = false;
         _tryStartRubberBand = false;
diff --git a/Source/Editor/Viewport/MainEditorGizmoViewport.cs b/Source/Editor/Viewport/MainEditorGizmoViewport.cs
index a1dc61b4a..587942844 100644
--- a/Source/Editor/Viewport/MainEditorGizmoViewport.cs
+++ b/Source/Editor/Viewport/MainEditorGizmoViewport.cs
@@ -494,8 +494,7 @@ namespace FlaxEditor.Viewport
         {
             base.OnLostFocus();
 
-            if (_rubberBandSelector.StopRubberBand())
-                EndMouseCapture();
+            _rubberBandSelector.StopRubberBand();
         }
 
         /// <inheritdoc />
@@ -503,8 +502,7 @@ namespace FlaxEditor.Viewport
         {
             base.OnMouseLeave();
 
-            if (_rubberBandSelector.StopRubberBand())
-                EndMouseCapture();
+            _rubberBandSelector.StopRubberBand();
         }
 
         /// <summary>
@@ -620,10 +618,7 @@ namespace FlaxEditor.Viewport
         {
             base.OnLeftMouseButtonDown();
 
-            if (_rubberBandSelector.TryStartingRubberBandSelection())
-            {
-                StartMouseCapture();
-            }
+            _rubberBandSelector.TryStartingRubberBandSelection();
         }
 
         /// <inheritdoc />
@@ -636,8 +631,6 @@ namespace FlaxEditor.Viewport
             // Select rubberbanded rect actor nodes or pick with gizmo
             if (!_rubberBandSelector.ReleaseRubberBandSelection())
             {
-                EndMouseCapture();
-
                 // Try to pick something with the current gizmo
                 Gizmos.Active?.Pick();
             }
@@ -648,6 +641,22 @@ namespace FlaxEditor.Viewport
             base.OnLeftMouseButtonUp();
         }
 
+        /// <inheritdoc />
+        public override bool OnMouseUp(Float2 location, MouseButton button)
+        {
+            if (base.OnMouseUp(location, button))
+                return true;
+
+            // Handle mouse going up when using rubber band with mouse capture that click up outside the view
+            if (button == MouseButton.Left && !new Rectangle(Float2.Zero, Size).Contains(ref location))
+            {
+                _rubberBandSelector.ReleaseRubberBandSelection();
+                return true;
+            }
+
+            return false;
+        }
+
         /// <inheritdoc />
         public override DragDropEffect OnDragEnter(ref Float2 location, DragData data)
         {