Skip to content

Commit 2aa59d8

Browse files
authored
Add JWT retryable exception type (#107)
1 parent c61af09 commit 2aa59d8

File tree

7 files changed

+147
-68
lines changed

7 files changed

+147
-68
lines changed
Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
package io.scalecube.security.environment;
2+
3+
import java.util.HashMap;
4+
import java.util.Map;
5+
import java.util.function.Supplier;
6+
import org.junit.jupiter.api.extension.BeforeAllCallback;
7+
import org.junit.jupiter.api.extension.ExtensionContext;
8+
import org.junit.jupiter.api.extension.ExtensionContext.Namespace;
9+
import org.junit.jupiter.api.extension.ParameterContext;
10+
import org.junit.jupiter.api.extension.ParameterResolutionException;
11+
import org.junit.jupiter.api.extension.ParameterResolver;
12+
13+
public class IntegrationEnvironmentFixture
14+
implements BeforeAllCallback, ExtensionContext.Store.CloseableResource, ParameterResolver {
15+
16+
private static final Map<Class<?>, Supplier<?>> PARAMETERS_TO_RESOLVE = new HashMap<>();
17+
18+
private static VaultEnvironment vaultEnvironment;
19+
20+
@Override
21+
public void beforeAll(ExtensionContext context) {
22+
context
23+
.getRoot()
24+
.getStore(Namespace.GLOBAL)
25+
.getOrComputeIfAbsent(
26+
this.getClass(),
27+
key -> {
28+
vaultEnvironment = VaultEnvironment.start();
29+
return this;
30+
});
31+
32+
PARAMETERS_TO_RESOLVE.put(VaultEnvironment.class, () -> vaultEnvironment);
33+
}
34+
35+
@Override
36+
public void close() {
37+
if (vaultEnvironment != null) {
38+
vaultEnvironment.close();
39+
}
40+
}
41+
42+
@Override
43+
public boolean supportsParameter(
44+
ParameterContext parameterContext, ExtensionContext extensionContext)
45+
throws ParameterResolutionException {
46+
Class<?> type = parameterContext.getParameter().getType();
47+
return PARAMETERS_TO_RESOLVE.keySet().stream().anyMatch(type::isAssignableFrom);
48+
}
49+
50+
@Override
51+
public Object resolveParameter(
52+
ParameterContext parameterContext, ExtensionContext extensionContext)
53+
throws ParameterResolutionException {
54+
Class<?> type = parameterContext.getParameter().getType();
55+
return PARAMETERS_TO_RESOLVE.get(type).get();
56+
}
57+
}

tests/src/test/java/io/scalecube/security/environment/VaultEnvironment.java

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,13 @@ public static Throwable getRootCause(Throwable throwable) {
211211
return throwable;
212212
}
213213

214+
public String newServiceToken() {
215+
String keyName = createIdentityKey(); // oidc/key
216+
String roleName = createIdentityRole(keyName); // oidc/role
217+
String clientToken = login(); // onboard entity with policy
218+
return generateIdentityToken(clientToken, roleName);
219+
}
220+
214221
@Override
215222
public void close() {
216223
vault.stop();
Lines changed: 29 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package io.scalecube.security.tokens.jwt;
22

33
import static io.scalecube.security.environment.VaultEnvironment.getRootCause;
4+
import static org.hamcrest.CoreMatchers.instanceOf;
5+
import static org.hamcrest.MatcherAssert.assertThat;
6+
import static org.hamcrest.core.StringStartsWith.startsWith;
47
import static org.junit.jupiter.api.Assertions.assertNotNull;
58
import static org.junit.jupiter.api.Assertions.assertTrue;
69
import static org.junit.jupiter.api.Assertions.fail;
@@ -9,34 +12,20 @@
912
import static org.mockito.Mockito.when;
1013

1114
import io.jsonwebtoken.Locator;
15+
import io.scalecube.security.environment.IntegrationEnvironmentFixture;
1216
import io.scalecube.security.environment.VaultEnvironment;
1317
import java.security.Key;
1418
import java.time.Duration;
1519
import java.util.concurrent.TimeUnit;
16-
import org.junit.jupiter.api.AfterAll;
17-
import org.junit.jupiter.api.Assertions;
18-
import org.junit.jupiter.api.BeforeAll;
1920
import org.junit.jupiter.api.Test;
21+
import org.junit.jupiter.api.extension.ExtendWith;
2022

23+
@ExtendWith(IntegrationEnvironmentFixture.class)
2124
public class JsonwebtokenResolverTests {
2225

23-
private static VaultEnvironment vaultEnvironment;
24-
25-
@BeforeAll
26-
static void beforeAll() {
27-
vaultEnvironment = VaultEnvironment.start();
28-
}
29-
30-
@AfterAll
31-
static void afterAll() {
32-
if (vaultEnvironment != null) {
33-
vaultEnvironment.close();
34-
}
35-
}
36-
3726
@Test
38-
void testResolveTokenSuccessfully() throws Exception {
39-
final var token = generateToken();
27+
void testResolveTokenSuccessfully(VaultEnvironment vaultEnvironment) throws Exception {
28+
final var token = vaultEnvironment.newServiceToken();
4029

4130
final var jwtToken =
4231
new JsonwebtokenResolver(
@@ -50,13 +39,13 @@ void testResolveTokenSuccessfully() throws Exception {
5039
.get(3, TimeUnit.SECONDS);
5140

5241
assertNotNull(jwtToken, "jwtToken");
53-
Assertions.assertTrue(jwtToken.header().size() > 0, "jwtToken.header: " + jwtToken.header());
54-
Assertions.assertTrue(jwtToken.payload().size() > 0, "jwtToken.payload: " + jwtToken.payload());
42+
assertTrue(jwtToken.header().size() > 0, "jwtToken.header: " + jwtToken.header());
43+
assertTrue(jwtToken.payload().size() > 0, "jwtToken.payload: " + jwtToken.payload());
5544
}
5645

5746
@Test
58-
void testJwksKeyLocatorThrowsError() {
59-
final var token = generateToken();
47+
void testJwksKeyLocatorThrowsError(VaultEnvironment vaultEnvironment) {
48+
final var token = vaultEnvironment.newServiceToken();
6049

6150
Locator<Key> keyLocator = mock(Locator.class);
6251
when(keyLocator.locate(any())).thenThrow(new RuntimeException("Cannot get key"));
@@ -66,16 +55,25 @@ void testJwksKeyLocatorThrowsError() {
6655
fail("Expected exception");
6756
} catch (Exception e) {
6857
final var ex = getRootCause(e);
69-
assertNotNull(ex);
70-
assertNotNull(ex.getMessage());
71-
assertTrue(ex.getMessage().startsWith("Cannot get key"), "Exception: " + ex);
58+
assertThat(ex, instanceOf(RuntimeException.class));
59+
assertThat(ex.getMessage(), startsWith("Cannot get key"));
7260
}
7361
}
7462

75-
private static String generateToken() {
76-
String keyName = vaultEnvironment.createIdentityKey(); // oidc/key
77-
String roleName = vaultEnvironment.createIdentityRole(keyName); // oidc/role
78-
String clientToken = vaultEnvironment.login(); // onboard entity with policy
79-
return vaultEnvironment.generateIdentityToken(clientToken, roleName);
63+
@Test
64+
void testJwksKeyLocatorThrowsRetryableError(VaultEnvironment vaultEnvironment) {
65+
final var token = vaultEnvironment.newServiceToken();
66+
67+
Locator<Key> keyLocator = mock(Locator.class);
68+
when(keyLocator.locate(any())).thenThrow(new JwtUnavailableException("JWKS timeout"));
69+
70+
try {
71+
new JsonwebtokenResolver(keyLocator).resolve(token).get(3, TimeUnit.SECONDS);
72+
fail("Expected exception");
73+
} catch (Exception e) {
74+
final var ex = getRootCause(e);
75+
assertThat(ex, instanceOf(JwtUnavailableException.class));
76+
assertThat(ex.getMessage(), startsWith("JWKS timeout"));
77+
}
8078
}
8179
}

tests/src/test/java/io/scalecube/security/vault/VaultServiceTokenTests.java

Lines changed: 15 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,14 @@
22

33
import static io.scalecube.security.environment.VaultEnvironment.getRootCause;
44
import static java.util.concurrent.CompletableFuture.completedFuture;
5+
import static org.hamcrest.MatcherAssert.assertThat;
6+
import static org.hamcrest.core.StringStartsWith.startsWith;
57
import static org.junit.jupiter.api.Assertions.assertNotNull;
68
import static org.junit.jupiter.api.Assertions.assertTrue;
79
import static org.junit.jupiter.api.Assertions.fail;
810
import static org.testcontainers.shaded.org.apache.commons.lang3.RandomStringUtils.randomAlphabetic;
911

12+
import io.scalecube.security.environment.IntegrationEnvironmentFixture;
1013
import io.scalecube.security.environment.VaultEnvironment;
1114
import io.scalecube.security.tokens.jwt.JsonwebtokenResolver;
1215
import io.scalecube.security.tokens.jwt.JwksKeyLocator;
@@ -17,29 +20,15 @@
1720
import java.util.Map;
1821
import java.util.concurrent.ExecutionException;
1922
import java.util.concurrent.TimeUnit;
20-
import org.junit.jupiter.api.AfterAll;
21-
import org.junit.jupiter.api.Assertions;
22-
import org.junit.jupiter.api.BeforeAll;
2323
import org.junit.jupiter.api.Test;
24+
import org.junit.jupiter.api.extension.ExtendWith;
2425

26+
@ExtendWith(IntegrationEnvironmentFixture.class)
2527
public class VaultServiceTokenTests {
2628

27-
private static VaultEnvironment vaultEnvironment;
28-
29-
@BeforeAll
30-
static void beforeAll() {
31-
vaultEnvironment = VaultEnvironment.start();
32-
}
33-
34-
@AfterAll
35-
static void afterAll() {
36-
if (vaultEnvironment != null) {
37-
vaultEnvironment.close();
38-
}
39-
}
40-
4129
@Test
42-
void testGetServiceTokenUsingWrongCredentials() throws Exception {
30+
void testGetServiceTokenUsingWrongCredentials(VaultEnvironment vaultEnvironment)
31+
throws Exception {
4332
final var serviceTokenSupplier =
4433
new VaultServiceTokenSupplier.Builder()
4534
.vaultAddress(vaultEnvironment.vaultAddr())
@@ -54,14 +43,12 @@ void testGetServiceTokenUsingWrongCredentials() throws Exception {
5443
} catch (ExecutionException e) {
5544
final var ex = getRootCause(e);
5645
assertNotNull(ex);
57-
assertNotNull(ex.getMessage());
58-
assertTrue(
59-
ex.getMessage().contains("Failed to get service token, status=403"), "Exception: " + ex);
46+
assertThat(ex.getMessage(), startsWith("Failed to get service token, status=403"));
6047
}
6148
}
6249

6350
@Test
64-
void testGetNonExistingServiceToken() throws Exception {
51+
void testGetNonExistingServiceToken(VaultEnvironment vaultEnvironment) throws Exception {
6552
final var nonExistingServiceRole = "non-existing-role-" + System.currentTimeMillis();
6653

6754
final var serviceTokenSupplier =
@@ -78,14 +65,12 @@ void testGetNonExistingServiceToken() throws Exception {
7865
} catch (ExecutionException e) {
7966
final var ex = getRootCause(e);
8067
assertNotNull(ex);
81-
assertNotNull(ex.getMessage());
82-
assertTrue(
83-
ex.getMessage().contains("Failed to get service token, status=400"), "Exception: " + ex);
68+
assertThat(ex.getMessage(), startsWith("Failed to get service token, status=400"));
8469
}
8570
}
8671

8772
@Test
88-
void testGetServiceTokenByWrongServiceRole() throws Exception {
73+
void testGetServiceTokenByWrongServiceRole(VaultEnvironment vaultEnvironment) throws Exception {
8974
final var now = System.currentTimeMillis();
9075
final var serviceRole1 = "role1-" + now;
9176
final var serviceRole2 = "role2-" + now;
@@ -122,14 +107,12 @@ void testGetServiceTokenByWrongServiceRole() throws Exception {
122107
} catch (ExecutionException e) {
123108
final var ex = getRootCause(e);
124109
assertNotNull(ex);
125-
assertNotNull(ex.getMessage());
126-
assertTrue(
127-
ex.getMessage().contains("Failed to get service token, status=400"), "Exception: " + ex);
110+
assertThat(ex.getMessage(), startsWith("Failed to get service token, status=400"));
128111
}
129112
}
130113

131114
@Test
132-
void testGetServiceTokenSuccessfully() throws Exception {
115+
void testGetServiceTokenSuccessfully(VaultEnvironment vaultEnvironment) throws Exception {
133116
final var now = System.currentTimeMillis();
134117
final var serviceRole = "role-" + now;
135118
final var tags = Map.of("type", "ops", "ns", "develop");
@@ -164,8 +147,8 @@ void testGetServiceTokenSuccessfully() throws Exception {
164147
.get(3, TimeUnit.SECONDS);
165148

166149
assertNotNull(jwtToken, "jwtToken");
167-
Assertions.assertTrue(jwtToken.header().size() > 0, "jwtToken.header: " + jwtToken.header());
168-
Assertions.assertTrue(jwtToken.payload().size() > 0, "jwtToken.payload: " + jwtToken.payload());
150+
assertTrue(jwtToken.header().size() > 0, "jwtToken.header: " + jwtToken.header());
151+
assertTrue(jwtToken.payload().size() > 0, "jwtToken.payload: " + jwtToken.payload());
169152
}
170153

171154
private static String toQualifiedName(String role, Map<String, String> tags) {

tokens/src/main/java/io/scalecube/security/tokens/jwt/JwksKeyLocator.java

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import java.net.http.HttpRequest;
1818
import java.net.http.HttpResponse;
1919
import java.net.http.HttpResponse.BodyHandlers;
20+
import java.net.http.HttpTimeoutException;
2021
import java.security.Key;
2122
import java.security.KeyFactory;
2223
import java.security.PublicKey;
@@ -55,13 +56,11 @@ protected Key locate(JwsHeader header) {
5556
kid -> {
5657
final var key = findKeyById(computeKeyList(), kid);
5758
if (key == null) {
58-
throw new RuntimeException("Cannot find key by kid: " + kid);
59+
throw new JwtUnavailableException("Cannot find key by kid: " + kid);
5960
}
6061
return new CachedKey(key, System.currentTimeMillis() + keyTtl);
6162
})
6263
.key();
63-
} catch (Exception ex) {
64-
throw new JwtTokenException(ex);
6564
} finally {
6665
tryCleanup();
6766
}
@@ -77,8 +76,13 @@ private JwkInfoList computeKeyList() {
7776
.send(
7877
HttpRequest.newBuilder(jwksUri).GET().timeout(requestTimeout).build(),
7978
BodyHandlers.ofInputStream());
80-
} catch (Exception e) {
81-
throw new RuntimeException("Failed to retrive jwk keys", e);
79+
} catch (HttpTimeoutException e) {
80+
throw new JwtUnavailableException("Failed to retrive jwk keys", e);
81+
} catch (IOException e) {
82+
throw new RuntimeException(e);
83+
} catch (InterruptedException e) {
84+
Thread.currentThread().interrupt();
85+
throw new RuntimeException(e);
8286
}
8387

8488
final var statusCode = httpResponse.statusCode();

tokens/src/main/java/io/scalecube/security/tokens/jwt/JwtTokenException.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@
22

33
import java.util.StringJoiner;
44

5+
/**
6+
* Generic exception type for JWT token resolution errors. Used as part {@link JwtTokenResolver}
7+
* mechanism, and responsible to abstract token resolution problems.
8+
*/
59
public class JwtTokenException extends RuntimeException {
610

711
public JwtTokenException(String message) {
Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
package io.scalecube.security.tokens.jwt;
2+
3+
/**
4+
* Special JWT exception type indicating transient error during token resolution. For example such
5+
* transient errors are:
6+
*
7+
* <ul>
8+
* <li>Key Rotation: JWKS endpoints often implement key rotation policies where keys are
9+
* periodically changed for security reasons. If the JWT was issued with a "kid" that
10+
* corresponds to a key that has since been rotated out, that key won't be available in the
11+
* JWKS anymore.
12+
* <li>Network or Server Issues: if the JWKS URI is temporarily down, inaccessible, or
13+
* experiencing issues, cleint might not be able to retrieve the keys, or the list of keys
14+
* might be incomplete or outdated.
15+
* </ul>
16+
*/
17+
public class JwtUnavailableException extends JwtTokenException {
18+
19+
public JwtUnavailableException(String message) {
20+
super(message);
21+
}
22+
23+
public JwtUnavailableException(String message, Throwable cause) {
24+
super(message, cause);
25+
}
26+
}

0 commit comments

Comments
 (0)