From ab9290772b46ac2087eeb9f9e5fdfaf848fa03a5 Mon Sep 17 00:00:00 2001
From: smoogipoo <smoogipoo@smgi.me>
Date: Tue, 8 Jun 2021 17:54:54 +0900
Subject: [PATCH] Fix a similar case with online play sub-screens

---
 .../Navigation/TestSceneScreenNavigation.cs   | 28 +++++++++++++++++++
 .../Screens/OnlinePlay/OnlinePlayScreen.cs    |  5 +++-
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/osu.Game.Tests/Visual/Navigation/TestSceneScreenNavigation.cs b/osu.Game.Tests/Visual/Navigation/TestSceneScreenNavigation.cs
index 6fc19b95b9..52401d32e5 100644
--- a/osu.Game.Tests/Visual/Navigation/TestSceneScreenNavigation.cs
+++ b/osu.Game.Tests/Visual/Navigation/TestSceneScreenNavigation.cs
@@ -9,16 +9,19 @@ using osu.Framework.Screens;
 using osu.Framework.Testing;
 using osu.Game.Beatmaps;
 using osu.Game.Graphics.UserInterface;
+using osu.Game.Online.Multiplayer;
 using osu.Game.Overlays;
 using osu.Game.Overlays.Mods;
 using osu.Game.Overlays.Toolbar;
 using osu.Game.Rulesets.Mods;
 using osu.Game.Rulesets.Osu.Mods;
+using osu.Game.Screens.OnlinePlay.Components;
 using osu.Game.Screens.Play;
 using osu.Game.Screens.Ranking;
 using osu.Game.Screens.Select;
 using osu.Game.Screens.Select.Options;
 using osu.Game.Tests.Beatmaps.IO;
+using osu.Game.Tests.Visual.Multiplayer;
 using osuTK;
 using osuTK.Input;
 
@@ -306,6 +309,18 @@ namespace osu.Game.Tests.Visual.Navigation
             AddAssert("Toolbar is hidden", () => Game.Toolbar.State.Value == Visibility.Hidden);
         }
 
+        [Test]
+        public void TestPushMatchSubScreenAndPressBackButtonImmediately()
+        {
+            TestMultiplayer multiplayer = null;
+
+            PushAndConfirm(() => multiplayer = new TestMultiplayer());
+
+            AddStep("open room", () => multiplayer.OpenNewRoom());
+            AddStep("press back button", () => Game.ChildrenOfType<BackButton>().First().Action());
+            AddWaitStep("wait two frames", 2);
+        }
+
         private void pushEscape() =>
             AddStep("Press escape", () => InputManager.Key(Key.Escape));
 
@@ -330,5 +345,18 @@ namespace osu.Game.Tests.Visual.Navigation
 
             protected override bool DisplayStableImportPrompt => false;
         }
+
+        private class TestMultiplayer : Screens.OnlinePlay.Multiplayer.Multiplayer
+        {
+            [Cached(typeof(MultiplayerClient))]
+            public readonly TestMultiplayerClient Client;
+
+            public TestMultiplayer()
+            {
+                Client = new TestMultiplayerClient((TestMultiplayerRoomManager)RoomManager);
+            }
+
+            protected override RoomManager CreateRoomManager() => new TestMultiplayerRoomManager();
+        }
     }
 }
diff --git a/osu.Game/Screens/OnlinePlay/OnlinePlayScreen.cs b/osu.Game/Screens/OnlinePlay/OnlinePlayScreen.cs
index 90e499c67f..e418d36d40 100644
--- a/osu.Game/Screens/OnlinePlay/OnlinePlayScreen.cs
+++ b/osu.Game/Screens/OnlinePlay/OnlinePlayScreen.cs
@@ -253,7 +253,10 @@ namespace osu.Game.Screens.OnlinePlay
 
         public override bool OnBackButton()
         {
-            if ((screenStack.CurrentScreen as IOnlinePlaySubScreen)?.OnBackButton() == true)
+            if (!(screenStack.CurrentScreen is IOnlinePlaySubScreen onlineSubScreen))
+                return false;
+
+            if (((Drawable)onlineSubScreen).IsLoaded && onlineSubScreen.AllowBackButton && onlineSubScreen.OnBackButton())
                 return true;
 
             if (screenStack.CurrentScreen != null && !(screenStack.CurrentScreen is LoungeSubScreen))