From 1a5f59aebc76c286dc0f144bbbcc91d2966874e7 Mon Sep 17 00:00:00 2001 From: Sam Sternberg Date: Tue, 28 Apr 2026 08:26:56 -0400 Subject: [PATCH 1/4] feat(auth): Add JSON object context support for Conditional Data Access - Add typed getContext() (String), getContextAsObject(), getContextAsMap() to Credentials; @SuppressWarnings scoped to Credentials where the invariant is guaranteed by the typed setContext overloads - Remove unchecked cast and @SuppressWarnings from Utils.generateBearerToken by using the new typed accessors - Assert getContextAsObject/getContextAsMap behaviour in CredentialsTests Co-Authored-By: Claude Sonnet 4.6 (1M context) --- .../java/com/skyflow/config/Credentials.java | 11 +++++++++- src/main/java/com/skyflow/utils/Utils.java | 21 +++++++++---------- .../com/skyflow/config/CredentialsTests.java | 3 +++ 3 files changed, 23 insertions(+), 12 deletions(-) diff --git a/src/main/java/com/skyflow/config/Credentials.java b/src/main/java/com/skyflow/config/Credentials.java index c2594ef6..85c0bfde 100644 --- a/src/main/java/com/skyflow/config/Credentials.java +++ b/src/main/java/com/skyflow/config/Credentials.java @@ -33,10 +33,19 @@ public void setRoles(ArrayList roles) { this.roles = roles; } - public Object getContext() { + public String getContext() { + return context instanceof String ? (String) context : null; + } + + public Object getContextAsObject() { return context; } + @SuppressWarnings("unchecked") + public Map getContextAsMap() { + return context instanceof Map ? (Map) context : null; + } + public void setContext(String context) { this.context = context; } diff --git a/src/main/java/com/skyflow/utils/Utils.java b/src/main/java/com/skyflow/utils/Utils.java index b33b08c1..2fe26aff 100644 --- a/src/main/java/com/skyflow/utils/Utils.java +++ b/src/main/java/com/skyflow/utils/Utils.java @@ -47,28 +47,27 @@ 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) { BearerToken.BearerTokenBuilder builder = BearerToken.builder() .setCredentials(new File(credentials.getPath())) .setRoles(credentials.getRoles()); - Object ctx = credentials.getContext(); - if (ctx instanceof String) { - builder.setCtx((String) ctx); - } else if (ctx instanceof Map) { - builder.setCtx((Map) ctx); + Map ctxMap = credentials.getContextAsMap(); + if (ctxMap != null) { + builder.setCtx(ctxMap); + } else { + builder.setCtx(credentials.getContext()); } return builder.build().getBearerToken(); } else if (credentials.getCredentialsString() != null) { BearerToken.BearerTokenBuilder builder = BearerToken.builder() .setCredentials(credentials.getCredentialsString()) .setRoles(credentials.getRoles()); - Object ctx = credentials.getContext(); - if (ctx instanceof String) { - builder.setCtx((String) ctx); - } else if (ctx instanceof Map) { - builder.setCtx((Map) ctx); + Map ctxMap = credentials.getContextAsMap(); + if (ctxMap != null) { + builder.setCtx(ctxMap); + } else { + builder.setCtx(credentials.getContext()); } return builder.build().getBearerToken(); } else { diff --git a/src/test/java/com/skyflow/config/CredentialsTests.java b/src/test/java/com/skyflow/config/CredentialsTests.java index a9a9153f..5b6688f8 100644 --- a/src/test/java/com/skyflow/config/CredentialsTests.java +++ b/src/test/java/com/skyflow/config/CredentialsTests.java @@ -281,6 +281,9 @@ public void testValidMapContextInCredentials() { ctxMap.put("user_id", "user_12345"); credentials.setContext(ctxMap); Validations.validateCredentials(credentials); + Assert.assertNull(credentials.getContext()); + Assert.assertEquals(ctxMap, credentials.getContextAsObject()); + Assert.assertEquals(ctxMap, credentials.getContextAsMap()); } catch (SkyflowException e) { Assert.fail(INVALID_EXCEPTION_THROWN); } From 73cceb6b5f0221a6b4dcd6a225ef839a2d434b7b Mon Sep 17 00:00:00 2001 From: Sam Sternberg Date: Tue, 28 Apr 2026 08:30:29 -0400 Subject: [PATCH 2/4] feat(auth): Add builder-level Map context validation and tests - Extract validateCtxMapKeys() from validateCredentials into a public helper on Validations so builders can call it directly - Add throws SkyflowException + validateCtxMapKeys() to setCtx(Map) in BearerToken and SignedDataTokens builders for early validation - Add nested map, credentials-string, and invalid-key tests for both BearerToken and SignedDataTokens builders Co-Authored-By: Claude Sonnet 4.6 (1M context) --- .../serviceaccount/util/BearerToken.java | 7 +- .../serviceaccount/util/SignedDataTokens.java | 7 +- .../utils/validations/Validations.java | 13 +++ .../serviceaccount/util/BearerTokenTests.java | 75 ++++++++++++++++++ .../util/SignedDataTokensTests.java | 79 +++++++++++++++++++ 5 files changed, 179 insertions(+), 2 deletions(-) diff --git a/src/main/java/com/skyflow/serviceaccount/util/BearerToken.java b/src/main/java/com/skyflow/serviceaccount/util/BearerToken.java index 9e3a6d63..84c11c4a 100644 --- a/src/main/java/com/skyflow/serviceaccount/util/BearerToken.java +++ b/src/main/java/com/skyflow/serviceaccount/util/BearerToken.java @@ -17,6 +17,8 @@ import com.skyflow.utils.logger.LogUtil; import io.jsonwebtoken.Jwts; +import com.skyflow.utils.validations.Validations; + import java.io.File; import java.io.FileNotFoundException; import java.io.FileReader; @@ -220,7 +222,10 @@ public BearerTokenBuilder setCtx(String ctx) { return this; } - public BearerTokenBuilder setCtx(Map ctx) { + public BearerTokenBuilder setCtx(Map ctx) throws SkyflowException { + if (ctx != null && !ctx.isEmpty()) { + Validations.validateCtxMapKeys(ctx); + } this.ctx = ctx; 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 0ce14007..98f098c9 100644 --- a/src/main/java/com/skyflow/serviceaccount/util/SignedDataTokens.java +++ b/src/main/java/com/skyflow/serviceaccount/util/SignedDataTokens.java @@ -13,6 +13,8 @@ import com.skyflow.utils.logger.LogUtil; import io.jsonwebtoken.Jwts; +import com.skyflow.utils.validations.Validations; + import java.io.File; import java.io.FileNotFoundException; import java.io.FileReader; @@ -199,7 +201,10 @@ public SignedDataTokensBuilder setCtx(String ctx) { return this; } - public SignedDataTokensBuilder setCtx(Map ctx) { + public SignedDataTokensBuilder setCtx(Map ctx) throws SkyflowException { + if (ctx != null && !ctx.isEmpty()) { + Validations.validateCtxMapKeys(ctx); + } this.ctx = ctx; return this; } diff --git a/src/main/java/com/skyflow/utils/validations/Validations.java b/src/main/java/com/skyflow/utils/validations/Validations.java index e1f18795..bc937e2f 100644 --- a/src/main/java/com/skyflow/utils/validations/Validations.java +++ b/src/main/java/com/skyflow/utils/validations/Validations.java @@ -45,6 +45,19 @@ public class Validations { private Validations() { } + public static void validateCtxMapKeys(Map ctxMap) throws SkyflowException { + Pattern ctxKeyPattern = Pattern.compile(Constants.CONTEXT_KEY_REGEX); + 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)); + } + } + } + public static void validateVaultConfig(VaultConfig vaultConfig) throws SkyflowException { String vaultId = vaultConfig.getVaultId(); String clusterId = vaultConfig.getClusterId(); diff --git a/src/test/java/com/skyflow/serviceaccount/util/BearerTokenTests.java b/src/test/java/com/skyflow/serviceaccount/util/BearerTokenTests.java index ecd38e84..d4bbab4e 100644 --- a/src/test/java/com/skyflow/serviceaccount/util/BearerTokenTests.java +++ b/src/test/java/com/skyflow/serviceaccount/util/BearerTokenTests.java @@ -249,4 +249,79 @@ public void testInvalidTokenURIInCredentialsForCredentials() throws SkyflowExcep Assert.assertEquals(ErrorMessage.InvalidTokenUri.getMessage(), e.getMessage()); } } + + @Test + public void testBearerTokenBuilderWithNestedMapContext() { + try { + File file = new File(credentialsFilePath); + Map nestedContext = new HashMap<>(); + Map user = new HashMap<>(); + user.put("role", "admin"); + user.put("level", 5); + nestedContext.put("user", user); + nestedContext.put("project_id", "proj_123"); + BearerToken.builder().setCredentials(file).setCtx(nestedContext).build(); + } catch (Exception e) { + Assert.fail(INVALID_EXCEPTION_THROWN); + } + } + + @Test + public void testBearerTokenBuilderWithMapContextAndCredentialsString() { + try { + Map mapContext = new HashMap<>(); + mapContext.put("role", "admin"); + BearerToken.builder().setCredentials(credentialsString).setCtx(mapContext).build(); + } catch (Exception e) { + Assert.fail(INVALID_EXCEPTION_THROWN); + } + } + + @Test + public void testBearerTokenWithInvalidCtxMapKeyContainingHyphen() { + try { + Map mapContext = new HashMap<>(); + mapContext.put("invalid-key", "value"); + BearerToken.builder().setCtx(mapContext); + Assert.fail(EXCEPTION_NOT_THROWN); + } catch (SkyflowException e) { + Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode()); + Assert.assertEquals( + Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), "invalid-key"), + e.getMessage() + ); + } + } + + @Test + public void testBearerTokenWithInvalidCtxMapKeyContainingDot() { + try { + Map mapContext = new HashMap<>(); + mapContext.put("invalid.key", "value"); + BearerToken.builder().setCtx(mapContext); + Assert.fail(EXCEPTION_NOT_THROWN); + } catch (SkyflowException e) { + Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode()); + Assert.assertEquals( + Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), "invalid.key"), + e.getMessage() + ); + } + } + + @Test + public void testBearerTokenWithInvalidCtxMapKeyContainingSpace() { + try { + Map mapContext = new HashMap<>(); + mapContext.put("invalid key", "value"); + BearerToken.builder().setCtx(mapContext); + Assert.fail(EXCEPTION_NOT_THROWN); + } catch (SkyflowException e) { + Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode()); + Assert.assertEquals( + Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), "invalid key"), + e.getMessage() + ); + } + } } diff --git a/src/test/java/com/skyflow/serviceaccount/util/SignedDataTokensTests.java b/src/test/java/com/skyflow/serviceaccount/util/SignedDataTokensTests.java index 5e2cbe60..e578fe02 100644 --- a/src/test/java/com/skyflow/serviceaccount/util/SignedDataTokensTests.java +++ b/src/test/java/com/skyflow/serviceaccount/util/SignedDataTokensTests.java @@ -238,4 +238,83 @@ public void testSignedDataTokenResponse() { Assert.fail(INVALID_EXCEPTION_THROWN); } } + + @Test + public void testSignedDataTokensBuilderWithNestedMapContext() { + try { + File file = new File(credentialsFilePath); + Map nestedContext = new HashMap<>(); + Map user = new HashMap<>(); + user.put("role", "admin"); + user.put("level", 5); + nestedContext.put("user", user); + nestedContext.put("project_id", "proj_123"); + SignedDataTokens.builder() + .setCredentials(file).setCtx(nestedContext).setDataTokens(dataTokens).setTimeToLive(ttl) + .build(); + } catch (Exception e) { + Assert.fail(INVALID_EXCEPTION_THROWN); + } + } + + @Test + public void testSignedDataTokensBuilderWithMapContextAndCredentialsString() { + try { + Map mapContext = new HashMap<>(); + mapContext.put("role", "admin"); + SignedDataTokens.builder() + .setCredentials(credentialsString).setCtx(mapContext).setDataTokens(dataTokens).setTimeToLive(ttl) + .build(); + } catch (Exception e) { + Assert.fail(INVALID_EXCEPTION_THROWN); + } + } + + @Test + public void testSignedDataTokensWithInvalidCtxMapKeyContainingHyphen() { + try { + Map mapContext = new HashMap<>(); + mapContext.put("invalid-key", "value"); + SignedDataTokens.builder().setCtx(mapContext); + Assert.fail(EXCEPTION_NOT_THROWN); + } catch (SkyflowException e) { + Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode()); + Assert.assertEquals( + Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), "invalid-key"), + e.getMessage() + ); + } + } + + @Test + public void testSignedDataTokensWithInvalidCtxMapKeyContainingDot() { + try { + Map mapContext = new HashMap<>(); + mapContext.put("invalid.key", "value"); + SignedDataTokens.builder().setCtx(mapContext); + Assert.fail(EXCEPTION_NOT_THROWN); + } catch (SkyflowException e) { + Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode()); + Assert.assertEquals( + Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), "invalid.key"), + e.getMessage() + ); + } + } + + @Test + public void testSignedDataTokensWithInvalidCtxMapKeyContainingSpace() { + try { + Map mapContext = new HashMap<>(); + mapContext.put("invalid key", "value"); + SignedDataTokens.builder().setCtx(mapContext); + Assert.fail(EXCEPTION_NOT_THROWN); + } catch (SkyflowException e) { + Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode()); + Assert.assertEquals( + Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), "invalid key"), + e.getMessage() + ); + } + } } From 8828c89d5ba57b30f3141a04f67a8bd54b26fbd9 Mon Sep 17 00:00:00 2001 From: Sam Sternberg Date: Tue, 28 Apr 2026 08:31:41 -0400 Subject: [PATCH 3/4] fix: use getContextAsObject() in validateCredentials for Map context check getContext() now returns String (null for Map), so validateCredentials was silently skipping Map context validation. Switch to getContextAsObject() to restore instanceof Map branching. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- src/main/java/com/skyflow/utils/validations/Validations.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/com/skyflow/utils/validations/Validations.java b/src/main/java/com/skyflow/utils/validations/Validations.java index bc937e2f..d118b2d9 100644 --- a/src/main/java/com/skyflow/utils/validations/Validations.java +++ b/src/main/java/com/skyflow/utils/validations/Validations.java @@ -175,7 +175,7 @@ public static void validateCredentials(Credentials credentials) throws SkyflowEx String credentialsString = credentials.getCredentialsString(); String token = credentials.getToken(); String apiKey = credentials.getApiKey(); - Object context = credentials.getContext(); + Object context = credentials.getContextAsObject(); ArrayList roles = credentials.getRoles(); if (path != null) nonNullMembers++; From 949f5b946fdf41014a1b2b8df121ceb6512fe5c5 Mon Sep 17 00:00:00 2001 From: Sam Sternberg Date: Tue, 28 Apr 2026 08:32:21 -0400 Subject: [PATCH 4/4] docs: update getContext() return type note in README Co-Authored-By: Claude Sonnet 4.6 (1M context) --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 3f8a9adb..f77f13a9 100644 --- a/README.md +++ b/README.md @@ -2817,7 +2817,7 @@ ctx.put("department", "finance"); credentials.setContext(ctx); ``` -> **Note:** `getContext()` returns `Object` — callers should use `instanceof` if they need to inspect the type. +> **Note:** `getContext()` returns the context as a `String` (or `null` if a Map was set). Use `getContextAsMap()` to retrieve a Map context, or `getContextAsObject()` to retrieve either as the underlying `Object`. Context map keys must contain only alphanumeric characters and underscores (`[a-zA-Z0-9_]`). Invalid keys will throw a `SkyflowException`.