SK-2503: Add support for custom tokenUri

src/error/codes/index.ts

@@ -17,6 +17,7 @@ const SKYFLOW_ERROR_CODE = {
     INVALID_PARSED_CREDENTIALS_STRING: { http_code: 400, message: errorMessages.INVALID_PARSED_CREDENTIALS_STRING },
     INVALID_KEY: { http_code: 400, message: errorMessages.INVALID_KEY },
     INVALID_CREDENTIALS_FILE_PATH: { http_code: 400, message: errorMessages.INVALID_CREDENTIALS_FILE_PATH },
+    INVALID_TOKEN_URI: { http_code: 400, message: errorMessages.INVALID_TOKEN_URI },
 
     INVALID_BEARER_TOKEN_WITH_ID: { http_code: 400, message: errorMessages.INVALID_BEARER_TOKEN_WITH_ID },
     INVALID_PARSED_CREDENTIALS_STRING_WITH_ID: { http_code: 400, message: errorMessages.INVALID_PARSED_CREDENTIALS_STRING_WITH_ID },

src/error/messages/index.ts

@@ -22,6 +22,7 @@ const errorMessages = {
     INVALID_CREDENTIAL_FILE_PATH: `${errorPrefix} Initialization failed. Invalid credentials. Expected file path to be a string.`,
 
     INVALID_CREDENTIALS_FILE_PATH: `${errorPrefix} Initialization failed. Invalid skyflow credentials. Expected file path to exists.`,
+    INVALID_TOKEN_URI: `${errorPrefix} Initialization failed. Invalid Skyflow credentials. The token URI must be a string and a valid URL.`,
     INVALID_KEY: `${errorPrefix} Initialization failed. Invalid skyflow credentials. Specify a valid api key.`,
     INVALID_PARSED_CREDENTIALS_STRING: `${errorPrefix} Initialization failed. Invalid skyflow credentials. Specify a valid credentials string.`,
     INVALID_BEARER_TOKEN: `${errorPrefix} Initialization failed. Invalid skyflow credentials. Specify a valid token.`,

src/service-account/index.ts

@@ -1,7 +1,7 @@
 import * as fs from 'fs';
 import jwt from "jsonwebtoken";
 import { V1GetAuthTokenRequest, V1GetAuthTokenResponse } from '../ _generated_/rest/api';
-import { getBaseUrl, LogLevel, MessageType, parameterizedString, printLog } from '../utils';
+import { getBaseUrl, isValidURL, LogLevel, MessageType, parameterizedString, printLog } from '../utils';
 import Client from './client';
 import logs from '../utils/logs';
 import SkyflowError from '../error';
@@ -13,6 +13,7 @@ export type BearerTokenOptions = {
     ctx?: string | Record<string, any>,
     roleIDs?: string[],
     logLevel?: LogLevel,
+    tokenUri?: string,
 }
 
 export type GenerateTokenOptions = {
@@ -29,6 +30,7 @@ export type SignedDataTokensOptions = {
     timeToLive?: number,
     ctx?: string | Record<string, any>,
     logLevel?: LogLevel,
+    tokenUri?: string
 }
 
 export type TokenResponse = {
@@ -98,6 +100,17 @@ function getToken(credentials, options?: BearerTokenOptions): Promise<TokenRespo
                 printLog(logs.errorLogs.NOT_A_VALID_JSON, MessageType.ERROR, options?.logLevel);
                 reject(new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_JSON_FORMAT));
             }
+
+            if (Object.prototype.hasOwnProperty.call(options, 'tokenUri')) {
+                if (options?.tokenUri === undefined || typeof options.tokenUri !== 'string' || !isValidURL(options.tokenUri)) {
+                    throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI);
+                }
+            }
+
+            if (options?.tokenUri) {
+                credentialsObj.tokenURI = options.tokenUri;
+            }
+
             const expiryTime = Math.floor(Date.now() / 1000) + 3600;
             const claims = {
                 iss: credentialsObj.clientID,
@@ -231,6 +244,16 @@ function getSignedTokens(credentials, options: SignedDataTokensOptions): Promise
                 reject(new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_JSON_FORMAT));
             }
 
+            if (Object.prototype.hasOwnProperty.call(options, 'tokenUri')) {  
+                if (options?.tokenUri === undefined || typeof options.tokenUri !== 'string' || !isValidURL(options.tokenUri)) {  
+                    throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI);
+                }
+            }
+
+            if (options?.tokenUri) {
+                credentialsObj.tokenURI = options.tokenUri;
+            }
+
             let expiryTime;
             if (options?.timeToLive && options?.timeToLive !== null) {
                 expiryTime = Math.floor(Date.now() / 1000) + options?.timeToLive;

src/utils/index.ts

@@ -288,6 +288,7 @@ export async function getToken(credentials: Credentials, logLevel?: LogLevel): P
             roleIDs: stringCred.roles,
             ctx: stringCred.context,
             logLevel,
+            tokenUri: stringCred.tokenUri,
         });
     }
 
@@ -298,6 +299,7 @@ export async function getToken(credentials: Credentials, logLevel?: LogLevel): P
             roleIDs: pathCred.roles,
             ctx: pathCred.context,
             logLevel,
+            tokenUri: pathCred.tokenUri,
         });
     }
 

src/utils/validations/index.ts

@@ -172,6 +172,11 @@ export const validateCredentialsWithId = (credentials: Credentials, type: string
         if (pathCred.context !== undefined && (typeof pathCred.context !== 'string' && typeof pathCred.context !== 'object')) {
             throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_CONTEXT, [type, typeId, id]);
         }
+        if(Object.prototype.hasOwnProperty.call(pathCred, 'tokenUri')) {
+            if (pathCred.tokenUri === undefined || typeof pathCred.tokenUri !== 'string' || !isValidURL(pathCred.tokenUri)) {
+                throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI, [type, typeId, id]);
+            }
+        }
     }
 
     // Validate StringCredentials
@@ -187,6 +192,11 @@ export const validateCredentialsWithId = (credentials: Credentials, type: string
         if (stringCred.context !== undefined && (typeof stringCred.context !== 'string' && typeof stringCred.context !== 'object')) {
             throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_CONTEXT, [type, typeId, id]);
         }
+        if (Object.prototype.hasOwnProperty.call(stringCred, 'tokenUri')) {
+            if (stringCred.tokenUri === undefined || typeof stringCred.tokenUri !== 'string' || !isValidURL(stringCred.tokenUri)) {
+                throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI, [type, typeId, id]);
+            }
+        }
     }
 
     // Validate ApiKeyCredentials
@@ -298,6 +308,12 @@ export const validateSkyflowCredentials = (credentials: Credentials, logLevel: L
         if (pathCred.context !== undefined && (typeof pathCred.context !== 'string' && typeof pathCred.context !== 'object')) {
             throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_CONTEXT);
         }
+
+        if(Object.prototype.hasOwnProperty.call(pathCred, 'tokenUri')) {
+            if (pathCred.tokenUri === undefined || typeof pathCred.tokenUri !== 'string' || !isValidURL(pathCred.tokenUri)) {
+                throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI);
+            }
+        }
     }
 
     // Validate StringCredentials
@@ -314,6 +330,11 @@ export const validateSkyflowCredentials = (credentials: Credentials, logLevel: L
         if (stringCred.context !== undefined && (typeof stringCred.context !== 'string' && typeof stringCred.context !== 'object')) {
             throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_CONTEXT);
         }
+        if (Object.prototype.hasOwnProperty.call(stringCred, 'tokenUri')) {
+            if (stringCred.tokenUri === undefined || typeof stringCred.tokenUri !== 'string' || !isValidURL(stringCred.tokenUri)) {
+                throw new SkyflowError(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI);
+            }
+        }
     }
 
     // Validate ApiKeyCredentials

src/vault/config/credentials/index.ts

@@ -12,12 +12,14 @@ export interface PathCredentials {
   path: string;
   roles?: Array<string>;
   context?: string | Record<string, any>;
+  tokenUri?: string;
 }
 
 export interface StringCredentials {
   credentialsString: string;
   roles?: Array<string>;
   context?: string | Record<string, any>
+  tokenUri?: string;
 }
 
 export interface ApiKeyCredentials {

test/service-account/token.test.js

@@ -336,6 +336,13 @@ describe('Signed Data Token Generation Test', () => {
 
 describe('getToken Tests', () => {
     let mockClient;
+    const validCredentials = {
+        clientID: "test-client-id",
+        keyID: "test-key-id",
+        tokenURI: "https://test-token-uri.com",
+        privateKey: "KEY",
+        data: "DATA",
+    };
     const credentials = {
         clientID: "test-client-id",
         keyID: "test-key-id",
@@ -413,4 +420,74 @@ describe('getToken Tests', () => {
             expect(err).toBeDefined();
         }
     });
+
+    test("should use tokenUri from options if provided and valid", async () => {
+        const validCredsString = JSON.stringify(validCredentials);
+        const validTokenOptions = { tokenUri: "https://override-token-uri.com" };
+        const getBaseUrlSpy = jest.spyOn(require('../../src/utils'), 'getBaseUrl');
+        await getToken(validCredsString, validTokenOptions);
+        expect(getBaseUrlSpy).toHaveBeenCalledWith(validTokenOptions.tokenUri);
+    });
+
+    test("should throw error if tokenUri in options is invalid", async () => {
+        const validCredsString = JSON.stringify(validCredentials);
+        const invalidOptions = { tokenUri: "not-a-valid-url" };
+        await expect(getToken(validCredsString, invalidOptions)).rejects.toThrow();
+    });
+});
+
+
+describe('getToken and getSignedTokens tokenUri override tests', () => {
+    const validCreds = {
+        clientID: "test-client-id",
+        keyID: "test-key-id",
+        tokenURI: "https://original-token-uri.com",
+        privateKey: "KEY",
+        data: "DATA",
+    };
+
+    const validCredsString = JSON.stringify(validCreds);
+
+    const validSignedTokenOptions = {
+        dataTokens: ['datatoken1'],
+        tokenUri: "https://override-token-uri.com"
+    };
+
+    const validTokenOptions = {
+        tokenUri: "https://override-token-uri.com"
+    };
+
+    beforeEach(() => {
+        jest.spyOn(jwt, 'sign').mockReturnValue('mocked_token');
+    });
+
+    afterEach(() => {
+        jest.restoreAllMocks();
+    });
+
+    test('getToken uses tokenUri from options if provided', async () => {
+        const getBaseUrlSpy = jest.spyOn(require('../../src/utils'), 'getBaseUrl');
+        await getToken(validCredsString, validTokenOptions);
+        expect(getBaseUrlSpy).toHaveBeenCalledWith(validTokenOptions.tokenUri);
+    });
+
+    test('generateSignedDataTokensFromCreds uses tokenUri from options if provided', async () => {
+        let capturedClaims = null;
+        jest.spyOn(jwt, 'sign').mockImplementation((claims, key, opts) => {
+            capturedClaims = claims;
+            return 'mocked_token';
+        });
+        await generateSignedDataTokensFromCreds(validCredsString, validSignedTokenOptions);
+        expect(capturedClaims.aud).toBe(validSignedTokenOptions.tokenUri);
+    });
+
+    test('getToken throws error if tokenUri in options is invalid', async () => {
+        const invalidOptions = { tokenUri: "not-a-valid-url" };
+        await expect(getToken(validCredsString, invalidOptions)).rejects.toThrow();
+    });
+
+    test('generateSignedDataTokensFromCreds throws error if tokenUri in options is invalid', async () => {
+        const invalidOptions = { dataTokens: ['datatoken1'], tokenUri: "not-a-valid-url" };
+        await expect(generateSignedDataTokensFromCreds(validCredsString, invalidOptions)).rejects.toThrow();
+    });
 });

test/utils/validations.test.js

@@ -4046,4 +4046,140 @@ describe('validateCredentialsWithId', () => {
     expect(() => validateCredentialsWithId(null, type, typeId, id))
       .toThrow(SKYFLOW_ERROR_CODE.INVALID_CREDENTIALS_WITH_ID);
   });
-});
+});
+
+describe('validateCredentialsWithId and validateSkyflowCredentials - tokenUri validation', () => {
+  const type = 'vault';
+  const typeId = 'vault_id';
+  const id = 'test-id';
+
+  const validUrl = 'https://valid.url/token';
+
+  test('validateCredentialsWithId: should throw error if tokenUri is present but not a string (PathCredentials)', () => {
+    const credentials = {
+      path: '/valid/path',
+      tokenUri: 123
+    };
+    jest.spyOn(fs, 'existsSync').mockReturnValue(true);
+    expect(() => validateCredentialsWithId(credentials, type, typeId, id))
+      .toThrow(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI);
+  });
+
+  test('validateCredentialsWithId: should throw error if tokenUri is present but not a valid URL (PathCredentials)', () => {
+    const credentials = {
+      path: '/valid/path',
+      tokenUri: 'not-a-url'
+    };
+    jest.spyOn(fs, 'existsSync').mockReturnValue(true);
+    expect(() => validateCredentialsWithId(credentials, type, typeId, id))
+      .toThrow(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI);
+  });
+
+  test('validateCredentialsWithId: should accept valid tokenUri (PathCredentials)', () => {
+    const credentials = {
+      path: '/valid/path',
+      tokenUri: validUrl
+    };
+    jest.spyOn(fs, 'existsSync').mockReturnValue(true);
+    expect(() => validateCredentialsWithId(credentials, type, typeId, id)).not.toThrow();
+  });
+
+  test('validateCredentialsWithId: should throw error if tokenUri is present but not a string (StringCredentials)', () => {
+    const credentials = {
+      credentialsString: JSON.stringify({ clientID: 'c', keyID: 'k' }),
+      tokenUri: 123
+    };
+    expect(() => validateCredentialsWithId(credentials, type, typeId, id))
+      .toThrow(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI);
+  });
+
+  test('validateCredentialsWithId: should throw error if tokenUri is present but not a valid URL (StringCredentials)', () => {
+    const credentials = {
+      credentialsString: JSON.stringify({ clientID: 'c', keyID: 'k' }),
+      tokenUri: 'not-a-url'
+    };
+    expect(() => validateCredentialsWithId(credentials, type, typeId, id))
+      .toThrow(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI);
+  });
+
+  test('validateCredentialsWithId: should accept valid tokenUri (StringCredentials)', () => {
+    const credentials = {
+      credentialsString: JSON.stringify({ clientID: 'c', keyID: 'k' }),
+      tokenUri: validUrl
+    };
+    expect(() => validateCredentialsWithId(credentials, type, typeId, id)).not.toThrow();
+  });
+
+  test('validateCredentialsWithId: should accept valid tokenUri (TokenCredentials)', () => {
+    jest.spyOn(require('../../src/utils/jwt-utils'), 'isExpired').mockReturnValue(false);
+    const credentials = {
+      token: 'valid-token',
+      tokenUri: validUrl
+    };
+    expect(() => validateCredentialsWithId(credentials, type, typeId, id)).not.toThrow();
+  });
+
+  test('validateSkyflowCredentials: should throw error if tokenUri is present but not a string (PathCredentials)', () => {
+    const credentials = {
+      path: '/valid/path',
+      tokenUri: 123
+    };
+    jest.spyOn(fs, 'existsSync').mockReturnValue(true);
+    expect(() => validateSkyflowCredentials(credentials))
+      .toThrow(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI);
+  });
+
+  test('validateSkyflowCredentials: should throw error if tokenUri is present but not a valid URL (PathCredentials)', () => {
+    const credentials = {
+      path: '/valid/path',
+      tokenUri: 'not-a-url'
+    };
+    jest.spyOn(fs, 'existsSync').mockReturnValue(true);
+    expect(() => validateSkyflowCredentials(credentials))
+      .toThrow(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI);
+  });
+
+  test('validateSkyflowCredentials: should accept valid tokenUri (PathCredentials)', () => {
+    const credentials = {
+      path: '/valid/path',
+      tokenUri: validUrl
+    };
+    jest.spyOn(fs, 'existsSync').mockReturnValue(true);
+    expect(() => validateSkyflowCredentials(credentials)).not.toThrow();
+  });
+
+  test('validateSkyflowCredentials: should throw error if tokenUri is present but not a string (StringCredentials)', () => {
+    const credentials = {
+      credentialsString: JSON.stringify({ clientID: 'c', keyID: 'k' }),
+      tokenUri: 123
+    };
+    expect(() => validateSkyflowCredentials(credentials))
+      .toThrow(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI);
+  });
+
+  test('validateSkyflowCredentials: should throw error if tokenUri is present but not a valid URL (StringCredentials)', () => {
+    const credentials = {
+      credentialsString: JSON.stringify({ clientID: 'c', keyID: 'k' }),
+      tokenUri: 'not-a-url'
+    };
+    expect(() => validateSkyflowCredentials(credentials))
+      .toThrow(SKYFLOW_ERROR_CODE.INVALID_TOKEN_URI);
+  });
+
+  test('validateSkyflowCredentials: should accept valid tokenUri (StringCredentials)', () => {
+    const credentials = {
+      credentialsString: JSON.stringify({ clientID: 'c', keyID: 'k' }),
+      tokenUri: validUrl
+    };
+    expect(() => validateSkyflowCredentials(credentials)).not.toThrow();
+  });
+
+  test('validateSkyflowCredentials: should accept valid tokenUri (TokenCredentials)', () => {
+    jest.spyOn(require('../../src/utils/jwt-utils'), 'isExpired').mockReturnValue(false);
+    const credentials = {
+      token: 'valid-token',
+      tokenUri: validUrl
+    };
+    expect(() => validateSkyflowCredentials(credentials)).not.toThrow();
+  });
+});