Skip to content

Commit a79cb21

Browse files
saileshwar-skyflowsaileshwar-skyflow
authored
SK-2505: Add support for custom tokenUri (#271)
* SK-2505: add support for custom token uri
1 parent 6c28b9a commit a79cb21

File tree

10 files changed

+1275
-28
lines changed

10 files changed

+1275
-28
lines changed

‎src/main/java/com/skyflow/config/Credentials.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,13 @@ public class Credentials {
99
private String credentialsString;
1010
private String token;
1111
private String apiKey;
12+
private String tokenUri;
1213

1314
public Credentials() {
1415
this.path = null;
1516
this.context = null;
1617
this.credentialsString = null;
18+
this.tokenUri = null;
1719
}
1820

1921
public String getPath() {
@@ -63,4 +65,12 @@ public String getApiKey() {
6365
public void setApiKey(String apiKey) {
6466
this.apiKey = apiKey;
6567
}
68+
69+
public String getTokenUri() {
70+
return tokenUri;
71+
}
72+
73+
public void setTokenUri(String tokenUri) {
74+
this.tokenUri = tokenUri;
75+
}
6676
}

‎src/main/java/com/skyflow/errors/ErrorMessage.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ public enum ErrorMessage {
4444
MissingClientId("%s0 Initialization failed. Unable to read client ID in credentials. Verify your client ID."),
4545
MissingKeyId("%s0 Initialization failed. Unable to read key ID in credentials. Verify your key ID."),
4646
MissingTokenUri("%s0 Initialization failed. Unable to read token URI in credentials. Verify your token URI."),
47-
InvalidTokenUri("%s0 Initialization failed. Token URI in not a valid URL in credentials. Verify your token URI."),
47+
InvalidTokenUri("%s0 Initialization failed. Invalid Skyflow credentials. The token URI must be a string and a valid URL."),
4848
JwtInvalidFormat("%s0 Initialization failed. Invalid private key format. Verify your credentials."),
4949
InvalidAlgorithm("%s0 Initialization failed. Invalid algorithm to parse private key. Specify valid algorithm."),
5050
InvalidKeySpec("%s0 Initialization failed. Unable to parse RSA private key. Verify your credentials."),

‎src/main/java/com/skyflow/serviceaccount/util/BearerToken.java‎

Lines changed: 28 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -34,21 +34,23 @@ public class BearerToken {
3434
private final String ctx;
3535
private final ArrayList<String> roles;
3636
private final String credentialsType;
37+
private final String tokenUri;
3738

3839
private BearerToken(BearerTokenBuilder builder) {
3940
this.credentialsFile = builder.credentialsFile;
4041
this.credentialsString = builder.credentialsString;
4142
this.ctx = builder.ctx;
4243
this.roles = builder.roles;
4344
this.credentialsType = builder.credentialsType;
45+
this.tokenUri = builder.tokenUri;
4446
}
4547

4648
public static BearerTokenBuilder builder() {
4749
return new BearerTokenBuilder();
4850
}
4951

5052
private static V1GetAuthTokenResponse generateBearerTokenFromCredentials(
51-
File credentialsFile, String context, ArrayList<String> roles
53+
File credentialsFile, String context, ArrayList<String> roles, String overrideTokenUri
5254
) throws SkyflowException {
5355
LogUtil.printInfoLog(InfoLogs.GENERATE_BEARER_TOKEN_FROM_CREDENTIALS_TRIGGERED.getLog());
5456
try {
@@ -58,7 +60,7 @@ private static V1GetAuthTokenResponse generateBearerTokenFromCredentials(
5860
}
5961
FileReader reader = new FileReader(String.valueOf(credentialsFile));
6062
JsonObject serviceAccountCredentials = JsonParser.parseReader(reader).getAsJsonObject();
61-
return getBearerTokenFromCredentials(serviceAccountCredentials, context, roles);
63+
return getBearerTokenFromCredentials(serviceAccountCredentials, context, roles, overrideTokenUri);
6264
} catch (JsonSyntaxException e) {
6365
LogUtil.printErrorLog(ErrorLogs.INVALID_CREDENTIALS_FILE_FORMAT.getLog());
6466
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), Utils.parameterizedString(
@@ -71,7 +73,7 @@ private static V1GetAuthTokenResponse generateBearerTokenFromCredentials(
7173
}
7274

7375
private static V1GetAuthTokenResponse generateBearerTokenFromCredentialString(
74-
String credentials, String context, ArrayList<String> roles
76+
String credentials, String context, ArrayList<String> roles, String overrideTokenUri
7577
) throws SkyflowException {
7678
LogUtil.printInfoLog(InfoLogs.GENERATE_BEARER_TOKEN_FROM_CREDENTIALS_STRING_TRIGGERED.getLog());
7779
try {
@@ -80,7 +82,7 @@ private static V1GetAuthTokenResponse generateBearerTokenFromCredentialString(
8082
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.InvalidCredentials.getMessage());
8183
}
8284
JsonObject serviceAccountCredentials = JsonParser.parseString(credentials).getAsJsonObject();
83-
return getBearerTokenFromCredentials(serviceAccountCredentials, context, roles);
85+
return getBearerTokenFromCredentials(serviceAccountCredentials, context, roles, overrideTokenUri);
8486
} catch (JsonSyntaxException e) {
8587
LogUtil.printErrorLog(ErrorLogs.INVALID_CREDENTIALS_STRING_FORMAT.getLog());
8688
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(),
@@ -89,7 +91,7 @@ private static V1GetAuthTokenResponse generateBearerTokenFromCredentialString(
8991
}
9092

9193
private static V1GetAuthTokenResponse getBearerTokenFromCredentials(
92-
JsonObject credentials, String context, ArrayList<String> roles
94+
JsonObject credentials, String context, ArrayList<String> roles, String overrideTokenUri
9395
) throws SkyflowException {
9496
try {
9597
JsonElement privateKey = credentials.get("privateKey");
@@ -111,17 +113,19 @@ private static V1GetAuthTokenResponse getBearerTokenFromCredentials(
111113
}
112114

113115
JsonElement tokenURI = credentials.get("tokenURI");
114-
if (tokenURI == null) {
116+
if (tokenURI == null && overrideTokenUri == null) {
115117
LogUtil.printErrorLog(ErrorLogs.TOKEN_URI_IS_REQUIRED.getLog());
116118
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.MissingTokenUri.getMessage());
117119
}
118120

121+
String finalTokenUri = (overrideTokenUri != null) ? overrideTokenUri : tokenURI.getAsString();
122+
119123
PrivateKey pvtKey = Utils.getPrivateKeyFromPem(privateKey.getAsString());
120124
String signedUserJWT = getSignedToken(
121-
clientID.getAsString(), keyID.getAsString(), tokenURI.getAsString(), pvtKey, context
125+
clientID.getAsString(), keyID.getAsString(), finalTokenUri, pvtKey, context
122126
);
123127

124-
String basePath = Utils.getBaseURL(tokenURI.getAsString());
128+
String basePath = Utils.getBaseURL(finalTokenUri);
125129
API_CLIENT_BUILDER.url(basePath);
126130
ApiClient apiClient = API_CLIENT_BUILDER.token("token").build();
127131
AuthenticationClient authenticationApi = apiClient.authentication();
@@ -174,10 +178,10 @@ public synchronized String getBearerToken() throws SkyflowException {
174178
V1GetAuthTokenResponse response;
175179
String accessToken = null;
176180
if (this.credentialsFile != null && Objects.equals(this.credentialsType, "FILE")) {
177-
response = generateBearerTokenFromCredentials(this.credentialsFile, this.ctx, this.roles);
181+
response = generateBearerTokenFromCredentials(this.credentialsFile, this.ctx, this.roles, this.tokenUri);
178182
accessToken = response.getAccessToken().get();
179183
} else if (this.credentialsString != null && Objects.equals(this.credentialsType, "STRING")) {
180-
response = generateBearerTokenFromCredentialString(this.credentialsString, this.ctx, this.roles);
184+
response = generateBearerTokenFromCredentialString(this.credentialsString, this.ctx, this.roles, this.tokenUri);
181185
accessToken = response.getAccessToken().get();
182186
}
183187
LogUtil.printInfoLog(InfoLogs.GET_BEARER_TOKEN_SUCCESS.getLog());
@@ -191,6 +195,7 @@ public static class BearerTokenBuilder {
191195
private String ctx;
192196
private ArrayList<String> roles;
193197
private String credentialsType;
198+
private String tokenUri;
194199

195200
private BearerTokenBuilder() {
196201
}
@@ -221,6 +226,19 @@ public BearerTokenBuilder setRoles(ArrayList<String> roles) {
221226
return this;
222227
}
223228

229+
public BearerTokenBuilder setTokenUri(String tokenUri) throws SkyflowException {
230+
if (tokenUri != null && !tokenUri.isEmpty()) {
231+
try {
232+
new java.net.URL(tokenUri);
233+
this.tokenUri = tokenUri;
234+
} catch (MalformedURLException e) {
235+
LogUtil.printErrorLog(ErrorLogs.INVALID_TOKEN_URI.getLog());
236+
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.InvalidTokenUri.getMessage());
237+
}
238+
}
239+
return this;
240+
}
241+
224242
public BearerToken build() {
225243
return new BearerToken(this);
226244
}

‎src/main/java/com/skyflow/serviceaccount/util/SignedDataTokens.java‎

Lines changed: 32 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import java.io.File;
1717
import java.io.FileNotFoundException;
1818
import java.io.FileReader;
19+
import java.net.MalformedURLException;
1920
import java.security.PrivateKey;
2021
import java.util.ArrayList;
2122
import java.util.Date;
@@ -29,6 +30,7 @@ public class SignedDataTokens {
2930
private final String ctx;
3031
private final ArrayList<String> dataTokens;
3132
private final Integer timeToLive;
33+
private final String tokenUri;
3234

3335
private SignedDataTokens(SignedDataTokensBuilder builder) {
3436
this.credentialsFile = builder.credentialsFile;
@@ -37,14 +39,15 @@ private SignedDataTokens(SignedDataTokensBuilder builder) {
3739
this.ctx = builder.ctx;
3840
this.dataTokens = builder.dataTokens;
3941
this.timeToLive = builder.timeToLive;
42+
this.tokenUri = builder.tokenUri;
4043
}
4144

4245
public static SignedDataTokensBuilder builder() {
4346
return new SignedDataTokensBuilder();
4447
}
4548

4649
private static List<SignedDataTokenResponse> generateSignedTokenFromCredentialsFile(
47-
File credentialsFile, ArrayList<String> dataTokens, Integer timeToLive, String context
50+
File credentialsFile, ArrayList<String> dataTokens, Integer timeToLive, String context, String overrideTokenUri
4851
) throws SkyflowException {
4952
LogUtil.printInfoLog(InfoLogs.GENERATE_SIGNED_TOKENS_FROM_CREDENTIALS_FILE_TRIGGERED.getLog());
5053
List<SignedDataTokenResponse> responseToken;
@@ -55,7 +58,7 @@ private static List<SignedDataTokenResponse> generateSignedTokenFromCredentialsF
5558
}
5659
FileReader reader = new FileReader(String.valueOf(credentialsFile));
5760
JsonObject serviceAccountCredentials = JsonParser.parseReader(reader).getAsJsonObject();
58-
responseToken = generateSignedTokensFromCredentials(serviceAccountCredentials, dataTokens, timeToLive, context);
61+
responseToken = generateSignedTokensFromCredentials(serviceAccountCredentials, dataTokens, timeToLive, context, overrideTokenUri);
5962
} catch (JsonSyntaxException e) {
6063
LogUtil.printErrorLog(ErrorLogs.INVALID_CREDENTIALS_FILE_FORMAT.getLog());
6164
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), Utils.parameterizedString(
@@ -69,7 +72,7 @@ private static List<SignedDataTokenResponse> generateSignedTokenFromCredentialsF
6972
}
7073

7174
private static List<SignedDataTokenResponse> generateSignedTokensFromCredentialsString(
72-
String credentials, ArrayList<String> dataTokens, Integer timeToLive, String context
75+
String credentials, ArrayList<String> dataTokens, Integer timeToLive, String context, String overrideTokenUri
7376
) throws SkyflowException {
7477
LogUtil.printInfoLog(InfoLogs.GENERATE_SIGNED_TOKENS_FROM_CREDENTIALS_STRING_TRIGGERED.getLog());
7578
List<SignedDataTokenResponse> responseToken;
@@ -79,7 +82,7 @@ private static List<SignedDataTokenResponse> generateSignedTokensFromCredentials
7982
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.InvalidCredentials.getMessage());
8083
}
8184
JsonObject serviceAccountCredentials = JsonParser.parseString(credentials).getAsJsonObject();
82-
responseToken = generateSignedTokensFromCredentials(serviceAccountCredentials, dataTokens, timeToLive, context);
85+
responseToken = generateSignedTokensFromCredentials(serviceAccountCredentials, dataTokens, timeToLive, context, overrideTokenUri);
8386
} catch (JsonSyntaxException e) {
8487
LogUtil.printErrorLog(ErrorLogs.INVALID_CREDENTIALS_STRING_FORMAT.getLog());
8588
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(),
@@ -89,7 +92,7 @@ private static List<SignedDataTokenResponse> generateSignedTokensFromCredentials
8992
}
9093

9194
private static List<SignedDataTokenResponse> generateSignedTokensFromCredentials(
92-
JsonObject credentials, ArrayList<String> dataTokens, Integer timeToLive, String context
95+
JsonObject credentials, ArrayList<String> dataTokens, Integer timeToLive, String context, String overrideTokenUri
9396
) throws SkyflowException {
9497
List<SignedDataTokenResponse> signedDataTokens = null;
9598
try {
@@ -112,7 +115,7 @@ private static List<SignedDataTokenResponse> generateSignedTokensFromCredentials
112115
}
113116
PrivateKey pvtKey = Utils.getPrivateKeyFromPem(privateKey.getAsString());
114117
signedDataTokens = getSignedToken(
115-
clientID.getAsString(), keyID.getAsString(), pvtKey, dataTokens, timeToLive, context);
118+
clientID.getAsString(), keyID.getAsString(), pvtKey, dataTokens, timeToLive, context, overrideTokenUri);
116119
} catch (RuntimeException e) {
117120
LogUtil.printErrorLog(ErrorLogs.SIGNED_DATA_TOKENS_REJECTED.getLog());
118121
throw new SkyflowException(e);
@@ -122,7 +125,7 @@ private static List<SignedDataTokenResponse> generateSignedTokensFromCredentials
122125

123126
private static List<SignedDataTokenResponse> getSignedToken(
124127
String clientID, String keyID, PrivateKey pvtKey,
125-
ArrayList<String> dataTokens, Integer timeToLive, String context
128+
ArrayList<String> dataTokens, Integer timeToLive, String context, String overrideTokenUri
126129
) {
127130
final Date createdDate = new Date();
128131
final Date expirationDate;
@@ -133,6 +136,11 @@ private static List<SignedDataTokenResponse> getSignedToken(
133136
expirationDate = new Date(createdDate.getTime() + 60000); // Valid for 60 seconds
134137
}
135138

139+
String finalTokenUri = null;
140+
if (overrideTokenUri != null && !overrideTokenUri.isEmpty()) {
141+
finalTokenUri = overrideTokenUri;
142+
}
143+
136144
List<SignedDataTokenResponse> list = new ArrayList<>();
137145
for (String dataToken : dataTokens) {
138146
String eachSignedDataToken = Jwts.builder()
@@ -142,6 +150,7 @@ private static List<SignedDataTokenResponse> getSignedToken(
142150
.claim("sub", clientID)
143151
.claim("ctx", context)
144152
.claim("tok", dataToken)
153+
.claim("aud", finalTokenUri)
145154
.expiration(expirationDate)
146155
.signWith(pvtKey, Jwts.SIG.RS256)
147156
.compact();
@@ -155,9 +164,9 @@ public synchronized List<SignedDataTokenResponse> getSignedDataTokens() throws S
155164
LogUtil.printInfoLog(InfoLogs.GET_SIGNED_DATA_TOKENS_TRIGGERED.getLog());
156165
List<SignedDataTokenResponse> signedToken = new ArrayList<>();
157166
if (this.credentialsFile != null && Objects.equals(this.credentialsType, "FILE")) {
158-
signedToken = generateSignedTokenFromCredentialsFile(this.credentialsFile, this.dataTokens, this.timeToLive, this.ctx);
167+
signedToken = generateSignedTokenFromCredentialsFile(this.credentialsFile, this.dataTokens, this.timeToLive, this.ctx, this.tokenUri);
159168
} else if (this.credentialsString != null && Objects.equals(this.credentialsType, "STRING")) {
160-
signedToken = generateSignedTokensFromCredentialsString(this.credentialsString, this.dataTokens, this.timeToLive, this.ctx);
169+
signedToken = generateSignedTokensFromCredentialsString(this.credentialsString, this.dataTokens, this.timeToLive, this.ctx, this.tokenUri);
161170
}
162171
LogUtil.printInfoLog(InfoLogs.GET_SIGNED_DATA_TOKEN_SUCCESS.getLog());
163172
return signedToken;
@@ -170,6 +179,7 @@ public static class SignedDataTokensBuilder {
170179
private String credentialsString;
171180
private String ctx;
172181
private String credentialsType;
182+
private String tokenUri;
173183

174184
private SignedDataTokensBuilder() {
175185
}
@@ -205,6 +215,19 @@ public SignedDataTokensBuilder setTimeToLive(Integer timeToLive) {
205215
return this;
206216
}
207217

218+
public SignedDataTokensBuilder setTokenUri(String tokenUri) throws SkyflowException {
219+
if (tokenUri != null && !tokenUri.isEmpty()) {
220+
try {
221+
new java.net.URL(tokenUri);
222+
this.tokenUri = tokenUri;
223+
} catch (MalformedURLException e) {
224+
LogUtil.printErrorLog(ErrorLogs.INVALID_TOKEN_URI.getLog());
225+
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.InvalidTokenUri.getMessage());
226+
}
227+
}
228+
return this;
229+
}
230+
208231
public SignedDataTokens build() {
209232
return new SignedDataTokens(this);
210233
}

‎src/main/java/com/skyflow/utils/Utils.java‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -49,19 +49,23 @@ public static String getVaultURL(String clusterId, Env env) {
4949

5050
public static String generateBearerToken(Credentials credentials) throws SkyflowException {
5151
if (credentials.getPath() != null) {
52-
return BearerToken.builder()
52+
BearerToken.BearerTokenBuilder builder = BearerToken.builder()
5353
.setCredentials(new File(credentials.getPath()))
5454
.setRoles(credentials.getRoles())
55-
.setCtx(credentials.getContext())
56-
.build()
57-
.getBearerToken();
55+
.setCtx(credentials.getContext());
56+
if (credentials.getTokenUri() != null) {
57+
builder.setTokenUri(credentials.getTokenUri());
58+
}
59+
return builder.build().getBearerToken();
5860
} else if (credentials.getCredentialsString() != null) {
59-
return BearerToken.builder()
61+
BearerToken.BearerTokenBuilder builder = BearerToken.builder()
6062
.setCredentials(credentials.getCredentialsString())
6163
.setRoles(credentials.getRoles())
62-
.setCtx(credentials.getContext())
63-
.build()
64-
.getBearerToken();
64+
.setCtx(credentials.getContext());
65+
if (credentials.getTokenUri() != null) {
66+
builder.setTokenUri(credentials.getTokenUri());
67+
}
68+
return builder.build().getBearerToken();
6569
} else {
6670
return credentials.getToken();
6771
}

‎src/main/java/com/skyflow/utils/validations/Validations.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,7 @@ public static void validateCredentials(Credentials credentials) throws SkyflowEx
154154
String token = credentials.getToken();
155155
String apiKey = credentials.getApiKey();
156156
String context = credentials.getContext();
157+
String tokenUri = credentials.getTokenUri();
157158
ArrayList<String> roles = credentials.getRoles();
158159

159160
if (path != null) nonNullMembers++;
@@ -212,6 +213,11 @@ public static void validateCredentials(Credentials credentials) throws SkyflowEx
212213
LogUtil.printErrorLog(ErrorLogs.EMPTY_OR_NULL_CONTEXT.getLog());
213214
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.EmptyContext.getMessage());
214215
}
216+
217+
if (tokenUri != null && isInvalidURL(tokenUri)) {
218+
LogUtil.printErrorLog(ErrorLogs.INVALID_TOKEN_URI.getLog());
219+
throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.InvalidTokenUri.getMessage());
220+
}
215221
}
216222

217223
public static void validateDetokenizeRequest(DetokenizeRequest detokenizeRequest) throws SkyflowException {

0 commit comments

Comments
 (0)