diff --git a/src/api/TgBotApiImpl.cpp b/src/api/TgBotApiImpl.cpp index 335cfff8..29020ebf 100644 --- a/src/api/TgBotApiImpl.cpp +++ b/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(); + 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(std::move(message), how); - const auto botCommand = ext->get(); - 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: " diff --git a/src/api/components/Restart.cpp b/src/api/components/Restart.cpp index 99a1c3c8..7513c0d6 100644 --- a/src/api/components/Restart.cpp +++ b/src/api/components/Restart.cpp @@ -1,9 +1,10 @@ +#include #include -#include -#include +#include #include #include -#include +#include +#include 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 restartBuf = {0}; int count = 0; std::vector 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()->id, - message->get()}, - true); + std::array 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(Locale::Default)), Strings::RESTARTING_BOT)); - // Call exeve - execve(exe, argv, myEnviron.data()); + execve(argv[0], argv, myEnviron.data()); } \ No newline at end of file diff --git a/src/command_modules/support/restartfmt_parser/restartfmt_parser.cpp b/src/command_modules/support/restartfmt_parser/restartfmt_parser.cpp index 03525e39..51fedfb2 100644 --- a/src/command_modules/support/restartfmt_parser/restartfmt_parser.cpp +++ b/src/command_modules/support/restartfmt_parser/restartfmt_parser.cpp @@ -4,78 +4,78 @@ #include #include #include +#include #include +#include #include "utils/Env.hpp" -std::optional 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 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::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; } \ No newline at end of file diff --git a/src/command_modules/support/restartfmt_parser/restartfmt_parser.hpp b/src/command_modules/support/restartfmt_parser/restartfmt_parser.hpp index ebbf72ff..c2dd0094 100644 --- a/src/command_modules/support/restartfmt_parser/restartfmt_parser.hpp +++ b/src/command_modules/support/restartfmt_parser/restartfmt_parser.hpp @@ -1,37 +1,32 @@ #include +#include #include #include -#include #include #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 fromString(const std::string& string, - bool withPrefix = false); - - // parse the string from environment variables - static std::optional 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:") + diff --git a/src/include/api/TgBotApiImpl.hpp b/src/include/api/TgBotApiImpl.hpp index 54169104..fdb22469 100644 --- a/src/include/api/TgBotApiImpl.hpp +++ b/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; diff --git a/src/main.cpp b/src/main.cpp index bada4d69..844a8528 100644 --- a/src/main.cpp +++ b/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()); diff --git a/src/utils/Env.hpp b/src/utils/Env.hpp index 79838a72..7ed4a635 100644 --- a/src/utils/Env.hpp +++ b/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; };