Skip to content

Commit db3d027

Browse files
CaideyipiCaideyipi
committed
Pipe: use session context in legacy receiver loaders (#18111)
* Pipe: use session context in legacy receiver loaders * Update DeletionLoader.java * Fix legacy pipe receiver i18n message
1 parent 427d402 commit db3d027

6 files changed

Lines changed: 142 additions & 66 deletions

File tree

‎integration-test/src/test/java/org/apache/iotdb/pipe/it/single/IoTDBLegacyPipeReceiverSecurityIT.java‎

Lines changed: 107 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,10 @@
2222
import org.apache.iotdb.common.rpc.thrift.TSStatus;
2323
import org.apache.iotdb.commons.client.property.ThriftClientProperty;
2424
import org.apache.iotdb.commons.conf.IoTDBConstant;
25+
import org.apache.iotdb.commons.path.MeasurementPath;
2526
import org.apache.iotdb.commons.pipe.sink.client.IoTDBSyncClient;
27+
import org.apache.iotdb.db.pipe.sink.payload.legacy.PipeData;
28+
import org.apache.iotdb.db.storageengine.dataregion.modification.Deletion;
2629
import org.apache.iotdb.it.env.EnvFactory;
2730
import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper;
2831
import org.apache.iotdb.it.framework.IoTDBTestRunner;
@@ -35,22 +38,34 @@
3538
import org.apache.iotdb.service.rpc.thrift.TSyncIdentityInfo;
3639
import org.apache.iotdb.service.rpc.thrift.TSyncTransportMetaInfo;
3740

41+
import org.apache.tsfile.utils.ReadWriteIOUtils;
3842
import org.junit.AfterClass;
3943
import org.junit.Assert;
4044
import org.junit.BeforeClass;
4145
import org.junit.Test;
4246
import org.junit.experimental.categories.Category;
4347
import org.junit.runner.RunWith;
4448

49+
import java.io.ByteArrayOutputStream;
50+
import java.io.DataOutputStream;
4551
import java.io.File;
4652
import java.nio.ByteBuffer;
4753
import java.nio.charset.StandardCharsets;
54+
import java.sql.Connection;
55+
import java.sql.ResultSet;
56+
import java.sql.SQLException;
57+
import java.sql.Statement;
4858
import java.time.ZoneId;
4959

5060
@RunWith(IoTDBTestRunner.class)
5161
@Category({LocalStandaloneIT.class})
5262
public class IoTDBLegacyPipeReceiverSecurityIT {
5363

64+
private static final String LEGACY_PIPE_USER = "pipeHack";
65+
private static final String LEGACY_PIPE_PASSWORD = "StrngPsWd@623451";
66+
private static final String LEGACY_DATABASE = "root.legacy_poc";
67+
private static final String LEGACY_TIMESERIES = LEGACY_DATABASE + ".d1.s1";
68+
5469
@BeforeClass
5570
public static void setUp() {
5671
EnvFactory.getEnv().initClusterEnvironment();
@@ -99,11 +114,101 @@ public void testRejectPathTraversalFileNameInLegacyTransportFile() throws Except
99114
}
100115
}
101116

117+
@Test
118+
public void testLegacyPipeDataDeleteUsesAuthenticatedUserPermission() throws Exception {
119+
prepareLegacyPipePrivilegeEscalationData();
120+
assertDirectDeleteDeniedForLegacyPipeUser();
121+
122+
final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0);
123+
try (final IoTDBSyncClient client =
124+
new IoTDBSyncClient(
125+
new ThriftClientProperty.Builder().build(),
126+
dataNode.getIp(),
127+
dataNode.getPort(),
128+
false,
129+
null,
130+
null)) {
131+
final TSOpenSessionResp openSessionResp =
132+
client.openSession(createOpenSessionReq(LEGACY_PIPE_USER, LEGACY_PIPE_PASSWORD));
133+
Assert.assertEquals(
134+
TSStatusCode.SUCCESS_STATUS.getStatusCode(), openSessionResp.getStatus().getCode());
135+
136+
try {
137+
final TSStatus handshakeStatus =
138+
client.handshake(
139+
new TSyncIdentityInfo(
140+
"legacyPipePrivilege", System.currentTimeMillis(), "UNKNOWN", ""));
141+
Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), handshakeStatus.getCode());
142+
143+
final TSStatus status =
144+
client.sendPipeData(ByteBuffer.wrap(createDeletionPipeDataPayload()));
145+
146+
Assert.assertEquals(TSStatusCode.PIPESERVER_ERROR.getStatusCode(), status.getCode());
147+
} finally {
148+
client.closeSession(new TSCloseSessionReq(openSessionResp.getSessionId()));
149+
}
150+
}
151+
152+
assertLegacyPocRowCount(2);
153+
}
154+
155+
private void prepareLegacyPipePrivilegeEscalationData() throws SQLException {
156+
try (final Connection connection = EnvFactory.getEnv().getConnection();
157+
final Statement statement = connection.createStatement()) {
158+
statement.execute("CREATE DATABASE " + LEGACY_DATABASE);
159+
statement.execute(
160+
"CREATE TIMESERIES " + LEGACY_TIMESERIES + " WITH DATATYPE=INT64,ENCODING=PLAIN");
161+
statement.execute("INSERT INTO root.legacy_poc.d1(time,s1) VALUES (1,1),(2,2)");
162+
statement.execute("CREATE USER " + LEGACY_PIPE_USER + " '" + LEGACY_PIPE_PASSWORD + "'");
163+
statement.execute("GRANT USE_PIPE ON root.** TO USER " + LEGACY_PIPE_USER);
164+
}
165+
}
166+
167+
private void assertDirectDeleteDeniedForLegacyPipeUser() throws SQLException {
168+
try (final Connection connection =
169+
EnvFactory.getEnv().getConnection(LEGACY_PIPE_USER, LEGACY_PIPE_PASSWORD);
170+
final Statement statement = connection.createStatement()) {
171+
final SQLException exception =
172+
Assert.assertThrows(
173+
SQLException.class,
174+
() -> statement.execute("DELETE FROM " + LEGACY_TIMESERIES + " WHERE time <= 1"));
175+
Assert.assertTrue(exception.getMessage().contains("WRITE_DATA"));
176+
}
177+
}
178+
179+
private byte[] createDeletionPipeDataPayload() throws Exception {
180+
final ByteArrayOutputStream byteStream = new ByteArrayOutputStream();
181+
try (final DataOutputStream stream = new DataOutputStream(byteStream)) {
182+
stream.writeByte(PipeData.PipeDataType.DELETION.getType());
183+
stream.writeLong(1L);
184+
ReadWriteIOUtils.write(LEGACY_DATABASE, stream);
185+
new Deletion(new MeasurementPath(LEGACY_TIMESERIES), 0, Long.MIN_VALUE, 1)
186+
.serializeWithoutFileOffset(stream);
187+
}
188+
return byteStream.toByteArray();
189+
}
190+
191+
private void assertLegacyPocRowCount(final int expectedCount) throws SQLException {
192+
try (final Connection connection = EnvFactory.getEnv().getConnection();
193+
final Statement statement = connection.createStatement();
194+
final ResultSet resultSet = statement.executeQuery("SELECT s1 FROM root.legacy_poc.d1")) {
195+
int actualCount = 0;
196+
while (resultSet.next()) {
197+
++actualCount;
198+
}
199+
Assert.assertEquals(expectedCount, actualCount);
200+
}
201+
}
202+
102203
private TSOpenSessionReq createOpenSessionReq() {
204+
return createOpenSessionReq("root", "root");
205+
}
206+
207+
private TSOpenSessionReq createOpenSessionReq(final String username, final String password) {
103208
final TSOpenSessionReq req = new TSOpenSessionReq();
104209
req.setClient_protocol(TSProtocolVersion.IOTDB_SERVICE_PROTOCOL_V3);
105-
req.setUsername("root");
106-
req.setPassword("root");
210+
req.setUsername(username);
211+
req.setPassword(password);
107212
req.setZoneId(ZoneId.systemDefault().toString());
108213
req.putToConfiguration("version", IoTDBConstant.ClientVersion.V_1_0.toString());
109214
req.putToConfiguration("sql_dialect", "tree");

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/legacy/IoTDBLegacyPipeReceiverAgent.java‎

Lines changed: 16 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -22,21 +22,14 @@
2222

2323
import org.apache.iotdb.common.rpc.thrift.TSStatus;
2424
import org.apache.iotdb.commons.conf.CommonDescriptor;
25-
import org.apache.iotdb.commons.exception.IllegalPathException;
26-
import org.apache.iotdb.commons.path.PartialPath;
2725
import org.apache.iotdb.commons.pipe.receiver.PipeReceiverFilePathUtils;
2826
import org.apache.iotdb.commons.utils.FileUtils;
29-
import org.apache.iotdb.db.auth.AuthorityChecker;
30-
import org.apache.iotdb.db.conf.IoTDBDescriptor;
3127
import org.apache.iotdb.db.pipe.sink.payload.legacy.PipeData;
3228
import org.apache.iotdb.db.pipe.sink.payload.legacy.TsFilePipeData;
3329
import org.apache.iotdb.db.protocol.session.SessionManager;
3430
import org.apache.iotdb.db.queryengine.common.SessionInfo;
35-
import org.apache.iotdb.db.queryengine.plan.Coordinator;
3631
import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher;
3732
import org.apache.iotdb.db.queryengine.plan.analyze.schema.ISchemaFetcher;
38-
import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult;
39-
import org.apache.iotdb.db.queryengine.plan.statement.metadata.DatabaseSchemaStatement;
4033
import org.apache.iotdb.pipe.api.exception.PipeException;
4134
import org.apache.iotdb.rpc.RpcUtils;
4235
import org.apache.iotdb.rpc.TSStatusCode;
@@ -53,7 +46,6 @@
5346
import java.io.RandomAccessFile;
5447
import java.nio.ByteBuffer;
5548
import java.nio.file.Paths;
56-
import java.time.ZoneId;
5749
import java.util.Map;
5850
import java.util.Objects;
5951
import java.util.concurrent.ConcurrentHashMap;
@@ -76,8 +68,6 @@ public class IoTDBLegacyPipeReceiverAgent {
7668
private final Map<Long, Map<String, Long>> connectionIdToStartIndexRecord =
7769
new ConcurrentHashMap<>();
7870

79-
private final Map<String, String> registeredDatabase = new ConcurrentHashMap<>();
80-
8171
// The sync connectionId is unique in one IoTDB instance.
8272
private final AtomicLong connectionIdGenerator = new AtomicLong();
8373

@@ -118,12 +108,6 @@ public TSStatus handshake(
118108
new File(getFileDataDir(identityInfo)).mkdirs();
119109
}
120110
createConnection(identityInfo);
121-
if (!StringUtils.isEmpty(identityInfo.getDatabase())
122-
&& !registerDatabase(identityInfo.getDatabase(), partitionFetcher, schemaFetcher)) {
123-
return RpcUtils.getStatus(
124-
TSStatusCode.PIPESERVER_ERROR,
125-
String.format("Auto register database %s error.", identityInfo.getDatabase()));
126-
}
127111
return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS, "");
128112
}
129113

@@ -138,42 +122,6 @@ private void createConnection(final SyncIdentityInfo identityInfo) {
138122
connectionIdToIdentityInfoMap.put(connectionId, identityInfo);
139123
}
140124

141-
private boolean registerDatabase(
142-
String database, IPartitionFetcher partitionFetcher, ISchemaFetcher schemaFetcher) {
143-
if (registeredDatabase.containsKey(database)) {
144-
return true;
145-
}
146-
try {
147-
DatabaseSchemaStatement statement =
148-
new DatabaseSchemaStatement(DatabaseSchemaStatement.DatabaseSchemaStatementType.CREATE);
149-
statement.setDatabasePath(new PartialPath(database));
150-
long queryId = SessionManager.getInstance().requestQueryId();
151-
ExecutionResult result =
152-
Coordinator.getInstance()
153-
.executeForTreeModel(
154-
statement,
155-
queryId,
156-
new SessionInfo(0, AuthorityChecker.SUPER_USER, ZoneId.systemDefault(), ""),
157-
"",
158-
partitionFetcher,
159-
schemaFetcher,
160-
IoTDBDescriptor.getInstance().getConfig().getQueryTimeoutThreshold(),
161-
false);
162-
if (result.status.code != TSStatusCode.SUCCESS_STATUS.getStatusCode()
163-
&& result.status.code != TSStatusCode.DATABASE_ALREADY_EXISTS.getStatusCode()) {
164-
LOGGER.error(
165-
"Create Database error, statement: {}, result status : {}.", statement, result.status);
166-
return false;
167-
}
168-
} catch (IllegalPathException e) {
169-
LOGGER.error("Parse database PartialPath {} error", database, e);
170-
return false;
171-
}
172-
173-
registeredDatabase.put(database, "");
174-
return true;
175-
}
176-
177125
/**
178126
* Receive {@link PipeData} and load it into IoTDB Engine.
179127
*
@@ -182,7 +130,12 @@ private boolean registerDatabase(
182130
* @throws TException The connection between the sender and the receiver has not been established
183131
* by {@link IoTDBLegacyPipeReceiverAgent#handshake}
184132
*/
185-
public TSStatus transportPipeData(ByteBuffer buff) throws TException {
133+
public TSStatus transportPipeData(final ByteBuffer buff) throws TException {
134+
return transportPipeData(buff, null);
135+
}
136+
137+
public TSStatus transportPipeData(final ByteBuffer buff, final SessionInfo sessionInfo)
138+
throws TException {
186139
// step1. check connection
187140
SyncIdentityInfo identityInfo = getCurrentSyncIdentityInfo();
188141
if (identityInfo == null) {
@@ -217,7 +170,7 @@ public TSStatus transportPipeData(ByteBuffer buff) throws TException {
217170
pipeData.getPipeDataType(),
218171
pipeData);
219172
try {
220-
pipeData.createLoader().load();
173+
pipeData.createLoader().load(sessionInfo == null ? getCurrentSessionInfo() : sessionInfo);
221174
LOGGER.info(
222175
"Load pipeData with serialize number {} successfully.", pipeData.getSerialNumber());
223176
} catch (PipeException e) {
@@ -243,6 +196,15 @@ private SyncIdentityInfo getCurrentSyncIdentityInfo() {
243196
}
244197
}
245198

199+
private SessionInfo getCurrentSessionInfo() {
200+
final SessionManager sessionManager = SessionManager.getInstance();
201+
if (!sessionManager.checkLogin(sessionManager.getCurrSession())) {
202+
throw new PipeException("Legacy pipe receiver requires a logged-in session.");
203+
}
204+
205+
return sessionManager.getSessionInfo(sessionManager.getCurrSession());
206+
}
207+
246208
/**
247209
* handle when successfully receive tsFilePipeData. Rename .patch file and reset tsFilePipeData's
248210
* path.

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/legacy/loader/DeletionLoader.java‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
package org.apache.iotdb.db.pipe.receiver.protocol.legacy.loader;
2121

22+
import org.apache.iotdb.common.rpc.thrift.TSStatus;
2223
import org.apache.iotdb.commons.conf.CommonDescriptor;
2324
import org.apache.iotdb.db.auth.AuthorityChecker;
2425
import org.apache.iotdb.db.conf.IoTDBDescriptor;
@@ -37,7 +38,6 @@
3738
import org.slf4j.Logger;
3839
import org.slf4j.LoggerFactory;
3940

40-
import java.time.ZoneId;
4141
import java.util.Collections;
4242

4343
/** This loader is used to load deletion plan. */
@@ -52,19 +52,25 @@ public DeletionLoader(Deletion deletion) {
5252
}
5353

5454
@Override
55-
public void load() throws PipeException {
55+
public void load(final SessionInfo sessionInfo) throws PipeException {
5656
if (CommonDescriptor.getInstance().getConfig().isReadOnly()) {
5757
throw new PipeException("storage engine readonly");
5858
}
5959
try {
6060
Statement statement = generateStatement();
61+
TSStatus authorityStatus =
62+
AuthorityChecker.checkAuthority(statement, sessionInfo.getUserName());
63+
if (authorityStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
64+
throw new PipeException(authorityStatus.getMessage());
65+
}
66+
6167
long queryId = SessionManager.getInstance().requestQueryId();
6268
ExecutionResult result =
6369
Coordinator.getInstance()
6470
.executeForTreeModel(
6571
statement,
6672
queryId,
67-
new SessionInfo(0, AuthorityChecker.SUPER_USER, ZoneId.systemDefault(), ""),
73+
sessionInfo,
6874
"",
6975
PARTITION_FETCHER,
7076
SCHEMA_FETCHER,

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/legacy/loader/ILoader.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
package org.apache.iotdb.db.pipe.receiver.protocol.legacy.loader;
2121

22+
import org.apache.iotdb.db.queryengine.common.SessionInfo;
2223
import org.apache.iotdb.db.queryengine.plan.analyze.ClusterPartitionFetcher;
2324
import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher;
2425
import org.apache.iotdb.db.queryengine.plan.analyze.schema.ClusterSchemaFetcher;
@@ -34,5 +35,5 @@ public interface ILoader {
3435

3536
ISchemaFetcher SCHEMA_FETCHER = ClusterSchemaFetcher.getInstance();
3637

37-
void load();
38+
void load(SessionInfo sessionInfo);
3839
}

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/legacy/loader/TsFileLoader.java‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121

2222
import org.apache.iotdb.commons.exception.IllegalPathException;
2323
import org.apache.iotdb.commons.path.PartialPath;
24-
import org.apache.iotdb.db.auth.AuthorityChecker;
2524
import org.apache.iotdb.db.conf.IoTDBDescriptor;
2625
import org.apache.iotdb.db.exception.load.LoadFileException;
2726
import org.apache.iotdb.db.protocol.session.SessionManager;
@@ -36,7 +35,6 @@
3635
import org.slf4j.LoggerFactory;
3736

3837
import java.io.File;
39-
import java.time.ZoneId;
4038

4139
/** This loader is used to load tsFiles. If .mods file exists, it will be loaded as well. */
4240
public class TsFileLoader implements ILoader {
@@ -52,22 +50,23 @@ public TsFileLoader(File tsFile, String database) {
5250
}
5351

5452
@Override
55-
public void load() {
53+
public void load(final SessionInfo sessionInfo) {
5654
try {
5755
LoadTsFileStatement statement = LoadTsFileStatement.createUnchecked(tsFile.getAbsolutePath());
5856
statement.setDeleteAfterLoad(true);
5957
statement.setConvertOnTypeMismatch(true);
6058
statement.setDatabaseLevel(parseSgLevel());
6159
statement.setVerifySchema(true);
62-
statement.setAutoCreateDatabase(false);
60+
statement.setAutoCreateDatabase(
61+
IoTDBDescriptor.getInstance().getConfig().isAutoCreateSchemaEnabled());
6362

6463
long queryId = SessionManager.getInstance().requestQueryId();
6564
ExecutionResult result =
6665
Coordinator.getInstance()
6766
.executeForTreeModel(
6867
statement,
6968
queryId,
70-
new SessionInfo(0, AuthorityChecker.SUPER_USER, ZoneId.systemDefault(), ""),
69+
sessionInfo,
7170
"",
7271
PARTITION_FETCHER,
7372
SCHEMA_FETCHER,

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2776,7 +2776,10 @@ public TSStatus sendPipeData(final ByteBuffer buff) throws TException {
27762776
if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
27772777
return status;
27782778
}
2779-
return PipeDataNodeAgent.receiver().legacy().transportPipeData(buff);
2779+
return PipeDataNodeAgent.receiver()
2780+
.legacy()
2781+
.transportPipeData(
2782+
buff, SESSION_MANAGER.getSessionInfo(SESSION_MANAGER.getCurrSession()));
27802783
} finally {
27812784
SESSION_MANAGER.updateIdleTime();
27822785
}

0 commit comments

Comments
 (0)