From d31225c55dc84b19af19af25cc5958789a97c172 Mon Sep 17 00:00:00 2001
From: Nikolay Shestakov <tesseract@ydb.tech>
Date: Mon, 5 Feb 2024 15:36:55 +0500
Subject: [PATCH] Fix the processing of incorrect SourceId (#1555)

* Fix the processing of incorrect SourceId

* fix
---
 .../persqueue/ut/partition_chooser_ut.cpp     | 25 +++++++++++++++++++
 ...ion_chooser_impl__abstract_chooser_actor.h |  9 +++++--
 ...artition_chooser_impl__old_chooser_actor.h |  4 ++-
 ...partition_chooser_impl__sm_chooser_actor.h |  4 ++-
 .../partition_chooser_impl__table_helper.h    |  2 +-
 5 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/ydb/core/persqueue/ut/partition_chooser_ut.cpp b/ydb/core/persqueue/ut/partition_chooser_ut.cpp
index 17eefd280848..7961aa8b04f4 100644
--- a/ydb/core/persqueue/ut/partition_chooser_ut.cpp
+++ b/ydb/core/persqueue/ut/partition_chooser_ut.cpp
@@ -621,6 +621,19 @@ Y_UNIT_TEST(TPartitionChooserActor_SplitMergeEnabled_PreferedPartition_OtherPart
     AssertTable(server, "A_Source_10", 0, 13);
 }
 
+Y_UNIT_TEST(TPartitionChooserActor_SplitMergeEnabled_BadSourceId_Test) {
+    NPersQueue::TTestServer server = CreateServer();
+
+    auto config = CreateConfig0(true);
+    AddPartition(config, 0, {}, {});
+    CreatePQTabletMock(server, 0, ETopicPartitionStatus::Active);
+
+    auto r = ChoosePartition(server, config, "base64:a***");
+
+    UNIT_ASSERT(r->Error);
+}
+
+
 Y_UNIT_TEST(TPartitionChooserActor_SplitMergeDisabled_NewSourceId_Test) {
     NPersQueue::TTestServer server = CreateServer();
 
@@ -692,4 +705,16 @@ Y_UNIT_TEST(TPartitionChooserActor_SplitMergeDisabled_PreferedPartition_Inactive
     UNIT_ASSERT(r->Error);
 }
 
+Y_UNIT_TEST(TPartitionChooserActor_SplitMergeDisabled_BadSourceId_Test) {
+    NPersQueue::TTestServer server = CreateServer();
+
+    auto config = CreateConfig0(false);
+    AddPartition(config, 0, {}, {});
+
+    auto r = ChoosePartition(server, config, "base64:a***");
+
+    UNIT_ASSERT(r->Error);
+}
+
+
 }
diff --git a/ydb/core/persqueue/writer/partition_chooser_impl__abstract_chooser_actor.h b/ydb/core/persqueue/writer/partition_chooser_impl__abstract_chooser_actor.h
index 9f8f4be0df09..4b385f779cda 100644
--- a/ydb/core/persqueue/writer/partition_chooser_impl__abstract_chooser_actor.h
+++ b/ydb/core/persqueue/writer/partition_chooser_impl__abstract_chooser_actor.h
@@ -54,8 +54,13 @@ class TAbstractPartitionChooserActor: public TActorBootstrapped<TDerived> {
         return TActor<TDerived>::SelfId();
     }
 
-    void Initialize(const NActors::TActorContext& ctx) {
-        TableHelper.Initialize(ctx, SourceId);
+    [[nodiscard]] bool Initialize(const NActors::TActorContext& ctx) {
+        if (TableHelper.Initialize(ctx, SourceId)) {
+            return true;
+        }
+        StartIdle();        
+        TThis::ReplyError(ErrorCode::BAD_REQUEST, "Bad SourceId", ctx);
+        return false;
     }
 
     void PassAway() {
diff --git a/ydb/core/persqueue/writer/partition_chooser_impl__old_chooser_actor.h b/ydb/core/persqueue/writer/partition_chooser_impl__old_chooser_actor.h
index f8d9d35e6eed..33e1f3dd2786 100644
--- a/ydb/core/persqueue/writer/partition_chooser_impl__old_chooser_actor.h
+++ b/ydb/core/persqueue/writer/partition_chooser_impl__old_chooser_actor.h
@@ -37,7 +37,9 @@ class TPartitionChooserActor: public TAbstractPartitionChooserActor<TPartitionCh
     }
 
     void Bootstrap(const TActorContext& ctx) {
-        TThis::Initialize(ctx);
+        if (!TThis::Initialize(ctx)) {
+            return;
+        }
         TThis::InitTable(ctx);
     }
 
diff --git a/ydb/core/persqueue/writer/partition_chooser_impl__sm_chooser_actor.h b/ydb/core/persqueue/writer/partition_chooser_impl__sm_chooser_actor.h
index 1bc599d87238..2342ae49f308 100644
--- a/ydb/core/persqueue/writer/partition_chooser_impl__sm_chooser_actor.h
+++ b/ydb/core/persqueue/writer/partition_chooser_impl__sm_chooser_actor.h
@@ -38,7 +38,9 @@ class TSMPartitionChooserActor: public TAbstractPartitionChooserActor<TSMPartiti
     }
 
     void Bootstrap(const TActorContext& ctx) {
-        TThis::Initialize(ctx);
+        if (!TThis::Initialize(ctx)) {
+            return;
+        }
         BoundaryPartition = ChoosePartitionSync();
 
         if (TThis::SourceId) {
diff --git a/ydb/core/persqueue/writer/partition_chooser_impl__table_helper.h b/ydb/core/persqueue/writer/partition_chooser_impl__table_helper.h
index 3dee14a904c2..0268b7d88e67 100644
--- a/ydb/core/persqueue/writer/partition_chooser_impl__table_helper.h
+++ b/ydb/core/persqueue/writer/partition_chooser_impl__table_helper.h
@@ -40,7 +40,7 @@ class TTableHelper {
         return SeqNo_;
     }
 
-    bool Initialize(const TActorContext& ctx, const TString& sourceId) {
+    [[nodiscard]] bool Initialize(const TActorContext& ctx, const TString& sourceId) {
         const auto& pqConfig = AppData(ctx)->PQConfig;
 
         TableGeneration = pqConfig.GetTopicsAreFirstClassCitizen() ? ESourceIdTableGeneration::PartitionMapping