From 4e81e2ece63f3b2c3c526ddac24f7b01b58a1536 Mon Sep 17 00:00:00 2001 From: zhangdong Date: Thu, 13 Feb 2025 12:09:37 +0800 Subject: [PATCH] [enhance](auth)The priority of attributes is higher than session variable (#47185) The current priority is that session variable is higher than user property, which is incorrect because session variable can be set freely by ordinary users and will bypass the restrictions set by administrators Scope of Impact: - getInsertTimeoutS - getQueryTimeoutS - getMaxExecMemByte --- .../org/apache/doris/qe/ConnectContext.java | 65 ++++++++++++++++- .../java/org/apache/doris/qe/Coordinator.java | 4 +- .../apache/doris/qe/ConnectContextTest.java | 70 +++++++++++++++++++ .../account_p0/test_property_session.groovy | 56 +++++++++++++++ 4 files changed, 190 insertions(+), 5 deletions(-) create mode 100644 regression-test/suites/account_p0/test_property_session.groovy diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java index c160d7e77b916f..6a645be18e6135 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java @@ -76,6 +76,7 @@ import com.google.common.collect.Maps; import com.google.common.collect.Sets; import io.netty.util.concurrent.FastThreadLocal; +import org.apache.commons.lang3.StringUtils; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.json.JSONObject; @@ -1028,15 +1029,75 @@ public String getCurrentConnectedFEIp() { public int getExecTimeout() { if (executor != null && executor.isSyncLoadKindStmt()) { // particular for insert stmt, we can expand other type of timeout in the same way - return Math.max(sessionVariable.getInsertTimeoutS(), sessionVariable.getQueryTimeoutS()); + return Math.max(getInsertTimeoutS(), getQueryTimeoutS()); } else if (executor != null && executor.isAnalyzeStmt()) { return sessionVariable.getAnalyzeTimeoutS(); } else { // normal query stmt - return sessionVariable.getQueryTimeoutS(); + return getQueryTimeoutS(); } } + /** + * First, retrieve from the user's attributes. If not, retrieve from the session variable + * + * @return insertTimeoutS + */ + public int getInsertTimeoutS() { + int userInsertTimeout = getInsertTimeoutSFromProperty(); + if (userInsertTimeout > 0) { + return userInsertTimeout; + } + return sessionVariable.getInsertTimeoutS(); + } + + private int getInsertTimeoutSFromProperty() { + if (env == null || env.getAuth() == null || StringUtils.isEmpty(getQualifiedUser())) { + return 0; + } + return env.getAuth().getInsertTimeout(getQualifiedUser()); + } + + /** + * First, retrieve from the user's attributes. If not, retrieve from the session variable + * + * @return queryTimeoutS + */ + public int getQueryTimeoutS() { + int userQueryTimeout = getQueryTimeoutSFromProperty(); + if (userQueryTimeout > 0) { + return userQueryTimeout; + } + return sessionVariable.getQueryTimeoutS(); + } + + private int getQueryTimeoutSFromProperty() { + if (env == null || env.getAuth() == null || StringUtils.isEmpty(getQualifiedUser())) { + return 0; + } + return env.getAuth().getQueryTimeout(getQualifiedUser()); + } + + /** + * First, retrieve from the user's attributes. If not, retrieve from the session variable + * + * @return maxExecMemByte + */ + public long getMaxExecMemByte() { + long userLimit = getMaxExecMemByteFromProperty(); + if (userLimit > 0) { + return userLimit; + } + return sessionVariable.getMaxExecMemByte(); + } + + private long getMaxExecMemByteFromProperty() { + if (env == null || env.getAuth() == null || StringUtils.isEmpty(getQualifiedUser())) { + return 0L; + } + return env.getAuth().getExecMemLimit(getQualifiedUser()); + } + public void setResultAttachedInfo(Map resultAttachedInfo) { this.resultAttachedInfo = resultAttachedInfo; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java index 58631524024721..05594f9c235f92 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java @@ -403,9 +403,7 @@ private void setFromUserProperty(ConnectContext connectContext) { this.queryOptions.setResourceLimit(resourceLimit); } // set exec mem limit - long maxExecMemByte = connectContext.getSessionVariable().getMaxExecMemByte(); - long memLimit = maxExecMemByte > 0 ? maxExecMemByte : - Env.getCurrentEnv().getAuth().getExecMemLimit(qualifiedUser); + long memLimit = connectContext.getMaxExecMemByte(); if (memLimit > 0) { // overwrite the exec_mem_limit from session variable; this.queryOptions.setMemLimit(memLimit); diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java index 4bb30b14a632be..6ad9532a29784b 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java @@ -23,6 +23,7 @@ import org.apache.doris.mysql.privilege.Auth; import org.apache.doris.thrift.TUniqueId; +import mockit.Expectations; import mockit.Mocked; import org.junit.Assert; import org.junit.Before; @@ -202,4 +203,73 @@ public void testThreadLocal() { Assert.assertNotNull(ConnectContext.get()); Assert.assertEquals(ctx, ConnectContext.get()); } + + @Test + public void testGetMaxExecMemByte() { + ConnectContext context = new ConnectContext(); + context.setQualifiedUser("a"); + context.setEnv(env); + long sessionValue = 2097153L; + long propertyValue = 2097154L; + // only session + context.getSessionVariable().setMaxExecMemByte(sessionValue); + long result = context.getMaxExecMemByte(); + Assert.assertEquals(sessionValue, result); + // has property + new Expectations() { + { + auth.getExecMemLimit(anyString); + minTimes = 0; + result = propertyValue; + } + }; + result = context.getMaxExecMemByte(); + Assert.assertEquals(propertyValue, result); + } + + @Test + public void testGetQueryTimeoutS() { + ConnectContext context = new ConnectContext(); + context.setQualifiedUser("a"); + context.setEnv(env); + int sessionValue = 1; + int propertyValue = 2; + // only session + context.getSessionVariable().setQueryTimeoutS(sessionValue); + long result = context.getQueryTimeoutS(); + Assert.assertEquals(sessionValue, result); + // has property + new Expectations() { + { + auth.getQueryTimeout(anyString); + minTimes = 0; + result = propertyValue; + } + }; + result = context.getQueryTimeoutS(); + Assert.assertEquals(propertyValue, result); + } + + @Test + public void testInsertQueryTimeoutS() { + ConnectContext context = new ConnectContext(); + context.setQualifiedUser("a"); + context.setEnv(env); + int sessionValue = 1; + int propertyValue = 2; + // only session + context.getSessionVariable().setInsertTimeoutS(sessionValue); + long result = context.getInsertTimeoutS(); + Assert.assertEquals(sessionValue, result); + // has property + new Expectations() { + { + auth.getInsertTimeout(anyString); + minTimes = 0; + result = propertyValue; + } + }; + result = context.getInsertTimeoutS(); + Assert.assertEquals(propertyValue, result); + } } diff --git a/regression-test/suites/account_p0/test_property_session.groovy b/regression-test/suites/account_p0/test_property_session.groovy new file mode 100644 index 00000000000000..57b2dad747b123 --- /dev/null +++ b/regression-test/suites/account_p0/test_property_session.groovy @@ -0,0 +1,56 @@ +// 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. + +import org.junit.Assert; + +suite("test_property_session") { + String suiteName = "test_property_session" + String userName = "${suiteName}_user" + String pwd = 'C123_567p' + sql """drop user if exists ${userName}""" + sql """CREATE USER '${userName}' IDENTIFIED BY '${pwd}'""" + + //cloud-mode + if (isCloudMode()) { + def clusters = sql " SHOW CLUSTERS; " + assertTrue(!clusters.isEmpty()) + def validCluster = clusters[0][0] + sql """GRANT USAGE_PRIV ON CLUSTER `${validCluster}` TO ${userName}"""; + } + sql """GRANT select_PRIV ON *.*.* TO ${userName}"""; + connect(user=userName, password="${pwd}", url=context.config.jdbcUrl) { + sql """ + set query_timeout=1; + """ + test { + sql """ + select sleep(3); + """ + exception "timeout" + } + } + + // the priority of property should be higher than session + sql """set property for '${userName}' 'query_timeout' = '10';""" + connect(user=userName, password="${pwd}", url=context.config.jdbcUrl) { + sql """ + select sleep(3); + """ + } + + sql """drop user if exists ${userName}""" +}