openai-api/communicate.hpp

#ifndef COMMUNICATE_HPP
#define COMMUNICATE_HPP

#include <algorithm>
#include <mutex>
#include <string>
#include <vector>

#if defined(min)
#undef min
#endif

enum FrameType : int16_t {
    TYPE_REQUEST = 0,
    TYPE_RESPONSE_SUCCESS,
    TYPE_RESPONSE_ERROR,
    TYPE_OUT_OF_LIMIT,
};

struct FrameData {
    ~FrameData()
    {
        len = 0;
        delete[] data;
    }
    FrameType type;
    char* data{};
    int len{};
    int16_t protk{};
    int16_t coptk{};
};

class CMutBuffer
{
public:
    CMutBuffer() = default;

public:
    void push(const char* data, int len)
    {
        std::lock_guard<std::mutex> lock(mutex_);
        buffer_.insert(buffer_.end(), data, data + len);
    }
    int index_of(const char* data, int len, int start_pos = 0)
    {
        std::lock_guard<std::mutex> lock(mutex_);
        if (start_pos < 0 || start_pos >= static_cast<int>(buffer_.size()) || len <= 0) {
            return -1;
        }
        auto it = std::search(buffer_.begin() + start_pos, buffer_.end(), data, data + len);
        if (it != buffer_.end()) {
            return std::distance(buffer_.begin(), it);
        }
        return -1;
    }
    const char* get_data() const
    {
        return buffer_.data();
    }
    int get_len() const
    {
        return static_cast<int>(buffer_.size());
    }
    void remove_of(int start_pos, int len)
    {
        std::lock_guard<std::mutex> lock(mutex_);
        if (start_pos < 0 || start_pos >= static_cast<int>(buffer_.size()) || len <= 0) {
            return;
        }
        auto end_pos = std::min(start_pos + len, static_cast<int>(buffer_.size()));
        buffer_.erase(buffer_.begin() + start_pos, buffer_.begin() + end_pos);
    }
    void clear()
    {
        std::lock_guard<std::mutex> lock(mutex_);
        buffer_.clear();
    }

private:
    std::vector<char> buffer_;
    std::mutex mutex_;
};

/*
【TCP 数据协议 】
    header 2 char: 0xFF 0xFE
    type   2 char:
    len    4 char:
    data    xxxxx:
    protk  2 char:
    coptk  2 char:
    tail   2 char: 0xFF 0xFF
*/
inline FrameData* com_parse(CMutBuffer& buffer)
{
    FrameData* r = nullptr;
    constexpr char header[] = {0xFF, 0xFE};
    constexpr char tail[] = {0xFF, 0xFF};

    int find = buffer.index_of(header, sizeof(header));
    if (find < 0) {
        return r;
    }

    int len{};
    int16_t type{};
    int16_t protk{};
    int16_t coptk{};

    std::memcpy(&type, buffer.get_data() + find + sizeof(header), sizeof(type));
    std::memcpy(&len, buffer.get_data() + find + sizeof(header) + sizeof(type), sizeof(len));

    if (len < 1) {
        return r;
    }

    int tail_index = sizeof(header) + sizeof(type) + sizeof(len) + sizeof(protk) + sizeof(coptk) + len;
    if (std::memcmp(buffer.get_data() + tail_index, tail, sizeof(tail)) != 0) {
        return r;
    }

    std::memcpy(&protk, buffer.get_data() + find + sizeof(header) + sizeof(type) + sizeof(len) + len, sizeof(protk));
    std::memcpy(&coptk, buffer.get_data() + find + sizeof(header) + sizeof(type) + sizeof(len) + sizeof(protk) + len, sizeof(coptk));

    r = new FrameData();
    r->type = static_cast<FrameType>(type);
    r->len = len;
    r->data = new char[len];
    std::memcpy(r->data, buffer.get_data() + find + sizeof(header) + sizeof(type) + sizeof(len), len);
    r->protk = protk;
    r->coptk = coptk;

    buffer.remove_of(0, tail_index + 2);
    return r;
}

/*
【TCP 数据协议 】
    header 2 char: 0xFF 0xFE
    type   2 char:
    len    4 char:
    data    xxxxx:
    protk  2 char:
    coptk  2 char:
    tail   2 char: 0xFF 0xFF
*/
inline bool com_pack(FrameData* data, char** out_buf, int& len)
{
    if (data == nullptr) {
        return false;
    }
    if (data->data == nullptr) {
        data->len = 0;
    }

    constexpr char header[] = {0xFF, 0xFE};
    constexpr char tail[] = {0xFF, 0xFF};

    len = sizeof(header) + sizeof(data->type) + sizeof(data->len) + sizeof(data->protk) + sizeof(data->coptk) + data->len +
          sizeof(tail);
    *out_buf = new char[len];
    std::memset(*out_buf, 0, len);
    std::memcpy(*out_buf, header, sizeof(header));
    std::memcpy(*out_buf + sizeof(header), &data->type, sizeof(data->type));
    std::memcpy(*out_buf + sizeof(header) + sizeof(data->type), &data->len, sizeof(data->len));
    if (data->data != nullptr) {
        std::memcpy(*out_buf + sizeof(header) + sizeof(data->type) + sizeof(data->len), data->data, data->len);
    }
    std::memcpy(*out_buf + sizeof(header) + sizeof(data->type) + sizeof(data->len) + data->len, &data->protk,
                sizeof(data->protk));
    std::memcpy(*out_buf + sizeof(header) + sizeof(data->type) + sizeof(data->len) + data->len + sizeof(data->protk),
                &data->coptk, sizeof(data->coptk));
    std::memcpy(*out_buf + sizeof(header) + sizeof(data->type) + sizeof(data->len) + sizeof(data->protk) + sizeof(data->coptk) +
                    data->len,
                tail, sizeof(tail));
    return true;
};

#endif   // COMMUNICATE_HPP