From 71eb7502951283584c8af632a56d21eae0217a97 Mon Sep 17 00:00:00 2001
From: Chenfeng Bao <chenfeng.bao@d2l.com>
Date: Wed, 23 Oct 2024 09:01:46 -0700
Subject: [PATCH] add per-page timeout for passwordless auth (#492)

We have an overall timeout of 15 seconds for passwordless auth, but also
we should not get stuck on a single page for anywhere near that long.
---
 src/D2L.Bmx/OktaAuthenticator.cs | 102 ++++++++++++++++++-------------
 1 file changed, 60 insertions(+), 42 deletions(-)

diff --git a/src/D2L.Bmx/OktaAuthenticator.cs b/src/D2L.Bmx/OktaAuthenticator.cs
index 66e95d59..b58bf51b 100644
--- a/src/D2L.Bmx/OktaAuthenticator.cs
+++ b/src/D2L.Bmx/OktaAuthenticator.cs
@@ -114,51 +114,10 @@ private bool TryAuthenticateFromCache(
 		string user,
 		string browserPath
 	) {
-		await using var browser = await browserLauncher.LaunchAsync( browserPath );
-
-		using var cancellationTokenSource = new CancellationTokenSource( TimeSpan.FromSeconds( 15 ) );
-		var sessionIdTcs = new TaskCompletionSource<string?>( TaskCreationOptions.RunContinuationsAsynchronously );
-		cancellationTokenSource.Token.Register( () => sessionIdTcs.TrySetCanceled() );
 		string? sessionId = null;
 
 		try {
-			using var page = await browser.NewPageAsync().WaitAsync( cancellationTokenSource.Token );
-			int attempt = 1;
-
-			page.Load += ( _, _ ) => _ = GetSessionCookieAsync();
-			await page.GoToAsync( orgUrl.AbsoluteUri ).WaitAsync( cancellationTokenSource.Token );
-			sessionId = await sessionIdTcs.Task;
-
-			async Task GetSessionCookieAsync() {
-				var url = new Uri( page.Url );
-				if( url.Host == orgUrl.Host ) {
-					string title = await page.GetTitleAsync().WaitAsync( cancellationTokenSource.Token );
-					// DSSO can sometimes takes more than one attempt.
-					// If the path is '/' with 'sign in' in the title, it means DSSO is not available and we should stop retrying.
-					if( title.Contains( "sign in", StringComparison.OrdinalIgnoreCase ) ) {
-						if( attempt < 3 && url.AbsolutePath != "/" ) {
-							attempt++;
-							await page.GoToAsync( orgUrl.AbsoluteUri ).WaitAsync( cancellationTokenSource.Token );
-						} else {
-							if( BmxEnvironment.IsDebug ) {
-								if( url.AbsolutePath == "/" ) {
-									consoleWriter.WriteWarning(
-										"Okta passwordless authentication is not available."
-									);
-								} else {
-									consoleWriter.WriteWarning( "Okta passwordless authentication failed" );
-								}
-							}
-							sessionIdTcs.SetResult( null );
-						}
-						return;
-					}
-				}
-				var cookies = await page.GetCookiesAsync( orgUrl.AbsoluteUri ).WaitAsync( cancellationTokenSource.Token );
-				if( Array.Find( cookies, c => c.Name == "sid" )?.Value is string sid ) {
-					sessionIdTcs.SetResult( sid );
-				}
-			}
+			sessionId = await GetSessionIdFromBrowserAsync( browserPath, orgUrl );
 		} catch( TaskCanceledException ) {
 			if( BmxEnvironment.IsDebug ) {
 				consoleWriter.WriteWarning( "Okta passwordless authentication timed out." );
@@ -189,6 +148,65 @@ The provided Okta user '{providedLogin}' does not match the system configured pa
 		return oktaAuthenticatedClient;
 	}
 
+	private async Task<string?> GetSessionIdFromBrowserAsync( string browserPath, Uri orgUrl ) {
+		await using var browser = await browserLauncher.LaunchAsync( browserPath );
+
+		var sessionIdTcs = new TaskCompletionSource<string?>( TaskCreationOptions.RunContinuationsAsynchronously );
+
+		// cancel if the total time exceeds 15 seconds, including all page loads and retries
+		using var cancellationTokenSource = new CancellationTokenSource( TimeSpan.FromSeconds( 15 ) );
+		cancellationTokenSource.Token.Register( () => sessionIdTcs.TrySetCanceled() );
+
+		// cancel if we're stuck on a single page for 3 seconds
+		using var pageTimer = new System.Timers.Timer( TimeSpan.FromSeconds( 3 ) ) { AutoReset = false };
+		pageTimer.Elapsed += ( _, _ ) => cancellationTokenSource.Cancel();
+		pageTimer.Start();
+
+		using var page = await browser.NewPageAsync().WaitAsync( cancellationTokenSource.Token );
+		int attempt = 1;
+
+		page.Load += ( _, _ ) => _ = OnPageLoadAsync();
+		await page.GoToAsync( orgUrl.AbsoluteUri ).WaitAsync( cancellationTokenSource.Token );
+		return await sessionIdTcs.Task;
+
+		async Task OnPageLoadAsync() {
+			// reset the 3-sec per-page timer on every page load
+			lock( pageTimer ) {
+				pageTimer.Stop();
+				pageTimer.Start();
+			}
+
+			var url = new Uri( page.Url );
+			if( url.Host == orgUrl.Host ) {
+				string title = await page.GetTitleAsync().WaitAsync( cancellationTokenSource.Token );
+				// DSSO can sometimes takes more than one attempt.
+				// If the path is '/' with 'sign in' in the title, it means DSSO is not available and we should stop retrying.
+				if( title.Contains( "sign in", StringComparison.OrdinalIgnoreCase ) ) {
+					if( attempt < 3 && url.AbsolutePath != "/" ) {
+						attempt++;
+						await page.GoToAsync( orgUrl.AbsoluteUri ).WaitAsync( cancellationTokenSource.Token );
+					} else {
+						if( BmxEnvironment.IsDebug ) {
+							if( url.AbsolutePath == "/" ) {
+								consoleWriter.WriteWarning(
+									"Okta passwordless authentication is not available."
+								);
+							} else {
+								consoleWriter.WriteWarning( "Okta passwordless authentication failed" );
+							}
+						}
+						sessionIdTcs.SetResult( null );
+					}
+					return;
+				}
+			}
+			var cookies = await page.GetCookiesAsync( orgUrl.AbsoluteUri ).WaitAsync( cancellationTokenSource.Token );
+			if( Array.Find( cookies, c => c.Name == "sid" )?.Value is string sid ) {
+				sessionIdTcs.SetResult( sid );
+			}
+		}
+	}
+
 	private async Task<IOktaAuthenticatedClient> GetPasswordAuthenticatedClientAsync( Uri orgUrl, string user ) {
 		string password = consolePrompter.PromptPassword();