From f7ba5ee94b8e8a488789542b48379295ca7d3285 Mon Sep 17 00:00:00 2001
From: Jordan Brown <Cyberboss@users.noreply.github.com>
Date: Sun, 12 Jan 2020 13:46:11 -0500
Subject: [PATCH] Add support for console cancellation during setup wizard

---
 src/Tgstation.Server.Host/Core/SetupWizard.cs | 73 ++++++++++++-------
 src/Tgstation.Server.Host/IO/Console.cs       | 41 ++++++++++-
 src/Tgstation.Server.Host/IO/IConsole.cs      |  5 ++
 3 files changed, 91 insertions(+), 28 deletions(-)

diff --git a/src/Tgstation.Server.Host/Core/SetupWizard.cs b/src/Tgstation.Server.Host/Core/SetupWizard.cs
index bc3a0759b8..dff6446387 100644
--- a/src/Tgstation.Server.Host/Core/SetupWizard.cs
+++ b/src/Tgstation.Server.Host/Core/SetupWizard.cs
@@ -638,41 +638,60 @@ namespace Tgstation.Server.Host.Core
 			}
 
 			var userConfigFileName = String.Format(CultureInfo.InvariantCulture, "appsettings.{0}.json", hostingEnvironment.EnvironmentName);
-			var exists = await ioManager.FileExists(userConfigFileName, cancellationToken).ConfigureAwait(false);
 
-			bool shouldRunBasedOnAutodetect;
-			if (exists)
+			async Task HandleSetupCancel()
 			{
-				var bytes = await ioManager.ReadAllBytes(userConfigFileName, cancellationToken).ConfigureAwait(false);
-				var contents = Encoding.UTF8.GetString(bytes);
-				var existingConfigIsEmpty = String.IsNullOrWhiteSpace(contents) || contents.Trim() == "{}";
-				logger.LogTrace("Configuration json detected. Empty: {0}", existingConfigIsEmpty);
-				shouldRunBasedOnAutodetect = existingConfigIsEmpty;
-			}
-			else
-			{
-				shouldRunBasedOnAutodetect = true;
-				logger.LogTrace("No configuration json detected");
+				await console.WriteAsync(String.Empty, true, default).ConfigureAwait(false);
+				await console.WriteAsync("Aborting setup!", true, default).ConfigureAwait(false);
 			}
 
-			if (!shouldRunBasedOnAutodetect)
-			{
-				if (forceRun)
+			// Link passed cancellationToken with cancel key press
+			Task finalTask = Task.CompletedTask;
+			using (var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, console.CancelKeyPress))
+			using ((cancellationToken = cts.Token).Register(() => finalTask = HandleSetupCancel()))
+				try
 				{
-					logger.LogTrace("Asking user to bypass due to force run request...");
-					await console.WriteAsync(String.Format(CultureInfo.InvariantCulture, "The configuration settings are requesting the setup wizard be run, but you already appear to have a configuration file ({0})!", userConfigFileName), true, cancellationToken).ConfigureAwait(false);
+					var exists = await ioManager.FileExists(userConfigFileName, cancellationToken).ConfigureAwait(false);
 
-					forceRun = await PromptYesNo("Continue running setup wizard? (y/n): ", cancellationToken).ConfigureAwait(false);
+					bool shouldRunBasedOnAutodetect;
+					if (exists)
+					{
+						var bytes = await ioManager.ReadAllBytes(userConfigFileName, cancellationToken).ConfigureAwait(false);
+						var contents = Encoding.UTF8.GetString(bytes);
+						var existingConfigIsEmpty = String.IsNullOrWhiteSpace(contents) || contents.Trim() == "{}";
+						logger.LogTrace("Configuration json detected. Empty: {0}", existingConfigIsEmpty);
+						shouldRunBasedOnAutodetect = existingConfigIsEmpty;
+					}
+					else
+					{
+						shouldRunBasedOnAutodetect = true;
+						logger.LogTrace("No configuration json detected");
+					}
+
+					if (!shouldRunBasedOnAutodetect)
+					{
+						if (forceRun)
+						{
+							logger.LogTrace("Asking user to bypass due to force run request...");
+							await console.WriteAsync(String.Format(CultureInfo.InvariantCulture, "The configuration settings are requesting the setup wizard be run, but you already appear to have a configuration file ({0})!", userConfigFileName), true, cancellationToken).ConfigureAwait(false);
+
+							forceRun = await PromptYesNo("Continue running setup wizard? (y/n): ", cancellationToken).ConfigureAwait(false);
+						}
+
+						if (!forceRun)
+							return false;
+					}
+
+					// flush the logs to prevent console conflicts
+					await asyncDelayer.Delay(TimeSpan.FromSeconds(1), cancellationToken).ConfigureAwait(false);
+
+					await RunWizard(userConfigFileName, cancellationToken).ConfigureAwait(false);
+				}
+				finally
+				{
+					await finalTask.ConfigureAwait(false);
 				}
 
-				if (!forceRun)
-					return false;
-			}
-
-			// flush the logs to prevent console conflicts
-			await asyncDelayer.Delay(TimeSpan.FromSeconds(1), cancellationToken).ConfigureAwait(false);
-
-			await RunWizard(userConfigFileName, cancellationToken).ConfigureAwait(false);
 			return true;
 		}
 	}
diff --git a/src/Tgstation.Server.Host/IO/Console.cs b/src/Tgstation.Server.Host/IO/Console.cs
index 9f1fceba80..0d90d54c25 100644
--- a/src/Tgstation.Server.Host/IO/Console.cs
+++ b/src/Tgstation.Server.Host/IO/Console.cs
@@ -6,11 +6,50 @@ using System.Threading.Tasks;
 namespace Tgstation.Server.Host.IO
 {
 	/// <inheritdoc />
-	sealed class Console : IConsole
+	sealed class Console : IConsole, IDisposable
 	{
 		/// <inheritdoc />
 		public bool Available => Environment.UserInteractive;
 
+		/// <inheritdoc />
+		public CancellationToken CancelKeyPress => cancelKeyCts.Token;
+
+		/// <summary>
+		/// The <see cref="CancellationTokenSource"/> for <see cref="CancelKeyPress"/>.
+		/// </summary>
+		readonly CancellationTokenSource cancelKeyCts;
+
+		/// <summary>
+		/// If the <see cref="Console"/> was disposed;
+		/// </summary>
+		bool disposed;
+
+		/// <summary>
+		/// Initializes a new instance of the <see cref="Console"/> <see langword="class"/>.
+		/// </summary>
+		public Console()
+		{
+			cancelKeyCts = new CancellationTokenSource();
+			System.Console.CancelKeyPress += (sender, e) =>
+			{
+				lock (cancelKeyCts)
+				{
+					if (!disposed)
+						cancelKeyCts.Cancel();
+				}
+			};
+		}
+
+		/// <inheritdoc />
+		public void Dispose()
+		{
+			lock (cancelKeyCts)
+			{
+				cancelKeyCts.Dispose();
+				disposed = true;
+			}
+		}
+
 		void CheckAvailable()
 		{
 			if (!Available)
diff --git a/src/Tgstation.Server.Host/IO/IConsole.cs b/src/Tgstation.Server.Host/IO/IConsole.cs
index 119b19d5dc..405ffb425b 100644
--- a/src/Tgstation.Server.Host/IO/IConsole.cs
+++ b/src/Tgstation.Server.Host/IO/IConsole.cs
@@ -13,6 +13,11 @@ namespace Tgstation.Server.Host.IO
 		/// </summary>
 		bool Available { get; }
 
+		/// <summary>
+		/// Gets a <see cref="CancellationToken"/> that triggers if Crtl+C or an equivalent is pressed.
+		/// </summary>
+		CancellationToken CancelKeyPress { get; }
+
 		/// <summary>
 		/// Write some <paramref name="text"/> to the <see cref="IConsole"/>
 		/// </summary>