From 60b9fc122bbc32c66279a9cee079f22c5f5612fa Mon Sep 17 00:00:00 2001 From: Slim Bouguerra Date: Sun, 1 Oct 2023 07:46:28 -0700 Subject: [PATCH] Cleanly reject session authorization queries for unsupported clients Adds SESSION_AUTHORIZATION to ClientCapabilities, so that the client side support for this feature can be identified using the headers sent to Trino. --- .../io/trino/client/ClientCapabilities.java | 4 ++- .../ResetSessionAuthorizationTask.java | 5 ++++ .../SetSessionAuthorizationTask.java | 5 ++++ .../test/java/io/trino/tests/TestServer.java | 29 +++++++++++++++++++ 4 files changed, 42 insertions(+), 1 deletion(-) diff --git a/client/trino-client/src/main/java/io/trino/client/ClientCapabilities.java b/client/trino-client/src/main/java/io/trino/client/ClientCapabilities.java index 809e2c43e581..ba55f3872ea4 100644 --- a/client/trino-client/src/main/java/io/trino/client/ClientCapabilities.java +++ b/client/trino-client/src/main/java/io/trino/client/ClientCapabilities.java @@ -23,5 +23,7 @@ public enum ClientCapabilities // time(p) without time zone // interval X(p1) to Y(p2) // When this capability is not set, the server returns datetime types with precision = 3 - PARAMETRIC_DATETIME; + PARAMETRIC_DATETIME, + // Whether clients support the session authorization set/reset feature + SESSION_AUTHORIZATION; } diff --git a/core/trino-main/src/main/java/io/trino/execution/ResetSessionAuthorizationTask.java b/core/trino-main/src/main/java/io/trino/execution/ResetSessionAuthorizationTask.java index 6bc1b2e66867..8af7bf504227 100644 --- a/core/trino-main/src/main/java/io/trino/execution/ResetSessionAuthorizationTask.java +++ b/core/trino-main/src/main/java/io/trino/execution/ResetSessionAuthorizationTask.java @@ -16,6 +16,7 @@ import com.google.common.util.concurrent.ListenableFuture; import com.google.inject.Inject; import io.trino.Session; +import io.trino.client.ClientCapabilities; import io.trino.execution.warnings.WarningCollector; import io.trino.spi.TrinoException; import io.trino.sql.tree.Expression; @@ -26,6 +27,7 @@ import static com.google.common.util.concurrent.Futures.immediateFuture; import static io.trino.spi.StandardErrorCode.GENERIC_USER_ERROR; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static java.util.Objects.requireNonNull; public class ResetSessionAuthorizationTask @@ -53,6 +55,9 @@ public ListenableFuture execute( WarningCollector warningCollector) { Session session = stateMachine.getSession(); + if (!session.getClientCapabilities().contains(ClientCapabilities.SESSION_AUTHORIZATION.toString())) { + throw new TrinoException(NOT_SUPPORTED, "RESET SESSION AUTHORIZATION not supported by client"); + } session.getTransactionId().ifPresent(transactionId -> { if (!transactionManager.getTransactionInfo(transactionId).isAutoCommitContext()) { throw new TrinoException(GENERIC_USER_ERROR, "Can't reset authorization user in the middle of a transaction"); diff --git a/core/trino-main/src/main/java/io/trino/execution/SetSessionAuthorizationTask.java b/core/trino-main/src/main/java/io/trino/execution/SetSessionAuthorizationTask.java index 99afede9118c..b351549a23ee 100644 --- a/core/trino-main/src/main/java/io/trino/execution/SetSessionAuthorizationTask.java +++ b/core/trino-main/src/main/java/io/trino/execution/SetSessionAuthorizationTask.java @@ -16,6 +16,7 @@ import com.google.common.util.concurrent.ListenableFuture; import com.google.inject.Inject; import io.trino.Session; +import io.trino.client.ClientCapabilities; import io.trino.execution.warnings.WarningCollector; import io.trino.security.AccessControl; import io.trino.spi.TrinoException; @@ -31,6 +32,7 @@ import static com.google.common.base.Preconditions.checkState; import static com.google.common.util.concurrent.Futures.immediateFuture; import static io.trino.spi.StandardErrorCode.GENERIC_USER_ERROR; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static java.util.Objects.requireNonNull; public class SetSessionAuthorizationTask @@ -60,6 +62,9 @@ public ListenableFuture execute( WarningCollector warningCollector) { Session session = stateMachine.getSession(); + if (!session.getClientCapabilities().contains(ClientCapabilities.SESSION_AUTHORIZATION.toString())) { + throw new TrinoException(NOT_SUPPORTED, "SET SESSION AUTHORIZATION not supported by client"); + } Identity originalIdentity = session.getOriginalIdentity(); // Set authorization user in the middle of a transaction is disallowed by the SQL spec session.getTransactionId().ifPresent(transactionId -> { diff --git a/testing/trino-tests/src/test/java/io/trino/tests/TestServer.java b/testing/trino-tests/src/test/java/io/trino/tests/TestServer.java index 1769da31ffc5..825fc08f3587 100644 --- a/testing/trino-tests/src/test/java/io/trino/tests/TestServer.java +++ b/testing/trino-tests/src/test/java/io/trino/tests/TestServer.java @@ -68,6 +68,7 @@ import static io.trino.SystemSessionProperties.MAX_HASH_PARTITION_COUNT; import static io.trino.SystemSessionProperties.QUERY_MAX_MEMORY; import static io.trino.client.ClientCapabilities.PATH; +import static io.trino.client.ClientCapabilities.SESSION_AUTHORIZATION; import static io.trino.client.ProtocolHeaders.TRINO_HEADERS; import static io.trino.spi.StandardErrorCode.INCOMPATIBLE_CLIENT; import static io.trino.testing.TestingSession.testSessionBuilder; @@ -291,6 +292,34 @@ public void testSetPathSupportByClient() } } + @Test + public void testSetSessionSupportByClient() + { + try (TestingTrinoClient testingClient = new TestingTrinoClient(server, testSessionBuilder().setClientCapabilities(Set.of()).build())) { + assertThatThrownBy(() -> testingClient.execute("SET SESSION AUTHORIZATION userA")) + .hasMessage("SET SESSION AUTHORIZATION not supported by client"); + } + + try (TestingTrinoClient testingClient = new TestingTrinoClient(server, testSessionBuilder().setClientCapabilities(Set.of( + SESSION_AUTHORIZATION.name())).build())) { + testingClient.execute("SET SESSION AUTHORIZATION userA"); + } + } + + @Test + public void testResetSessionSupportByClient() + { + try (TestingTrinoClient testingClient = new TestingTrinoClient(server, testSessionBuilder().setClientCapabilities(Set.of()).build())) { + assertThatThrownBy(() -> testingClient.execute("RESET SESSION AUTHORIZATION")) + .hasMessage("RESET SESSION AUTHORIZATION not supported by client"); + } + + try (TestingTrinoClient testingClient = new TestingTrinoClient(server, testSessionBuilder().setClientCapabilities(Set.of( + SESSION_AUTHORIZATION.name())).build())) { + testingClient.execute("RESET SESSION AUTHORIZATION"); + } + } + private void checkVersionOnError(String query, @Language("RegExp") String proofOfOrigin) { QueryResults queryResults = postQuery(request -> request