Skip to content

SK-2505: Add support for custom tokenUri #271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
saileshwar-skyflow merged 3 commits into from
Jan 27, 2026
Merged

SK-2505: Add support for custom tokenUri #271

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
dae60ff
SK-2505: add support for custom token uri
saileshwar-skyflow Jan 23, 2026
9aa9d68
SK-2505: add unit tests
saileshwar-skyflow Jan 23, 2026
d828aac
SK-2505: increase tests coverage
saileshwar-skyflow Jan 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions src/main/java/com/skyflow/config/Credentials.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,13 @@ public class Credentials {
private String credentialsString;
private String token;
private String apiKey;
private String tokenUri;

public Credentials() {
this.path = null;
this.context = null;
this.credentialsString = null;
this.tokenUri = null;
}

public String getPath() {
Expand Down Expand Up @@ -63,4 +65,12 @@ public String getApiKey() {
public void setApiKey(String apiKey) {
this.apiKey = apiKey;
}

public String getTokenUri() {
return tokenUri;
}

public void setTokenUri(String tokenUri) {
this.tokenUri = tokenUri;
}
}
2 changes: 1 addition & 1 deletion src/main/java/com/skyflow/errors/ErrorMessage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ public enum ErrorMessage {
MissingClientId("%s0 Initialization failed. Unable to read client ID in credentials. Verify your client ID."),
MissingKeyId("%s0 Initialization failed. Unable to read key ID in credentials. Verify your key ID."),
MissingTokenUri("%s0 Initialization failed. Unable to read token URI in credentials. Verify your token URI."),
InvalidTokenUri("%s0 Initialization failed. Token URI in not a valid URL in credentials. Verify your token URI."),
InvalidTokenUri("%s0 Initialization failed. Invalid Skyflow credentials. The token URI must be a string and a valid URL."),
JwtInvalidFormat("%s0 Initialization failed. Invalid private key format. Verify your credentials."),
InvalidAlgorithm("%s0 Initialization failed. Invalid algorithm to parse private key. Specify valid algorithm."),
InvalidKeySpec("%s0 Initialization failed. Unable to parse RSA private key. Verify your credentials."),
Expand Down
38 changes: 28 additions & 10 deletions src/main/java/com/skyflow/serviceaccount/util/BearerToken.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,21 +34,23 @@ public class BearerToken {
private final String ctx;
private final ArrayList<String> roles;
private final String credentialsType;
private final String tokenUri;

private BearerToken(BearerTokenBuilder builder) {
this.credentialsFile = builder.credentialsFile;
this.credentialsString = builder.credentialsString;
this.ctx = builder.ctx;
this.roles = builder.roles;
this.credentialsType = builder.credentialsType;
this.tokenUri = builder.tokenUri;
}

public static BearerTokenBuilder builder() {
return new BearerTokenBuilder();
}

private static V1GetAuthTokenResponse generateBearerTokenFromCredentials(
File credentialsFile, String context, ArrayList<String> roles
File credentialsFile, String context, ArrayList<String> roles, String overrideTokenUri
) throws SkyflowException {
LogUtil.printInfoLog(InfoLogs.GENERATE_BEARER_TOKEN_FROM_CREDENTIALS_TRIGGERED.getLog());
try {
Expand All @@ -58,7 +60,7 @@ private static V1GetAuthTokenResponse generateBearerTokenFromCredentials(
}
FileReader reader = new FileReader(String.valueOf(credentialsFile));
JsonObject serviceAccountCredentials = JsonParser.parseReader(reader).getAsJsonObject();
return getBearerTokenFromCredentials(serviceAccountCredentials, context, roles);
return getBearerTokenFromCredentials(serviceAccountCredentials, context, roles, overrideTokenUri);
} catch (JsonSyntaxException e) {
LogUtil.printErrorLog(ErrorLogs.INVALID_CREDENTIALS_FILE_FORMAT.getLog());
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), Utils.parameterizedString(
Expand All @@ -71,7 +73,7 @@ private static V1GetAuthTokenResponse generateBearerTokenFromCredentials(
}

private static V1GetAuthTokenResponse generateBearerTokenFromCredentialString(
String credentials, String context, ArrayList<String> roles
String credentials, String context, ArrayList<String> roles, String overrideTokenUri
) throws SkyflowException {
LogUtil.printInfoLog(InfoLogs.GENERATE_BEARER_TOKEN_FROM_CREDENTIALS_STRING_TRIGGERED.getLog());
try {
Expand All @@ -80,7 +82,7 @@ private static V1GetAuthTokenResponse generateBearerTokenFromCredentialString(
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.InvalidCredentials.getMessage());
}
JsonObject serviceAccountCredentials = JsonParser.parseString(credentials).getAsJsonObject();
return getBearerTokenFromCredentials(serviceAccountCredentials, context, roles);
return getBearerTokenFromCredentials(serviceAccountCredentials, context, roles, overrideTokenUri);
} catch (JsonSyntaxException e) {
LogUtil.printErrorLog(ErrorLogs.INVALID_CREDENTIALS_STRING_FORMAT.getLog());
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(),
Expand All @@ -89,7 +91,7 @@ private static V1GetAuthTokenResponse generateBearerTokenFromCredentialString(
}

private static V1GetAuthTokenResponse getBearerTokenFromCredentials(
JsonObject credentials, String context, ArrayList<String> roles
JsonObject credentials, String context, ArrayList<String> roles, String overrideTokenUri
) throws SkyflowException {
try {
JsonElement privateKey = credentials.get("privateKey");
Expand All @@ -111,17 +113,19 @@ private static V1GetAuthTokenResponse getBearerTokenFromCredentials(
}

JsonElement tokenURI = credentials.get("tokenURI");
if (tokenURI == null) {
if (tokenURI == null && overrideTokenUri == null) {
LogUtil.printErrorLog(ErrorLogs.TOKEN_URI_IS_REQUIRED.getLog());
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.MissingTokenUri.getMessage());
}

String finalTokenUri = (overrideTokenUri != null) ? overrideTokenUri : tokenURI.getAsString();

PrivateKey pvtKey = Utils.getPrivateKeyFromPem(privateKey.getAsString());
String signedUserJWT = getSignedToken(
clientID.getAsString(), keyID.getAsString(), tokenURI.getAsString(), pvtKey, context
clientID.getAsString(), keyID.getAsString(), finalTokenUri, pvtKey, context
);

String basePath = Utils.getBaseURL(tokenURI.getAsString());
String basePath = Utils.getBaseURL(finalTokenUri);
API_CLIENT_BUILDER.url(basePath);
ApiClient apiClient = API_CLIENT_BUILDER.token("token").build();
AuthenticationClient authenticationApi = apiClient.authentication();
Expand Down Expand Up @@ -174,10 +178,10 @@ public synchronized String getBearerToken() throws SkyflowException {
V1GetAuthTokenResponse response;
String accessToken = null;
if (this.credentialsFile != null && Objects.equals(this.credentialsType, "FILE")) {
response = generateBearerTokenFromCredentials(this.credentialsFile, this.ctx, this.roles);
response = generateBearerTokenFromCredentials(this.credentialsFile, this.ctx, this.roles, this.tokenUri);
accessToken = response.getAccessToken().get();
} else if (this.credentialsString != null && Objects.equals(this.credentialsType, "STRING")) {
response = generateBearerTokenFromCredentialString(this.credentialsString, this.ctx, this.roles);
response = generateBearerTokenFromCredentialString(this.credentialsString, this.ctx, this.roles, this.tokenUri);
accessToken = response.getAccessToken().get();
}
LogUtil.printInfoLog(InfoLogs.GET_BEARER_TOKEN_SUCCESS.getLog());
Expand All @@ -191,6 +195,7 @@ public static class BearerTokenBuilder {
private String ctx;
private ArrayList<String> roles;
private String credentialsType;
private String tokenUri;

private BearerTokenBuilder() {
}
Expand Down Expand Up @@ -221,6 +226,19 @@ public BearerTokenBuilder setRoles(ArrayList<String> roles) {
return this;
}

public BearerTokenBuilder setTokenUri(String tokenUri) throws SkyflowException {
if (tokenUri != null && !tokenUri.isEmpty()) {
try {
new java.net.URL(tokenUri);
this.tokenUri = tokenUri;
} catch (MalformedURLException e) {
LogUtil.printErrorLog(ErrorLogs.INVALID_TOKEN_URI.getLog());
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.InvalidTokenUri.getMessage());
}
}
return this;
}

public BearerToken build() {
return new BearerToken(this);
}
Expand Down
41 changes: 32 additions & 9 deletions src/main/java/com/skyflow/serviceaccount/util/SignedDataTokens.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.net.MalformedURLException;
import java.security.PrivateKey;
import java.util.ArrayList;
import java.util.Date;
Expand All @@ -29,6 +30,7 @@ public class SignedDataTokens {
private final String ctx;
private final ArrayList<String> dataTokens;
private final Integer timeToLive;
private final String tokenUri;

private SignedDataTokens(SignedDataTokensBuilder builder) {
this.credentialsFile = builder.credentialsFile;
Expand All @@ -37,14 +39,15 @@ private SignedDataTokens(SignedDataTokensBuilder builder) {
this.ctx = builder.ctx;
this.dataTokens = builder.dataTokens;
this.timeToLive = builder.timeToLive;
this.tokenUri = builder.tokenUri;
}

public static SignedDataTokensBuilder builder() {
return new SignedDataTokensBuilder();
}

private static List<SignedDataTokenResponse> generateSignedTokenFromCredentialsFile(
File credentialsFile, ArrayList<String> dataTokens, Integer timeToLive, String context
File credentialsFile, ArrayList<String> dataTokens, Integer timeToLive, String context, String overrideTokenUri
) throws SkyflowException {
LogUtil.printInfoLog(InfoLogs.GENERATE_SIGNED_TOKENS_FROM_CREDENTIALS_FILE_TRIGGERED.getLog());
List<SignedDataTokenResponse> responseToken;
Expand All @@ -55,7 +58,7 @@ private static List<SignedDataTokenResponse> generateSignedTokenFromCredentialsF
}
FileReader reader = new FileReader(String.valueOf(credentialsFile));
JsonObject serviceAccountCredentials = JsonParser.parseReader(reader).getAsJsonObject();
responseToken = generateSignedTokensFromCredentials(serviceAccountCredentials, dataTokens, timeToLive, context);
responseToken = generateSignedTokensFromCredentials(serviceAccountCredentials, dataTokens, timeToLive, context, overrideTokenUri);
} catch (JsonSyntaxException e) {
LogUtil.printErrorLog(ErrorLogs.INVALID_CREDENTIALS_FILE_FORMAT.getLog());
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), Utils.parameterizedString(
Expand All @@ -69,7 +72,7 @@ private static List<SignedDataTokenResponse> generateSignedTokenFromCredentialsF
}

private static List<SignedDataTokenResponse> generateSignedTokensFromCredentialsString(
String credentials, ArrayList<String> dataTokens, Integer timeToLive, String context
String credentials, ArrayList<String> dataTokens, Integer timeToLive, String context, String overrideTokenUri
) throws SkyflowException {
LogUtil.printInfoLog(InfoLogs.GENERATE_SIGNED_TOKENS_FROM_CREDENTIALS_STRING_TRIGGERED.getLog());
List<SignedDataTokenResponse> responseToken;
Expand All @@ -79,7 +82,7 @@ private static List<SignedDataTokenResponse> generateSignedTokensFromCredentials
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.InvalidCredentials.getMessage());
}
JsonObject serviceAccountCredentials = JsonParser.parseString(credentials).getAsJsonObject();
responseToken = generateSignedTokensFromCredentials(serviceAccountCredentials, dataTokens, timeToLive, context);
responseToken = generateSignedTokensFromCredentials(serviceAccountCredentials, dataTokens, timeToLive, context, overrideTokenUri);
} catch (JsonSyntaxException e) {
LogUtil.printErrorLog(ErrorLogs.INVALID_CREDENTIALS_STRING_FORMAT.getLog());
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(),
Expand All @@ -89,7 +92,7 @@ private static List<SignedDataTokenResponse> generateSignedTokensFromCredentials
}

private static List<SignedDataTokenResponse> generateSignedTokensFromCredentials(
JsonObject credentials, ArrayList<String> dataTokens, Integer timeToLive, String context
JsonObject credentials, ArrayList<String> dataTokens, Integer timeToLive, String context, String overrideTokenUri
) throws SkyflowException {
List<SignedDataTokenResponse> signedDataTokens = null;
try {
Expand All @@ -112,7 +115,7 @@ private static List<SignedDataTokenResponse> generateSignedTokensFromCredentials
}
PrivateKey pvtKey = Utils.getPrivateKeyFromPem(privateKey.getAsString());
signedDataTokens = getSignedToken(
clientID.getAsString(), keyID.getAsString(), pvtKey, dataTokens, timeToLive, context);
clientID.getAsString(), keyID.getAsString(), pvtKey, dataTokens, timeToLive, context, overrideTokenUri);
} catch (RuntimeException e) {
LogUtil.printErrorLog(ErrorLogs.SIGNED_DATA_TOKENS_REJECTED.getLog());
throw new SkyflowException(e);
Expand All @@ -122,7 +125,7 @@ private static List<SignedDataTokenResponse> generateSignedTokensFromCredentials

private static List<SignedDataTokenResponse> getSignedToken(
String clientID, String keyID, PrivateKey pvtKey,
ArrayList<String> dataTokens, Integer timeToLive, String context
ArrayList<String> dataTokens, Integer timeToLive, String context, String overrideTokenUri
) {
final Date createdDate = new Date();
final Date expirationDate;
Expand All @@ -133,6 +136,11 @@ private static List<SignedDataTokenResponse> getSignedToken(
expirationDate = new Date(createdDate.getTime() + 60000); // Valid for 60 seconds
}

String finalTokenUri = null;
if (overrideTokenUri != null && !overrideTokenUri.isEmpty()) {
finalTokenUri = overrideTokenUri;
}

List<SignedDataTokenResponse> list = new ArrayList<>();
for (String dataToken : dataTokens) {
String eachSignedDataToken = Jwts.builder()
Expand All @@ -142,6 +150,7 @@ private static List<SignedDataTokenResponse> getSignedToken(
.claim("sub", clientID)
.claim("ctx", context)
.claim("tok", dataToken)
.claim("aud", finalTokenUri)
.expiration(expirationDate)
.signWith(pvtKey, Jwts.SIG.RS256)
.compact();
Expand All @@ -155,9 +164,9 @@ public synchronized List<SignedDataTokenResponse> getSignedDataTokens() throws S
LogUtil.printInfoLog(InfoLogs.GET_SIGNED_DATA_TOKENS_TRIGGERED.getLog());
List<SignedDataTokenResponse> signedToken = new ArrayList<>();
if (this.credentialsFile != null && Objects.equals(this.credentialsType, "FILE")) {
signedToken = generateSignedTokenFromCredentialsFile(this.credentialsFile, this.dataTokens, this.timeToLive, this.ctx);
signedToken = generateSignedTokenFromCredentialsFile(this.credentialsFile, this.dataTokens, this.timeToLive, this.ctx, this.tokenUri);
} else if (this.credentialsString != null && Objects.equals(this.credentialsType, "STRING")) {
signedToken = generateSignedTokensFromCredentialsString(this.credentialsString, this.dataTokens, this.timeToLive, this.ctx);
signedToken = generateSignedTokensFromCredentialsString(this.credentialsString, this.dataTokens, this.timeToLive, this.ctx, this.tokenUri);
}
LogUtil.printInfoLog(InfoLogs.GET_SIGNED_DATA_TOKEN_SUCCESS.getLog());
return signedToken;
Expand All @@ -170,6 +179,7 @@ public static class SignedDataTokensBuilder {
private String credentialsString;
private String ctx;
private String credentialsType;
private String tokenUri;

private SignedDataTokensBuilder() {
}
Expand Down Expand Up @@ -205,6 +215,19 @@ public SignedDataTokensBuilder setTimeToLive(Integer timeToLive) {
return this;
}

public SignedDataTokensBuilder setTokenUri(String tokenUri) throws SkyflowException {
if (tokenUri != null && !tokenUri.isEmpty()) {
try {
new java.net.URL(tokenUri);
this.tokenUri = tokenUri;
} catch (MalformedURLException e) {
LogUtil.printErrorLog(ErrorLogs.INVALID_TOKEN_URI.getLog());
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.InvalidTokenUri.getMessage());
}
}
return this;
}

public SignedDataTokens build() {
return new SignedDataTokens(this);
}
Expand Down
20 changes: 12 additions & 8 deletions src/main/java/com/skyflow/utils/Utils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,19 +49,23 @@ public static String getVaultURL(String clusterId, Env env) {

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();
.setCtx(credentials.getContext());
if (credentials.getTokenUri() != null) {
builder.setTokenUri(credentials.getTokenUri());
}
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();
.setCtx(credentials.getContext());
if (credentials.getTokenUri() != null) {
builder.setTokenUri(credentials.getTokenUri());
}
return builder.build().getBearerToken();
} else {
return credentials.getToken();
}
Expand Down
6 changes: 6 additions & 0 deletions src/main/java/com/skyflow/utils/validations/Validations.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,7 @@ public static void validateCredentials(Credentials credentials) throws SkyflowEx
String token = credentials.getToken();
String apiKey = credentials.getApiKey();
String context = credentials.getContext();
String tokenUri = credentials.getTokenUri();
ArrayList<String> roles = credentials.getRoles();

if (path != null) nonNullMembers++;
Expand Down Expand Up @@ -212,6 +213,11 @@ public static void validateCredentials(Credentials credentials) throws SkyflowEx
LogUtil.printErrorLog(ErrorLogs.EMPTY_OR_NULL_CONTEXT.getLog());
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.EmptyContext.getMessage());
}

if (tokenUri != null && isInvalidURL(tokenUri)) {
LogUtil.printErrorLog(ErrorLogs.INVALID_TOKEN_URI.getLog());
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.InvalidTokenUri.getMessage());
}
}

public static void validateDetokenizeRequest(DetokenizeRequest detokenizeRequest) throws SkyflowException {
Expand Down
Loading
Loading