Subscription: support encrypted password auth in consumer builder

integration-test/src/test/java/org/apache/iotdb/subscription/it/triple/treemodel/regression/param/IoTDBEncryptedPasswordPullConsumerIT.java

@@ -0,0 +1,174 @@
+/*
+ * 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.
+ */
+
+package org.apache.iotdb.subscription.it.triple.treemodel.regression.param;
+
+import org.apache.iotdb.commons.utils.AuthUtils;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.MultiClusterIT2SubscriptionTreeRegressionMisc;
+import org.apache.iotdb.rpc.IoTDBConnectionException;
+import org.apache.iotdb.rpc.StatementExecutionException;
+import org.apache.iotdb.session.subscription.consumer.tree.SubscriptionTreePullConsumer;
+import org.apache.iotdb.subscription.it.triple.treemodel.regression.AbstractSubscriptionTreeRegressionIT;
+
+import org.apache.thrift.TException;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.file.metadata.enums.CompressionType;
+import org.apache.tsfile.file.metadata.enums.TSEncoding;
+import org.apache.tsfile.write.record.Tablet;
+import org.apache.tsfile.write.schema.IMeasurementSchema;
+import org.apache.tsfile.write.schema.MeasurementSchema;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Ignore;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+@Ignore("TODO: enable after encrypted password subscription IT stabilizes")
+@RunWith(IoTDBTestRunner.class)
+@Category({MultiClusterIT2SubscriptionTreeRegressionMisc.class})
+public class IoTDBEncryptedPasswordPullConsumerIT extends AbstractSubscriptionTreeRegressionIT {
+
+  private static final String DATABASE = "root.TestEncryptedPasswordPullConsumer";
+  private static final String DEVICE = DATABASE + ".d_0";
+  private static final String TOPIC_NAME = "TestEncryptedPasswordPullConsumerTopic";
+  private static final String USERNAME = "encrypted_user";
+  private static final String PASSWORD = "encrypted_user_123";
+  private static final String ENCRYPTED_PASSWORD = AuthUtils.encryptPassword(PASSWORD);
+  private static final String WRONG_ENCRYPTED_PASSWORD =
+      AuthUtils.encryptPassword("wrong_encrypted_user_123");
+
+  private static final List<IMeasurementSchema> SCHEMA_LIST = new ArrayList<>();
+
+  static {
+    SCHEMA_LIST.add(new MeasurementSchema("s_0", TSDataType.INT64));
+    SCHEMA_LIST.add(new MeasurementSchema("s_1", TSDataType.DOUBLE));
+  }
+
+  private SubscriptionTreePullConsumer consumer;
+
+  @Override
+  @Before
+  public void setUp() throws Exception {
+    super.setUp();
+    createDB(DATABASE);
+    createTopic_s(TOPIC_NAME, "root.**", null, null, false);
+    session_src.createTimeseries(
+        DEVICE + ".s_0", TSDataType.INT64, TSEncoding.GORILLA, CompressionType.LZ4);
+    session_src.createTimeseries(
+        DEVICE + ".s_1", TSDataType.DOUBLE, TSEncoding.TS_2DIFF, CompressionType.LZ4);
+    session_dest.createTimeseries(
+        DEVICE + ".s_0", TSDataType.INT64, TSEncoding.GORILLA, CompressionType.LZ4);
+    session_dest.createTimeseries(
+        DEVICE + ".s_1", TSDataType.DOUBLE, TSEncoding.TS_2DIFF, CompressionType.LZ4);
+    session_src.executeNonQueryStatement("create user " + USERNAME + " '" + PASSWORD + "'");
+    session_src.executeNonQueryStatement("grant read,write on root.** to user " + USERNAME);
+    assertTrue(subs.getTopic(TOPIC_NAME).isPresent());
+  }
+
+  @Override
+  @After
+  public void tearDown() throws Exception {
+    try {
+      if (consumer != null) {
+        consumer.close();
+      }
+    } catch (final Exception ignored) {
+    }
+    try {
+      subs.dropTopic(TOPIC_NAME);
+    } catch (final Exception ignored) {
+    }
+    try {
+      session_src.executeNonQueryStatement("drop user " + USERNAME);
+    } catch (final Exception ignored) {
+    }
+    dropDB(DATABASE);
+    super.tearDown();
+  }
+
+  @Test
+  public void testSubscribeWithEncryptedPassword()
+      throws TException,
+          IoTDBConnectionException,
+          IOException,
+          StatementExecutionException,
+          InterruptedException {
+    consumer = createConsumer("encrypted-password-group", ENCRYPTED_PASSWORD);
+
+    consumer.open();
+    consumer.subscribe(TOPIC_NAME);
+    assertEquals(1, subs.getSubscriptions().size(), "subscribe with encrypted password");
+
+    insertData(1706659200000L);
+    consume_data(consumer, session_dest);
+    check_count(
+        4,
+        "select count(s_0) from " + DEVICE + " where time >= 1706659200000",
+        "encrypted password consumption");
+  }
+
+  @Test
+  public void testSubscribeFailsWithWrongEncryptedPassword()
+      throws IoTDBConnectionException, StatementExecutionException {
+    consumer = createConsumer("wrong-encrypted-password-group", WRONG_ENCRYPTED_PASSWORD);
+
+    try {
+      consumer.open();
+      consumer.subscribe(TOPIC_NAME);
+      fail("subscribe should fail when encrypted password mismatches");
+    } catch (final Exception ignored) {
+      assertTrue(subs.getSubscriptions().isEmpty());
+    }
+  }
+
+  private SubscriptionTreePullConsumer createConsumer(
+      final String consumerGroupId, final String encryptedPassword) {
+    return new SubscriptionTreePullConsumer.Builder()
+        .host(SRC_HOST)
+        .port(SRC_PORT)
+        .username(USERNAME)
+        .password(PASSWORD)
+        .encryptedPassword(encryptedPassword)
+        .consumerId("consumer_" + consumerGroupId)
+        .consumerGroupId(consumerGroupId)
+        .buildPullConsumer();
+  }
+
+  private void insertData(long timestamp)
+      throws IoTDBConnectionException, StatementExecutionException {
+    final Tablet tablet = new Tablet(DEVICE, SCHEMA_LIST, 10);
+    for (int row = 0; row < 5; row++) {
+      final int rowIndex = tablet.getRowSize();
+      tablet.addTimestamp(rowIndex, timestamp);
+      tablet.addValue("s_0", rowIndex, row * 20L + row);
+      tablet.addValue("s_1", rowIndex, row + 2.45);
+      timestamp += row * 2000;
+    }
+    session_src.insertTablet(tablet);
+  }
+}

iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConfig.java

@@ -76,6 +76,10 @@ public String getPassword() {
     return getString(ConsumerConstant.PASSWORD_KEY);
   }
 
+  public String getEncryptedPassword() {
+    return getString(ConsumerConstant.ENCRYPTED_PASSWORD_KEY);
+  }
+
   public String getSqlDialect() {
     return getString(ConsumerConstant.SQL_DIALECT_KEY);
   }

iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConstant.java

@@ -34,6 +34,7 @@ public class ConsumerConstant {
 
   public static final String USERNAME_KEY = "username";
   public static final String PASSWORD_KEY = "password";
+  public static final String ENCRYPTED_PASSWORD_KEY = "encrypted-password";
 
   public static final String CONSUMER_ID_KEY = "consumer-id";
   public static final String CONSUMER_GROUP_ID_KEY = "group-id";

iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java

@@ -100,6 +100,7 @@ abstract class AbstractSubscriptionConsumer implements AutoCloseable {
 
   private final String username;
   private final String password;
+  private final String encryptedPassword;
 
   protected String consumerId;
   protected String consumerGroupId;
@@ -177,6 +178,7 @@ protected AbstractSubscriptionConsumer(final AbstractSubscriptionConsumerBuilder
 
     this.username = builder.username;
     this.password = builder.password;
+    this.encryptedPassword = builder.encryptedPassword;
 
     this.consumerId = builder.consumerId;
     this.consumerGroupId = builder.consumerGroupId;
@@ -207,6 +209,7 @@ protected AbstractSubscriptionConsumer(
                 (String)
                     properties.getOrDefault(
                         ConsumerConstant.PASSWORD_KEY, SessionConfig.DEFAULT_PASSWORD))
+            .encryptedPassword((String) properties.get(ConsumerConstant.ENCRYPTED_PASSWORD_KEY))
             .consumerId((String) properties.get(ConsumerConstant.CONSUMER_ID_KEY))
             .consumerGroupId((String) properties.get(ConsumerConstant.CONSUMER_GROUP_ID_KEY))
             .heartbeatIntervalMs(
@@ -387,6 +390,7 @@ protected abstract AbstractSubscriptionProvider constructSubscriptionProvider(
       final TEndPoint endPoint,
       final String username,
       final String password,
+      final String encryptedPassword,
       final String consumerId,
       final String consumerGroupId,
       final int thriftMaxFrameSize,
@@ -400,6 +404,7 @@ AbstractSubscriptionProvider constructProviderAndHandshake(final TEndPoint endPo
             endPoint,
             this.username,
             this.password,
+            this.encryptedPassword,
             this.consumerId,
             this.consumerGroupId,
             this.thriftMaxFrameSize,

iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumerBuilder.java

@@ -36,6 +36,7 @@ public class AbstractSubscriptionConsumerBuilder {
 
   protected String username = SessionConfig.DEFAULT_USER;
   protected String password = SessionConfig.DEFAULT_PASSWORD;
+  protected String encryptedPassword;
 
   protected String consumerId;
   protected String consumerGroupId;
@@ -76,6 +77,11 @@ public AbstractSubscriptionConsumerBuilder password(final String password) {
     return this;
   }
 
+  public AbstractSubscriptionConsumerBuilder encryptedPassword(final String encryptedPassword) {
+    this.encryptedPassword = encryptedPassword;
+    return this;
+  }
+
   public AbstractSubscriptionConsumerBuilder consumerId(@Nullable final String consumerId) {
     if (Objects.isNull(consumerId)) {
       return this;

iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java

@@ -90,6 +90,7 @@ public abstract class AbstractSubscriptionProvider {
 
   private final String username;
   private final String password;
+  private final String encryptedPassword;
   private final long heartbeatIntervalMs;
   private final int connectionTimeoutInMs;
 
@@ -105,6 +106,7 @@ protected AbstractSubscriptionProvider(
       final TEndPoint endPoint,
       final String username,
       final String password,
+      final String encryptedPassword,
       final String consumerId,
       final String consumerGroupId,
       final int thriftMaxFrameSize,
@@ -124,6 +126,7 @@ protected AbstractSubscriptionProvider(
     this.consumerGroupId = consumerGroupId;
     this.username = username;
     this.password = password;
+    this.encryptedPassword = encryptedPassword;
     this.heartbeatIntervalMs = heartbeatIntervalMs;
     this.connectionTimeoutInMs = connectionTimeoutInMs;
   }
@@ -175,6 +178,9 @@ synchronized void handshake() throws SubscriptionException, IoTDBConnectionExcep
     consumerAttributes.put(ConsumerConstant.CONSUMER_ID_KEY, consumerId);
     consumerAttributes.put(ConsumerConstant.USERNAME_KEY, username);
     consumerAttributes.put(ConsumerConstant.PASSWORD_KEY, password);
+    if (encryptedPassword != null) {
+      consumerAttributes.put(ConsumerConstant.ENCRYPTED_PASSWORD_KEY, encryptedPassword);
+    }
     consumerAttributes.put(ConsumerConstant.SQL_DIALECT_KEY, session.getSqlDialect());
     consumerAttributes.put(
         ConsumerConstant.HEARTBEAT_INTERVAL_MS_KEY, String.valueOf(heartbeatIntervalMs));

iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPullConsumerBuilder.java

@@ -58,6 +58,12 @@ public AbstractSubscriptionPullConsumerBuilder password(final String password) {
     return this;
   }
 
+  @Override
+  public AbstractSubscriptionPullConsumerBuilder encryptedPassword(final String encryptedPassword) {
+    super.encryptedPassword(encryptedPassword);
+    return this;
+  }
+
   @Override
   public AbstractSubscriptionPullConsumerBuilder consumerId(final String consumerId) {
     super.consumerId(consumerId);

iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPushConsumerBuilder.java

@@ -64,6 +64,12 @@ public AbstractSubscriptionPushConsumerBuilder password(final String password) {
     return this;
   }
 
+  @Override
+  public AbstractSubscriptionPushConsumerBuilder encryptedPassword(final String encryptedPassword) {
+    super.encryptedPassword(encryptedPassword);
+    return this;
+  }
+
   @Override
   public AbstractSubscriptionPushConsumerBuilder consumerId(final String consumerId) {
     super.consumerId(consumerId);

iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTableProvider.java

@@ -30,6 +30,7 @@ final class SubscriptionTableProvider extends AbstractSubscriptionProvider {
       final TEndPoint endPoint,
       final String username,
       final String password,
+      final String encryptedPassword,
       final String consumerId,
       final String consumerGroupId,
       final int thriftMaxFrameSize,
@@ -39,6 +40,7 @@ final class SubscriptionTableProvider extends AbstractSubscriptionProvider {
         endPoint,
         username,
         password,
+        encryptedPassword,
         consumerId,
         consumerGroupId,
         thriftMaxFrameSize,

iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumer.java

@@ -42,6 +42,7 @@ protected AbstractSubscriptionProvider constructSubscriptionProvider(
       final TEndPoint endPoint,
       final String username,
       final String password,
+      final String encryptedPassword,
       final String consumerId,
       final String consumerGroupId,
       final int thriftMaxFrameSize,
@@ -51,6 +52,7 @@ protected AbstractSubscriptionProvider constructSubscriptionProvider(
         endPoint,
         username,
         password,
+        encryptedPassword,
         consumerId,
         consumerGroupId,
         thriftMaxFrameSize,

iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java

@@ -56,6 +56,12 @@ public SubscriptionTablePullConsumerBuilder password(final String password) {
     return this;
   }
 
+  @Override
+  public SubscriptionTablePullConsumerBuilder encryptedPassword(final String encryptedPassword) {
+    super.encryptedPassword(encryptedPassword);
+    return this;
+  }
+
   @Override
   public SubscriptionTablePullConsumerBuilder consumerId(final String consumerId) {
     super.consumerId(consumerId);

iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePushConsumer.java

@@ -38,6 +38,7 @@ protected AbstractSubscriptionProvider constructSubscriptionProvider(
       final TEndPoint endPoint,
       final String username,
       final String password,
+      final String encryptedPassword,
       final String consumerId,
       final String consumerGroupId,
       final int thriftMaxFrameSize,
@@ -47,6 +48,7 @@ protected AbstractSubscriptionProvider constructSubscriptionProvider(
         endPoint,
         username,
         password,
+        encryptedPassword,
         consumerId,
         consumerGroupId,
         thriftMaxFrameSize,

iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePushConsumerBuilder.java

@@ -58,6 +58,12 @@ public SubscriptionTablePushConsumerBuilder password(final String password) {
     return this;
   }
 
+  @Override
+  public SubscriptionTablePushConsumerBuilder encryptedPassword(final String encryptedPassword) {
+    super.encryptedPassword(encryptedPassword);
+    return this;
+  }
+
   @Override
   public SubscriptionTablePushConsumerBuilder consumerId(final String consumerId) {
     super.consumerId(consumerId);