Pipe: use session context in legacy receiver loaders

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

@@ -22,7 +22,10 @@
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.commons.client.property.ThriftClientProperty;
 import org.apache.iotdb.commons.conf.IoTDBConstant;
+import org.apache.iotdb.commons.path.MeasurementPath;
 import org.apache.iotdb.commons.pipe.sink.client.IoTDBSyncClient;
+import org.apache.iotdb.db.pipe.sink.payload.legacy.PipeData;
+import org.apache.iotdb.db.storageengine.dataregion.modification.v1.Deletion;
 import org.apache.iotdb.isession.SessionConfig;
 import org.apache.iotdb.it.env.EnvFactory;
 import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper;
@@ -36,22 +39,34 @@
 import org.apache.iotdb.service.rpc.thrift.TSyncIdentityInfo;
 import org.apache.iotdb.service.rpc.thrift.TSyncTransportMetaInfo;
 
+import org.apache.tsfile.utils.ReadWriteIOUtils;
 import org.junit.AfterClass;
 import org.junit.Assert;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
 
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
 import java.io.File;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
 import java.time.ZoneId;
 
 @RunWith(IoTDBTestRunner.class)
 @Category({LocalStandaloneIT.class})
 public class IoTDBLegacyPipeReceiverSecurityIT {
 
+  private static final String LEGACY_PIPE_USER = "pipeHack";
+  private static final String LEGACY_PIPE_PASSWORD = "StrngPsWd@623451";
+  private static final String LEGACY_DATABASE = "root.legacy_poc";
+  private static final String LEGACY_TIMESERIES = LEGACY_DATABASE + ".d1.s1";
+
   @BeforeClass
   public static void setUp() {
     EnvFactory.getEnv().getConfig().getCommonConfig().setDatanodeMemoryProportion("3:3:1:1:1:0");
@@ -101,11 +116,101 @@ public void testRejectPathTraversalFileNameInLegacyTransportFile() throws Except
     }
   }
 
+  @Test
+  public void testLegacyPipeDataDeleteUsesAuthenticatedUserPermission() throws Exception {
+    prepareLegacyPipePrivilegeEscalationData();
+    assertDirectDeleteDeniedForLegacyPipeUser();
+
+    final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0);
+    try (final IoTDBSyncClient client =
+        new IoTDBSyncClient(
+            new ThriftClientProperty.Builder().build(),
+            dataNode.getIp(),
+            dataNode.getPort(),
+            false,
+            null,
+            null)) {
+      final TSOpenSessionResp openSessionResp =
+          client.openSession(createOpenSessionReq(LEGACY_PIPE_USER, LEGACY_PIPE_PASSWORD));
+      Assert.assertEquals(
+          TSStatusCode.SUCCESS_STATUS.getStatusCode(), openSessionResp.getStatus().getCode());
+
+      try {
+        final TSStatus handshakeStatus =
+            client.handshake(
+                new TSyncIdentityInfo(
+                    "legacyPipePrivilege", System.currentTimeMillis(), "UNKNOWN", ""));
+        Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), handshakeStatus.getCode());
+
+        final TSStatus status =
+            client.sendPipeData(ByteBuffer.wrap(createDeletionPipeDataPayload()));
+
+        Assert.assertEquals(TSStatusCode.PIPESERVER_ERROR.getStatusCode(), status.getCode());
+      } finally {
+        client.closeSession(new TSCloseSessionReq(openSessionResp.getSessionId()));
+      }
+    }
+
+    assertLegacyPocRowCount(2);
+  }
+
+  private void prepareLegacyPipePrivilegeEscalationData() throws SQLException {
+    try (final Connection connection = EnvFactory.getEnv().getConnection();
+        final Statement statement = connection.createStatement()) {
+      statement.execute("CREATE DATABASE " + LEGACY_DATABASE);
+      statement.execute(
+          "CREATE TIMESERIES " + LEGACY_TIMESERIES + " WITH DATATYPE=INT64,ENCODING=PLAIN");
+      statement.execute("INSERT INTO root.legacy_poc.d1(time,s1) VALUES (1,1),(2,2)");
+      statement.execute("CREATE USER " + LEGACY_PIPE_USER + " '" + LEGACY_PIPE_PASSWORD + "'");
+      statement.execute("GRANT SYSTEM ON root.** TO USER " + LEGACY_PIPE_USER);
+    }
+  }
+
+  private void assertDirectDeleteDeniedForLegacyPipeUser() throws SQLException {
+    try (final Connection connection =
+            EnvFactory.getEnv().getConnection(LEGACY_PIPE_USER, LEGACY_PIPE_PASSWORD);
+        final Statement statement = connection.createStatement()) {
+      final SQLException exception =
+          Assert.assertThrows(
+              SQLException.class,
+              () -> statement.execute("DELETE FROM " + LEGACY_TIMESERIES + " WHERE time <= 1"));
+      Assert.assertTrue(exception.getMessage().contains("WRITE_DATA"));
+    }
+  }
+
+  private byte[] createDeletionPipeDataPayload() throws Exception {
+    final ByteArrayOutputStream byteStream = new ByteArrayOutputStream();
+    try (final DataOutputStream stream = new DataOutputStream(byteStream)) {
+      stream.writeByte(PipeData.PipeDataType.DELETION.getType());
+      stream.writeLong(1L);
+      ReadWriteIOUtils.write(LEGACY_DATABASE, stream);
+      new Deletion(new MeasurementPath(LEGACY_TIMESERIES), 0, Long.MIN_VALUE, 1)
+          .serializeWithoutFileOffset(stream);
+    }
+    return byteStream.toByteArray();
+  }
+
+  private void assertLegacyPocRowCount(final int expectedCount) throws SQLException {
+    try (final Connection connection = EnvFactory.getEnv().getConnection();
+        final Statement statement = connection.createStatement();
+        final ResultSet resultSet = statement.executeQuery("SELECT s1 FROM root.legacy_poc.d1")) {
+      int actualCount = 0;
+      while (resultSet.next()) {
+        ++actualCount;
+      }
+      Assert.assertEquals(expectedCount, actualCount);
+    }
+  }
+
   private TSOpenSessionReq createOpenSessionReq() {
+    return createOpenSessionReq(SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD);
+  }
+
+  private TSOpenSessionReq createOpenSessionReq(final String username, final String password) {
     final TSOpenSessionReq req = new TSOpenSessionReq();
     req.setClient_protocol(TSProtocolVersion.IOTDB_SERVICE_PROTOCOL_V3);
-    req.setUsername(SessionConfig.DEFAULT_USER);
-    req.setPassword(SessionConfig.DEFAULT_PASSWORD);
+    req.setUsername(username);
+    req.setPassword(password);
     req.setZoneId(ZoneId.systemDefault().toString());
     req.putToConfiguration("version", IoTDBConstant.ClientVersion.V_1_0.toString());
     req.putToConfiguration("sql_dialect", "tree");

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

@@ -21,26 +21,18 @@
 package org.apache.iotdb.db.pipe.receiver.protocol.legacy;
 
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
-import org.apache.iotdb.commons.audit.UserEntity;
 import org.apache.iotdb.commons.conf.CommonDescriptor;
-import org.apache.iotdb.commons.exception.IllegalPathException;
 import org.apache.iotdb.commons.i18n.PipeMessages;
-import org.apache.iotdb.commons.path.PartialPath;
 import org.apache.iotdb.commons.pipe.receiver.PipeReceiverFilePathUtils;
 import org.apache.iotdb.commons.queryengine.common.SessionInfo;
 import org.apache.iotdb.commons.utils.FileUtils;
-import org.apache.iotdb.db.auth.AuthorityChecker;
-import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.i18n.DataNodeMiscMessages;
 import org.apache.iotdb.db.i18n.DataNodePipeMessages;
 import org.apache.iotdb.db.pipe.sink.payload.legacy.PipeData;
 import org.apache.iotdb.db.pipe.sink.payload.legacy.TsFilePipeData;
 import org.apache.iotdb.db.protocol.session.SessionManager;
-import org.apache.iotdb.db.queryengine.plan.Coordinator;
 import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher;
 import org.apache.iotdb.db.queryengine.plan.analyze.schema.ISchemaFetcher;
-import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult;
-import org.apache.iotdb.db.queryengine.plan.statement.metadata.DatabaseSchemaStatement;
 import org.apache.iotdb.pipe.api.exception.PipeException;
 import org.apache.iotdb.rpc.RpcUtils;
 import org.apache.iotdb.rpc.TSStatusCode;
@@ -57,7 +49,6 @@
 import java.io.RandomAccessFile;
 import java.nio.ByteBuffer;
 import java.nio.file.Paths;
-import java.time.ZoneId;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
@@ -80,8 +71,6 @@ public class IoTDBLegacyPipeReceiverAgent {
   private final Map<Long, Map<String, Long>> connectionIdToStartIndexRecord =
       new ConcurrentHashMap<>();
 
-  private final Map<String, String> registeredDatabase = new ConcurrentHashMap<>();
-
   // The sync connectionId is unique in one IoTDB instance.
   private final AtomicLong connectionIdGenerator = new AtomicLong();
 
@@ -124,12 +113,6 @@ public TSStatus handshake(
       new File(getFileDataDir(identityInfo)).mkdirs();
     }
     createConnection(identityInfo);
-    if (!StringUtils.isEmpty(identityInfo.getDatabase())
-        && !registerDatabase(identityInfo.getDatabase(), partitionFetcher, schemaFetcher)) {
-      return RpcUtils.getStatus(
-          TSStatusCode.PIPESERVER_ERROR,
-          String.format("Auto register database %s error.", identityInfo.getDatabase()));
-    }
     return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS, "");
   }
 
@@ -144,53 +127,6 @@ private void createConnection(final SyncIdentityInfo identityInfo) {
     connectionIdToIdentityInfoMap.put(connectionId, identityInfo);
   }
 
-  private boolean registerDatabase(
-      final String database,
-      final IPartitionFetcher partitionFetcher,
-      final ISchemaFetcher schemaFetcher) {
-    if (registeredDatabase.containsKey(database)) {
-      return true;
-    }
-    try {
-      final DatabaseSchemaStatement statement =
-          new DatabaseSchemaStatement(DatabaseSchemaStatement.DatabaseSchemaStatementType.CREATE);
-      statement.setDatabasePath(new PartialPath(database));
-      final long queryId = SessionManager.getInstance().requestQueryId();
-      final ExecutionResult result =
-          Coordinator.getInstance()
-              .executeForTreeModel(
-                  statement,
-                  queryId,
-                  new SessionInfo(
-                      0,
-                      new UserEntity(
-                          AuthorityChecker.SUPER_USER_ID,
-                          AuthorityChecker.SUPER_USER,
-                          IoTDBDescriptor.getInstance().getConfig().getInternalAddress()),
-                      ZoneId.systemDefault()),
-                  "",
-                  partitionFetcher,
-                  schemaFetcher,
-                  IoTDBDescriptor.getInstance().getConfig().getQueryTimeoutThreshold(),
-                  false,
-                  false);
-      if (result.status.code != TSStatusCode.SUCCESS_STATUS.getStatusCode()
-          && result.status.code != TSStatusCode.DATABASE_ALREADY_EXISTS.getStatusCode()) {
-        LOGGER.error(
-            DataNodePipeMessages.CREATE_DATABASE_ERROR_STATEMENT_RESULT_STATUS,
-            statement,
-            result.status);
-        return false;
-      }
-    } catch (final IllegalPathException e) {
-      LOGGER.error(DataNodePipeMessages.PARSE_DATABASE_PARTIALPATH_ERROR, database, e);
-      return false;
-    }
-
-    registeredDatabase.put(database, "");
-    return true;
-  }
-
   /**
    * Receive {@link PipeData} and load it into IoTDB Engine.
    *
@@ -200,6 +136,11 @@ private boolean registerDatabase(
    *     by {@link IoTDBLegacyPipeReceiverAgent#handshake}
    */
   public TSStatus transportPipeData(final ByteBuffer buff) throws TException {
+    return transportPipeData(buff, null);
+  }
+
+  public TSStatus transportPipeData(final ByteBuffer buff, final SessionInfo sessionInfo)
+      throws TException {
     // step1. check connection
     final SyncIdentityInfo identityInfo = getCurrentSyncIdentityInfo();
     if (identityInfo == null) {
@@ -235,7 +176,7 @@ public TSStatus transportPipeData(final ByteBuffer buff) throws TException {
         pipeData.getPipeDataType(),
         pipeData);
     try {
-      pipeData.createLoader().load();
+      pipeData.createLoader().load(sessionInfo == null ? getCurrentSessionInfo() : sessionInfo);
       LOGGER.info(
           DataNodePipeMessages.LOAD_PIPEDATA_WITH_SERIALIZE_NUMBER_SUCCESSFULLY,
           pipeData.getSerialNumber());
@@ -262,6 +203,15 @@ private SyncIdentityInfo getCurrentSyncIdentityInfo() {
     }
   }
 
+  private SessionInfo getCurrentSessionInfo() {
+    final SessionManager sessionManager = SessionManager.getInstance();
+    if (!sessionManager.checkLogin(sessionManager.getCurrSession())) {
+      throw new PipeException("Legacy pipe receiver requires a logged-in session.");
+    }
+
+    return sessionManager.getSessionInfoOfTreeModel(sessionManager.getCurrSession());
+  }
+
   /**
    * handle when successfully receive tsFilePipeData. Rename .patch file and reset tsFilePipeData's
    * path.

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

@@ -19,7 +19,7 @@
 
 package org.apache.iotdb.db.pipe.receiver.protocol.legacy.loader;
 
-import org.apache.iotdb.commons.audit.UserEntity;
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.commons.conf.CommonDescriptor;
 import org.apache.iotdb.commons.queryengine.common.SessionInfo;
 import org.apache.iotdb.db.auth.AuthorityChecker;
@@ -29,6 +29,7 @@
 import org.apache.iotdb.db.protocol.session.SessionManager;
 import org.apache.iotdb.db.queryengine.plan.Coordinator;
 import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult;
+import org.apache.iotdb.db.queryengine.plan.relational.security.TreeAccessCheckContext;
 import org.apache.iotdb.db.queryengine.plan.statement.Statement;
 import org.apache.iotdb.db.queryengine.plan.statement.crud.DeleteDataStatement;
 import org.apache.iotdb.db.queryengine.plan.statement.metadata.DeleteTimeSeriesStatement;
@@ -39,7 +40,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.time.ZoneId;
 import java.util.Collections;
 
 /** This loader is used to load deletion plan. */
@@ -54,25 +54,30 @@ public DeletionLoader(Deletion deletion) {
   }
 
   @Override
-  public void load() throws PipeException {
+  public void load(final SessionInfo sessionInfo) throws PipeException {
     if (CommonDescriptor.getInstance().getConfig().isReadOnly()) {
       throw new PipeException(DataNodePipeMessages.STORAGE_ENGINE_READONLY);
     }
     try {
       Statement statement = generateStatement();
+      TSStatus authorityStatus =
+          AuthorityChecker.checkAuthority(
+              statement,
+              new TreeAccessCheckContext(
+                  sessionInfo.getUserId(),
+                  sessionInfo.getUserName(),
+                  sessionInfo.getCliHostname()));
+      if (authorityStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+        throw new PipeException(authorityStatus.getMessage());
+      }
+
       long queryId = SessionManager.getInstance().requestQueryId();
       ExecutionResult result =
           Coordinator.getInstance()
               .executeForTreeModel(
                   statement,
                   queryId,
-                  new SessionInfo(
-                      0,
-                      new UserEntity(
-                          AuthorityChecker.SUPER_USER_ID,
-                          AuthorityChecker.SUPER_USER,
-                          IoTDBDescriptor.getInstance().getConfig().getInternalAddress()),
-                      ZoneId.systemDefault()),
+                  sessionInfo,
                   "",
                   PARTITION_FETCHER,
                   SCHEMA_FETCHER,

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

@@ -19,6 +19,7 @@
 
 package org.apache.iotdb.db.pipe.receiver.protocol.legacy.loader;
 
+import org.apache.iotdb.commons.queryengine.common.SessionInfo;
 import org.apache.iotdb.db.queryengine.plan.analyze.ClusterPartitionFetcher;
 import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher;
 import org.apache.iotdb.db.queryengine.plan.analyze.schema.ClusterSchemaFetcher;
@@ -34,5 +35,5 @@ public interface ILoader {
 
   ISchemaFetcher SCHEMA_FETCHER = ClusterSchemaFetcher.getInstance();
 
-  void load();
+  void load(SessionInfo sessionInfo);
 }