From 08254f398eb5904e0291c31a0b10bfa06474f91e Mon Sep 17 00:00:00 2001
From: Hanut Arora <36478891+ShellXploit@users.noreply.github.com>
Date: Sat, 20 Apr 2024 14:15:31 +0530
Subject: [PATCH] feat(cognito): :sparkles: Control JWT timestamps via props
 (#21)

---
 src/awsServices/cognito/cognito.ts            |  7 +-
 src/types/awsServices/cognito/cognito.ts      |  6 ++
 src/types/utils/getBaseJwtPayload.ts          |  2 +
 src/utils/getBaseJwtPayload.ts                | 19 +++--
 .../unit/awsServices/cognito/cognito.test.ts  | 74 +++++++++++++++++++
 5 files changed, 101 insertions(+), 7 deletions(-)

diff --git a/src/awsServices/cognito/cognito.ts b/src/awsServices/cognito/cognito.ts
index c21a982..554765a 100644
--- a/src/awsServices/cognito/cognito.ts
+++ b/src/awsServices/cognito/cognito.ts
@@ -14,7 +14,8 @@ export const getCognitoTokens = ({
 	asymmetricKeys,
 	user,
 	cognito = {},
-	userPool = {}
+	userPool = {},
+	jwtConfig = {}
 }: GetCognitoTokensProps): CognitoTokens => {
 	const userData = {
 		uuid: getUUID(),
@@ -33,7 +34,9 @@ export const getCognitoTokens = ({
 	const issuerUrl = getIssuerUrl(userPool);
 	const baseJwtPayload = getBaseJwtPayload({
 		issuer: issuerUrl,
-		subject: userData.uuid
+		subject: userData.uuid,
+		authTimeInEpoch: jwtConfig.authTimeInEpoch,
+		minutesToExpiry: jwtConfig.minutesToExpiry
 	});
 
 	const accessToken = getAccessToken(asymmetricKeys, {
diff --git a/src/types/awsServices/cognito/cognito.ts b/src/types/awsServices/cognito/cognito.ts
index a117fd5..de38f7b 100644
--- a/src/types/awsServices/cognito/cognito.ts
+++ b/src/types/awsServices/cognito/cognito.ts
@@ -31,11 +31,17 @@ interface UserPool {
 	region?: string;
 }
 
+interface JwtConfig {
+	authTimeInEpoch?: number;
+	minutesToExpiry?: number;
+}
+
 interface GetCognitoTokensProps {
 	asymmetricKeys: AsymmetricKeys;
 	user: User;
 	cognito?: Cognito;
 	userPool?: UserPool;
+	jwtConfig?: JwtConfig;
 }
 
 interface CognitoTokens {
diff --git a/src/types/utils/getBaseJwtPayload.ts b/src/types/utils/getBaseJwtPayload.ts
index 45e818c..049097d 100644
--- a/src/types/utils/getBaseJwtPayload.ts
+++ b/src/types/utils/getBaseJwtPayload.ts
@@ -1,6 +1,8 @@
 interface GetBaseJwtPayloadProps {
 	issuer: string;
 	subject: string;
+	authTimeInEpoch?: number;
+	minutesToExpiry?: number;
 }
 
 interface BaseJwtPayload {
diff --git a/src/utils/getBaseJwtPayload.ts b/src/utils/getBaseJwtPayload.ts
index 8222a77..a758fd6 100644
--- a/src/utils/getBaseJwtPayload.ts
+++ b/src/utils/getBaseJwtPayload.ts
@@ -6,9 +6,18 @@ import {
 
 const getBaseJwtPayload = ({
 	issuer,
-	subject
+	subject,
+	authTimeInEpoch,
+	minutesToExpiry = 60
 }: GetBaseJwtPayloadProps): BaseJwtPayload => {
-	const currentEpochTime = Math.floor(Date.now() / 1000);
+	const isEpochValid =
+		authTimeInEpoch && new Date(authTimeInEpoch * 1000).getTime() > 0;
+
+	const authTime = isEpochValid
+		? authTimeInEpoch
+		: Math.floor(Date.now() / 1000);
+
+	const secondsToExpiry = minutesToExpiry * 60;
 
 	return {
 		sub: subject,
@@ -16,9 +25,9 @@ const getBaseJwtPayload = ({
 		origin_jti: getUUID(),
 		event_id: getUUID(),
 		jti: getUUID(),
-		auth_time: currentEpochTime,
-		exp: currentEpochTime + 60 * 60,
-		iat: currentEpochTime
+		auth_time: authTime,
+		exp: authTime + secondsToExpiry,
+		iat: authTime
 	};
 };
 
diff --git a/tests/unit/awsServices/cognito/cognito.test.ts b/tests/unit/awsServices/cognito/cognito.test.ts
index e87fddb..8ae9f53 100644
--- a/tests/unit/awsServices/cognito/cognito.test.ts
+++ b/tests/unit/awsServices/cognito/cognito.test.ts
@@ -90,4 +90,78 @@ describe("Cognito", () => {
 			expect(error).not.toBeNull();
 		}
 	});
+
+	it("should return a valid tokens with custom jwt config - expiry", () => {
+		const asymmetricKeys = utils.getAsymmetricKeys();
+
+		const minutesToExpiry = 1;
+
+		const { id_token: idToken, access_token: accessToken } =
+			awsServices.cognito.getCognitoTokens({
+				asymmetricKeys,
+				user: {
+					emailId
+				},
+				jwtConfig: {
+					minutesToExpiry
+				}
+			});
+
+		let decodedToken: string | jwt.JwtPayload | null = null;
+		let authTime = 0;
+		let expiry = 0;
+		try {
+			decodedToken = jwt.decode(idToken) as jwt.JwtPayload;
+		} catch (error) {
+			expect(error).toBeNull();
+		}
+		expect(decodedToken).not.toBeNull();
+		authTime = decodedToken?.iat || 0;
+		expiry = decodedToken?.exp || 0;
+		expect(expiry - authTime).toBe(minutesToExpiry * 60);
+
+		try {
+			decodedToken = jwt.decode(accessToken) as jwt.JwtPayload;
+		} catch (error) {
+			expect(error).toBeNull();
+		}
+		expect(decodedToken).not.toBeNull();
+		authTime = decodedToken?.iat || 0;
+		expiry = decodedToken?.exp || 0;
+		expect(expiry - authTime).toBe(minutesToExpiry * 60);
+	});
+
+	it("should return a valid tokens with custom jwt config - auth time", () => {
+		const asymmetricKeys = utils.getAsymmetricKeys();
+
+		const authTimeInEpoch = 1000;
+
+		const { id_token: idToken, access_token: accessToken } =
+			awsServices.cognito.getCognitoTokens({
+				asymmetricKeys,
+				user: {
+					emailId
+				},
+				jwtConfig: {
+					authTimeInEpoch
+				}
+			});
+
+		let decodedToken: string | jwt.JwtPayload | null = null;
+		try {
+			decodedToken = jwt.decode(idToken) as jwt.JwtPayload;
+		} catch (error) {
+			expect(error).toBeNull();
+		}
+		expect(decodedToken).not.toBeNull();
+		expect(decodedToken?.iat).toBe(authTimeInEpoch);
+
+		try {
+			decodedToken = jwt.decode(accessToken) as jwt.JwtPayload;
+		} catch (error) {
+			expect(error).toBeNull();
+		}
+		expect(decodedToken).not.toBeNull();
+		expect(decodedToken?.iat).toBe(authTimeInEpoch);
+	});
 });