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..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 @@ -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,16 +195,33 @@ public void testAuthenticationFromMultipleThreadsWithLocallyStoredToken() @Test @Timeout(2) - public void testAuthenticationFromMultipleThreadsWithCachedToken() + public void testAuthenticationFromMultipleThreadsWithMemoryCachedToken() + { + KnownToken knownToken = KnownToken.memoryCached(); + testAuthenticationFromMultipleThreadsWithCachedToken(knownToken); + } + + @Test + @Timeout(10) + 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"))); + // 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, - () -> new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.memoryCached(), Duration.ofSeconds(1)) + 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) .collect(toImmutableList()); @@ -211,7 +230,7 @@ public void testAuthenticationFromMultipleThreadsWithCachedToken() 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); } 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`