From 5b56d07ccc374acbe410cd33ee0f8769ea19262e Mon Sep 17 00:00:00 2001
From: Leo Chen <coel@users.noreply.github.com>
Date: Mon, 15 Jan 2024 11:27:12 +1100
Subject: [PATCH] Use state if existing in session

---
 .../sdk/oauth2-flows/AuthorizationCode.spec.ts      | 13 +++++++++++++
 lib/sdk/oauth2-flows/AuthorizationCode.ts           |  7 ++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/lib/__tests__/sdk/oauth2-flows/AuthorizationCode.spec.ts b/lib/__tests__/sdk/oauth2-flows/AuthorizationCode.spec.ts
index fe152aa..bb228cb 100644
--- a/lib/__tests__/sdk/oauth2-flows/AuthorizationCode.spec.ts
+++ b/lib/__tests__/sdk/oauth2-flows/AuthorizationCode.spec.ts
@@ -96,6 +96,19 @@ describe('AuthorizationCode', () => {
       const state = searchParams.get('state');
       expect(state).toBe(expectedState);
     });
+
+    it('uses same state to generate authorization URL if existing in session', async () => {
+      const client = new AuthorizationCode(clientConfig, clientSecret);
+      const authURL = await client.createAuthorizationURL(sessionManager);
+      const searchParams = new URLSearchParams(authURL.search);
+      const firstState = searchParams.get('state');
+
+      const authURL2 = await client.createAuthorizationURL(sessionManager);
+      const searchParams2 = new URLSearchParams(authURL2.search);
+      const secondState = searchParams2.get('state');
+
+      expect(firstState).toBe(secondState);
+    });
   });
 
   describe('handleRedirectFromAuthDomain()', () => {
diff --git a/lib/sdk/oauth2-flows/AuthorizationCode.ts b/lib/sdk/oauth2-flows/AuthorizationCode.ts
index 5166658..ec920d9 100644
--- a/lib/sdk/oauth2-flows/AuthorizationCode.ts
+++ b/lib/sdk/oauth2-flows/AuthorizationCode.ts
@@ -36,7 +36,12 @@ export class AuthorizationCode extends AuthCodeAbstract {
     sessionManager: SessionManager,
     options: AuthURLOptions = {}
   ): Promise<URL> {
-    this.state = options.state ?? utilities.generateRandomString();
+    this.state =
+      options.state ??
+      ((await sessionManager.getSessionItem(
+        AuthorizationCode.STATE_KEY
+      )) as string) ??
+      utilities.generateRandomString();
     await sessionManager.setSessionItem(AuthorizationCode.STATE_KEY, this.state);
     const authURL = new URL(this.authorizationEndpoint);
     const authParams = this.generateAuthURLParams(options);