添加单元测试

CMakeLists.txt

@@ -2,6 +2,7 @@ cmake_minimum_required(VERSION 3.16)
 project(Transmission)
 
 add_subdirectory(src)
+add_subdirectory(test)
 
 add_executable(exec main.cpp
 )

src/CMakeLists.txt

@@ -17,7 +17,8 @@ add_library(Transmission
         Client.cpp
         Transmit.cpp
         CharBuff.cpp
-        "Exception/PrepareDataException.h"
+        Exception/PrepareDataException.h
 )
 
-target_link_libraries(Transmission protobuf proto)
+target_include_directories(Transmission PUBLIC ${CMAKE_CURRENT_BINARY_DIR})
+target_link_libraries(Transmission protobuf proto lzma)

src/Transmit.cpp

@@ -10,15 +10,17 @@
 #include <utility>
 #include <transmission.pb.h>
 
-#define DATA_STOP 0x20030507
-
 namespace Transmission {
+
+    const int32_t Transmit::DATA_STOP = 0x20030507;
+
     Transmit::Transmit(SOCKET fd, std::string ip) : ip(std::move(ip)), fd(fd) {
 
     }
 
     void Transmit::dataArrives(unsigned char *data, size_t size) {
         std::unique_lock<std::mutex> lock(dataMutex);
+        charBuff.put(data, 0, size);
         if (this->packSize < 0) {
             if (charBuff.getSize() >= 4) {
                 //读取下一个数据包的长度
@@ -36,18 +38,19 @@ namespace Transmission {
             //数据不够，等待填充数据
             return;
         }
-        int32_t magicNumber = (charBuff.get() << 0) + (charBuff.get() << 8) + (charBuff.get() << 16) + (charBuff.get() << 24);
-        if (magicNumber != DATA_STOP)
-            throw PrepareDataException("魔数错误");
 
         auto buff = std::shared_ptr<unsigned char>(new unsigned char[this->packSize], [](const unsigned char *p) {
             delete[] p;
         });
         charBuff.get(buff.get(), this->packSize);
+
+        int32_t magicNumber = (charBuff.get() << 0) + (charBuff.get() << 8) + (charBuff.get() << 16) + (charBuff.get() << 24);
+        if (magicNumber != DATA_STOP)
+            throw PrepareDataException("魔数错误");
         DataPackets dataPackets;
         if (!dataPackets.ParseFromArray(buff.get(), this->packSize))
             throw PrepareDataException("反序列化出错");
-
+        this->packSize = -1;
         switch (dataPackets.type()) {
             //无压缩
             case CompressAlgorithm::NOT: {
@@ -79,28 +82,33 @@ namespace Transmission {
 
     std::shared_ptr<unsigned char> Transmit::unLzma(const std::string &str, size_t &size) {
         lzma_stream stream = LZMA_STREAM_INIT;
-        lzma_ret ret = lzma_easy_encoder(&stream, LZMA_PRESET_DEFAULT, LZMA_CHECK_CRC64);
+        lzma_ret ret = lzma_stream_decoder(&stream, UINT64_MAX, LZMA_CONCATENATED);
         if (ret != LZMA_OK)
             throw PrepareDataException("LZMA初始化失败！");
         // 解压缩数据
         stream.next_in = reinterpret_cast<const uint8_t *>(str.data());
         stream.avail_in = str.size();
-        CharBuff buff((int64_t) lzma_stream_buffer_bound(str.size()));
+        CharBuff buff;
+        const int defaultSize = 1024;
         size_t lastSize = 0;
-        do {
+        while (true) {
-            size_t bufferSize = lzma_stream_buffer_bound(str.size());
+            //size_t bufferSize = defaultSize;
-            auto ptr = std::shared_ptr<unsigned char>(new unsigned char[bufferSize], [](const unsigned char *p) {
+            auto ptr = std::shared_ptr<unsigned char>(new unsigned char[defaultSize], [](const unsigned char *p) {
                 delete[] p;
             });
             stream.next_out = reinterpret_cast<uint8_t *>(ptr.get());
-            stream.avail_out = bufferSize;
+            stream.avail_out = defaultSize;
             ret = lzma_code(&stream, LZMA_FINISH);
+            if (ret == LZMA_STREAM_END) {
                 lastSize = stream.total_out - lastSize;
                 buff.put(ptr.get(), 0, (int64_t) lastSize);
-        } while (ret == LZMA_BUF_ERROR);
+                break;
-
+            }
-        if (ret != LZMA_OK && ret != LZMA_STREAM_END)
+            if (ret != LZMA_OK)
                 throw PrepareDataException("LZMA解压缩失败");
+            lastSize = stream.total_out - lastSize;
+            buff.put(ptr.get(), 0, (int64_t) lastSize);
+        }
 
         auto ptr = std::shared_ptr<unsigned char>(new unsigned char[buff.getSize()], [](const unsigned char *p) {
             delete[] p;
@@ -109,4 +117,82 @@ namespace Transmission {
         buff.get(ptr.get(), buff.getSize());
         return ptr;
     }
+
+    void Transmit::sendData(unsigned char *data, size_t size, int type) {
+        DataPackets dataPackets;
+        dataPackets.set_type(type);
+        std::string byteData(reinterpret_cast<char *>(data), size);
+        dataPackets.set_data(byteData);
+        if (size < 64) {
+            //小于64字节的数据包不压缩
+            pushData(dataPackets);
+        } else {
+            //大于64字节的数据使用lzma压缩后发送
+            pushLzmaData(dataPackets);
+        }
+    }
+
+    void Transmit::pushData(unsigned char *data, int32_t size) const {
+        ::send(this->fd, &size, sizeof(int32_t), 0);
+        ::send(this->fd, data, size, 0);
+        ::send(this->fd, &DATA_STOP, sizeof(DATA_STOP), 0);
+    }
+
+    void Transmit::pushData(DataPackets &data) {
+        data.set_algorithm(CompressAlgorithm::NOT);
+        size_t pack_size = data.ByteSizeLong();
+        auto *temp = new unsigned char[pack_size];
+        data.SerializeToArray(temp, (int) pack_size);
+        pushData(temp, pack_size);
+        delete[] temp;
+    }
+
+    void Transmit::pushLzmaData(DataPackets &data) {
+        // 初始化压缩流
+        lzma_stream stream = LZMA_STREAM_INIT;
+        lzma_ret ret = lzma_easy_encoder(&stream, 6, LZMA_CHECK_CRC64); // 压缩级别为 6，使用 CRC64 进行校验
+        if (ret != LZMA_OK)
+            throw PrepareDataException("LZMA初始化失败");
+
+
+        CharBuff buff;
+        size_t lastSize = 0;
+
+        stream.next_in = (const uint8_t *) data.data().data();
+        stream.avail_in = data.data().size();
+        //size_t maxCompressedSize = lzma_stream_buffer_bound(data.data().size());
+        const size_t maxCompressedSize = 1024;
+        while (true) {
+            auto ptr = std::shared_ptr<unsigned char>(new unsigned char[maxCompressedSize], [](const unsigned char *p) {
+                delete[] p;
+            });
+            stream.next_out = reinterpret_cast<uint8_t *>(ptr.get());
+            stream.avail_out = maxCompressedSize;
+            ret = lzma_code(&stream, LZMA_FINISH);
+            if (ret == LZMA_STREAM_END) {
+                lastSize = stream.total_out - lastSize;
+                buff.put(ptr.get(), 0, (int64_t) lastSize);
+                break;
+            }
+            if(ret != LZMA_OK)
+                throw PrepareDataException("LZMA压缩失败");
+            lastSize = stream.total_out - lastSize;
+            buff.put(ptr.get(), 0, (int64_t) lastSize);
+        }
+
+        auto ptr = std::shared_ptr<unsigned char>(new unsigned char[buff.getSize()], [](const unsigned char *p) {
+            delete[] p;
+        });
+        size_t size = buff.getSize();
+        buff.get(ptr.get(), buff.getSize());
+
+        data.set_algorithm(CompressAlgorithm::LZMA);
+        std::string byteData(reinterpret_cast<char *>(ptr.get()), size);
+        data.set_data(byteData);
+        size_t pack_size = data.ByteSizeLong();
+        auto *temp = new unsigned char[pack_size];
+        data.SerializeToArray(temp, (int) pack_size);
+        pushData(temp, pack_size);
+        delete[] temp;
+    }
 } // Transmission

src/Transmit.h

@@ -11,6 +11,7 @@
 #include <pcap/socket.h>
 #include <memory>
 #include "CharBuff.h"
+#include "transmission.pb.h"
 
 namespace Transmission {
     /// 传输层实现
@@ -23,6 +24,17 @@ namespace Transmission {
         int32_t packSize = -1;
         std::mutex dataMutex;
 
+    private:
+
+        /// 无压缩
+        void pushData(DataPackets &data);
+
+        /// lzma压缩后发送
+        void pushLzmaData(DataPackets &data);
+
+        /// 实际发送
+        virtual void pushData(unsigned char *data, int32_t size) const;
+
     protected:
 
         virtual std::shared_ptr<unsigned char> copyMem(const std::string &str);
@@ -31,11 +43,16 @@ namespace Transmission {
         virtual std::shared_ptr<unsigned char> unLzma(const std::string &str, size_t &size);
 
     public:
+        static const int32_t DATA_STOP;
+
         explicit Transmit(SOCKET fd, std::string ip);
 
         /// 收到来自网络的数据
         virtual void dataArrives(unsigned char *data, size_t size);
 
+        /// 发送数据
+        virtual void sendData(unsigned char *data, size_t size, int type = 0);
+
         /// 数据包就绪
         /// @param type 数据包类型
         /// @param data 解压缩后的数据包

test/CMakeLists.txt

@@ -0,0 +1,16 @@
+cmake_minimum_required(VERSION 3.16)
+project(Transmission)
+
+set(CMAKE_CXX_STANDARD 20)
+
+add_executable(Test
+        test_main.cpp
+)
+
+include_directories(../src)
+
+enable_testing()
+find_package(GTest REQUIRED)
+target_link_libraries(Test PRIVATE GTest::GTest GTest::Main Transmission absl_log_internal_check_op absl_log_internal_message)
+set(GTEST_LIB gtest gtest_main)
+add_test(NAME Test COMMAND Test)

test/test_main.cpp

@@ -0,0 +1,65 @@
+// 版权所有 (c) ling 保留所有权利。
+// 除非另行说明，否则仅允许在Transmission中使用此文件中的代码。
+// 
+// 由 ling 创建于 24-4-18.
+//
+#include <gtest/gtest.h>
+#include "Transmit.h"
+
+static bool isExec = false;
+static const char *shortStr = "Client Hello";
+static const char *LongStr = "Client HelloClient HelloClient HelloClient HelloClient HelloClient HelloClient HelloClient HelloClient HelloClient Hello";
+
+class Transmit : public Transmission::Transmit {
+private:
+    SOCKET fd;
+public:
+    explicit Transmit(SOCKET fd, std::string ip) : Transmission::Transmit(fd, ip) {
+        this->fd = fd;
+    }
+
+    void packetReady(int type, std::shared_ptr<unsigned char> data, size_t size) override {
+        if (type == 1)
+            ASSERT_TRUE(strcmp((char *) data.get(), shortStr) == 0);
+        else if (type == 2)
+            ASSERT_TRUE(strcmp((char *) data.get(), LongStr) == 0);
+        else
+            ASSERT_TRUE(false);
+        isExec = true;
+    }
+
+    void pushData(unsigned char *data, int32_t size) const override {
+        ::write(this->fd, &size, sizeof(int32_t));
+        ::write(this->fd, data, size);
+        ::write(this->fd, &DATA_STOP, sizeof(DATA_STOP));
+    }
+
+    void read() {
+        int32_t size = 0;
+        ::read(this->fd, &size, sizeof(size));
+        auto temp = new unsigned char[size];
+        ::read(this->fd, temp, size);
+        int32_t stop = 0;
+        ::read(this->fd, &stop, sizeof(stop));
+        dataArrives((unsigned char *) &size, sizeof(size));
+        dataArrives(temp, size);
+        dataArrives((unsigned char *) &stop, sizeof(stop));
+        delete[] temp;
+    }
+};
+
+TEST(Transmit测试, 1) {
+    int fds[2];
+    pipe(fds);
+    int read = fds[0];
+    int write = fds[1];
+    Transmit transmitRead(read, "");
+    Transmit transmitWrite(write, "");
+    transmitWrite.sendData((unsigned char *) shortStr, strlen(shortStr) + 1, 1);
+    transmitRead.read();
+    ASSERT_TRUE(isExec);
+    isExec = false;
+    transmitWrite.sendData((unsigned char *) LongStr, strlen(LongStr) + 1, 2);
+    transmitRead.read();
+    ASSERT_TRUE(isExec);
+}