JDBC external auth shared token #7309

s2lomon · 2021-03-16T22:50:27Z

This change allows sharing external authentication tokens
between different Connections, as long as they've been requested for
the same user.

Tokens are kept in ConcurrentHashMap, where jdbc user name is the key.

Different Connections - even for the same user, might be requested and used concurrently, I've decided to synchronize on resolving a token, to prevent spamming an end client with many web-browser pages opening. The idea is that whenever a user decide to share a token, only first request for a token will be served and all other will try to reuse token established by the previous request.

This option has been proposed only for jdbc, as it doesn't seem to have any meaning in cli.

client/trino-client/src/test/java/io/trino/client/auth/external/MockRedirectHandler.java

kokosing · 2021-03-17T09:03:18Z

client/trino-client/src/main/java/io/trino/client/auth/external/KnownToken.java

+            {
+                requireNonNull(tokenSource, "tokenSource is null");
+
+                knownToken = null;


client/trino-client/src/test/java/io/trino/client/auth/external/MockRedirectHandler.java

client/trino-cli/src/main/java/io/trino/cli/QueryRunner.java

client/trino-client/src/main/java/io/trino/client/auth/external/KnownToken.java

kokosing · 2021-03-17T09:15:50Z

client/trino-client/src/main/java/io/trino/client/auth/external/SharedToken.java

+        /*
+         * Operation for obtaining token is almost always bound to take a lot of time, as it will
+         * require user actions in a web-browser. The idea here is to trigger it only once, despite the number of
+         * concurrent invocations to setup a new token.


What if authentication flow failed somehow (server does not know it failed) and user would like to start a new one?

kokosing · 2021-03-17T09:33:47Z

client/trino-jdbc/src/main/java/io/trino/jdbc/ConnectionProperties.java

+    {
+        public ExternalAuthenticationSharedToken()
+        {
+            super("externalAuthenticationSharedToken", Optional.of("false"), NOT_REQUIRED, ALLOWED, BOOLEAN_CONVERTER);


It looks to me more about implementation than UX behaviour.

How about externalAuthenticationCache=NONE|MEMORY? Where NONE is default.

kokosing · 2021-03-17T09:36:26Z

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

+                .sleepOnRedirect(Duration.ofMillis(100));
+
+        ExternalAuthenticator authenticator = new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.shared("user-name"), Duration.ofSeconds(1));
+        List<Future<Request>> requests = executor.invokeAll(ImmutableList.<Callable<Request>>of(


Use IntStream and collect to list

Sorry, but since we need to generate Callable to change it to Future, doing this through stream requires construction like () -> () -> authenticator.authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""), which I find much less readable than simply repeating 4 times the same call.

kokosing · 2021-03-17T09:39:12Z

client/trino-client/src/main/java/io/trino/client/auth/external/SharedToken.java

+    }
+
+    @Override
+    public void setupToken(Supplier<Token> tokenSource)


I would make this method simply synchronized and same for getToken. Then you can use plain map.

client/trino-cli/src/main/java/io/trino/cli/QueryRunner.java

kokosing · 2021-03-23T09:37:19Z

client/trino-client/src/main/java/io/trino/client/auth/external/CashedToken.java

+class CashedToken
+        implements KnownToken
+{
+    private static final ReadWriteLock lock = new ReentrantReadWriteLock();


Can we have it as regular singleton, where you have single static instance and other fields are not static.

kokosing · 2021-03-23T09:37:52Z

client/trino-client/src/main/java/io/trino/client/auth/external/CashedToken.java

+ * many Connections will be actively using it. It's very important as obtaining the new token
+ * will take minutes, as it mostly requires user thinking time.
+ */
+class CashedToken


MemoryCachedKnownToken? It is not about the money....

kokosing · 2021-03-23T09:38:46Z

client/trino-client/src/main/java/io/trino/client/auth/external/CashedToken.java

+    {
+        //Try to lock and generate new token. If some other thread (Connection) has
+        //already obtained writeLock and is generating new token, then skipp this
+        //to block on getToken()


add space after // here and below

kokosing · 2021-03-23T09:39:34Z

client/trino-client/src/main/java/io/trino/client/auth/external/KnownToken.java

+
+    static KnownToken local()
+    {
+        return new KnownToken()


Please extract this as named class.

kokosing · 2021-03-23T09:49:36Z

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

+                    throw new RuntimeException(e);
+                }
+                catch (ExecutionException executionException) {
+                    if (executionException.getCause() != null) {


checkState, there should always be a cause.

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

kokosing · 2021-03-23T09:50:37Z

client/trino-jdbc/src/main/java/io/trino/jdbc/KnownTokenCache.java

+        @Override
+        KnownToken create()
+        {
+            return KnownToken.local();


inline this local() method.

I don't want to, as this would require declaring KnownToken classes instances with public modifiers etc. I want them to be hidden in the client-api and here used only through designated factory methods, without revealing their implementation etc.

kokosing · 2021-03-23T09:50:45Z

client/trino-jdbc/src/main/java/io/trino/jdbc/KnownTokenCache.java

+            return KnownToken.local();
+        }
+    },
+    CASHED {


kokosing · 2021-03-23T09:51:04Z

client/trino-jdbc/src/main/java/io/trino/jdbc/KnownTokenCache.java

+        @Override
+        KnownToken create()
+        {
+            return KnownToken.cashed();


inline this cached() (cachedInMemory()) method.

Just like above, I would have to make this implementation public and I don't think that's a good idea.

kokosing

Looks good.

client/trino-client/src/test/java/io/trino/client/auth/external/MockRedirectHandler.java

kokosing · 2021-03-24T19:51:05Z

client/trino-client/src/test/java/io/trino/client/auth/external/MockTokenPoller.java


 public final class MockTokenPoller
        implements TokenPoller
 {
-    private final Map<URI, Queue<TokenPollResult>> results = new HashMap<>();
+    private final Map<URI, BlockingDeque<TokenPollResult>> results = new HashMap<>();


Why this changed? Is this accessed by multiple threads?

Yes, in our new tests it is.

so results should also be thread safe. Isn't it?

kokosing · 2021-03-24T19:52:25Z

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

+        MockRedirectHandler redirectHandler = new MockRedirectHandler();
+
+        ExternalAuthenticator authenticator = new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.local(), Duration.ofSeconds(1));
+        List<Future<Request>> requests = times(4,


nit: each arg in separate line (please put 4 to the next line)

kokosing · 2021-03-24T19:54:11Z

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

+    public void testAuthenticationFromMultipleThreadsWithLocallyStoredToken()
+            throws Exception
+    {
+        ExecutorService executor = newFixedThreadPool(5);


Let's have a shared executor service per entire class. So we would avoid opening and closing it in each tests.

Use newCachedThreadPool(daemonThreadsNamed(getClass().getName() + "-%s"));

It will make the tests less separated and will add additional handling for the assertions. Right now since I'm blocking until all the tasks are finished, I can simply assert the end results. I can change the thread pool construction though.

kokosing · 2021-03-24T19:55:10Z

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

+                .sleepOnRedirect(Duration.ofMillis(10));
+
+        ExternalAuthenticator authenticator = new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.memoryCached(), Duration.ofSeconds(1));
+        List<Future<Request>> requests = times(4,


nit: args in separate lines

kokosing · 2021-03-24T19:55:44Z

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

+        assertThat(redirectHandler.getRedirectionCount()).isEqualTo(1);
+    }
+
+    static Stream<Callable<Request>> times(int times, Callable<Request> request)


private, please put it to the very bottom

kokosing · 2021-03-24T20:02:53Z

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

+            throws Exception
+    {
+        ExecutorService executor = newFixedThreadPool(5);
+        ExecutorService interruptableThreadPool = newFixedThreadPool(1);


Why do you need two pools?

To generate interrupted exception instead of cancellation exception (which happens why I simply cancel a thread)

kokosing · 2021-03-24T20:03:55Z

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

+                .collect(toImmutableList());
+
+        Thread.sleep(100);
+        interruptableThreadPool.shutdownNow();


did you mean interruptedAuthentication.cancel(true)?

No, as even with true for the interrupt, the end exception is Cancelation rather than Interrupted Exception and it seems to me that it's better for the case that we want to check here.

kokosing · 2021-03-24T20:05:05Z

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

+        ExternalAuthenticator authenticator = new ExternalAuthenticator(redirectHandler, (uri, duration) -> TokenPollResult.pending(uri), KnownToken.memoryCached(), Duration.ofMillis(1));
+        Future<Request> interruptedAuthentication = interruptableThreadPool.submit(
+                () -> authenticator.authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\"")));
+        List<Future<Request>> requests = times(2,


I think you need to sleep before you schedule anything here, to make sure interruptedAuthentication went sleep

I can add a very short sleep, just in case.

lukasz-walkiewicz · 2021-03-30T15:08:29Z

client/trino-client/src/main/java/io/trino/client/auth/external/LocalKnowToken.java

+ * and it's designed to use it in fully serialized manner.
+ */
+@NotThreadSafe
+class LocalKnowToken


lukasz-walkiewicz · 2021-03-30T15:21:29Z

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

+                .extracting(Future::get)
+                .extracting(Request::headers)
+                .extracting(headers -> headers.get(AUTHORIZATION))
+                .allMatch("Bearer valid-token"::equals);


containsOnly?

kokosing · 2021-03-31T19:50:49Z

client/trino-client/src/test/java/io/trino/client/auth/external/MockTokenPoller.java


 public final class MockTokenPoller
        implements TokenPoller
 {
-    private final Map<URI, Queue<TokenPollResult>> results = new HashMap<>();
+    private final Map<URI, BlockingDeque<TokenPollResult>> results = new HashMap<>();


so results should also be thread safe. Isn't it?

kokosing · 2021-03-31T19:51:17Z

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

+    public void shutDownThreadPool()
+            throws InterruptedException
+    {
+        executor.shutdown();


shutdownNow() is enough

kokosing · 2021-03-31T19:51:31Z

client/trino-client/src/test/java/io/trino/client/auth/external/TestExternalAuthenticator.java

+
+    @AfterClass(alwaysRun = true)
+    public void shutDownThreadPool()
+            throws InterruptedException


just throws Exception

kokosing · 2021-04-01T08:58:55Z

@dain @electrum would you like to take a look?

This change allows sharing external authentication tokens between different Connections. Each time when a new token is required, first Connection that needs it, will handle obtaining a new token when all the other Connections wait for this operation to finish. Token is kept in memmory, guarded by ReadWriteLock. To enable token cache use externalAuthenticationTokenCache=MEMORY Default value for externalAuthenticationTokenCache is NONE.

kokosing · 2021-04-09T09:47:50Z

CI hit: #7534

cla-bot bot added the cla-signed label Mar 16, 2021

s2lomon requested a review from kokosing March 16, 2021 22:50

s2lomon added the WIP label Mar 16, 2021

kokosing reviewed Mar 17, 2021

View reviewed changes

findepi changed the title ~~Jdbc external auth shared token~~ JDBC external auth shared token Mar 17, 2021

s2lomon force-pushed the jdbc-external-auth-shared-token branch from 22eb640 to b4eb4e6 Compare March 23, 2021 00:52

kokosing reviewed Mar 23, 2021

View reviewed changes

s2lomon force-pushed the jdbc-external-auth-shared-token branch from b4eb4e6 to 850779d Compare March 24, 2021 14:12

s2lomon removed the WIP label Mar 24, 2021

s2lomon requested review from kokosing and lukasz-walkiewicz March 24, 2021 14:16

s2lomon force-pushed the jdbc-external-auth-shared-token branch from 850779d to 80ce2f0 Compare March 24, 2021 15:46

kokosing reviewed Mar 24, 2021

View reviewed changes

lukasz-walkiewicz approved these changes Mar 30, 2021

View reviewed changes

Extract class MockRedirectHandler

629b0b6

s2lomon force-pushed the jdbc-external-auth-shared-token branch from 80ce2f0 to 55db7e4 Compare March 31, 2021 15:27

s2lomon requested a review from kokosing March 31, 2021 15:27

kokosing reviewed Mar 31, 2021

View reviewed changes

s2lomon force-pushed the jdbc-external-auth-shared-token branch from 55db7e4 to 7162450 Compare April 8, 2021 17:03

kokosing approved these changes Apr 9, 2021

View reviewed changes

kokosing merged commit ceb2d1b into trinodb:master Apr 9, 2021

kokosing mentioned this pull request Apr 9, 2021

Release notes for 356 #7533

Closed

9 tasks

phd3 mentioned this pull request Apr 22, 2021

Flaky testAuthenticationFromMultipleThreadsWithCachedTokenAfterAuthenticateFails #7697

Closed

martint added this to the 356 milestone Apr 26, 2021

JDBC external auth shared token #7309

JDBC external auth shared token #7309

Conversation

s2lomon commented Mar 16, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kokosing left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kokosing commented Apr 1, 2021

kokosing commented Apr 9, 2021

s2lomon commented Mar 16, 2021 •

edited

Loading