2019-06-26 16:11:22 -04:00
|
|
|
#include <csignal>
|
|
|
|
#include <misc/memtracker.h>
|
|
|
|
#include "CompressionHandler.h"
|
|
|
|
#define QLZ_COMPRESSION_LEVEL 1
|
|
|
|
#define QLZ_MEMORY_SAFE
|
|
|
|
#include "qlz/QuickLZ.h"
|
|
|
|
#include "buffers.h"
|
|
|
|
|
|
|
|
using namespace ts;
|
|
|
|
using namespace ts::connection;
|
|
|
|
using namespace std;
|
|
|
|
|
2020-01-26 20:21:39 -05:00
|
|
|
namespace ts::compression {
|
|
|
|
class thread_buffer {
|
|
|
|
public:
|
|
|
|
void* get_buffer(size_t size) {
|
|
|
|
if(size > 1024 * 1024 *5) /* we don't want to keep such big buffers in memory */
|
|
|
|
return malloc(size);
|
|
|
|
|
|
|
|
if(this->buffer_length < size) {
|
|
|
|
free(this->buffer_ptr);
|
|
|
|
|
|
|
|
size = std::max(size, (size_t) 1024);
|
|
|
|
this->buffer_ptr = malloc(size);
|
|
|
|
this->buffer_length = size;
|
|
|
|
}
|
|
|
|
return buffer_ptr;
|
|
|
|
}
|
|
|
|
|
|
|
|
void free_buffer(void* ptr) {
|
|
|
|
if(ptr == this->buffer_ptr) return;
|
|
|
|
free(ptr);
|
|
|
|
}
|
|
|
|
|
|
|
|
~thread_buffer() {
|
|
|
|
free(this->buffer_ptr);
|
|
|
|
}
|
|
|
|
private:
|
|
|
|
void* buffer_ptr{nullptr};
|
|
|
|
size_t buffer_length{0};
|
|
|
|
};
|
|
|
|
thread_local thread_buffer qlz_buffer{};
|
|
|
|
|
|
|
|
size_t qlz_decompressed_size(const void* payload, size_t payload_length) {
|
|
|
|
return qlz_size_decompressed((char*) payload) + 400;
|
|
|
|
}
|
|
|
|
|
|
|
|
bool qlz_decompress_payload(const void* payload, void* buffer, size_t* buffer_size) {
|
|
|
|
assert(payload != buffer);
|
|
|
|
|
|
|
|
qlz_state_decompress state{};
|
|
|
|
size_t data_length = qlz_decompress((char*) payload, (char*) buffer, &state);
|
|
|
|
if(data_length <= 0)
|
|
|
|
return false;
|
|
|
|
|
|
|
|
/* test for overflow */
|
|
|
|
if(data_length > *buffer_size) terminate();
|
|
|
|
*buffer_size = data_length;
|
|
|
|
return true;
|
|
|
|
}
|
|
|
|
|
|
|
|
size_t qlz_compressed_size(const void* payload, size_t payload_length) {
|
|
|
|
//// "Always allocate size + 400 bytes for the destination buffer when compressing." <= http://www.quicklz.com/manual.html
|
|
|
|
return max(min(payload_length * 2, (size_t) (payload_length + 400ULL)), (size_t) 24ULL); /* at least 12 bytes (QLZ header) */
|
|
|
|
}
|
|
|
|
|
|
|
|
bool qlz_compress_payload(const void* payload, size_t payload_length, void* buffer, size_t* buffer_length) {
|
|
|
|
assert(payload != buffer);
|
|
|
|
assert(*buffer_length >= qlz_compressed_size(payload, payload_length));
|
|
|
|
|
|
|
|
qlz_state_compress state{};
|
|
|
|
size_t compressed_length = qlz_compress(payload, (char*) buffer, payload_length, &state);
|
|
|
|
if(compressed_length > *buffer_length) terminate();
|
|
|
|
|
|
|
|
if(compressed_length <= 0)
|
|
|
|
return false;
|
|
|
|
*buffer_length = compressed_length;
|
|
|
|
return true;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2019-06-26 16:11:22 -04:00
|
|
|
bool CompressionHandler::compress(protocol::BasicPacket* packet, std::string &error) {
|
2020-01-23 20:49:59 -05:00
|
|
|
auto packet_payload = packet->data();
|
|
|
|
auto header_length = packet->length() - packet_payload.length();
|
2019-10-19 07:55:45 -04:00
|
|
|
|
2020-01-26 20:21:39 -05:00
|
|
|
size_t max_compressed_payload_size = compression::qlz_compressed_size(packet_payload.data_ptr(), packet_payload.length());
|
2020-01-23 20:49:59 -05:00
|
|
|
auto target_buffer = buffer::allocate_buffer(max_compressed_payload_size + header_length);
|
2019-06-26 16:11:22 -04:00
|
|
|
|
2020-01-26 20:21:39 -05:00
|
|
|
size_t compressed_size{max_compressed_payload_size};
|
|
|
|
if(!compression::qlz_compress_payload(packet_payload.data_ptr(), packet_payload.length(), &target_buffer[header_length], &compressed_size)) return false;
|
2019-06-26 16:11:22 -04:00
|
|
|
|
2020-01-23 20:49:59 -05:00
|
|
|
memcpy(target_buffer.data_ptr(), packet->buffer().data_ptr(), header_length);
|
2020-01-26 20:21:39 -05:00
|
|
|
packet->buffer(target_buffer.range(0, compressed_size + header_length));
|
2020-01-23 20:49:59 -05:00
|
|
|
return true;
|
2019-06-26 16:11:22 -04:00
|
|
|
}
|
|
|
|
|
|
|
|
bool CompressionHandler::decompress(protocol::BasicPacket* packet, std::string &error) {
|
2020-01-26 20:21:39 -05:00
|
|
|
auto expected_length = compression::qlz_decompressed_size(packet->data().data_ptr(), packet->data().length());
|
2020-01-23 20:49:59 -05:00
|
|
|
if(expected_length > this->max_packet_size){ //Max 16MB. (97% Compression!)
|
|
|
|
error = "Invalid packet size. (Calculated target length of " + to_string(expected_length) + ". Max length: " + to_string(this->max_packet_size) + ")";
|
|
|
|
return false;
|
|
|
|
}
|
|
|
|
auto header_length = packet->header().length() + packet->mac().length();
|
|
|
|
auto buffer = buffer::allocate_buffer(expected_length + header_length);
|
2019-06-26 16:11:22 -04:00
|
|
|
|
2020-01-26 20:21:39 -05:00
|
|
|
size_t compressed_size{expected_length};
|
2020-02-02 14:11:11 -05:00
|
|
|
if(!compression::qlz_decompress_payload(packet->data().data_ptr(), &buffer[header_length], &compressed_size)) return false;
|
2020-02-02 14:44:35 -05:00
|
|
|
memcpy(buffer.data_ptr(), packet->buffer().data_ptr(), header_length);
|
2020-01-26 20:21:39 -05:00
|
|
|
|
|
|
|
packet->buffer(buffer.range(0, compressed_size + header_length));
|
2020-01-23 20:49:59 -05:00
|
|
|
return true;
|
2019-06-26 16:11:22 -04:00
|
|
|
}
|
|
|
|
|
|
|
|
bool CompressionHandler::progressPacketIn(protocol::BasicPacket* packet, std::string &error) {
|
2020-01-23 20:49:59 -05:00
|
|
|
if(packet->isCompressed()) {
|
|
|
|
if(!this->decompress(packet, error)) return false;
|
|
|
|
packet->setCompressed(false);
|
|
|
|
}
|
|
|
|
return true;
|
2019-06-26 16:11:22 -04:00
|
|
|
}
|
|
|
|
|
|
|
|
bool CompressionHandler::progressPacketOut(protocol::BasicPacket* packet, std::string& error) {
|
2020-01-23 20:49:59 -05:00
|
|
|
if(packet->has_flag(protocol::PacketFlag::Compressed) && !packet->isCompressed()) {
|
|
|
|
if(!this->compress(packet, error)) return false;
|
|
|
|
packet->setCompressed(true);
|
|
|
|
}
|
|
|
|
return true;
|
2019-06-26 16:11:22 -04:00
|
|
|
}
|
|
|
|
|
|
|
|
CompressionHandler::CompressionHandler() { }
|
|
|
|
CompressionHandler::~CompressionHandler() { }
|