Skip to content

Keep generated RSA-key-pair for JWT token broker on heap #1661

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
merged 2 commits into from
May 28, 2025
Merged

Keep generated RSA-key-pair for JWT token broker on heap #1661

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f8ca476
Keep generated RSA-key-pair for JWT token broker on heap
snazy May 22, 2025
14029a0
update config docs
snazy May 23, 2025
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
14 changes: 5 additions & 9 deletions quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTRSAKeyPairTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@
import com.auth0.jwt.interfaces.DecodedJWT;
import io.quarkus.test.junit.QuarkusTest;
import jakarta.inject.Inject;
import java.nio.file.Path;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import org.apache.polaris.core.PolarisCallContext;
Expand All @@ -39,14 +38,14 @@
import org.apache.polaris.core.persistence.dao.entity.EntityResult;
import org.apache.polaris.core.persistence.dao.entity.PrincipalSecretsResult;
import org.apache.polaris.service.auth.JWTRSAKeyPair;
import org.apache.polaris.service.auth.KeyProvider;
import org.apache.polaris.service.auth.LocalRSAKeyProvider;
import org.apache.polaris.service.auth.PemUtils;
import org.apache.polaris.service.auth.TokenBroker;
import org.apache.polaris.service.auth.TokenRequestValidator;
import org.apache.polaris.service.auth.TokenResponse;
import org.apache.polaris.service.types.TokenType;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import org.mockito.Mockito;

@QuarkusTest
Expand All @@ -55,10 +54,8 @@ public class JWTRSAKeyPairTest {
@Inject protected PolarisConfigurationStore configurationStore;

@Test
public void testSuccessfulTokenGeneration(@TempDir Path tempDir) throws Exception {
Path privateFileLocation = tempDir.resolve("test-private.pem");
Path publicFileLocation = tempDir.resolve("test-public.pem");
PemUtils.generateKeyPair(privateFileLocation, publicFileLocation);
public void testSuccessfulTokenGeneration() throws Exception {
var keyPair = PemUtils.generateKeyPair();

final String clientId = "test-client-id";
final String scope = "PRINCIPAL_ROLE:TEST";
Expand All @@ -82,8 +79,8 @@ public void testSuccessfulTokenGeneration(@TempDir Path tempDir) throws Exceptio
Mockito.when(
metastoreManager.loadEntity(polarisCallContext, 0L, 1L, PolarisEntityType.PRINCIPAL))
.thenReturn(new EntityResult(principal));
TokenBroker tokenBroker =
new JWTRSAKeyPair(metastoreManager, 420, publicFileLocation, privateFileLocation);
KeyProvider provider = new LocalRSAKeyProvider(keyPair);
TokenBroker tokenBroker = new JWTRSAKeyPair(metastoreManager, 420, provider);
TokenResponse token =
tokenBroker.generateFromClientSecrets(
clientId,
Expand All @@ -95,7 +92,6 @@ public void testSuccessfulTokenGeneration(@TempDir Path tempDir) throws Exceptio
assertThat(token).isNotNull();
assertThat(token.getExpiresIn()).isEqualTo(420);

LocalRSAKeyProvider provider = new LocalRSAKeyProvider(publicFileLocation, privateFileLocation);
assertThat(provider.getPrivateKey()).isNotNull();
assertThat(provider.getPublicKey()).isNotNull();
JWTVerifier verifier =
Expand Down
6 changes: 2 additions & 4 deletions service/common/src/main/java/org/apache/polaris/service/auth/JWTRSAKeyPair.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
package org.apache.polaris.service.auth;

import com.auth0.jwt.algorithms.Algorithm;
import java.nio.file.Path;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
Expand All @@ -32,10 +31,9 @@ public class JWTRSAKeyPair extends JWTBroker {
public JWTRSAKeyPair(
PolarisMetaStoreManager metaStoreManager,
int maxTokenGenerationInSeconds,
Path publicKeyFile,
Path privateKeyFile) {
KeyProvider keyProvider) {
super(metaStoreManager, maxTokenGenerationInSeconds);
keyProvider = new LocalRSAKeyProvider(publicKeyFile, privateKeyFile);
this.keyProvider = keyProvider;
}

@Override
Expand Down
33 changes: 14 additions & 19 deletions service/common/src/main/java/org/apache/polaris/service/auth/JWTRSAKeyPairFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,6 @@
import io.smallrye.common.annotation.Identifier;
import jakarta.enterprise.context.ApplicationScoped;
import jakarta.inject.Inject;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.NoSuchAlgorithmException;
import java.time.Duration;
import java.util.concurrent.ConcurrentHashMap;
Expand Down Expand Up @@ -59,28 +56,26 @@ public TokenBroker apply(RealmContext realmContext) {
private JWTRSAKeyPair createTokenBroker(RealmContext realmContext) {
AuthenticationRealmConfiguration config = authenticationConfiguration.forRealm(realmContext);
Duration maxTokenGeneration = config.tokenBroker().maxTokenGeneration();
RSAKeyPairConfiguration keyPairConfiguration =
config.tokenBroker().rsaKeyPair().orElseGet(this::generateKeyPair);
KeyProvider keyProvider =
config
.tokenBroker()
.rsaKeyPair()
.map(this::fileSystemKeyPair)
.orElseGet(this::generateEphemeralKeyPair);
PolarisMetaStoreManager metaStoreManager =
metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext);
return new JWTRSAKeyPair(
metaStoreManager,
(int) maxTokenGeneration.toSeconds(),
keyPairConfiguration.publicKeyFile(),
keyPairConfiguration.privateKeyFile());
return new JWTRSAKeyPair(metaStoreManager, (int) maxTokenGeneration.toSeconds(), keyProvider);
}

private RSAKeyPairConfiguration generateKeyPair() {
private KeyProvider fileSystemKeyPair(RSAKeyPairConfiguration config) {
return LocalRSAKeyProvider.fromFiles(config.publicKeyFile(), config.privateKeyFile());
}

private KeyProvider generateEphemeralKeyPair() {
try {
Path privateFileLocation = Files.createTempFile("polaris-private", ".pem");
Path publicFileLocation = Files.createTempFile("polaris-public", ".pem");
PemUtils.generateKeyPair(privateFileLocation, publicFileLocation);
return new GeneratedKeyPair(privateFileLocation, publicFileLocation);
} catch (IOException | NoSuchAlgorithmException e) {
return new LocalRSAKeyProvider(PemUtils.generateKeyPair());
} catch (NoSuchAlgorithmException e) {
throw new RuntimeException(e);
}
}

private record GeneratedKeyPair(Path privateKeyFile, Path publicKeyFile)
implements RSAKeyPairConfiguration {}
}
61 changes: 39 additions & 22 deletions service/common/src/main/java/org/apache/polaris/service/auth/LocalRSAKeyProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,27 +18,55 @@
*/
package org.apache.polaris.service.auth;

import jakarta.annotation.Nonnull;
import java.io.IOException;
import java.nio.file.Path;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.PublicKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
* Class that can load public / private keys stored on localhost. Meant to be a simple
* implementation for now where a PEM file is loaded off disk.
*/
/** Holds a public / private key pair in memory. */
public class LocalRSAKeyProvider implements KeyProvider {

private static final Logger LOGGER = LoggerFactory.getLogger(LocalRSAKeyProvider.class);

private final Path publicKeyFileLocation;
private final Path privateKeyFileLocation;
private final PublicKey publicKey;
private final PrivateKey privateKey;

public LocalRSAKeyProvider(@Nonnull KeyPair keyPair) {
this(keyPair.getPublic(), keyPair.getPrivate());
}

public LocalRSAKeyProvider(@Nonnull PublicKey publicKey, @Nonnull PrivateKey privateKey) {
this.publicKey = publicKey;
this.privateKey = privateKey;
}

public LocalRSAKeyProvider(Path publicKeyFileLocation, Path privateKeyFileLocation) {
this.publicKeyFileLocation = publicKeyFileLocation;
this.privateKeyFileLocation = privateKeyFileLocation;
public static LocalRSAKeyProvider fromFiles(
@Nonnull Path publicKeyFile, @Nonnull Path privateKeyFile) {
return new LocalRSAKeyProvider(
readPublicKeyFile(publicKeyFile), readPrivateKeyFile(privateKeyFile));
}

private static PrivateKey readPrivateKeyFile(Path privateKeyFileLocation) {
try {
return PemUtils.readPrivateKeyFromFile(privateKeyFileLocation, "RSA");
} catch (IOException e) {
LOGGER.error("Unable to read private key from file {}", privateKeyFileLocation, e);
throw new RuntimeException(
"Unable to read private key from file " + privateKeyFileLocation, e);
}
}

private static PublicKey readPublicKeyFile(Path publicKeyFileLocation) {
try {
return PemUtils.readPublicKeyFromFile(publicKeyFileLocation, "RSA");
} catch (IOException e) {
LOGGER.error("Unable to read public key from file {}", publicKeyFileLocation, e);
throw new RuntimeException("Unable to read public key from file " + publicKeyFileLocation, e);
}
}

/**
Expand All @@ -48,12 +76,7 @@ public LocalRSAKeyProvider(Path publicKeyFileLocation, Path privateKeyFileLocati
*/
@Override
public PublicKey getPublicKey() {
try {
return PemUtils.readPublicKeyFromFile(publicKeyFileLocation, "RSA");
} catch (IOException e) {
LOGGER.error("Unable to read public key from file {}", publicKeyFileLocation, e);
throw new RuntimeException("Unable to read public key from file " + publicKeyFileLocation, e);
}
return publicKey;
Comment on lines 78 to +79
Copy link
Contributor

@dimas-b dimas-b May 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change will require server restarts for key rotation. I'm not sure if on-the-fly key rotation was envisioned as a supported use case before, but it might be worth flagging this behaviour change for visibility.

Copy link
Member Author

@snazy snazy May 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You'd have to keep the previous (public) key around and referenced by their id, as the auth0 RSAKeyProvider interface defines, which Polaris' KeyProvider does not. If you change the keys in a running Polaris system, all current tokens become immediately invalid.

Copy link
Member Author

@snazy snazy May 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Implementing the ability to rotate keys consistently across multiple Polaris instances is a bigger effort, and should be backed by a secrets manager.

Copy link
Contributor

@dimas-b dimas-b May 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not saying the old code was perfect, but I sense that it might have had a purpose 🤔

@collado-mike WDYT?

Copy link
Contributor

@dimas-b dimas-b May 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To be clear: I personally support merging this PR.

}

/**
Expand All @@ -63,12 +86,6 @@ public PublicKey getPublicKey() {
*/
@Override
public PrivateKey getPrivateKey() {
try {
return PemUtils.readPrivateKeyFromFile(privateKeyFileLocation, "RSA");
} catch (IOException e) {
LOGGER.error("Unable to read private key from file {}", privateKeyFileLocation, e);
throw new RuntimeException(
"Unable to read private key from file " + privateKeyFileLocation, e);
}
return privateKey;
}
}
18 changes: 13 additions & 5 deletions service/common/src/main/java/org/apache/polaris/service/auth/PemUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,23 +120,31 @@ public static PrivateKey readPrivateKeyFromFile(Path filepath, String algorithm)
return PemUtils.getPrivateKey(bytes, algorithm);
}

public static void generateKeyPair(Path privateFileLocation, Path publicFileLocation)
throws NoSuchAlgorithmException, IOException {
public static KeyPair generateKeyPair() throws NoSuchAlgorithmException {
KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA");
kpg.initialize(2048);
KeyPair kp = kpg.generateKeyPair();
return kpg.generateKeyPair();
}

public static void generateKeyPairFiles(Path privateFileLocation, Path publicFileLocation)
throws NoSuchAlgorithmException, IOException {
writeKeyPairFiles(generateKeyPair(), privateFileLocation, publicFileLocation);
}

public static void writeKeyPairFiles(
KeyPair keyPair, Path privateFileLocation, Path publicFileLocation) throws IOException {
try (BufferedWriter writer = Files.newBufferedWriter(privateFileLocation, UTF_8)) {
writer.write("-----BEGIN PRIVATE KEY-----");
writer.newLine();
writer.write(Base64.getMimeEncoder().encodeToString(kp.getPrivate().getEncoded()));
writer.write(Base64.getMimeEncoder().encodeToString(keyPair.getPrivate().getEncoded()));
writer.newLine();
writer.write("-----END PRIVATE KEY-----");
writer.newLine();
}
try (BufferedWriter writer = Files.newBufferedWriter(publicFileLocation, UTF_8)) {
writer.write("-----BEGIN PUBLIC KEY-----");
writer.newLine();
writer.write(Base64.getMimeEncoder().encodeToString(kp.getPublic().getEncoded()));
writer.write(Base64.getMimeEncoder().encodeToString(keyPair.getPublic().getEncoded()));
writer.newLine();
writer.write("-----END PUBLIC KEY-----");
writer.newLine();
Expand Down
75 changes: 75 additions & 0 deletions service/common/src/test/java/org/apache/polaris/service/auth/LocalRSAKeyProviderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/

package org.apache.polaris.service.auth;

import static org.assertj.core.api.InstanceOfAssertFactories.BYTE_ARRAY;

import java.nio.file.Path;
import java.security.PrivateKey;
import java.security.PublicKey;
import org.assertj.core.api.SoftAssertions;
import org.assertj.core.api.junit.jupiter.InjectSoftAssertions;
import org.assertj.core.api.junit.jupiter.SoftAssertionsExtension;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.api.io.TempDir;

@ExtendWith(SoftAssertionsExtension.class)
public class LocalRSAKeyProviderTest {
@InjectSoftAssertions SoftAssertions soft;

@Test
public void fromFiles(@TempDir Path tempDir) throws Exception {
var publicKeyFile = tempDir.resolve("public.key");
var privateKeyFile = tempDir.resolve("private.key");
PemUtils.generateKeyPairFiles(privateKeyFile, publicKeyFile);

var generatedPublicKey = PemUtils.readPublicKeyFromFile(publicKeyFile, "RSA");
var generatedPrivateKey = PemUtils.readPrivateKeyFromFile(privateKeyFile, "RSA");

var keyProvider = LocalRSAKeyProvider.fromFiles(publicKeyFile, privateKeyFile);
soft.assertThat(keyProvider)
.extracting(KeyProvider::getPublicKey)
.extracting(PublicKey::getEncoded, BYTE_ARRAY)
.containsExactly(generatedPublicKey.getEncoded());
soft.assertThat(keyProvider)
.extracting(KeyProvider::getPrivateKey)
.extracting(PrivateKey::getEncoded, BYTE_ARRAY)
.containsExactly(generatedPrivateKey.getEncoded());
}

@Test
public void onHeap() throws Exception {
var keyPair = PemUtils.generateKeyPair();

var generatedPublicKey = keyPair.getPublic();
var generatedPrivateKey = keyPair.getPrivate();

var keyProvider = new LocalRSAKeyProvider(keyPair);
soft.assertThat(keyProvider)
.extracting(KeyProvider::getPublicKey)
.extracting(PublicKey::getEncoded, BYTE_ARRAY)
.containsExactly(generatedPublicKey.getEncoded());
soft.assertThat(keyProvider)
.extracting(KeyProvider::getPrivateKey)
.extracting(PrivateKey::getEncoded, BYTE_ARRAY)
.containsExactly(generatedPrivateKey.getEncoded());
}
}
6 changes: 3 additions & 3 deletions site/content/in-dev/unreleased/configuration.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,10 +93,10 @@ read-only mode, as Polaris only reads the configuration file once, at startup.
| `polaris.features.realm-overrides."my-realm"."SKIP_CREDENTIAL_SUBSCOPING_INDIRECTION"` | `true` | "Override" realm features, here the skip credential subscoping indirection flag. |
| `polaris.authentication.authenticator.type` | `default` | Define the Polaris authenticator type. |
| `polaris.authentication.token-service.type` | `default` | Define the Polaris token service type. |
| `polaris.authentication.token-broker.type` | `rsa-key-pair` | Define the Polaris token broker type. |
| `polaris.authentication.token-broker.type` | `rsa-key-pair` | Define the Polaris token broker type. Also configure the location of the key files. For RSA: if the locations of the key files are not configured, an ephemeral key-pair will be created on each Polaris server instance startup, which breaks existing tokens after server restarts and is also incompatible with running multiple Polaris server instances. |
| `polaris.authentication.token-broker.max-token-generation` | `PT1H` | Define the max token generation policy on the token broker. |
| `polaris.authentication.token-broker.rsa-key-pair.public-key-file` | `/tmp/public.key` | Define the location of the public key file. |
| `polaris.authentication.token-broker.rsa-key-pair.private-key-file` | `/tmp/private.key` | Define the location of the private key file. |
| `polaris.authentication.token-broker.rsa-key-pair.private-key-file` | | Define the location of the RSA-256 private key file, if present the `public-key` file must be specified, too. |
| `polaris.authentication.token-broker.rsa-key-pair.public-key-file` | | Define the location of the RSA-256 public key file, if present the `private-key` file must be specified, too. |
| `polaris.authentication.token-broker.symmetric-key.secret` | `secret` | Define the secret of the symmetric key. |
| `polaris.authentication.token-broker.symmetric-key.file` | `/tmp/symmetric.key` | Define the location of the symmetric key file. |
| `polaris.storage.aws.access-key` | `accessKey` | Define the AWS S3 access key. If unset, the default credential provider chain will be used. |
Expand Down
Loading