From 34d2b8d05c16a92ba71ffe1d4e99ca0f41ac3028 Mon Sep 17 00:00:00 2001 From: Sam Sternberg Date: Mon, 30 Mar 2026 14:34:34 -0400 Subject: [PATCH] feat: add JSON object support for ctx field in bearer token and signed data token generation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend the Java SDK's bearer token and signed data token generation to accept a JSON object (Map) for the ctx field, in addition to the existing String type. This enables structured context for conditional data access policies where ctx object keys map to Skyflow CEL policy variables (e.g., request.context.role, request.context.department). Changes: - Credentials: widen context field from String to Object, add overloaded setContext(Map) - BearerToken/SignedDataTokens: widen ctx to Object, add overloaded setCtx(Map), conditionally include ctx in JWT claims - Utils: dispatch to correct setCtx overload based on context type - Validations: validate both String and Map context types, validate map keys match [a-zA-Z0-9_]+ for CEL compatibility - ErrorMessage/ErrorLogs: add InvalidContextType and InvalidContextMapKey - Tests: add Map-based context tests for Credentials, BearerToken, and SignedDataTokens (51 tests, all passing) - Samples: add JSON object context examples - README: document both string and object ctx patterns with CEL policy variable mapping Technical note: JJWT's .claim(String, Object) handles both types — String serializes as a JSON string, Map serializes as a JSON object in the JWT payload. No custom serialization needed. Resolves: SK-2679, DOCU-1438 Co-Authored-By: Claude Opus 4.6 (1M context) --- README.md | 163 +++++++----------- ...arerTokenGenerationWithContextExample.java | 66 ++++--- .../SignedTokenGenerationExample.java | 85 +++++---- .../java/com/skyflow/config/Credentials.java | 9 +- .../java/com/skyflow/errors/ErrorMessage.java | 2 + src/main/java/com/skyflow/logs/ErrorLogs.java | 2 + .../serviceaccount/util/BearerToken.java | 31 ++-- .../serviceaccount/util/SignedDataTokens.java | 31 ++-- src/main/java/com/skyflow/utils/Utils.java | 29 ++-- .../utils/validations/Validations.java | 32 +++- .../com/skyflow/config/CredentialsTests.java | 86 +++++++++ .../serviceaccount/util/BearerTokenTests.java | 27 +++ .../util/SignedDataTokensTests.java | 16 ++ 13 files changed, 377 insertions(+), 202 deletions(-) diff --git a/README.md b/README.md index 023bb288..3f8a9adb 100644 --- a/README.md +++ b/README.md @@ -2771,73 +2771,59 @@ public class BearerTokenGenerationExample { ## Generate bearer tokens with context -**Context-aware authorization** embeds context values into a bearer token during its generation and so you can reference those values in your policies. This enables more flexible access controls, such as helping you track end-user identity when making API calls using service accounts, and facilitates using signed data tokens during detokenization. . +**Context-aware authorization** embeds context values into a bearer token during its generation so you can reference those values in your policies. This enables more flexible access controls, such as helping you track end-user identity when making API calls using service accounts, and facilitates using signed data tokens during detokenization. A service account with the `context_id` identifier generates bearer tokens containing context information, represented as a JWT claim in a Skyflow-generated bearer token. Tokens generated from such service accounts include a `context_identifier` claim, are valid for 60 minutes, and can be used to make API calls to the Data and Management APIs, depending on the service account's permissions. -[Example](https://github.com/skyflowapi/skyflow-java/blob/main/samples/src/main/java/com/example/serviceaccount/BearerTokenGenerationWithContextExample.java): +The `setCtx()` method accepts either a **String** or a **`Map`**: + +**String context** — use when your policy references a single context value: ```java -import com.skyflow.errors.SkyflowException; -import com.skyflow.serviceaccount.util.BearerToken; +BearerToken token = BearerToken.builder() + .setCredentials(new File(filePath)) + .setCtx("user_12345") + .build(); +``` -import java.io.File; +**JSON object context** — use when your policy needs multiple context values for conditional data access. Each key in the `Map` maps to a Skyflow CEL policy variable under `request.context.*`: -/** - * Example program to generate a Bearer Token using Skyflow's BearerToken utility. - * The token is generated using two approaches: - * 1. By providing the credentials.json file path. - * 2. By providing the contents of credentials.json as a string. - */ -public class BearerTokenGenerationWithContextExample { - public static void main(String[] args) { - // Variable to store the generated Bearer Token - String bearerToken = null; +```java +Map ctx = new HashMap<>(); +ctx.put("role", "admin"); +ctx.put("department", "finance"); +ctx.put("user_id", "user_12345"); - // Approach 1: Generate Bearer Token by specifying the path to the credentials.json file - try { - // Replace with the full path to your credentials.json file - String filePath = ""; +BearerToken token = BearerToken.builder() + .setCredentials(new File(filePath)) + .setCtx(ctx) + .build(); +``` - // Create a BearerToken object using the file path - BearerToken token = BearerToken.builder() - .setCredentials(new File(filePath)) // Set credentials using a File object - .setCtx("abc") // Set context string (example: "abc") - .build(); // Build the BearerToken object +With the map above, your Skyflow policies can reference `request.context.role`, `request.context.department`, and `request.context.user_id` to make conditional access decisions. - // Retrieve the Bearer Token as a string - bearerToken = token.getBearerToken(); +You can also set context on `Credentials` for automatic token generation: - // Print the generated Bearer Token to the console - System.out.println(bearerToken); - } catch (SkyflowException e) { - // Handle exceptions specific to Skyflow operations - e.printStackTrace(); - } +```java +// String context +Credentials credentials = new Credentials(); +credentials.setPath("path/to/credentials.json"); +credentials.setContext("user_12345"); - // Approach 2: Generate Bearer Token by specifying the contents of credentials.json as a string - try { - // Replace with the actual contents of your credentials.json file - String fileContents = ""; +// Map context +Map ctx = new HashMap<>(); +ctx.put("role", "admin"); +ctx.put("department", "finance"); +credentials.setContext(ctx); +``` - // Create a BearerToken object using the file contents as a string - BearerToken token = BearerToken.builder() - .setCredentials(fileContents) // Set credentials using a string representation of the file - .setCtx("abc") // Set context string (example: "abc") - .build(); // Build the BearerToken object +> **Note:** `getContext()` returns `Object` — callers should use `instanceof` if they need to inspect the type. - // Retrieve the Bearer Token as a string - bearerToken = token.getBearerToken(); +Context map keys must contain only alphanumeric characters and underscores (`[a-zA-Z0-9_]`). Invalid keys will throw a `SkyflowException`. - // Print the generated Bearer Token to the console - System.out.println(bearerToken); - } catch (SkyflowException e) { - // Handle exceptions specific to Skyflow operations - e.printStackTrace(); - } - } -} -``` +[Full example](https://github.com/skyflowapi/skyflow-java/blob/main/samples/src/main/java/com/example/serviceaccount/BearerTokenGenerationWithContextExample.java) + +See Skyflow's [context-aware authorization](https://docs.skyflow.com) and [conditional data access](https://docs.skyflow.com) docs for policy variable syntax like `request.context.*`. ## Generate scoped bearer tokens @@ -2903,58 +2889,31 @@ with the private key of the service account credentials, which adds an additiona be detokenized by passing the signed data token and a bearer token generated from service account credentials. The service account must have appropriate permissions and context to detokenize the signed data tokens. -[Example](https://github.com/skyflowapi/skyflow-java/blob/main/samples/src/main/java/com/example/serviceaccount/SignedTokenGenerationExample.java): +The `setCtx()` method on `SignedDataTokensBuilder` also accepts either a **String** or a **`Map`**, using the same format as bearer tokens: ```java -import com.skyflow.errors.SkyflowException; -import com.skyflow.serviceaccount.util.SignedDataTokenResponse; -import com.skyflow.serviceaccount.util.SignedDataTokens; - -import java.io.File; -import java.util.ArrayList; -import java.util.List; - -public class SignedTokenGenerationExample { - public static void main(String[] args) { - List signedTokenValues; - // Generate Signed data token with context by specifying credentials.json file path - try { - String filePath = ""; - String context = "abc"; - ArrayList dataTokens = new ArrayList<>(); - dataTokens.add("YOUR_DATA_TOKEN_1"); - SignedDataTokens signedToken = SignedDataTokens.builder() - .setCredentials(new File(filePath)) - .setCtx(context) - .setTimeToLive(30) // in seconds - .setDataTokens(dataTokens) - .build(); - signedTokenValues = signedToken.getSignedDataTokens(); - System.out.println(signedTokenValues); - } catch (SkyflowException e) { - e.printStackTrace(); - } - - // Generate Signed data token with context by specifying credentials.json as string - try { - String fileContents = ""; - String context = "abc"; - ArrayList dataTokens = new ArrayList<>(); - dataTokens.add("YOUR_DATA_TOKEN_1"); - SignedDataTokens signedToken = SignedDataTokens.builder() - .setCredentials(fileContents) - .setCtx(context) - .setTimeToLive(30) // in seconds - .setDataTokens(dataTokens) - .build(); - signedTokenValues = signedToken.getSignedDataTokens(); - System.out.println(signedTokenValues); - } catch (SkyflowException e) { - e.printStackTrace(); - } - } -} -``` +// String context +SignedDataTokens signedToken = SignedDataTokens.builder() + .setCredentials(new File(filePath)) + .setCtx("user_12345") + .setTimeToLive(30) + .setDataTokens(dataTokens) + .build(); + +// JSON object context +Map ctx = new HashMap<>(); +ctx.put("role", "analyst"); +ctx.put("department", "research"); + +SignedDataTokens signedToken = SignedDataTokens.builder() + .setCredentials(new File(filePath)) + .setCtx(ctx) + .setTimeToLive(30) + .setDataTokens(dataTokens) + .build(); +``` + +[Full example](https://github.com/skyflowapi/skyflow-java/blob/main/samples/src/main/java/com/example/serviceaccount/SignedTokenGenerationExample.java) Response: diff --git a/samples/src/main/java/com/example/serviceaccount/BearerTokenGenerationWithContextExample.java b/samples/src/main/java/com/example/serviceaccount/BearerTokenGenerationWithContextExample.java index 3ed9e267..fcb2a407 100644 --- a/samples/src/main/java/com/example/serviceaccount/BearerTokenGenerationWithContextExample.java +++ b/samples/src/main/java/com/example/serviceaccount/BearerTokenGenerationWithContextExample.java @@ -4,57 +4,69 @@ import com.skyflow.serviceaccount.util.BearerToken; import java.io.File; +import java.util.HashMap; +import java.util.Map; /** * Example program to generate a Bearer Token using Skyflow's BearerToken utility. - * The token is generated using two approaches: - * 1. By providing the credentials.json file path. - * 2. By providing the contents of credentials.json as a string. + * The token is generated using three approaches: + * 1. By providing a string context. + * 2. By providing a JSON object context (Map) for conditional data access policies. + * 3. By providing the credentials as a string with context. */ public class BearerTokenGenerationWithContextExample { public static void main(String[] args) { - // Variable to store the generated Bearer Token String bearerToken = null; - // Approach 1: Generate Bearer Token by specifying the path to the credentials.json file + // Approach 1: Bearer token with string context + // Use a simple string identifier when your policy references a single context value. try { - // Replace with the full path to your credentials.json file String filePath = ""; - - // Create a BearerToken object using the file path BearerToken token = BearerToken.builder() - .setCredentials(new File(filePath)) // Set credentials using a File object - .setCtx("abc") // Set context string (example: "abc") - .build(); // Build the BearerToken object + .setCredentials(new File(filePath)) + .setCtx("user_12345") + .build(); - // Retrieve the Bearer Token as a string bearerToken = token.getBearerToken(); - - // Print the generated Bearer Token to the console - System.out.println(bearerToken); + System.out.println("Bearer token (string context): " + bearerToken); } catch (SkyflowException e) { - // Handle exceptions specific to Skyflow operations e.printStackTrace(); } - // Approach 2: Generate Bearer Token by specifying the contents of credentials.json as a string + // Approach 2: Bearer token with JSON object context + // Use a structured Map when your policy needs multiple context values. + // Each key maps to a Skyflow CEL policy variable under request.context.* + // For example, the map below enables policies like: + // request.context.role == "admin" && request.context.department == "finance" try { - // Replace with the actual contents of your credentials.json file - String fileContents = ""; + String filePath = ""; + Map ctx = new HashMap<>(); + ctx.put("role", "admin"); + ctx.put("department", "finance"); + ctx.put("user_id", "user_12345"); - // Create a BearerToken object using the file contents as a string BearerToken token = BearerToken.builder() - .setCredentials(fileContents) // Set credentials using a string representation of the file - .setCtx("abc") // Set context string (example: "abc") - .build(); // Build the BearerToken object + .setCredentials(new File(filePath)) + .setCtx(ctx) + .build(); - // Retrieve the Bearer Token as a string bearerToken = token.getBearerToken(); + System.out.println("Bearer token (object context): " + bearerToken); + } catch (SkyflowException e) { + e.printStackTrace(); + } + + // Approach 3: Bearer token with string context from credentials string + try { + String fileContents = ""; + BearerToken token = BearerToken.builder() + .setCredentials(fileContents) + .setCtx("user_12345") + .build(); - // Print the generated Bearer Token to the console - System.out.println(bearerToken); + bearerToken = token.getBearerToken(); + System.out.println("Bearer token (creds string): " + bearerToken); } catch (SkyflowException e) { - // Handle exceptions specific to Skyflow operations e.printStackTrace(); } } diff --git a/samples/src/main/java/com/example/serviceaccount/SignedTokenGenerationExample.java b/samples/src/main/java/com/example/serviceaccount/SignedTokenGenerationExample.java index 98f552f3..517d5cc8 100644 --- a/samples/src/main/java/com/example/serviceaccount/SignedTokenGenerationExample.java +++ b/samples/src/main/java/com/example/serviceaccount/SignedTokenGenerationExample.java @@ -6,68 +6,83 @@ import java.io.File; import java.util.ArrayList; +import java.util.HashMap; import java.util.List; +import java.util.Map; /** - * This example demonstrates how to generate Signed Data Tokens using two methods: - * 1. Specifying the path to a credentials JSON file. - * 2. Providing the credentials JSON as a string. - *

- * Signed data tokens are used to verify and securely transmit data with a specified context and TTL. + * This example demonstrates how to generate Signed Data Tokens using: + * 1. String context. + * 2. JSON object context (Map) for conditional data access policies. + * 3. Credentials string with context. */ public class SignedTokenGenerationExample { public static void main(String[] args) { - List signedTokenValues; // List to store signed data token responses + List signedTokenValues; - // Example 1: Generate Signed Data Token using a credentials file path + // Example 1: Signed data tokens with string context try { - // Step 1: Specify the path to the service account credentials JSON file - String filePath = ""; // Replace with the actual file path - - // Step 2: Set the context and create the list of data tokens to be signed - String context = "abc"; // Replace with your specific context (e.g., session identifier) + String filePath = ""; + String context = "user_12345"; ArrayList dataTokens = new ArrayList<>(); - dataTokens.add("YOUR_DATA_TOKEN_1"); // Replace with your actual data token(s) + dataTokens.add("YOUR_DATA_TOKEN_1"); - // Step 3: Build the SignedDataTokens object SignedDataTokens signedToken = SignedDataTokens.builder() - .setCredentials(new File(filePath)) // Provide the credentials file - .setCtx(context) // Set the context for the token - .setTimeToLive(30) // Set the TTL (in seconds) - .setDataTokens(dataTokens) // Set the data tokens to sign + .setCredentials(new File(filePath)) + .setCtx(context) + .setTimeToLive(30) + .setDataTokens(dataTokens) .build(); - // Step 4: Retrieve and print the signed data tokens signedTokenValues = signedToken.getSignedDataTokens(); - System.out.println("Signed Tokens (using file path): " + signedTokenValues); + System.out.println("Signed Tokens (string context): " + signedTokenValues); } catch (SkyflowException e) { - System.out.println("Error occurred while generating signed tokens using file path:"); e.printStackTrace(); } - // Example 2: Generate Signed Data Token using credentials JSON as a string + // Example 2: Signed data tokens with JSON object context + // Each key maps to a Skyflow CEL policy variable under request.context.* + // For example: request.context.role == "analyst" && request.context.department == "research" try { - // Step 1: Provide the contents of the credentials JSON file as a string - String fileContents = ""; // Replace with actual JSON content + String filePath = ""; + Map ctx = new HashMap<>(); + ctx.put("role", "analyst"); + ctx.put("department", "research"); + ctx.put("user_id", "user_67890"); + + ArrayList dataTokens = new ArrayList<>(); + dataTokens.add("YOUR_DATA_TOKEN_1"); + + SignedDataTokens signedToken = SignedDataTokens.builder() + .setCredentials(new File(filePath)) + .setCtx(ctx) + .setTimeToLive(30) + .setDataTokens(dataTokens) + .build(); - // Step 2: Set the context and create the list of data tokens to be signed - String context = "abc"; // Replace with your specific context + signedTokenValues = signedToken.getSignedDataTokens(); + System.out.println("Signed Tokens (object context): " + signedTokenValues); + } catch (SkyflowException e) { + e.printStackTrace(); + } + + // Example 3: Signed data tokens from credentials string + try { + String fileContents = ""; + String context = "user_12345"; ArrayList dataTokens = new ArrayList<>(); - dataTokens.add("YOUR_DATA_TOKEN_1"); // Replace with your actual data token(s) + dataTokens.add("YOUR_DATA_TOKEN_1"); - // Step 3: Build the SignedDataTokens object SignedDataTokens signedToken = SignedDataTokens.builder() - .setCredentials(fileContents) // Provide the credentials as a string - .setCtx(context) // Set the context for the token - .setTimeToLive(30) // Set the TTL (in seconds) - .setDataTokens(dataTokens) // Set the data tokens to sign + .setCredentials(fileContents) + .setCtx(context) + .setTimeToLive(30) + .setDataTokens(dataTokens) .build(); - // Step 4: Retrieve and print the signed data tokens signedTokenValues = signedToken.getSignedDataTokens(); - System.out.println("Signed Tokens (using credentials string): " + signedTokenValues); + System.out.println("Signed Tokens (creds string): " + signedTokenValues); } catch (SkyflowException e) { - System.out.println("Error occurred while generating signed tokens using credentials string:"); e.printStackTrace(); } } diff --git a/src/main/java/com/skyflow/config/Credentials.java b/src/main/java/com/skyflow/config/Credentials.java index f1865dc7..c2594ef6 100644 --- a/src/main/java/com/skyflow/config/Credentials.java +++ b/src/main/java/com/skyflow/config/Credentials.java @@ -1,11 +1,12 @@ package com.skyflow.config; import java.util.ArrayList; +import java.util.Map; public class Credentials { private String path; private ArrayList roles; - private String context; + private Object context; private String credentialsString; private String token; private String apiKey; @@ -32,7 +33,7 @@ public void setRoles(ArrayList roles) { this.roles = roles; } - public String getContext() { + public Object getContext() { return context; } @@ -40,6 +41,10 @@ public void setContext(String context) { this.context = context; } + public void setContext(Map context) { + this.context = context; + } + public String getCredentialsString() { return credentialsString; } diff --git a/src/main/java/com/skyflow/errors/ErrorMessage.java b/src/main/java/com/skyflow/errors/ErrorMessage.java index 8885a6c8..fc222522 100644 --- a/src/main/java/com/skyflow/errors/ErrorMessage.java +++ b/src/main/java/com/skyflow/errors/ErrorMessage.java @@ -34,6 +34,8 @@ public enum ErrorMessage { EmptyRoles("%s0 Initialization failed. Invalid roles. Specify at least one role."), EmptyRoleInRoles("%s0 Initialization failed. Invalid role. Specify a valid role."), EmptyContext("%s0 Initialization failed. Invalid context. Specify a valid context."), + InvalidContextType("%s0 Initialization failed. Invalid context type. Specify context as a String or Map."), + InvalidContextMapKey("%s0 Initialization failed. Invalid key '%s1' in context map. Keys must contain only alphanumeric characters and underscores."), // Bearer token generation FileNotFound("%s0 Initialization failed. Credential file not found at %s1. Verify the file path."), diff --git a/src/main/java/com/skyflow/logs/ErrorLogs.java b/src/main/java/com/skyflow/logs/ErrorLogs.java index eb5ea742..e6e8b304 100644 --- a/src/main/java/com/skyflow/logs/ErrorLogs.java +++ b/src/main/java/com/skyflow/logs/ErrorLogs.java @@ -25,6 +25,8 @@ public enum ErrorLogs { EMPTY_ROLES("Invalid credentials. Roles can not be empty."), EMPTY_OR_NULL_ROLE_IN_ROLES("Invalid credentials. Role can not be null or empty in roles at index %s1."), EMPTY_OR_NULL_CONTEXT("Invalid credentials. Context can not be empty."), + INVALID_CONTEXT_TYPE("Invalid credentials. Context must be a String or Map."), + INVALID_CONTEXT_MAP_KEY("Invalid credentials. Context map key '%s1' contains invalid characters."), // Bearer token generation INVALID_BEARER_TOKEN("Bearer token is invalid or expired."), diff --git a/src/main/java/com/skyflow/serviceaccount/util/BearerToken.java b/src/main/java/com/skyflow/serviceaccount/util/BearerToken.java index 21000f0f..9e3a6d63 100644 --- a/src/main/java/com/skyflow/serviceaccount/util/BearerToken.java +++ b/src/main/java/com/skyflow/serviceaccount/util/BearerToken.java @@ -24,6 +24,7 @@ import java.security.PrivateKey; import java.util.ArrayList; import java.util.Date; +import java.util.Map; import java.util.Objects; public class BearerToken { @@ -31,7 +32,7 @@ public class BearerToken { private static final ApiClientBuilder API_CLIENT_BUILDER = new ApiClientBuilder(); private final File credentialsFile; private final String credentialsString; - private final String ctx; + private final Object ctx; private final ArrayList roles; private final String credentialsType; @@ -48,7 +49,7 @@ public static BearerTokenBuilder builder() { } private static V1GetAuthTokenResponse generateBearerTokenFromCredentials( - File credentialsFile, String context, ArrayList roles + File credentialsFile, Object context, ArrayList roles ) throws SkyflowException { LogUtil.printInfoLog(InfoLogs.GENERATE_BEARER_TOKEN_FROM_CREDENTIALS_TRIGGERED.getLog()); try { @@ -71,7 +72,7 @@ private static V1GetAuthTokenResponse generateBearerTokenFromCredentials( } private static V1GetAuthTokenResponse generateBearerTokenFromCredentialString( - String credentials, String context, ArrayList roles + String credentials, Object context, ArrayList roles ) throws SkyflowException { LogUtil.printInfoLog(InfoLogs.GENERATE_BEARER_TOKEN_FROM_CREDENTIALS_STRING_TRIGGERED.getLog()); try { @@ -89,7 +90,7 @@ private static V1GetAuthTokenResponse generateBearerTokenFromCredentialString( } private static V1GetAuthTokenResponse getBearerTokenFromCredentials( - JsonObject credentials, String context, ArrayList roles + JsonObject credentials, Object context, ArrayList roles ) throws SkyflowException { try { JsonElement privateKey = credentials.get("privateKey"); @@ -144,19 +145,22 @@ private static V1GetAuthTokenResponse getBearerTokenFromCredentials( } private static String getSignedToken( - String clientID, String keyID, String tokenURI, PrivateKey pvtKey, String context + String clientID, String keyID, String tokenURI, PrivateKey pvtKey, Object context ) { final Date createdDate = new Date(); final Date expirationDate = new Date(createdDate.getTime() + (3600 * 1000)); - return Jwts.builder() + io.jsonwebtoken.JwtBuilder builder = Jwts.builder() .claim("iss", clientID) .claim("key", keyID) .claim("aud", tokenURI) .claim("sub", clientID) - .claim("ctx", context) - .expiration(expirationDate) - .signWith(pvtKey, Jwts.SIG.RS256) - .compact(); + .expiration(expirationDate); + + if (context != null) { + builder.claim("ctx", context); + } + + return builder.signWith(pvtKey, Jwts.SIG.RS256).compact(); } private static String getScopeUsingRoles(ArrayList roles) { @@ -188,7 +192,7 @@ public synchronized String getBearerToken() throws SkyflowException { public static class BearerTokenBuilder { private File credentialsFile; private String credentialsString; - private String ctx; + private Object ctx; private ArrayList roles; private String credentialsType; @@ -216,6 +220,11 @@ public BearerTokenBuilder setCtx(String ctx) { return this; } + public BearerTokenBuilder setCtx(Map ctx) { + this.ctx = ctx; + return this; + } + public BearerTokenBuilder setRoles(ArrayList roles) { this.roles = roles; return this; diff --git a/src/main/java/com/skyflow/serviceaccount/util/SignedDataTokens.java b/src/main/java/com/skyflow/serviceaccount/util/SignedDataTokens.java index 70a5a330..0ce14007 100644 --- a/src/main/java/com/skyflow/serviceaccount/util/SignedDataTokens.java +++ b/src/main/java/com/skyflow/serviceaccount/util/SignedDataTokens.java @@ -20,13 +20,14 @@ import java.util.ArrayList; import java.util.Date; import java.util.List; +import java.util.Map; import java.util.Objects; public class SignedDataTokens { private final File credentialsFile; private final String credentialsString; private final String credentialsType; - private final String ctx; + private final Object ctx; private final ArrayList dataTokens; private final Integer timeToLive; @@ -44,7 +45,7 @@ public static SignedDataTokensBuilder builder() { } private static List generateSignedTokenFromCredentialsFile( - File credentialsFile, ArrayList dataTokens, Integer timeToLive, String context + File credentialsFile, ArrayList dataTokens, Integer timeToLive, Object context ) throws SkyflowException { LogUtil.printInfoLog(InfoLogs.GENERATE_SIGNED_TOKENS_FROM_CREDENTIALS_FILE_TRIGGERED.getLog()); List responseToken; @@ -69,7 +70,7 @@ private static List generateSignedTokenFromCredentialsF } private static List generateSignedTokensFromCredentialsString( - String credentials, ArrayList dataTokens, Integer timeToLive, String context + String credentials, ArrayList dataTokens, Integer timeToLive, Object context ) throws SkyflowException { LogUtil.printInfoLog(InfoLogs.GENERATE_SIGNED_TOKENS_FROM_CREDENTIALS_STRING_TRIGGERED.getLog()); List responseToken; @@ -89,7 +90,7 @@ private static List generateSignedTokensFromCredentials } private static List generateSignedTokensFromCredentials( - JsonObject credentials, ArrayList dataTokens, Integer timeToLive, String context + JsonObject credentials, ArrayList dataTokens, Integer timeToLive, Object context ) throws SkyflowException { List signedDataTokens = null; try { @@ -122,7 +123,7 @@ private static List generateSignedTokensFromCredentials private static List getSignedToken( String clientID, String keyID, PrivateKey pvtKey, - ArrayList dataTokens, Integer timeToLive, String context + ArrayList dataTokens, Integer timeToLive, Object context ) { final Date createdDate = new Date(); final Date expirationDate; @@ -135,16 +136,19 @@ private static List getSignedToken( List list = new ArrayList<>(); for (String dataToken : dataTokens) { - String eachSignedDataToken = Jwts.builder() + io.jsonwebtoken.JwtBuilder builder = Jwts.builder() .claim("iss", "sdk") .claim("iat", (createdDate.getTime() / 1000)) .claim("key", keyID) .claim("sub", clientID) - .claim("ctx", context) .claim("tok", dataToken) - .expiration(expirationDate) - .signWith(pvtKey, Jwts.SIG.RS256) - .compact(); + .expiration(expirationDate); + + if (context != null) { + builder.claim("ctx", context); + } + + String eachSignedDataToken = builder.signWith(pvtKey, Jwts.SIG.RS256).compact(); SignedDataTokenResponse responseObject = new SignedDataTokenResponse(dataToken, eachSignedDataToken); list.add(responseObject); } @@ -168,7 +172,7 @@ public static class SignedDataTokensBuilder { private Integer timeToLive; private File credentialsFile; private String credentialsString; - private String ctx; + private Object ctx; private String credentialsType; private SignedDataTokensBuilder() { @@ -195,6 +199,11 @@ public SignedDataTokensBuilder setCtx(String ctx) { return this; } + public SignedDataTokensBuilder setCtx(Map ctx) { + this.ctx = ctx; + return this; + } + public SignedDataTokensBuilder setDataTokens(ArrayList dataTokens) { this.dataTokens = dataTokens; return this; diff --git a/src/main/java/com/skyflow/utils/Utils.java b/src/main/java/com/skyflow/utils/Utils.java index 165c6a80..b33b08c1 100644 --- a/src/main/java/com/skyflow/utils/Utils.java +++ b/src/main/java/com/skyflow/utils/Utils.java @@ -47,21 +47,30 @@ public static String getVaultURL(String clusterId, Env env) { return sb.toString(); } + @SuppressWarnings("unchecked") public static String generateBearerToken(Credentials credentials) throws SkyflowException { if (credentials.getPath() != null) { - return BearerToken.builder() + BearerToken.BearerTokenBuilder builder = BearerToken.builder() .setCredentials(new File(credentials.getPath())) - .setRoles(credentials.getRoles()) - .setCtx(credentials.getContext()) - .build() - .getBearerToken(); + .setRoles(credentials.getRoles()); + Object ctx = credentials.getContext(); + if (ctx instanceof String) { + builder.setCtx((String) ctx); + } else if (ctx instanceof Map) { + builder.setCtx((Map) ctx); + } + return builder.build().getBearerToken(); } else if (credentials.getCredentialsString() != null) { - return BearerToken.builder() + BearerToken.BearerTokenBuilder builder = BearerToken.builder() .setCredentials(credentials.getCredentialsString()) - .setRoles(credentials.getRoles()) - .setCtx(credentials.getContext()) - .build() - .getBearerToken(); + .setRoles(credentials.getRoles()); + Object ctx = credentials.getContext(); + if (ctx instanceof String) { + builder.setCtx((String) ctx); + } else if (ctx instanceof Map) { + builder.setCtx((Map) ctx); + } + return builder.build().getBearerToken(); } else { return credentials.getToken(); } diff --git a/src/main/java/com/skyflow/utils/validations/Validations.java b/src/main/java/com/skyflow/utils/validations/Validations.java index 1d078fa9..76ac4f93 100644 --- a/src/main/java/com/skyflow/utils/validations/Validations.java +++ b/src/main/java/com/skyflow/utils/validations/Validations.java @@ -162,7 +162,7 @@ public static void validateCredentials(Credentials credentials) throws SkyflowEx String credentialsString = credentials.getCredentialsString(); String token = credentials.getToken(); String apiKey = credentials.getApiKey(); - String context = credentials.getContext(); + Object context = credentials.getContext(); ArrayList roles = credentials.getRoles(); if (path != null) nonNullMembers++; @@ -217,9 +217,33 @@ public static void validateCredentials(Credentials credentials) throws SkyflowEx } } } - if (context != null && context.trim().isEmpty()) { - LogUtil.printErrorLog(ErrorLogs.EMPTY_OR_NULL_CONTEXT.getLog()); - throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.EmptyContext.getMessage()); + if (context != null) { + if (context instanceof String) { + String ctxStr = (String) context; + if (ctxStr.trim().isEmpty()) { + LogUtil.printErrorLog(ErrorLogs.EMPTY_OR_NULL_CONTEXT.getLog()); + throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.EmptyContext.getMessage()); + } + } else if (context instanceof Map) { + Map ctxMap = (Map) context; + if (ctxMap.isEmpty()) { + LogUtil.printErrorLog(ErrorLogs.EMPTY_OR_NULL_CONTEXT.getLog()); + throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.EmptyContext.getMessage()); + } + Pattern ctxKeyPattern = Pattern.compile("^[a-zA-Z0-9_]+$"); + for (Object key : ctxMap.keySet()) { + if (key == null || !ctxKeyPattern.matcher(key.toString()).matches()) { + String keyStr = key == null ? "null" : key.toString(); + LogUtil.printErrorLog(Utils.parameterizedString( + ErrorLogs.INVALID_CONTEXT_MAP_KEY.getLog(), keyStr)); + throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), + Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), keyStr)); + } + } + } else { + LogUtil.printErrorLog(ErrorLogs.INVALID_CONTEXT_TYPE.getLog()); + throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.InvalidContextType.getMessage()); + } } } diff --git a/src/test/java/com/skyflow/config/CredentialsTests.java b/src/test/java/com/skyflow/config/CredentialsTests.java index ff9fcfda..a9a9153f 100644 --- a/src/test/java/com/skyflow/config/CredentialsTests.java +++ b/src/test/java/com/skyflow/config/CredentialsTests.java @@ -10,6 +10,10 @@ import org.junit.Test; import java.util.ArrayList; +import java.util.HashMap; +import java.util.Map; + +import com.skyflow.utils.Utils; public class CredentialsTests { private static final String INVALID_EXCEPTION_THROWN = "Should not have thrown any exception"; @@ -266,4 +270,86 @@ public void testEmptyContextInCredentials() { } } + @Test + public void testValidMapContextInCredentials() { + try { + Credentials credentials = new Credentials(); + credentials.setPath(path); + Map ctxMap = new HashMap<>(); + ctxMap.put("role", "admin"); + ctxMap.put("department", "finance"); + ctxMap.put("user_id", "user_12345"); + credentials.setContext(ctxMap); + Validations.validateCredentials(credentials); + } catch (SkyflowException e) { + Assert.fail(INVALID_EXCEPTION_THROWN); + } + } + + @Test + public void testEmptyMapContextInCredentials() { + try { + Credentials credentials = new Credentials(); + credentials.setPath(path); + Map ctxMap = new HashMap<>(); + credentials.setContext(ctxMap); + Validations.validateCredentials(credentials); + Assert.fail(EXCEPTION_NOT_THROWN); + } catch (SkyflowException e) { + Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode()); + Assert.assertEquals(ErrorMessage.EmptyContext.getMessage(), e.getMessage()); + } + } + + @Test + public void testInvalidMapKeyInContextCredentials() { + try { + Credentials credentials = new Credentials(); + credentials.setPath(path); + Map ctxMap = new HashMap<>(); + ctxMap.put("valid_key", "value"); + ctxMap.put("invalid-key", "value"); + credentials.setContext(ctxMap); + Validations.validateCredentials(credentials); + Assert.fail(EXCEPTION_NOT_THROWN); + } catch (SkyflowException e) { + Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode()); + Assert.assertTrue(e.getMessage().contains("invalid-key")); + } + } + + @Test + public void testMapContextWithNestedObjects() { + try { + Credentials credentials = new Credentials(); + credentials.setPath(path); + Map nested = new HashMap<>(); + nested.put("level", 2); + Map ctxMap = new HashMap<>(); + ctxMap.put("role", "admin"); + ctxMap.put("metadata", nested); + credentials.setContext(ctxMap); + Validations.validateCredentials(credentials); + } catch (SkyflowException e) { + Assert.fail(INVALID_EXCEPTION_THROWN); + } + } + + @Test + public void testMapContextWithMixedValueTypes() { + try { + Credentials credentials = new Credentials(); + credentials.setPath(path); + Map ctxMap = new HashMap<>(); + ctxMap.put("role", "admin"); + ctxMap.put("level", 3); + ctxMap.put("active", true); + ctxMap.put("timestamp", "2025-12-25T10:30:00Z"); + credentials.setContext(ctxMap); + Validations.validateCredentials(credentials); + } catch (SkyflowException e) { + Assert.fail(INVALID_EXCEPTION_THROWN); + } + } + } diff --git a/src/test/java/com/skyflow/serviceaccount/util/BearerTokenTests.java b/src/test/java/com/skyflow/serviceaccount/util/BearerTokenTests.java index 4bfb697c..ecd38e84 100644 --- a/src/test/java/com/skyflow/serviceaccount/util/BearerTokenTests.java +++ b/src/test/java/com/skyflow/serviceaccount/util/BearerTokenTests.java @@ -10,6 +10,8 @@ import java.io.File; import java.util.ArrayList; +import java.util.HashMap; +import java.util.Map; public class BearerTokenTests { private static final String INVALID_EXCEPTION_THROWN = "Should not have thrown any exception"; @@ -51,7 +53,32 @@ public void testBearerTokenBuilderWithCredentialsString() { } catch (Exception e) { Assert.fail(INVALID_EXCEPTION_THROWN); } + } + + @Test + public void testBearerTokenBuilderWithMapContext() { + try { + Map ctxMap = new HashMap<>(); + ctxMap.put("role", "admin"); + ctxMap.put("department", "finance"); + ctxMap.put("user_id", "user_12345"); + File file = new File(credentialsFilePath); + BearerToken.builder().setCredentials(file).setCtx(ctxMap).build(); + } catch (Exception e) { + Assert.fail(INVALID_EXCEPTION_THROWN); + } + } + @Test + public void testBearerTokenBuilderWithMapContextFromString() { + try { + Map ctxMap = new HashMap<>(); + ctxMap.put("role", "analyst"); + ctxMap.put("level", 3); + BearerToken.builder().setCredentials(credentialsString).setCtx(ctxMap).build(); + } catch (Exception e) { + Assert.fail(INVALID_EXCEPTION_THROWN); + } } @Test diff --git a/src/test/java/com/skyflow/serviceaccount/util/SignedDataTokensTests.java b/src/test/java/com/skyflow/serviceaccount/util/SignedDataTokensTests.java index f431042d..5e2cbe60 100644 --- a/src/test/java/com/skyflow/serviceaccount/util/SignedDataTokensTests.java +++ b/src/test/java/com/skyflow/serviceaccount/util/SignedDataTokensTests.java @@ -10,6 +10,8 @@ import java.io.File; import java.util.ArrayList; +import java.util.HashMap; +import java.util.Map; public class SignedDataTokensTests { private static final String INVALID_EXCEPTION_THROWN = "Should not have thrown any exception"; @@ -57,7 +59,21 @@ public void testSignedDataTokensBuilderWithCredentialsString() { } catch (Exception e) { Assert.fail(INVALID_EXCEPTION_THROWN); } + } + @Test + public void testSignedDataTokensBuilderWithMapContext() { + try { + Map ctxMap = new HashMap<>(); + ctxMap.put("role", "admin"); + ctxMap.put("department", "finance"); + File file = new File(credentialsFilePath); + SignedDataTokens.builder() + .setCredentials(file).setCtx(ctxMap).setDataTokens(dataTokens).setTimeToLive(ttl) + .build(); + } catch (Exception e) { + Assert.fail(INVALID_EXCEPTION_THROWN); + } } @Test