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`. 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/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/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/main/java/com/skyflow/utils/validations/Validations.java b/src/main/java/com/skyflow/utils/validations/Validations.java index e1f18795..d118b2d9 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(); @@ -162,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++; 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); } 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() + ); + } + } }