From 69b3983b8d838ef3650dac7017b8ab740e0c3259 Mon Sep 17 00:00:00 2001 From: Soo Hwan Na Date: Wed, 12 Jun 2024 20:52:51 +0900 Subject: [PATCH] TgBot++: socket: Make uploadfile contain in-out paths --- src/socket/TgBotSocketClient.cpp | 35 ++++++------ src/socket/include/TgBotSocket_Export.hpp | 55 ++++++++++++------- .../interface/impl/bot/SocketDataHandler.cpp | 45 ++++++++------- .../impl/bot/TgBotSocketFileHelper.hpp | 25 ++++++--- .../impl/bot/TgBotSocketInterface.hpp | 6 +- 5 files changed, 101 insertions(+), 65 deletions(-) diff --git a/src/socket/TgBotSocketClient.cpp b/src/socket/TgBotSocketClient.cpp index cb24bcfc..7898cfe0 100644 --- a/src/socket/TgBotSocketClient.cpp +++ b/src/socket/TgBotSocketClient.cpp @@ -83,9 +83,6 @@ std::string_view AckTypeToStr(callback::AckType type) { } // namespace -// TODO: WTF is this -FileDataHelper::DataFromFileParam p; - struct ClientParser : TgBotSocketParser { explicit ClientParser(SocketInterfaceBase* interface) : TgBotSocketParser(interface) {} @@ -93,12 +90,11 @@ struct ClientParser : TgBotSocketParser { using callback::AckType; using callback::GenericAck; - callback::GetUptimeCallback callbackData = {}; - callback::GenericAck result{}; std::string resultText; switch (pkt.header.cmd) { case Command::CMD_GET_UPTIME_CALLBACK: { + callback::GetUptimeCallback callbackData = {}; pkt.data.assignTo(callbackData); LOG(INFO) << "Server replied: " << callbackData.uptime.data(); break; @@ -109,31 +105,39 @@ struct ClientParser : TgBotSocketParser { break; } case Command::CMD_UPLOAD_FILE_DRY_CALLBACK: { - pkt.data.assignTo(result); + callback::UploadFileDryCallback callbackData{}; + pkt.data.assignTo(callbackData); LOG(INFO) << "Response from server: " - << AckTypeToStr(result.result); - if (result.result != AckType::SUCCESS) { - LOG(ERROR) << "Reason: " << result.error_msg.data(); + << AckTypeToStr(callbackData.result); + if (callbackData.result != AckType::SUCCESS) { + LOG(ERROR) << "Reason: " << callbackData.error_msg.data(); } else { interface->closeSocketHandle(context); context = interface->createClientSocket().value(); LOG(INFO) << "Recreated client socket"; + auto params_in = callbackData.requestdata; + FileDataHelper::DataFromFileParam param; + param.destfilepath = params_in.destfilepath.data(); + param.filepath = params_in.srcfilepath.data(); + param.options = params_in.options; auto newPkt = FileDataHelper::DataFromFile< - FileDataHelper::UPLOAD_FILE>(p); + FileDataHelper::UPLOAD_FILE>(param); LOG(INFO) << "Sending the actual file content again..."; interface->writeToSocket(context, newPkt->toSocketData()); onNewBuffer(context); } break; } - case Command::CMD_GENERIC_ACK: - pkt.data.assignTo(result); + case Command::CMD_GENERIC_ACK: { + callback::GenericAck callbackData{}; + pkt.data.assignTo(callbackData); LOG(INFO) << "Response from server: " - << AckTypeToStr(result.result); - if (result.result != AckType::SUCCESS) { - LOG(ERROR) << "Reason: " << result.error_msg.data(); + << AckTypeToStr(callbackData.result); + if (callbackData.result != AckType::SUCCESS) { + LOG(ERROR) << "Reason: " << callbackData.error_msg.data(); } break; + } default: LOG(ERROR) << "Unhandled callback of command: " << static_cast(pkt.header.cmd); @@ -246,7 +250,6 @@ int main(int argc, char** argv) { params.options.overwrite = true; pkt = FileDataHelper::DataFromFile( params); - p = params; break; } case Command::CMD_DOWNLOAD_FILE: { diff --git a/src/socket/include/TgBotSocket_Export.hpp b/src/socket/include/TgBotSocket_Export.hpp index 89d5df38..670c61a5 100644 --- a/src/socket/include/TgBotSocket_Export.hpp +++ b/src/socket/include/TgBotSocket_Export.hpp @@ -63,7 +63,8 @@ struct TGSOCKET_ATTR_PACKED PacketHeader { // 3: Uploadfile has a sha256sum check, std::array conversions // 4: Move CMD_UPLOAD_FILE_DRY to internal namespace // 5: Use the packed attribute for structs - constexpr static int DATA_VERSION = 5; + // 6: Make CMD_UPLOAD_FILE_DRY_CALLBACK return sperate callback, and add srcpath to UploadFile + constexpr static int DATA_VERSION = 6; constexpr static int64_t MAGIC_VALUE = MAGIC_VALUE_BASE + DATA_VERSION; int64_t magic = MAGIC_VALUE; ///< Magic value to verify the packet @@ -125,6 +126,10 @@ struct Packet { } }; +using PathStringArray = std::array; +using MessageStringArray = std::array; +using SHA256StringArray = std::array; + namespace data { enum class FileType { @@ -146,7 +151,7 @@ enum class CtrlSpamBlock { struct TGSOCKET_ATTR_PACKED WriteMsgToChatId { ChatId chat; // destination chatid - std::array message; // Msg to send + MessageStringArray message; // Msg to send }; struct TGSOCKET_ATTR_PACKED ObserveChatId { @@ -158,7 +163,7 @@ struct TGSOCKET_ATTR_PACKED ObserveChatId { struct TGSOCKET_ATTR_PACKED SendFileToChatId { ChatId chat; // Destination ChatId FileType fileType; // File type for file - std::array filePath; // Path to file (local) + PathStringArray filePath; // Path to file (local) }; struct TGSOCKET_ATTR_PACKED ObserveAllChats { @@ -170,28 +175,35 @@ struct TGSOCKET_ATTR_PACKED DeleteControllerById { int controller_id; }; -struct TGSOCKET_ATTR_PACKED UploadFile { - std::array filepath{}; // Destination file name - std::array - sha256_hash{}; // SHA256 hash of the file +struct TGSOCKET_ATTR_PACKED UploadFileDry { + PathStringArray destfilepath{}; // Destination file name + PathStringArray srcfilepath{}; // Source file name (This is not used on the remote, used if dry=true) + SHA256StringArray sha256_hash{}; // SHA256 hash of the file // Returns AckType::ERROR_COMMAND_IGNORED on options failure struct TGSOCKET_ATTR_PACKED Options { - bool overwrite = false; // Overwrite existing file if exists - bool hash_ignore = - false; // Ignore hash, always upload file even if the same file - // exists and the hash matches. Depends on overwrite flag - // for actual overwriting - bool dry_run = - false; // If true, just check the hashes and return result. + // Overwrite existing file if exists + bool overwrite = false; + + // Ignore hash, always upload file even if the same file + // exists and the hash matches. Depends on overwrite flag + // for actual overwriting + bool hash_ignore = false; + + // If true, just check the hashes and return result. + bool dry_run = false; } options; +}; + +struct TGSOCKET_ATTR_PACKED UploadFile : public UploadFileDry { + using Options = UploadFileDry::Options; uint8_t buf[]; // Buffer }; struct TGSOCKET_ATTR_PACKED DownloadFile { - std::array filepath{}; // Path to file (in remote) - std::array destfilename{}; // Destination file name - uint8_t buf[]; // Buffer + PathStringArray filepath{}; // Path to file (in remote) + PathStringArray destfilename{}; // Destination file name + uint8_t buf[]; // Buffer }; } // namespace data @@ -213,7 +225,7 @@ enum class AckType { struct TGSOCKET_ATTR_PACKED GenericAck { AckType result; // result type // Error message, only valid when result type is not SUCCESS - std::array error_msg{}; + MessageStringArray error_msg{}; // Create a new instance of the Generic Ack, error. explicit GenericAck(AckType result, const std::string& errorMsg) @@ -227,6 +239,10 @@ struct TGSOCKET_ATTR_PACKED GenericAck { static GenericAck ok() { return GenericAck(AckType::SUCCESS, "OK"); } }; +struct TGSOCKET_ATTR_PACKED UploadFileDryCallback : public GenericAck { + data::UploadFileDry requestdata; +}; + } // namespace callback } // namespace TgBotSocket @@ -242,7 +258,7 @@ ASSERT_SIZE(ObserveChatId, 9); ASSERT_SIZE(SendFileToChatId, 268); ASSERT_SIZE(ObserveAllChats, 1); ASSERT_SIZE(DeleteControllerById, 4); -ASSERT_SIZE(UploadFile, 291); +ASSERT_SIZE(UploadFile, 547); ASSERT_SIZE(DownloadFile, 512); ASSERT_SIZE(PacketHeader, 24); } // namespace TgBotSocket::data @@ -250,6 +266,7 @@ ASSERT_SIZE(PacketHeader, 24); namespace TgBotSocket::callback { ASSERT_SIZE(GetUptimeCallback, 21); ASSERT_SIZE(GenericAck, 260); +ASSERT_SIZE(UploadFileDryCallback, 807); } // namespace TgBotSocket::callback #undef ASSERT_SIZE diff --git a/src/socket/interface/impl/bot/SocketDataHandler.cpp b/src/socket/interface/impl/bot/SocketDataHandler.cpp index f0bab642..8b1e5fb1 100644 --- a/src/socket/interface/impl/bot/SocketDataHandler.cpp +++ b/src/socket/interface/impl/bot/SocketDataHandler.cpp @@ -196,25 +196,29 @@ GenericAck SocketInterfaceTgBot::handle_DeleteControllerById(const void* ptr) { } GenericAck SocketInterfaceTgBot::handle_UploadFile( - const void* ptr, TgBotSocket::PacketHeader::length_type len, bool dry) { + const void* ptr, TgBotSocket::PacketHeader::length_type len) { + if (!FileDataHelper::DataToFile(ptr, len)) { + return GenericAck(AckType::ERROR_RUNTIME_ERROR, "Failed to write file"); + } + return GenericAck::ok(); +} + +UploadFileDryCallback SocketInterfaceTgBot::handle_UploadFileDry( + const void* ptr, TgBotSocket::PacketHeader::length_type len) { bool ret = false; + const auto f = static_cast(ptr); + UploadFileDryCallback callback; + callback.requestdata = *f; - if (dry) { - ret = FileDataHelper::DataToFile(ptr, - len); - if (!ret) { - return GenericAck(AckType::ERROR_COMMAND_IGNORED, - "Options verification failed"); - } + ret = FileDataHelper::DataToFile(ptr, len); + if (!ret) { + copyTo(callback.error_msg, "Options verification failed"); + callback.result = AckType::ERROR_COMMAND_IGNORED; } else { - ret = FileDataHelper::DataToFile(ptr, len); - if (!ret) { - return GenericAck(AckType::ERROR_RUNTIME_ERROR, - "Failed to write file"); - } + copyTo(callback.error_msg, "OK"); + callback.result = AckType::SUCCESS; } - - return GenericAck::ok(); + return callback; } bool SocketInterfaceTgBot::handle_DownloadFile(SocketConnContext ctx, @@ -252,7 +256,7 @@ bool SocketInterfaceTgBot::handle_GetUptime(SocketConnContext ctx, void SocketInterfaceTgBot::handle_CommandPacket(SocketConnContext ctx, TgBotSocket::Packet pkt) { const void* ptr = pkt.data.get(); - std::variant ret; + std::variant ret; switch (pkt.header.cmd) { case Command::CMD_WRITE_MSG_TO_CHAT_ID: @@ -277,10 +281,10 @@ void SocketInterfaceTgBot::handle_CommandPacket(SocketConnContext ctx, ret = handle_GetUptime(ctx, ptr); break; case Command::CMD_UPLOAD_FILE: - ret = handle_UploadFile(ptr, pkt.header.data_size, false); + ret = handle_UploadFile(ptr, pkt.header.data_size); break; case Command::CMD_UPLOAD_FILE_DRY: - ret = handle_UploadFile(ptr, pkt.header.data_size, true); + ret = handle_UploadFileDry(ptr, pkt.header.data_size); break; case Command::CMD_DOWNLOAD_FILE: ret = handle_DownloadFile(ctx, ptr); @@ -305,10 +309,9 @@ void SocketInterfaceTgBot::handle_CommandPacket(SocketConnContext ctx, break; } case Command::CMD_UPLOAD_FILE_DRY: { - // Need a special cmd here - GenericAck result = std::get(ret); + const auto result = std::get(ret); Packet ackpkt(Command::CMD_UPLOAD_FILE_DRY_CALLBACK, &result, - sizeof(GenericAck)); + sizeof(UploadFileDryCallback)); LOG(INFO) << "Sending CMD_UPLOAD_FILE_DRY ack: " << std::boolalpha << (result.result == AckType::SUCCESS); interface->writeToSocket(ctx, ackpkt.toSocketData()); diff --git a/src/socket/interface/impl/bot/TgBotSocketFileHelper.hpp b/src/socket/interface/impl/bot/TgBotSocketFileHelper.hpp index 13d20763..fbfffa41 100644 --- a/src/socket/interface/impl/bot/TgBotSocketFileHelper.hpp +++ b/src/socket/interface/impl/bot/TgBotSocketFileHelper.hpp @@ -4,6 +4,7 @@ #include #include +#include #include #include #include @@ -18,6 +19,10 @@ template inline void copyTo(std::array& arr_in, const char* buf) { strncpy(arr_in.data(), buf, size); } +template +inline void copyTo(std::array& arr_in, std::array buf) { + copyTo(arr_in, buf.data()); +} namespace FileDataHelper { using len_t = TgBotSocket::PacketHeader::length_type; @@ -81,6 +86,7 @@ std::optional DataFromFile( // Implementation starts using TgBotSocket::data::DownloadFile; using TgBotSocket::data::UploadFile; +using TgBotSocket::data::UploadFileDry; struct HashContainer { std::array m_data; }; @@ -151,8 +157,8 @@ inline std::ostream& operator<<(std::ostream& self, const HashContainer& data) { template <> bool FileDataHelper::DataToFile( const void* ptr, len_t len) { - const auto* data = static_cast(ptr); - const char* filename = data->filepath.data(); + const auto* data = static_cast(ptr); + const char* filename = data->destfilepath.data(); std::error_code errc; bool exists = false; @@ -202,7 +208,7 @@ bool FileDataHelper::DataToFile(const void* ptr, len_t len) { const auto* data = static_cast(ptr); - return writeFileCommon(data->filepath.data(), &data->buf[0], + return writeFileCommon(data->destfilepath.data(), &data->buf[0], len - sizeof(UploadFile)); } @@ -231,7 +237,7 @@ FileDataHelper::DataFromFile( // The front bytes of the buffer is UploadFile, hence cast it auto* uploadFile = static_cast(resultPointer.get()); // Copy destination file name info to the buffer - copyTo(uploadFile->filepath, params.destfilepath.string().c_str()); + copyTo(uploadFile->destfilepath, params.destfilepath.string().c_str()); // Copy source file data to the buffer memcpy(&uploadFile->buf[0], result->data.get(), result->size); // Calculate SHA256 hash @@ -261,11 +267,14 @@ FileDataHelper::DataFromFile( return std::nullopt; } // Create result packet buffer - auto resultPointer = createMemory(sizeof(UploadFile)); + auto resultPointer = createMemory(sizeof(UploadFileDry)); // The front bytes of the buffer is UploadFile, hence cast it - auto* uploadFile = static_cast(resultPointer.get()); + auto* uploadFile = static_cast(resultPointer.get()); // Copy destination file name info to the buffer - copyTo(uploadFile->filepath, params.destfilepath.string().c_str()); + copyTo(uploadFile->destfilepath, params.destfilepath.string().c_str()); + // Copy source file name to the buffer + copyTo(uploadFile->srcfilepath, params.filepath.string().c_str()); + // Calculate SHA256 hash SHA256(static_cast(result->data.get()), result->size, hash.m_data.data()); @@ -277,7 +286,7 @@ FileDataHelper::DataFromFile( uploadFile->options.dry_run = true; return TgBotSocket::Packet{TgBotSocket::Command::CMD_UPLOAD_FILE_DRY, - resultPointer.get(), sizeof(UploadFile)}; + resultPointer.get(), sizeof(UploadFileDry)}; } // Server -> diff --git a/src/socket/interface/impl/bot/TgBotSocketInterface.hpp b/src/socket/interface/impl/bot/TgBotSocketInterface.hpp index d3488758..2ef7e882 100644 --- a/src/socket/interface/impl/bot/TgBotSocketInterface.hpp +++ b/src/socket/interface/impl/bot/TgBotSocketInterface.hpp @@ -9,8 +9,10 @@ #include #include "TgBotPacketParser.hpp" +#include "TgBotSocket_Export.hpp" using TgBotSocket::callback::GenericAck; +using TgBotSocket::callback::UploadFileDryCallback; #ifdef WINDOWS_BUILD #include "impl/SocketWindows.hpp" @@ -59,7 +61,9 @@ struct SocketInterfaceTgBot : SingleThreadCtrlRunnable, static GenericAck handle_ObserveAllChats(const void* ptr); static GenericAck handle_DeleteControllerById(const void* ptr); static GenericAck handle_UploadFile( - const void* ptr, TgBotSocket::PacketHeader::length_type len, bool dry); + const void* ptr, TgBotSocket::PacketHeader::length_type len); + static UploadFileDryCallback handle_UploadFileDry( + const void* ptr, TgBotSocket::PacketHeader::length_type len); // These have their own ack handlers bool handle_GetUptime(SocketConnContext ctx, const void* ptr);