From 697dbec2b03656b87da35724c2f563cb8fa737e8 Mon Sep 17 00:00:00 2001 From: Sasha Sheikin Date: Wed, 18 Mar 2026 14:13:28 +0100 Subject: [PATCH 1/2] Allow to store externalAuthenticationToken in SYSTEM Cache (in dotfiles), so it can be reused between processes, without requiring authentication. This may be helpful in case different CLI are run from the same machine. --- .../client/auth/external/KnownToken.java | 5 + .../auth/external/SystemCachedKnownToken.java | 151 +++++++++++++ .../io/trino/client/uri/KnownTokenCache.java | 7 + .../external/TestExternalAuthenticator.java | 20 +- .../external/TestSystemCachedKnownToken.java | 212 ++++++++++++++++++ docs/src/main/sphinx/client/jdbc.md | 10 +- 6 files changed, 399 insertions(+), 6 deletions(-) create mode 100644 client/trino-client/src/main/java/io/trino/client/auth/external/SystemCachedKnownToken.java create mode 100644 client/trino-client/src/test/java/io/trino/client/auth/external/TestSystemCachedKnownToken.java diff --git a/client/trino-client/src/main/java/io/trino/client/auth/external/KnownToken.java b/client/trino-client/src/main/java/io/trino/client/auth/external/KnownToken.java index 240af9075763..01e0d7304713 100644 --- a/client/trino-client/src/main/java/io/trino/client/auth/external/KnownToken.java +++ b/client/trino-client/src/main/java/io/trino/client/auth/external/KnownToken.java @@ -31,4 +31,9 @@ static KnownToken memoryCached() { return MemoryCachedKnownToken.INSTANCE; } + + static KnownToken systemCached() + { + return SystemCachedKnownToken.INSTANCE; + } } diff --git a/client/trino-client/src/main/java/io/trino/client/auth/external/SystemCachedKnownToken.java b/client/trino-client/src/main/java/io/trino/client/auth/external/SystemCachedKnownToken.java new file mode 100644 index 000000000000..e14089501ce7 --- /dev/null +++ b/client/trino-client/src/main/java/io/trino/client/auth/external/SystemCachedKnownToken.java @@ -0,0 +1,151 @@ +/* + * Licensed 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 io.trino.client.auth.external; + +import com.google.common.collect.ImmutableSet; +import dev.failsafe.Failsafe; +import dev.failsafe.RetryPolicy; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.file.FileAlreadyExistsException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Duration; +import java.util.Optional; +import java.util.function.Supplier; + +import static java.nio.file.attribute.PosixFilePermission.OWNER_READ; +import static java.nio.file.attribute.PosixFilePermission.OWNER_WRITE; +import static java.time.temporal.ChronoUnit.MILLIS; +import static java.util.Objects.requireNonNull; + +/** + * This KnownToken instance persists the token to ~/.trino/.token on the filesystem, + * allowing it to be reused across separate CLI invocations. + * A lock file (~/.trino/.token.lck) is used to coordinate token acquisition + * across processes — its atomic creation acts as a cross-process tryLock. + * The implementation is similar to MemoryCachedKnownToken, but with LOCK_FILE acting as a Lock + */ +class SystemCachedKnownToken + implements KnownToken +{ + private static final Path DEFAULT_TRINO_DIR = Path.of(System.getProperty("user.home"), ".trino"); + // duration for the user to authenticate within IDP. It involves clicking and typing from a real person within a browser, so should be counted in minutes. + private static final Duration DEFAULT_LOCK_MAX_WAIT = Duration.ofMinutes(10); + + public static final SystemCachedKnownToken INSTANCE = new SystemCachedKnownToken(DEFAULT_TRINO_DIR); + + private final Path trinoDir; + private final Path tokenFile; + private final Path lockFile; + private final Duration lockMaxWait; + + SystemCachedKnownToken(Path trinoDir) + { + this(trinoDir, DEFAULT_LOCK_MAX_WAIT); + } + + SystemCachedKnownToken(Path trinoDir, Duration lockMaxWait) + { + this.trinoDir = requireNonNull(trinoDir, "trinoDir is null"); + this.tokenFile = trinoDir.resolve(".token"); + this.lockFile = trinoDir.resolve(".token.lck"); + this.lockMaxWait = requireNonNull(lockMaxWait, "lockMaxWait is null"); + } + + @Override + public Optional getToken() + { + // Wait while lock file exists, mimicking readLock blocking while writeLock is held + boolean lockFileExists = Failsafe.with(RetryPolicy.builder() + .handleResultIf(Boolean.TRUE::equals) + .withMaxAttempts(-1) + .withDelay(100, 1000, MILLIS) + .withMaxDuration(lockMaxWait) + .build()) + .get(() -> Files.exists(lockFile)); + if (lockFileExists) { + throw new IllegalStateException("Lock file " + lockFile + " for System Cached token still exists after waiting " + lockMaxWait + ". " + + "It may be created by another concurrent authentication, which is still in progress. " + + "If it's not a case and another transaction was abandoned - please remove the lock file manually and retry authentication."); + } + + if (!Files.exists(tokenFile)) { + return Optional.empty(); + } + try { + String content = Files.readString(tokenFile).trim(); + if (content.isEmpty()) { + return Optional.empty(); + } + return Optional.of(new Token(content)); + } + catch (IOException e) { + throw new UncheckedIOException("Failed to read token from " + tokenFile, e); + } + } + + @Override + public void setupToken(Supplier> tokenSource) + { + try { + Files.createDirectories(trinoDir); + } + catch (IOException e) { + throw new UncheckedIOException("Failed to create directory " + trinoDir, e); + } + + // Atomically create the lock file. If it already exists, another process + // is obtaining a token — skip, just like MemoryCachedKnownToken's tryLock. + try { + Files.createFile(lockFile); + } + catch (FileAlreadyExistsException e) { + return; + } + catch (IOException e) { + throw new UncheckedIOException("Failed to create lock file " + lockFile, e); + } + + try { + // Clear token before obtaining new one, as it might fail leaving old invalid token. + Files.deleteIfExists(tokenFile); + Optional token = tokenSource.get(); + token.ifPresent(this::writeTokenToFile); + } + catch (IOException e) { + throw new UncheckedIOException("Failed to update token file " + tokenFile, e); + } + finally { + try { + Files.deleteIfExists(lockFile); + } + catch (IOException e) { + throw new UncheckedIOException("Failed to delete lock file " + lockFile, e); + } + } + } + + private void writeTokenToFile(Token token) + { + try { + Files.writeString(tokenFile, token.token()); + Files.setPosixFilePermissions(tokenFile, ImmutableSet.of(OWNER_READ, OWNER_WRITE)); + } + catch (IOException e) { + throw new UncheckedIOException("Failed to write token to " + tokenFile, e); + } + } +} diff --git a/client/trino-client/src/main/java/io/trino/client/uri/KnownTokenCache.java b/client/trino-client/src/main/java/io/trino/client/uri/KnownTokenCache.java index c4ea61fa04e0..ce0c4576692d 100644 --- a/client/trino-client/src/main/java/io/trino/client/uri/KnownTokenCache.java +++ b/client/trino-client/src/main/java/io/trino/client/uri/KnownTokenCache.java @@ -30,6 +30,13 @@ KnownToken create() { return KnownToken.memoryCached(); } + }, + SYSTEM { + @Override + KnownToken create() + { + return KnownToken.systemCached(); + } }; abstract KnownToken create(); diff --git a/client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java b/client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java index 3cde942e784f..147044ebf4c2 100644 --- a/client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java +++ b/client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java @@ -25,10 +25,12 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestInstance; import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.api.parallel.Execution; import java.net.URI; import java.net.URISyntaxException; +import java.nio.file.Path; import java.time.Duration; import java.util.ArrayList; import java.util.List; @@ -193,7 +195,21 @@ public void testAuthenticationFromMultipleThreadsWithLocallyStoredToken() @Test @Timeout(2) - public void testAuthenticationFromMultipleThreadsWithCachedToken() + public void testAuthenticationFromMultipleThreadsWithMemoryCachedToken() + { + KnownToken knownToken = KnownToken.memoryCached(); + testAuthenticationFromMultipleThreadsWithCachedToken(knownToken); + } + + @Test + @Timeout(2) + public void testAuthenticationFromMultipleThreadsWithSystemCachedToken(@TempDir Path tempDir) + { + KnownToken knownToken = new SystemCachedKnownToken(tempDir); + testAuthenticationFromMultipleThreadsWithCachedToken(knownToken); + } + + private static void testAuthenticationFromMultipleThreadsWithCachedToken(KnownToken knownToken) { MockTokenPoller tokenPoller = new MockTokenPoller() .withResult(URI.create("http://token.uri"), successful(new Token("valid-token"))); @@ -202,7 +218,7 @@ public void testAuthenticationFromMultipleThreadsWithCachedToken() List> requests = times( 2, - () -> new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.memoryCached(), Duration.ofSeconds(1)) + () -> new ExternalAuthenticator(redirectHandler, tokenPoller, knownToken, Duration.ofSeconds(1)) .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""))) .map(executor::submit) .collect(toImmutableList()); diff --git a/client/trino-client/src/test/java/io/trino/client/auth/external/TestSystemCachedKnownToken.java b/client/trino-client/src/test/java/io/trino/client/auth/external/TestSystemCachedKnownToken.java new file mode 100644 index 000000000000..ec27466568f6 --- /dev/null +++ b/client/trino-client/src/test/java/io/trino/client/auth/external/TestSystemCachedKnownToken.java @@ -0,0 +1,212 @@ +/* + * Licensed 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 io.trino.client.auth.external; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Duration; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicBoolean; + +import static java.nio.file.attribute.PosixFilePermission.OWNER_READ; +import static java.nio.file.attribute.PosixFilePermission.OWNER_WRITE; +import static java.util.concurrent.Executors.newSingleThreadExecutor; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +@Timeout(30) +final class TestSystemCachedKnownToken +{ + @Test + void testGetTokenReturnsEmptyWhenNoTokenFile(@TempDir Path tempDir) + { + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(tempDir); + assertThat(knownToken.getToken()).isEmpty(); + } + + @Test + void testGetTokenReturnsTokenFromFile(@TempDir Path tempDir) + throws IOException + { + Files.writeString(tempDir.resolve(".token"), "test-token-value"); + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(tempDir); + assertThat(knownToken.getToken()) + .isPresent() + .hasValueSatisfying(token -> assertThat(token.token()).isEqualTo("test-token-value")); + } + + @Test + void testGetTokenReturnsEmptyForEmptyFile(@TempDir Path tempDir) + throws IOException + { + Files.writeString(tempDir.resolve(".token"), " \n "); + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(tempDir); + assertThat(knownToken.getToken()).isEmpty(); + } + + @Test + void testSetupTokenWritesTokenToFile(@TempDir Path tempDir) + { + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(tempDir); + knownToken.setupToken(() -> Optional.of(new Token("new-token"))); + assertThat(knownToken.getToken()) + .isPresent() + .hasValueSatisfying(token -> assertThat(token.token()).isEqualTo("new-token")); + } + + @Test + void testSetupTokenCreatesDirectory(@TempDir Path tempDir) + { + Path nestedDir = tempDir.resolve("nested").resolve("dir"); + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(nestedDir); + knownToken.setupToken(() -> Optional.of(new Token("token"))); + assertThat(Files.isDirectory(nestedDir)).isTrue(); + assertThat(knownToken.getToken()).isPresent(); + } + + @Test + void testSetupTokenClearsOldTokenBeforeObtainingNew(@TempDir Path tempDir) + throws IOException + { + Files.writeString(tempDir.resolve(".token"), "old-token"); + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(tempDir); + knownToken.setupToken(() -> Optional.of(new Token("new-token"))); + assertThat(knownToken.getToken()) + .isPresent() + .hasValueSatisfying(token -> assertThat(token.token()).isEqualTo("new-token")); + } + + @Test + void testSetupTokenClearsTokenWhenSourceReturnsEmpty(@TempDir Path tempDir) + throws IOException + { + Files.writeString(tempDir.resolve(".token"), "old-token"); + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(tempDir); + knownToken.setupToken(Optional::empty); + assertThat(knownToken.getToken()).isEmpty(); + } + + @Test + void testSetupTokenClearsTokenWhenSourceFails(@TempDir Path tempDir) + throws IOException + { + Files.writeString(tempDir.resolve(".token"), "old-token"); + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(tempDir); + + assertThatThrownBy(() -> knownToken.setupToken(() -> { + throw new RuntimeException("Auth is expected to fail"); + })).hasMessage("Auth is expected to fail"); + + // Old token should be cleared, lock file should be cleaned up + assertThat(knownToken.getToken()).isEmpty(); + assertThat(Files.exists(tempDir.resolve(".token.lck"))).isFalse(); + } + + @Test + void testSetupTokenRemovesLockFileAfterSuccess(@TempDir Path tempDir) + { + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(tempDir); + knownToken.setupToken(() -> Optional.of(new Token("token"))); + assertThat(Files.exists(tempDir.resolve(".token.lck"))).isFalse(); + } + + @Test + void testSetupTokenSkipsWhenLockFileAlreadyExists(@TempDir Path tempDir) + throws IOException + { + Files.writeString(tempDir.resolve(".token"), "existing-token"); + Files.createFile(tempDir.resolve(".token.lck")); + + AtomicBoolean tokenSourceCalled = new AtomicBoolean(false); + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(tempDir); + + knownToken.setupToken(() -> { + tokenSourceCalled.set(true); + return Optional.of(new Token("should-not-be-written")); + }); + + // Token source should not have been invoked + assertThat(tokenSourceCalled.get()).isFalse(); + // Original token file should be unchanged + assertThat(Files.readString(tempDir.resolve(".token"))).isEqualTo("existing-token"); + } + + @Test + void testSetupTokenSetsOwnerOnlyPermissions(@TempDir Path tempDir) + throws IOException + { + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(tempDir); + + knownToken.setupToken(() -> Optional.of(new Token("secret-token"))); + + assertThat(Files.getPosixFilePermissions(tempDir.resolve(".token"))) + .containsExactlyInAnyOrder(OWNER_READ, OWNER_WRITE); + } + + @Test + void testGetTokenThrowsWhenLockFileExistsAfterMaxWait(@TempDir Path tempDir) + throws IOException + { + Files.createFile(tempDir.resolve(".token.lck")); + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(tempDir, Duration.ofMillis(1001)); + + assertThatThrownBy(knownToken::getToken) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("Lock file") + .hasMessageContaining("still exists after waiting"); + } + + @Test + void testGetTokenWaitsForLockFileToBeRemoved(@TempDir Path tempDir) + throws Exception + { + Files.createFile(tempDir.resolve(".token.lck")); + Files.writeString(tempDir.resolve(".token"), "the-token"); + + SystemCachedKnownToken knownToken = new SystemCachedKnownToken(tempDir); + CountDownLatch getTokenStarted = new CountDownLatch(1); + + ExecutorService executor = newSingleThreadExecutor(); + try { + Future> future = executor.submit(() -> { + getTokenStarted.countDown(); + return knownToken.getToken(); + }); + + getTokenStarted.await(); + // Give getToken time to enter the retry loop + Thread.sleep(300); + assertThat(future.isDone()).isFalse(); + + // Remove lock file — getToken should now complete + Files.delete(tempDir.resolve(".token.lck")); + + Optional result = future.get(); + assertThat(result) + .isPresent() + .hasValueSatisfying(token -> assertThat(token.token()).isEqualTo("the-token")); + } + finally { + executor.shutdownNow(); + } + } +} diff --git a/docs/src/main/sphinx/client/jdbc.md b/docs/src/main/sphinx/client/jdbc.md index 2ce05df9c2fb..a085a63c77cb 100644 --- a/docs/src/main/sphinx/client/jdbc.md +++ b/docs/src/main/sphinx/client/jdbc.md @@ -249,10 +249,12 @@ may not be specified using both methods. - Allows the sharing of external authentication tokens between different connections for the same authenticated user until the cache is invalidated, such as when a client is restarted or when the classloader reloads the JDBC - driver. This is disabled by default, with a value of `NONE`. To enable, set - the value to `MEMORY`. If the JDBC driver is used in a shared mode by - different users, the first registered token is stored and authenticates all - users. + driver. This is disabled by default, with a value of `NONE`. Set the value + to `MEMORY` to cache the token in memory within the same process. Set the + value to `SYSTEM` to persist the token to the filesystem (`~/.trino/`), + allowing it to be reused across separate CLI or JDBC processes. If the JDBC + driver is used in a shared mode by different users, the first registered + token is stored and authenticates all users. * - `disableCompression` - Whether HTTP compression should be disabled. Defaults to `false`. * - `disallowLocalRedirect` From 2f148933b7edc91744d9269d5d25a40290ff093e Mon Sep 17 00:00:00 2001 From: Sasha Sheikin Date: Wed, 25 Mar 2026 11:10:25 +0100 Subject: [PATCH 2/2] Improve testAuthenticationFromMultipleThreadsWithCachedToken Storing tokens cached within system deals with filesystem, where delays are much more significant than in-memory. Increasing to 100 concurrent transaction increases the overall execution time, thus, also increasing timeout for a test method. --- .../auth/external/TestExternalAuthenticator.java | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java b/client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java index 147044ebf4c2..e69b72d38131 100644 --- a/client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java +++ b/client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java @@ -202,7 +202,7 @@ public void testAuthenticationFromMultipleThreadsWithMemoryCachedToken() } @Test - @Timeout(2) + @Timeout(10) public void testAuthenticationFromMultipleThreadsWithSystemCachedToken(@TempDir Path tempDir) { KnownToken knownToken = new SystemCachedKnownToken(tempDir); @@ -212,12 +212,15 @@ public void testAuthenticationFromMultipleThreadsWithSystemCachedToken(@TempDir private static void testAuthenticationFromMultipleThreadsWithCachedToken(KnownToken knownToken) { MockTokenPoller tokenPoller = new MockTokenPoller() - .withResult(URI.create("http://token.uri"), successful(new Token("valid-token"))); + // will be cached and reused + .withResult(URI.create("http://token.uri"), successful(new Token("first-token"))) + // should never be emitted + .withResult(URI.create("http://token.uri"), successful(new Token("second-token"))); MockRedirectHandler redirectHandler = new MockRedirectHandler() .sleepOnRedirect(Duration.ofSeconds(1)); List> requests = times( - 2, + 100, () -> new ExternalAuthenticator(redirectHandler, tokenPoller, knownToken, Duration.ofSeconds(1)) .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""))) .map(executor::submit) @@ -227,7 +230,7 @@ private static void testAuthenticationFromMultipleThreadsWithCachedToken(KnownTo assertion.requests() .extracting(Request::headers) .extracting(headers -> headers.get(AUTHORIZATION)) - .containsOnly("Bearer valid-token"); + .containsOnly("Bearer first-token"); assertion.assertThatNoExceptionsHasBeenThrown(); assertThat(redirectHandler.getRedirectionCount()).isEqualTo(1); }