Restart cmd: Fix multiple restarts

src/api/TgBotApiImpl.cpp

@@ -28,6 +28,7 @@
 
 #include "StringResLoader.hpp"
 #include "api/CommandModule.hpp"
+#include "api/MessageExt.hpp"
 
 bool TgBotApiImpl::validateValidArgs(const DynModule* module,
                                      MessageExt::Ptr& message) {
@@ -74,6 +75,16 @@ bool TgBotApiImpl::validateValidArgs(const DynModule* module,
     return true;
 }
 
+bool TgBotApiImpl::isMyCommand(const MessageExt::Ptr& message) const {
+    const auto botCommand = message->get<MessageAttrs::BotCommand>();
+    const auto target = botCommand.target;
+    if (target != getBotUser()->username && !target.empty()) {
+        DLOG(INFO) << "Ignore mismatched target: " << std::quoted(target);
+        return false;
+    }
+    return true;
+}
+
 bool TgBotApiImpl::authorized(const MessageExt::Ptr& message,
                               const std::string_view commandName,
                               AuthContext::Flags flags) const {
@@ -116,17 +127,14 @@ void TgBotApiImpl::commandHandler(const std::string& command,
                                   Message::Ptr message) {
     // Find the module first.
     const auto* module = (*kModuleLoader)[command];
-    static const std::string myName = getBotUser()->username;
 
     // Create MessageExt object.
     SplitMessageText how = module->_module->valid_args.enabled
                                ? module->_module->valid_args.split_type
                                : SplitMessageText::None;
     auto ext = std::make_shared<MessageExt>(std::move(message), how);
-    const auto botCommand = ext->get<MessageAttrs::BotCommand>();
-    const auto target = botCommand.target;
-    if (target != myName && !target.empty()) {
-        DLOG(INFO) << "Ignore mismatched target: " << std::quoted(target);
+
+    if (!isMyCommand(ext)) {
         return;
     }
 
@@ -135,6 +143,7 @@ void TgBotApiImpl::commandHandler(const std::string& command,
         LOG(INFO) << "Command module is unloaded: " << module->_module->name;
         return;
     }
+
     if (module->_module->function == nullptr) {
         // Just in case.
         LOG(ERROR) << "Command module does not have a function: "

src/api/components/Restart.cpp

@@ -1,9 +1,10 @@
+#include <Env.hpp>
 #include <api/MessageExt.hpp>
-#include <api/components/Restart.hpp>
-#include <restartfmt_parser.hpp>
+#include <api/components/ModuleManagement.hpp>
 #include <api/components/OnAnyMessage.hpp>
 #include <api/components/OnCallbackQuery.hpp>
-#include <api/components/ModuleManagement.hpp>
+#include <api/components/Restart.hpp>
+#include <restartfmt_parser.hpp>
 
 TgBotApiImpl::RestartCommand::RestartCommand(TgBotApiImpl::Ptr api)
     : _api(api) {
@@ -16,31 +17,40 @@ void TgBotApiImpl::RestartCommand::commandFunction(MessageExt::Ptr message) {
     if (!_api->authorized(message, "restart", AuthContext::Flags::None)) {
         return;
     }
+    if (!_api->isMyCommand(message)) {
+        return;
+    }
+    if (RestartFmt::isRestartedByThisMessage(message)) {
+        DLOG(INFO) << "Avoiding restart loop";
+        return;
+    }
 
     // typical chatid:int32_max
-    std::array<char, RestartFmt::MAX_KNOWN_LENGTH> restartBuf = {0};
     int count = 0;
     std::vector<char *> myEnviron;
 
-    const auto status = RestartFmt::handleMessage(_api);
-    LOG_IF(ERROR, status.code() == absl::StatusCode::kInvalidArgument)
-        << "Failed to handle restart message: " << status;
-    if (status.ok()) {
-        return;
-    }
-
     // Get the size of environment buffer
     for (; environ[count] != nullptr; ++count);
     // Get ready to insert 1 more to the environment buffer
     myEnviron.resize(count + 2);
     // Copy the environment buffer to the new buffer
     memcpy(myEnviron.data(), environ, count * sizeof(char *));
+    // Check if the environ RESTART exists already
+    auto [be, en] = std::ranges::remove_if(myEnviron, [](const char *env) {
+        return env != nullptr
+                   ? std::string_view(env).starts_with(RestartFmt::ENV_VAR_NAME)
+                   : false;
+    });
+    // If the environ RESTART exists, remove it
+    if (std::distance(be, en) != 0) {
+        DLOG(INFO) << "Removing existing RESTART=";
+        myEnviron.erase(be, en);
+    }
 
-    // Set the restart command env to the environment buffer
-    const auto restartEnv =
-        RestartFmt::toString({message->get<MessageAttrs::Chat>()->id,
-                              message->get<MessageAttrs::MessageId>()},
-                             true);
+    std::array<char, RestartFmt::MAX_KNOWN_LENGTH> restartBuf = {0};
+    std::string restartEnv =
+        fmt::format("{}={}", RestartFmt::ENV_VAR_NAME,
+                    RestartFmt::Type{message->message()}.to_string());
     strncpy(restartBuf.data(), restartEnv.c_str(), restartEnv.size());
 
     // Append the restart env to the environment buffer
@@ -73,6 +83,5 @@ void TgBotApiImpl::RestartCommand::commandFunction(MessageExt::Ptr message) {
                    message->get_or<MessageAttrs::Locale>(Locale::Default)),
                Strings::RESTARTING_BOT));
 
-    // Call exeve
-    execve(exe, argv, myEnviron.data());
+    execve(argv[0], argv, myEnviron.data());
 }

src/command_modules/support/restartfmt_parser/restartfmt_parser.cpp

@@ -4,78 +4,78 @@
 #include <absl/strings/str_cat.h>
 #include <absl/strings/str_split.h>
 #include <absl/strings/strip.h>
+#include <fmt/format.h>
 
 #include <TryParseStr.hpp>
+#include <stdexcept>
 
 #include "utils/Env.hpp"
 
-std::optional<RestartFmt::data_type> RestartFmt::fromString(
-    const std::string& string, bool withPrefix) {
-    std::string_view view = string;
-    ChatId parsedChatId = 0;
-    MessageId parsedMessageId = 0;
-    int parsedMessageThreadId = 0;
-
-    if (withPrefix && !absl::ConsumePrefix(&view, "RESTART=")) {
-        LOG(ERROR) << "Invalid format for RESTART=" << string
-                   << " (Doesn't have the expected prefix)";
-        return std::nullopt;
-    }
+RestartFmt::Type::Type(const std::string_view string) {
     std::vector<std::string> parts =
-        absl::StrSplit(view, ":", absl::SkipEmpty());
+        absl::StrSplit(string, ":", absl::SkipEmpty());
     if (parts.size() != 3) {
         LOG(ERROR) << "Invalid format for RESTART=" << string
                    << " (Expected three parts)";
-        return std::nullopt;
+        throw std::invalid_argument("Invalid format for RESTART");
     }
-    if (try_parse(parts[0], &parsedChatId) &&
-        try_parse(parts[1], &parsedMessageId) &&
-        try_parse(parts[2], &parsedMessageThreadId)) {
-        return Type{parsedChatId, parsedMessageId, parsedMessageThreadId};
-    } else {
+    if (!(try_parse(parts[0], &chat_id) && try_parse(parts[1], &message_id) &&
+          try_parse(parts[2], &message_thread_id))) {
         LOG(ERROR) << "Invalid format for RESTART=" << string
                    << " (Failed to parse chat_id and message_id)";
-        return std::nullopt;
+        throw std::invalid_argument("Invalid format for RESTART");
     }
 }
 
-std::optional<RestartFmt::data_type> RestartFmt::fromEnvVar() {
-    const auto restartStr = Env()[ENV_VAR_NAME];
-    if (!restartStr.has()) {
-        LOG(WARNING) << "No " << ENV_VAR_NAME << " environment variable found";
-        return std::nullopt;
-    }
-    return fromString(restartStr.get(), false);
+RestartFmt::Type::Type(const Message::Ptr& message)
+    : chat_id(message->chat->id),
+      message_id(message->messageId),
+      message_thread_id(message->messageThreadId) {}
+
+bool RestartFmt::Type::operator==(const Type& other) const {
+    return chat_id == other.chat_id && message_id == other.message_id &&
+           message_thread_id == other.message_thread_id;
 }
 
-std::string RestartFmt::toString(const data_type& data, bool withPrefix) {
-    std::string dataString = absl::StrCat(std::to_string(data.first), ":",
-                                          std::to_string(data.second), ":",
-                                          std::to_string(data.third));
+std::string RestartFmt::Type::to_string() const {
+    return fmt::format("{}:{}:{}", chat_id, message_id, message_thread_id);
+}
 
-    if (withPrefix) {
-        return absl::StrCat("RESTART=", dataString);
-    } else {
-        return dataString;
+bool RestartFmt::checkEnvAndVerifyRestart(TgBotApi::CPtr api) {
+    // Check if the environment variable is set and valid
+    LOG(INFO) << "RestartFmt::checkEnvAndVerifyRestart";
+    auto env = Env()[ENV_VAR_NAME];
+    if (!env.has()) {
+        LOG(INFO) << fmt::format("ENV_VAR {} is not set", ENV_VAR_NAME);
+        return false;
     }
+    std::string value = env.get();
+    LOG(INFO) << fmt::format("GETENV {}: is set to {}", ENV_VAR_NAME, value);
+    LOG(INFO) << "Restart successful";
+    try {
+        Type t{value};
+        api->sendReplyMessage(t.chat_id, t.message_id, t.message_thread_id,
+                              "Restart success!");
+    } catch (const std::invalid_argument& ex) {
+        LOG(ERROR) << "Invalid format for RESTART=" << value << ": "
+                   << ex.what();
+        return false;
+    }
+    return true;
 }
 
-absl::Status RestartFmt::handleMessage(TgBotApi::CPtr api) {
-    if (const auto env = Env()[ENV_VAR_NAME]; env.has()) {
-        const auto v = RestartFmt::fromEnvVar();
-        if (!v) {
-            LOG(ERROR) << "Invalid restart command format: " << env;
-            Env()[ENV_VAR_NAME].clear();
-            return absl::InvalidArgumentError(
-                "Invalid restart variable format");
-        } else {
-            LOG(INFO) << "Restart successful";
-            api->sendReplyMessage(v->first, v->second, v->third,
-                                  "Restart success!");
-            Env()[ENV_VAR_NAME].clear();
-            return absl::OkStatus();
-        }
+bool RestartFmt::isRestartedByThisMessage(const MessageExt::Ptr& message) {
+    auto env = Env()[ENV_VAR_NAME];
+    if (!env.has()) {
+        LOG(INFO) << fmt::format("ENV_VAR {} is not set", ENV_VAR_NAME);
+        return false;
+    }
+    std::string value = env.get();
+    LOG(INFO) << fmt::format("GETENV {}: is set to {}", ENV_VAR_NAME, value);
+    if (Type{message->message()} == Type{value}) {
+        LOG(INFO) << "bot is restarted by this message";
+        return true;
     }
-    // Could get here when the environment variable is not set
-    return absl::UnavailableError("Environment variable not set");
+    LOG(INFO) << "bot is not restarted by this message";
+    return false;
 }

src/command_modules/support/restartfmt_parser/restartfmt_parser.hpp

@@ -1,37 +1,32 @@
 
 #include <absl/status/status.h>
 
+#include <api/MessageExt.hpp>
 #include <api/TgBotApi.hpp>
 #include <limits>
-#include <optional>
 #include <string>
 
 #include "Types.h"
 
 class RestartFmt {
    public:
-    using MessageThreadId = int;
-
     struct Type {
-        ChatId first;
-        MessageId second;
-        MessageThreadId third;
+        ChatId chat_id{};
+        MessageId message_id{};
+        MessageThreadId message_thread_id{};
+
+        explicit Type(const std::string_view string);
+        explicit Type(const Message::Ptr& message);
+        std::string to_string() const;
+        bool operator==(const Type& other) const;
     };
     using data_type = Type;
 
-    // function to parse the given string and populate chat_id and message_id
-    static std::optional<data_type> fromString(const std::string& string,
-                                               bool withPrefix = false);
-
-    // parse the string from environment variables
-    static std::optional<data_type> fromEnvVar();
-
-    // function to convert chat_id and message_id to a string
-    static std::string toString(const data_type& data, bool withPrefix = false);
+    static bool checkEnvAndVerifyRestart(TgBotApi::CPtr api);
+    static bool isRestartedByThisMessage(const MessageExt::Ptr& message);
 
-    static absl::Status handleMessage(TgBotApi::CPtr api);
+    constexpr static std::string_view ENV_VAR_NAME = "RESTART";
 
-    constexpr static const char* ENV_VAR_NAME = "RESTART";
     // typical chatid:int32_max
     constexpr static size_t MAX_KNOWN_LENGTH =
         sizeof("RESTART=-00000000000:") +

src/include/api/TgBotApiImpl.hpp

@@ -455,9 +455,11 @@ class TgBotApiImpl : public TgBotApi {
     [[nodiscard]] EventBroadcaster& getEvents() { return _bot.getEvents(); }
     [[nodiscard]] const Api& getApi() const { return _bot.getApi(); }
 
-    bool authorized(const MessageExt::Ptr& message,
-                    const std::string_view commandName,
-                    AuthContext::Flags flags) const;
+    [[nodiscard]] bool authorized(const MessageExt::Ptr& message,
+                                  const std::string_view commandName,
+                                  AuthContext::Flags flags) const;
+
+    [[nodiscard]] bool isMyCommand(const MessageExt::Ptr& message) const;
 
     class Async;
     Bot _bot;

src/main.cpp

@@ -530,10 +530,8 @@ int main(int argc, char** argv) {
         }
     }
 
-#ifndef WINDOWS_BUILD
-    LOG_IF(WARNING, !RestartFmt::handleMessage(api).ok())
+    LOG_IF(WARNING, !RestartFmt::checkEnvAndVerifyRestart(api))
         << "Failed to handle restart message";
-#endif
 
     LOG(INFO) << "Subsystems initialized, bot started: " << argv[0];
     LOG(INFO) << fmt::format("Starting took {}", startupDp.get());

src/utils/Env.hpp

@@ -37,6 +37,25 @@ class TgBotUtils_API Env {
             return _key;
         }
 
+        ValueEntry& operator=(const ValueEntry& other) {
+            if (this == &other) {
+                return *this;
+            }
+            *this = other.get();
+            return *this;
+        }
+
+        // This is a ValueEntry, we are only playing with values...
+        ValueEntry& operator=(ValueEntry&& other) {
+            *this = std::move(other.get());
+            other.clear();
+            return *this;
+        }
+
+        bool operator==(const std::string_view other) const {
+            return get() == other;
+        }
+
         // Append a string to the current value.
         // Note: This will not overwrite the existing value.
         // To overwrite, use operator=
@@ -46,9 +65,6 @@ class TgBotUtils_API Env {
             return *this;
         }
 
-        // Do not allow copying of this element outside the function use.
-        NO_COPY_CTOR(ValueEntry);
-        NO_MOVE_CTOR(ValueEntry);
         ValueEntry() = delete;
         ~ValueEntry() = default;
     };