From c87aa30a1e79f90882310a1c33560deaf878df81 Mon Sep 17 00:00:00 2001
From: Vinayak Mishra <viks@vnykmshr.com>
Date: Fri, 3 Jul 2026 15:57:12 +0545
Subject: [PATCH] validateAndResolveCtx: port ctx validation from go/python
 SDKs

---
 src/error/codes/index.ts           |  2 ++
 src/error/messages/index.ts        |  2 ++
 src/service-account/index.ts       | 28 +++++++++++++++++++++++++---
 test/service-account/token.test.js | 29 +++++++++++++++++++++++++++++
 4 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/src/error/codes/index.ts b/src/error/codes/index.ts
index 672bd33a..64db75d2 100644
--- a/src/error/codes/index.ts
+++ b/src/error/codes/index.ts
@@ -64,6 +64,8 @@ const SKYFLOW_ERROR_CODE = {
 
     INVALID_ROLES_KEY_TYPE: { httpCode: 400, message: errorMessages.INVALID_ROLES_KEY_TYPE },
     INVALID_CONTEXT: { httpCode: 400, message: errorMessages.INVALID_CONTEXT },
+    INVALID_CTX_TYPE: { httpCode: 400, message: errorMessages.INVALID_CTX_TYPE },
+    INVALID_CTX_MAP_KEY: { httpCode: 400, message: errorMessages.INVALID_CTX_MAP_KEY },
     EMPTY_ROLES: { httpCode: 400, message: errorMessages.EMPTY_ROLES },
 
     INVALID_JSON_FORMAT: { httpCode: 400, message: errorMessages.INVALID_JSON_FORMAT },
diff --git a/src/error/messages/index.ts b/src/error/messages/index.ts
index 2b634c84..6f163d5a 100644
--- a/src/error/messages/index.ts
+++ b/src/error/messages/index.ts
@@ -64,6 +64,8 @@ const errorMessages = {
 
     INVALID_ROLES_KEY_TYPE: `${errorPrefix} Validation error. Invalid roles. Specify roles as an array.`,
     INVALID_CONTEXT: `${errorPrefix} Validation error. Invalid context. Specify context as a string.`,
+    INVALID_CTX_TYPE: `${errorPrefix} Validation error. Invalid context type. Specify context as a string, boolean, number, or JSON object.`,
+    INVALID_CTX_MAP_KEY: `${errorPrefix} Validation error. Invalid context map key '%s1'. Map keys must match pattern /^[a-zA-Z0-9_]+$/.`,
     EMPTY_ROLES: `${errorPrefix} Validation error. Invalid roles. Specify at least one role.`,
 
     INVALID_JSON_FORMAT: `${errorPrefix} Validation error. Credentials is not in valid JSON format. Verify the credentials.`,
diff --git a/src/service-account/index.ts b/src/service-account/index.ts
index a68903e0..1ae4c8cb 100644
--- a/src/service-account/index.ts
+++ b/src/service-account/index.ts
@@ -9,6 +9,28 @@ import SKYFLOW_ERROR_CODE from '../error/codes';
 import { ServiceAccountResponseError } from '../vault/types';
 import { WithRawResponse } from '../ _generated_/rest/core';
 
+const CTX_KEY_REGEX = /^[a-zA-Z0-9_]+$/;
+
+function validateAndResolveCtx(ctx?: string | Record<string, any> | boolean | number): any {
+    if (ctx === undefined || ctx === null) return undefined;
+    if (typeof ctx === 'string') return ctx === '' ? undefined : ctx;
+    if (typeof ctx === 'boolean' || typeof ctx === 'number') return ctx;
+    if (typeof ctx === 'object') {
+        if (Array.isArray(ctx)) {
+            throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_CTX_TYPE);
+        }
+        const keys = Object.keys(ctx);
+        if (keys.length === 0) return undefined;
+        for (const key of keys) {
+            if (!CTX_KEY_REGEX.test(key)) {
+                throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_CTX_MAP_KEY, [key]);
+            }
+        }
+        return ctx;
+    }
+    throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_CTX_TYPE);
+}
+
 function normalizeTokenOptions(options?: BearerTokenOptions): BearerTokenOptions | undefined {
     if (!options) return options;
     if (options.roleIDs !== undefined && options.roleIds === undefined) {
@@ -132,7 +154,7 @@ function getToken(credentials, options?: BearerTokenOptions): Promise<TokenRespo
                 aud: credentialsObj.tokenUri,
                 exp: expiryTime,
                 sub: credentialsObj.clientId,
-                ...(options && options.ctx ? { ctx: options.ctx } : {}),
+                ctx: validateAndResolveCtx(options?.ctx),
             };
             if (claims.iss == null) {
                 printLog(logs.errorLogs.CLIENT_ID_NOT_FOUND, MessageType.ERROR, options?.logLevel);
@@ -286,7 +308,7 @@ function getSignedTokens(credentials, options: SignedDataTokensOptions): Promise
                     exp: expiryTime,
                     sub: credentialsObj.clientId,
                     tok: token,
-                    ...(options?.ctx ? { ctx: options.ctx } : {}),
+                    ctx: validateAndResolveCtx(options?.ctx),
                 };
 
                 if (claims.key == null) {
@@ -395,4 +417,4 @@ function normalizeCredentials(obj: any): any {
     };
 }
 
-export { generateBearerToken, generateBearerTokenFromCreds, generateSignedDataTokens, generateSignedDataTokensFromCreds, getToken, successResponse, failureResponse };
\ No newline at end of file
+export { generateBearerToken, generateBearerTokenFromCreds, generateSignedDataTokens, generateSignedDataTokensFromCreds, getToken, successResponse, failureResponse, validateAndResolveCtx };
\ No newline at end of file
diff --git a/test/service-account/token.test.js b/test/service-account/token.test.js
index 9d73568f..d9963a1f 100644
--- a/test/service-account/token.test.js
+++ b/test/service-account/token.test.js
@@ -1146,3 +1146,32 @@ describe('service-account branch and line coverage', () => {
         ).rejects.toBeDefined();
     });
 });
+
+describe('ctx validation', () => {
+    const creds = JSON.stringify(validCredentials);
+
+    test('accepts string ctx', async () => {
+        await expect(getToken(creds, { ctx: 'test-ctx' })).resolves.toBeDefined();
+    });
+    test('accepts boolean ctx', async () => {
+        await expect(getToken(creds, { ctx: true })).resolves.toBeDefined();
+    });
+    test('accepts number ctx', async () => {
+        await expect(getToken(creds, { ctx: 42 })).resolves.toBeDefined();
+    });
+    test('accepts valid object ctx', async () => {
+        await expect(getToken(creds, { ctx: { key1: 'val' } })).resolves.toBeDefined();
+    });
+    test('omits empty string ctx from claims', async () => {
+        await expect(getToken(creds, { ctx: '' })).resolves.toBeDefined();
+    });
+    test('rejects array ctx', async () => {
+        await expect(getToken(creds, { ctx: ['bad'] })).rejects.toThrow(SKYFLOW_ERROR_CODE.INVALID_CTX_TYPE);
+    });
+    test('rejects ctx with invalid map key', async () => {
+        await expect(getToken(creds, { ctx: { 'bad key': 'val' } })).rejects.toThrow(SKYFLOW_ERROR_CODE.INVALID_CTX_MAP_KEY);
+    });
+    test('rejects function ctx', async () => {
+        await expect(getToken(creds, { ctx: () => {} })).rejects.toThrow(SKYFLOW_ERROR_CODE.INVALID_CTX_TYPE);
+    });
+});