Check permissions when store walle task

7 files changed, 91 insertions, 10 deletions

diff --git a/ydb/core/cms/cms.cpp b/ydb/core/cms/cms.cpp
index 45452c9a21c..8bac7eb17a5 100644
--- a/ydb/core/cms/cms.cpp
+++ b/ydb/core/cms/cms.cpp
@@ -1908,8 +1908,6 @@ void TCms::Handle(TEvCms::TEvWalleRemoveTaskRequest::TPtr &ev, const TActorConte
 void TCms::Handle(TEvCms::TEvStoreWalleTask::TPtr &ev, const TActorContext &ctx)
 {
     auto event = ev->Get();
-    State->WalleTasks.emplace(event->Task.TaskId, event->Task);
-    State->WalleRequests.emplace(event->Task.RequestId, event->Task.TaskId);
 
     auto handle = new IEventHandle(ev->Sender, SelfId(), new TEvCms::TEvWalleTaskStored(event->Task.TaskId), 0, ev->Cookie);
     Execute(CreateTxStoreWalleTask(event->Task, std::move(ev->Release()), handle), ctx);
diff --git a/ydb/core/cms/cms.h b/ydb/core/cms/cms.h
index 25cb8ef024b..dd02706dbb4 100644
--- a/ydb/core/cms/cms.h
+++ b/ydb/core/cms/cms.h
@@ -71,6 +71,7 @@ struct TEvCms {
         EvSetConfigResponse,
         EvSetMarkerResponse,
         EvResetMarkerResponse,
+        EvStoreWalleTaskFailed,
 
         EvEnd
     };
@@ -199,6 +200,23 @@ struct TEvCms {
         }
     };
 
+    struct TEvStoreWalleTaskFailed : public TEventLocal<TEvStoreWalleTaskFailed, EvStoreWalleTaskFailed> {
+        TString TaskId;
+        TString Reason;
+
+        TEvStoreWalleTaskFailed(TString id, TString reason)
+            : TaskId(id)
+            , Reason(reason)
+        {
+        }
+
+        TString ToString() const override
+        {
+            return Sprintf("%s { Task: %s Reason %s}", ToStringHeader().data(), TaskId.data(), Reason.data());
+        }
+    };
+
+
     struct TEvWalleTaskStored : public TEventLocal<TEvWalleTaskStored, EvWalleTaskStored> {
         TString TaskId;
 
diff --git a/ydb/core/cms/cms_tx_store_walle_task.cpp b/ydb/core/cms/cms_tx_store_walle_task.cpp
index f05828477a4..d6ca017f564 100644
--- a/ydb/core/cms/cms_tx_store_walle_task.cpp
+++ b/ydb/core/cms/cms_tx_store_walle_task.cpp
@@ -10,8 +10,7 @@ class TCms::TTxStoreWalleTask : public TTransactionBase<TCms> {
 public:
     TTxStoreWalleTask(TCms *self, const TWalleTaskInfo &task, THolder<IEventBase> req, TAutoPtr<IEventHandle> resp)
         : TBase(self)
-        , TaskId(task.TaskId)
-        , RequestId(task.RequestId)
+        , Task(task)
         , Request(std::move(req))
         , Response(std::move(resp))
     {
@@ -22,14 +21,30 @@ public:
     bool Execute(TTransactionContext &txc, const TActorContext &ctx) override
     {
         LOG_DEBUG(ctx, NKikimrServices::CMS, "TTxStoreWalleTask Execute");
+        
+        for (auto &perm : Task.Permissions) {
+            if (Self->State->Permissions.find(perm) == Self->State->Permissions.end()) {
+
+                Response.Reset(new IEventHandle(Response->Recipient, Response->Sender,
+                               new TEvCms::TEvStoreWalleTaskFailed(Task.TaskId,
+                                                                   TStringBuilder() << "There are no stored permissions for this task. "
+                                                                                    << "Maybe cleanup ran before task been stored. "
+                                                                                    << "Try request again"),
+                               0, Response->Cookie));
+                return true;
+            }
+        }
+
+        Self->State->WalleTasks.emplace(Task.TaskId, Task);
+        Self->State->WalleRequests.emplace(Task.RequestId, Task.TaskId);
 
         NIceDb::TNiceDb db(txc.DB);
-        auto row = db.Table<Schema::WalleTask>().Key(TaskId);
-        row.Update(NIceDb::TUpdate<Schema::WalleTask::RequestID>(RequestId));
+        auto row = db.Table<Schema::WalleTask>().Key(Task.TaskId);
+        row.Update(NIceDb::TUpdate<Schema::WalleTask::RequestID>(Task.RequestId));
 
         Self->AuditLog(ctx, TStringBuilder() << "Store wall-e task"
-            << ": id# " << TaskId
-            << ", requestId# " << RequestId);
+            << ": id# " << Task.TaskId
+            << ", requestId# " << Task.RequestId);
 
         return true;
     }
@@ -41,8 +56,7 @@ public:
     }
 
 private:
-    TString TaskId;
-    TString RequestId;
+    TWalleTaskInfo Task; 
     THolder<IEventBase> Request;
     TAutoPtr<IEventHandle> Response;
 };
diff --git a/ydb/core/cms/cms_ut.cpp b/ydb/core/cms/cms_ut.cpp
index 67e2b46cfb2..7356ae865de 100644
--- a/ydb/core/cms/cms_ut.cpp
+++ b/ydb/core/cms/cms_ut.cpp
@@ -1224,6 +1224,39 @@ Y_UNIT_TEST_SUITE(TCmsTest) {
             }
         }
     }
+
+    Y_UNIT_TEST(WalleCleanupTest)
+    {
+        TCmsTestEnv env(8);
+
+        TAutoPtr<NCms::TEvCms::TEvPermissionRequest> event = new NCms::TEvCms::TEvPermissionRequest;
+        event->Record.SetUser(WALLE_CMS_USER);
+        event->Record.SetPartialPermissionAllowed(true);
+        event->Record.SetDryRun(false);
+        event->Record.SetSchedule(false);
+
+        AddActions(event, MakeAction(TAction::RESTART_SERVICES, env.GetNodeId(3), 600000000, "storage"));
+
+        NKikimrCms::TPermissionResponse res;
+        res = env.CheckPermissionRequest(event, TStatus::ALLOW);
+
+        // Check that permission is stored
+        env.CheckListPermissions(WALLE_CMS_USER, 1);
+
+        // Adbance time to run cleanup
+        env.AdvanceCurrentTime(TDuration::Minutes(3));
+        env.RestartCms();
+        
+        // TODO:: перенести внутрь TCmsTestEnv
+        TAutoPtr<TEvCms::TEvStoreWalleTask> event_store = new TEvCms::TEvStoreWalleTask;
+        event_store->Task.TaskId = "walle-test-task-1";
+        event_store->Task.RequestId = res.GetRequestId();
+
+        for (auto &permission : res.GetPermissions())
+            event_store->Task.Permissions.insert(permission.GetId());
+
+        env.CheckWalleStoreTaskIsFailed(event_store.Release());
+    }
 }
 
 } // NCmsTest
diff --git a/ydb/core/cms/cms_ut_common.cpp b/ydb/core/cms/cms_ut_common.cpp
index 71414763706..5f6706f1744 100644
--- a/ydb/core/cms/cms_ut_common.cpp
+++ b/ydb/core/cms/cms_ut_common.cpp
@@ -746,6 +746,18 @@ TCmsTestEnv::CheckRequest(const TString &user,
     return rec;
 }
 
+
+void TCmsTestEnv::CheckWalleStoreTaskIsFailed(NCms::TEvCms::TEvStoreWalleTask* req)
+{ 
+    TString TaskId = req->Task.TaskId;
+    SendToPipe(CmsId, Sender, req, 0, GetPipeConfigWithRetries());
+    
+    TAutoPtr<IEventHandle> handle;
+    auto reply = GrabEdgeEventRethrow<TEvCms::TEvStoreWalleTaskFailed>(handle, TDuration::Seconds(30));
+    UNIT_ASSERT(reply);
+    UNIT_ASSERT_VALUES_EQUAL(reply->TaskId, TaskId);
+}
+
 void TCmsTestEnv::CheckWalleCreateTask(TAutoPtr<NCms::TEvCms::TEvWalleCreateTaskRequest> req,
                                        NKikimrCms::TStatus::ECode code)
 {
diff --git a/ydb/core/cms/cms_ut_common.h b/ydb/core/cms/cms_ut_common.h
index 588617d9b88..b4fa71c8b27 100644
--- a/ydb/core/cms/cms_ut_common.h
+++ b/ydb/core/cms/cms_ut_common.h
@@ -270,6 +270,8 @@ public:
         return CheckRequest(user, id, dry, NKikimrCms::MODE_MAX_AVAILABILITY, res, count);
     }
 
+    void CheckWalleStoreTaskIsFailed(NCms::TEvCms::TEvStoreWalleTask* req);
+
     template<typename... Ts>
     void CheckWalleCreateTask(const TString &id,
                               const TString &action,
diff --git a/ydb/core/cms/walle_create_task_adapter.cpp b/ydb/core/cms/walle_create_task_adapter.cpp
index 687a1d6eb0e..46a9fa52061 100644
--- a/ydb/core/cms/walle_create_task_adapter.cpp
+++ b/ydb/core/cms/walle_create_task_adapter.cpp
@@ -205,6 +205,10 @@ private:
         ctx.Send(Cms, request.Release());
     }
 
+    void Handle(TEvCms::TEvStoreWalleTaskFailed::TPtr &ev, const TActorContext &ctx) { 
+        ReplyWithErrorAndDie(TStatus::ERROR_TEMP, ev.Get()->Get()->Reason, ctx);
+    }
+
     void Finish(const TActorContext& ctx)
     {
         ReplyAndDie(Response, ctx);