ldb/kvs_rw_tools.c

#include "kvstore.h"
#include "kvs_rw_tools.h"
#include <arpa/inet.h>

int kvs_need(const uint8_t *p, const uint8_t *end, size_t n) {
	return (p + n <= end) ? 0 : -1;
}

// 注意u8类型不需要ntoh或者hton
int kvs_read_u8(const uint8_t **pp, const uint8_t *end, uint8_t *out) {
	const uint8_t *p = *pp;
	if (kvs_need(p, end, 1) < 0) return -1;
	*out = *p;
	*pp = p + 1;
	return 0;
}

int kvs_read_u16(const uint8_t **pp, const uint8_t *end, uint16_t *out) {
	const uint8_t *p = *pp;
	if (kvs_need(p, end, 2) < 0) return -1;
	uint16_t v;
	memcpy(&v, p, 2);
	*out = ntohs(v);
	*pp = p + 2;
	return 0;
}

int kvs_read_u32(const uint8_t **pp, const uint8_t *end, uint32_t *out) {
	const uint8_t *p = *pp;
	if (kvs_need(p, end, 4) < 0) return -1;
	uint32_t v;
	memcpy(&v, p, 4);
	*out = ntohl(v);
	*pp = p + 4;
	return 0;
}

int kvs_write_u8(uint8_t **pp, const uint8_t *end, uint8_t v) {
	uint8_t *p = *pp;
	if (kvs_need(p, end, 1) < 0) return -1;
	*p = v;
	*pp = p + 1;
	return 0;
}

int kvs_write_u16(uint8_t **pp, const uint8_t *end, uint16_t v) {
	uint8_t *p = *pp;
	if (kvs_need(p, end, 2) < 0) return -1;
	uint16_t be = htons(v);
	memcpy(p, &be, 2);
	*pp = p + 2;
	return 0;
}

int kvs_write_u32(uint8_t **pp, const uint8_t *end, uint32_t v) {
	uint8_t *p = *pp;
	if (kvs_need(p, end, 4) < 0) return -1;
	uint32_t be = htonl(v);
	memcpy(p, &be, 4);
	*pp = p + 4;
	return 0;
}

// return: -1 fail, 0 half, >0 consumed
int kvs_parse_one_cmd(const uint8_t *request, int request_length, kvs_req_t *req_out){
	if (!request || request_length <= 0 || !req_out) return -1;
	
	req_out->op = 0;
    req_out->argc = 0;
    req_out->args = NULL;

	const uint8_t *p = request;
    const uint8_t *end = request + (size_t)request_length;

	// OP + ARGC 
    if (kvs_need(p, end, 2)) {
        return 0; // NEED_MORE
    }

	uint8_t op = 0, argc = 0;
	if (kvs_read_u8(&p, end, &op) < 0) return -1;
	if (kvs_read_u8(&p, end, &argc) < 0) return -1;

	if (argc > KVS_MAX_ARGC) return -1;

	// 先扫描一遍确认整条命令数据都在 buffer 里
	const uint8_t *scan = p;
	uint32_t lens[KVS_MAX_ARGC];
	if (argc > 0) {
		for (uint8_t i = 0; i < argc; i++) {
			if (kvs_need(scan, end, 4)) {
				return 0; // NEED_MORE
			}
			uint32_t alen = 0;
			if (kvs_read_u32(&scan, end, &alen) < 0) return -1;

			// 防御：单个参数长度限制
			if (alen > KVS_MAX_ARGLEN) return -1;

			// 防御：scan + alen 越界 / 半包
			if (kvs_need(scan, end, (size_t)alen)) {
				return 0; // NEED_MORE
			}
			lens[i] = alen;
			scan += alen;
        }
	}

	size_t total_len = (size_t)(scan - request);
    if (total_len > KVS_MAX_CMD_BYTES) return -1;

	req_out->op = op;
    req_out->argc = argc;

	if (argc == 0) {
        return (int)total_len;
    }

	kvs_arg_t *args = (kvs_arg_t *)kvs_malloc((size_t)argc * sizeof(kvs_arg_t));
    if (!args) {
        kvs_free_request(req_out);
        return -1;
    }
	memset(args, 0, (size_t)argc * sizeof(kvs_arg_t));

	for (uint8_t i = 0; i < argc; i++) {
        uint32_t alen = 0;
        if (kvs_read_u32(&p, end, &alen) < 0) {
            kvs_free(args);
            kvs_free_request(req_out);
            return -1;
        }

        // alen 与 lens[i] 应当一致（扫描时读过），不一致说明解析器/输入异常
        if (alen != lens[i]) {
            kvs_free(args);
            kvs_free_request(req_out);
            return -1;
        }

        args[i].len = alen;
        args[i].data = p;     // 直接指向输入 buffer（零拷贝）
        p += alen;
    }

	
    req_out->args = args;

    return (int)(p - request);
}

void kvs_free_request(kvs_req_t *req) {
    if (!req) return;
    if (req->args) {
        kvs_free(req->args);
        req->args = NULL;
    }
    req->op = 0;
    req->argc = 0;
}

int kvs_build_one_rsp(const kvs_rsp_t *results, uint8_t *response, size_t response_cap){
	if (!results || !response) return -1;

	const uint8_t *end = response + response_cap;
    uint8_t *p = response;

    // 计算所需长度：1 + 1 + 4 + dlen
    // 注意防止 size_t 溢出
    size_t need = 1u + 1u + 4u + (size_t)results->dlen;
    if (need > response_cap) return -1;

    if (kvs_write_u8(&p, end, results->op) < 0) return -1;
    if (kvs_write_u8(&p, end, results->status) < 0) return -1;
    if (kvs_write_u32(&p, end, results->dlen) < 0) return -1;

    if (results->dlen > 0) {
        if (!results->data) return -1; // 有长度却没指针，视为错误
        if (kvs_need(p, end, (size_t)results->dlen) < 0) return -1;
        memcpy(p, results->data, results->dlen);
        p += results->dlen;
    }

    return (int)(p - response);
}