Skip to content

Commit 06d9f3b

Browse files
authored
Add UT for query auth check
1 parent 4513e55 commit 06d9f3b

File tree

3 files changed

+256
-3
lines changed

3 files changed

+256
-3
lines changed

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/AccessControlImpl.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,6 @@ public void checkCanShowOrDescTable(String userName, QualifiedObjectName tableNa
8686

8787
@Override
8888
public void checkUserHasMaintainPrivilege(String userName) {
89-
authChecker.checkMaintainPrivilege(userName);
89+
authChecker.checkGlobalPrivilege(userName, TableModelPrivilege.MAINTAIN);
9090
}
9191
}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/ITableAuthChecker.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,10 +66,11 @@ void checkTablePrivilege(
6666
void checkTableVisibility(String userName, QualifiedObjectName tableName);
6767

6868
/**
69-
* Check if user has global maintain privilege
69+
* Check if user has the specified global privilege
7070
*
7171
* @param userName name of user
72+
* @param privilege specified global privilege to be checked
7273
* @throws AccessDeniedException if not allowed
7374
*/
74-
void checkMaintainPrivilege(String userName);
75+
void checkGlobalPrivilege(String userName, TableModelPrivilege privilege);
7576
}
Lines changed: 252 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,252 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.iotdb.db.queryengine.plan.relational.analyzer;
21+
22+
import org.apache.iotdb.commons.exception.auth.AccessDeniedException;
23+
import org.apache.iotdb.db.protocol.session.IClientSession;
24+
import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
25+
import org.apache.iotdb.db.queryengine.common.SessionInfo;
26+
import org.apache.iotdb.db.queryengine.plan.execution.config.TableConfigTaskVisitor;
27+
import org.apache.iotdb.db.queryengine.plan.relational.metadata.QualifiedObjectName;
28+
import org.apache.iotdb.db.queryengine.plan.relational.security.AccessControlImpl;
29+
import org.apache.iotdb.db.queryengine.plan.relational.security.ITableAuthChecker;
30+
import org.apache.iotdb.db.queryengine.plan.relational.security.TableModelPrivilege;
31+
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Statement;
32+
import org.apache.iotdb.db.queryengine.plan.relational.sql.parser.SqlParser;
33+
34+
import org.junit.Test;
35+
import org.mockito.Mockito;
36+
37+
import java.time.ZoneId;
38+
import java.util.Collections;
39+
40+
import static org.apache.iotdb.db.queryengine.execution.warnings.WarningCollector.NOOP;
41+
import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestMatadata.DB1;
42+
import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestMatadata.TABLE1;
43+
import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestUtils.QUERY_ID;
44+
import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestUtils.TEST_MATADATA;
45+
import static org.junit.Assert.assertEquals;
46+
import static org.junit.Assert.fail;
47+
import static org.mockito.ArgumentMatchers.any;
48+
import static org.mockito.ArgumentMatchers.eq;
49+
50+
public class AuthTest {
51+
52+
private final SqlParser sqlParser = new SqlParser();
53+
54+
private final ZoneId zoneId = ZoneId.systemDefault();
55+
56+
private final QualifiedObjectName testdbTable1 = new QualifiedObjectName(DB1, TABLE1);
57+
58+
private final String userRoot = "root";
59+
60+
private final String user1 = "user1";
61+
62+
private final String user2 = "user2";
63+
64+
@Test
65+
public void testQueryRelatedAuth() {
66+
String sql = String.format("SELECT * FROM %s.%s", DB1, TABLE1);
67+
ITableAuthChecker authChecker = Mockito.mock(ITableAuthChecker.class);
68+
// user `root` always succeed
69+
Mockito.doNothing().when(authChecker).checkTablePrivilege(eq(userRoot), any(), any());
70+
// user `user1` don't hava testdb.table1's SELECT privilege
71+
String errorMsg =
72+
String.format(
73+
"%s doesn't have %s privilege on TABLE %s.%s",
74+
user1, TableModelPrivilege.SELECT, DB1, TABLE1);
75+
Mockito.doThrow(new AccessDeniedException(errorMsg))
76+
.when(authChecker)
77+
.checkTablePrivilege(eq(user1), eq(testdbTable1), eq(TableModelPrivilege.SELECT));
78+
// user `user2` has testdb.table1's SELECT privilege
79+
Mockito.doNothing()
80+
.when(authChecker)
81+
.checkTablePrivilege(eq(user2), eq(testdbTable1), eq(TableModelPrivilege.SELECT));
82+
83+
String userName = userRoot;
84+
try {
85+
analyzeSQL(sql, userName, authChecker);
86+
} catch (Exception e) {
87+
fail(e.getMessage());
88+
}
89+
90+
userName = user1;
91+
try {
92+
analyzeSQL(sql, userName, authChecker);
93+
fail("user1 should be denied");
94+
} catch (AccessDeniedException e) {
95+
assertEquals(AccessDeniedException.PREFIX + errorMsg, e.getMessage());
96+
} catch (Exception e) {
97+
fail("Unexpected exception : " + e.getMessage());
98+
}
99+
100+
userName = user2;
101+
try {
102+
analyzeSQL(sql, userName, authChecker, DB1);
103+
} catch (Exception e) {
104+
fail(e.getMessage());
105+
}
106+
}
107+
108+
@Test
109+
public void testDatabaseManagementRelatedAuth() {
110+
111+
// create database
112+
String databaseTest1 = "test1";
113+
String sql = String.format("CREATE DATABASE %s", databaseTest1);
114+
ITableAuthChecker authChecker = Mockito.mock(ITableAuthChecker.class);
115+
// user `root` always succeed
116+
Mockito.doNothing().when(authChecker).checkDatabasePrivilege(eq(userRoot), any(), any());
117+
// user `user1` don't hava test1's CREATE privilege
118+
String errorMsg =
119+
String.format(
120+
"%s doesn't have %s privilege on DATABASE %s",
121+
user1, TableModelPrivilege.CREATE, databaseTest1);
122+
Mockito.doThrow(new AccessDeniedException(errorMsg))
123+
.when(authChecker)
124+
.checkDatabasePrivilege(eq(user1), eq(databaseTest1), eq(TableModelPrivilege.CREATE));
125+
// user `user2` has test1's CREATE privilege
126+
Mockito.doNothing()
127+
.when(authChecker)
128+
.checkDatabasePrivilege(eq(user2), eq(databaseTest1), eq(TableModelPrivilege.CREATE));
129+
130+
String userName = userRoot;
131+
try {
132+
analyzeConfigTask(sql, userName, authChecker);
133+
} catch (Exception e) {
134+
fail(e.getMessage());
135+
}
136+
137+
userName = user1;
138+
try {
139+
analyzeConfigTask(sql, userName, authChecker);
140+
fail("user1 should be denied");
141+
} catch (AccessDeniedException e) {
142+
assertEquals(AccessDeniedException.PREFIX + errorMsg, e.getMessage());
143+
} catch (Exception e) {
144+
fail("Unexpected exception : " + e.getMessage());
145+
}
146+
147+
userName = user2;
148+
try {
149+
analyzeConfigTask(sql, userName, authChecker);
150+
} catch (Exception e) {
151+
fail(e.getMessage());
152+
}
153+
154+
// use database
155+
String databaseTest2 = "test2";
156+
// user `root` always succeed
157+
Mockito.doNothing().when(authChecker).checkDatabaseVisibility(eq(userRoot), any());
158+
// user `user1` can't see DATABASE test1
159+
errorMsg = String.format("%s has no privileges on DATABASE %s", user1, databaseTest1);
160+
Mockito.doThrow(new AccessDeniedException(errorMsg))
161+
.when(authChecker)
162+
.checkDatabaseVisibility(eq(user1), eq(databaseTest1));
163+
// user `user1` can see DATABASE test2
164+
Mockito.doNothing().when(authChecker).checkDatabaseVisibility(eq(user1), eq(databaseTest2));
165+
166+
sql = String.format("USE %s", databaseTest1);
167+
userName = userRoot;
168+
try {
169+
analyzeConfigTask(sql, userName, authChecker);
170+
} catch (Exception e) {
171+
fail(e.getMessage());
172+
}
173+
174+
userName = user1;
175+
try {
176+
analyzeConfigTask(sql, userName, authChecker);
177+
fail("user1 should be denied");
178+
} catch (AccessDeniedException e) {
179+
assertEquals(AccessDeniedException.PREFIX + errorMsg, e.getMessage());
180+
} catch (Exception e) {
181+
fail("Unexpected exception : " + e.getMessage());
182+
}
183+
184+
sql = String.format("USE %s", databaseTest2);
185+
try {
186+
analyzeConfigTask(sql, userName, authChecker);
187+
} catch (Exception e) {
188+
fail(e.getMessage());
189+
}
190+
191+
// TODO show databases
192+
193+
// TODO show databases
194+
195+
// TODO drop database
196+
197+
// TODO alter database
198+
199+
}
200+
201+
private void analyzeSQL(String sql, String userName, ITableAuthChecker authChecker) {
202+
analyzeSQL(sql, userName, authChecker, null);
203+
}
204+
205+
private void analyzeSQL(
206+
String sql,
207+
String userName,
208+
ITableAuthChecker authChecker,
209+
String databaseNameInSessionInfo) {
210+
Statement statement = sqlParser.createStatement(sql, zoneId);
211+
SessionInfo session =
212+
new SessionInfo(
213+
0, userName, zoneId, databaseNameInSessionInfo, IClientSession.SqlDialect.TABLE);
214+
StatementAnalyzerFactory statementAnalyzerFactory =
215+
new StatementAnalyzerFactory(TEST_MATADATA, sqlParser, new AccessControlImpl(authChecker));
216+
MPPQueryContext context = new MPPQueryContext(sql, QUERY_ID, 0, session, null, null);
217+
Analyzer analyzer =
218+
new Analyzer(
219+
context,
220+
session,
221+
statementAnalyzerFactory,
222+
Collections.emptyList(),
223+
Collections.emptyMap(),
224+
NOOP);
225+
analyzer.analyze(statement);
226+
}
227+
228+
private void analyzeConfigTask(String sql, String userName, ITableAuthChecker authChecker) {
229+
Statement statement = sqlParser.createStatement(sql, zoneId);
230+
SessionInfo session =
231+
new SessionInfo(0, userName, zoneId, null, IClientSession.SqlDialect.TABLE);
232+
MPPQueryContext context = new MPPQueryContext(sql, QUERY_ID, 0, session, null, null);
233+
234+
statement.accept(
235+
new TableConfigTaskVisitor(
236+
Mockito.mock(IClientSession.class), TEST_MATADATA, new AccessControlImpl(authChecker)),
237+
context);
238+
}
239+
240+
private void analyzeConfigTask(
241+
String sql, String userName, ITableAuthChecker authChecker, IClientSession clientSession) {
242+
Statement statement = sqlParser.createStatement(sql, zoneId);
243+
SessionInfo session =
244+
new SessionInfo(0, userName, zoneId, null, IClientSession.SqlDialect.TABLE);
245+
MPPQueryContext context = new MPPQueryContext(sql, QUERY_ID, 0, session, null, null);
246+
247+
statement.accept(
248+
new TableConfigTaskVisitor(
249+
clientSession, TEST_MATADATA, new AccessControlImpl(authChecker)),
250+
context);
251+
}
252+
}

0 commit comments

Comments
 (0)