Skip to content

Commit

Permalink
Fix racing when loading new JWKs from multiple threads
Browse files Browse the repository at this point in the history
The change ensure the mutation of JWKs is done in a single thread and
visible to all other threads, which in turn ensures validation to be
correctly performed concurrently.

Relates: elastic#88023
  • Loading branch information
ywangd committed Jul 25, 2022
1 parent be7c741 commit d91f12b
Showing 1 changed file with 23 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ boolean isEmpty() {
final List<String> allowedJwksAlgsPkc;
final List<String> allowedJwksAlgsHmac;
DelegatedAuthorizationSupport delegatedAuthorizationSupport = null;
ContentAndJwksAlgs contentAndJwksAlgsPkc;
volatile ContentAndJwksAlgs contentAndJwksAlgsPkc;
ContentAndJwksAlgs contentAndJwksAlgsHmac;
final URI jwkSetPathUri;

Expand Down Expand Up @@ -616,14 +616,13 @@ private void validateSignature(
primaryException
);

this.jwkSetLoader.load(ActionListener.wrap(newContentAndJwksAlgs -> {
if (Arrays.equals(this.contentAndJwksAlgsPkc.sha256, newContentAndJwksAlgs.sha256)) {
this.jwkSetLoader.reload(ActionListener.wrap(isUpdated -> {
if (false == isUpdated) {
// No change in JWKSet
logger.debug("Reloaded same PKC JWKs, can't retry verify JWT token=[{}]", tokenPrincipal);
listener.onFailure(primaryException);
return;
}
this.contentAndJwksAlgsPkc = newContentAndJwksAlgs;
// If all PKC JWKs were replaced, all PKC JWT cache entries need to be invalidated.
// Enhancement idea: Use separate caches for PKC vs HMAC JWKs, so only PKC entries get invalidated.
// Enhancement idea: When some JWKs are retained (ex: rotation), only invalidate for removed JWKs.
Expand Down Expand Up @@ -663,11 +662,31 @@ public void usageStats(final ActionListener<Map<String, Object>> listener) {
private class JwkSetLoader {
private final AtomicReference<ListenableFuture<ContentAndJwksAlgs>> reloadFutureRef = new AtomicReference<>();

/**
* Load the JWK sets and pass its content to the specified listener.
*/
void load(final ActionListener<ContentAndJwksAlgs> listener) {
final ListenableFuture<ContentAndJwksAlgs> future = this.getFuture();
future.addListener(listener);
}

/**
* Reload the JWK sets, compare to existing JWK sets and update it to the reloaded value if
* they are different. The listener is called with false if the reloaded content is the same
* as the existing one or true if they are different.
*/
void reload(final ActionListener<Boolean> listener) {
load(ActionListener.wrap(newContentAndJwksAlgs -> {
if (Arrays.equals(JwtRealm.this.contentAndJwksAlgsPkc.sha256, newContentAndJwksAlgs.sha256)) {
// No change in JWKSet
listener.onResponse(false);
} else {
JwtRealm.this.contentAndJwksAlgsPkc = newContentAndJwksAlgs;
listener.onResponse(true);
}
}, listener::onFailure));
}

private ListenableFuture<ContentAndJwksAlgs> getFuture() {
for (;;) {
final ListenableFuture<ContentAndJwksAlgs> existingFuture = this.reloadFutureRef.get();
Expand Down

0 comments on commit d91f12b

Please sign in to comment.