feat(cognito): ✨ Control JWT timestamps via props (#21)

ShellXploit · Apr 20, 2024 · 08254f3 · 08254f3
1 parent abe12aa
commit 08254f3
Show file tree

Hide file tree

Showing 5 changed files with 101 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
@@ -6,19 +6,28 @@ 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,
 		iss: issuer,
 		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
@@ -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);
+	});
 });