Skip to content

Commit

Permalink
socket: Add possiblity to skip enc/dec
Browse files Browse the repository at this point in the history
- TODO: Replace with valid tokens, I have to think if it is required to
- Update submodule to fix compile
  • Loading branch information
Royna2544 committed Dec 16, 2024
1 parent 1ee4ff3 commit 93ab993
Show file tree
Hide file tree
Showing 4 changed files with 50 additions and 18 deletions.
46 changes: 32 additions & 14 deletions src/socket/bot/PacketParser.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -232,25 +232,37 @@ std::optional<Packet> readPacket(const TgBotSocket::Context& context) {
LOG(ERROR) << "While reading data, failed";
return std::nullopt;
}
packet.data = *data;
if (!decryptPacket(packet)) {
return std::nullopt;
}
return packet;
}

bool Socket_API decryptPacket(TgBotSocket::Packet& packet) {
auto& header = packet.header;
auto& data = packet.data;

std::string_view session_token(header.session_token.data(),
header.session_token.size());
if (session_token.empty()) {
LOG(WARNING) << "No session token provided";
return std::nullopt;
if (header.session_token ==
TgBotSocket::Packet::Header::session_token_type{}) {
LOG(WARNING) << "No session token provided, decryption will be skipped.";
return true;
}
if (packet.header.hmac !=
HMAC::compute(static_cast<const uint8_t*>(data->get()), data->size(),
HMAC::compute(static_cast<const uint8_t*>(data.get()), data.size(),
session_token)) {
LOG(ERROR) << "HMAC mismatch";
return std::nullopt;
return false;
}
packet.data = decrypt_payload(header.session_token, data.value(),
packet.header.init_vector);
packet.data =
decrypt_payload(header.session_token, data, packet.header.init_vector);
if (!static_cast<bool>(packet.data)) {
LOG(ERROR) << "Decryption failed";
return std::nullopt;
return false;
}
return packet;
return true;
}

Packet Socket_API
Expand All @@ -266,12 +278,18 @@ createPacket(const Command command, const void* data,
if (data != nullptr && length > 0) {
packet.data.resize(length);
packet.data.assignFrom(data, length);
packet.data = encrypt_payload(sessionToken, packet.data,
packet.header.init_vector);

if (sessionToken != Packet::Header::session_token_type{}) {
packet.data = encrypt_payload(sessionToken, packet.data,
packet.header.init_vector);
packet.header.hmac =
HMAC::compute(static_cast<const uint8_t*>(packet.data.get()),
packet.header.data_size, sessionToken.data());
} else {
LOG(WARNING)
<< "No session token provided, encryption will be skipped";
}
packet.header.data_size = packet.data.size();
packet.header.hmac =
HMAC::compute(static_cast<const uint8_t*>(packet.data.get()),
packet.header.data_size, sessionToken.data());
}
return packet;
}
Expand Down
12 changes: 12 additions & 0 deletions src/socket/bot/PacketParser.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,18 @@ namespace TgBotSocket {
std::optional<Packet> Socket_API
readPacket(const TgBotSocket::Context& context);


/**
* @brief Decrypts a packet using the provided context.
*
* This function attempts to decrypt the given packet using the provided context.
* If successful, it returns `true`. If decryption fails, it returns `false`.
*
* @param packet The packet to decrypt
* @return `true` if the packet was successfully decrypted; otherwise, `false`.
*/
bool Socket_API decryptPacket(TgBotSocket::Packet& packet);

/**
* @brief Creates a packet with the given command and data.
*
Expand Down
8 changes: 5 additions & 3 deletions tests/SocketDataHandlerTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ class SocketDataHandlerTest : public ::testing::Test {
SharedMalloc packetData;
TgBotSocket::Packet::Header recv_header;

EXPECT_TRUE(TgBotSocket::decryptPacket(pkt));
EXPECT_CALL(*_mockImpl, write(_))
.WillOnce(DoAll(SaveArg<0>(&packetData), Return(true)));
mockInterface->handlePacket(*_mockImpl, std::move(pkt));
Expand Down Expand Up @@ -313,7 +314,8 @@ TEST_F(SocketDataHandlerTest, TestCmdUploadFileDryExistsOptSaidNo) {
TEST_F(SocketDataHandlerTest, TestCmdUploadFileOK) {
// Prepare file contents
const auto filemem = createFileMem();
SharedMalloc mem(sizeof(TgBotSocket::data::UploadFile) + filemem.size());
SharedMalloc mem(sizeof(TgBotSocket::data::UploadFileMeta) +
filemem.size());
auto* uploadfile = static_cast<TgBotSocket::data::UploadFile*>(mem.get());
uploadfile->srcfilepath = {"sourcefile"};
uploadfile->destfilepath = {"destinationfile"};
Expand All @@ -322,11 +324,11 @@ TEST_F(SocketDataHandlerTest, TestCmdUploadFileOK) {
uploadfile->options.overwrite = true;
uploadfile->options.dry_run = false;
mem.assignTo(filemem.get(), filemem.size(),
sizeof(TgBotSocket::data::UploadFile));
sizeof(TgBotSocket::data::UploadFileMeta));

// Set expectations
TgBotSocket::Packet pkt = TgBotSocket::createPacket(
TgBotSocket::Command::CMD_UPLOAD_FILE_DRY, mem.get(), mem.size(),
TgBotSocket::Command::CMD_UPLOAD_FILE, mem.get(), mem.size(),
TgBotSocket::PayloadType::Binary, {});
EXPECT_CALL(*_mockVFS, writeFile(FSP(uploadfile->destfilepath.data()), _,
filemem.size()))
Expand Down

0 comments on commit 93ab993

Please sign in to comment.