path: root/common/src/tfe_cmsg.cpp

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <arpa/inet.h>
#include <errno.h>
#include <pthread.h>

#include "tfe_types.h"
#include "tfe_utils.h"
#include "tfe_cmsg.h"

struct tfe_cmsg_tlv
{
	uint16_t type;
	uint16_t length;
	union
	{
		uint8_t value_as_uint8[0];
		uint16_t value_as_uint16[0];
		uint32_t value_as_uint32[0];
		unsigned char value_as_string[0];
	};
} __attribute__((packed));

struct tfe_cmsg
{
	uint8_t flag;
	uint8_t ref;
	pthread_rwlock_t rwlock;
	uint16_t nr_tlvs;
	struct tfe_cmsg_tlv* tlvs[TFE_CMSG_TLV_NR_MAX];
	uint16_t size;
} __attribute__((packed));

struct tfe_cmsg_serialize_header
{
	uint8_t __magic__[2]; /* Must be 0x4d, 0x5a */
	uint16_t nr_tlvs;
	struct tfe_cmsg_tlv tlvs[0];
} __attribute__((packed));

struct tfe_cmsg* tfe_cmsg_init()
{
	struct tfe_cmsg *cmsg = ALLOC(struct tfe_cmsg, 1);
	cmsg->size = sizeof(struct tfe_cmsg_serialize_header);

	pthread_rwlock_init(&(cmsg->rwlock), NULL);

	ATOMIC_ZERO(&cmsg->flag);
	ATOMIC_ZERO(&cmsg->ref);
	ATOMIC_INC(&cmsg->ref);
	return cmsg;
}

void tfe_cmsg_destroy(struct tfe_cmsg *cmsg)
{
	if(cmsg != NULL)
	{
		if ((__sync_sub_and_fetch(&cmsg->ref, 1) != 0))
			return;
		pthread_rwlock_wrlock(&cmsg->rwlock);
		for(int i = 0; i < TFE_CMSG_TLV_NR_MAX; i++)
		{
			FREE(&(cmsg->tlvs[i]));
		}
		pthread_rwlock_unlock(&cmsg->rwlock);
		pthread_rwlock_destroy(&cmsg->rwlock);
		FREE(&cmsg);
	}
}

void tfe_cmsg_dup(struct tfe_cmsg *cmsg)
{
	if (cmsg == NULL)
		return;
	ATOMIC_INC(&cmsg->ref);
}

void tfe_cmsg_set_flag(struct tfe_cmsg *cmsg, uint8_t flag)
{
	if (cmsg == NULL)
		return;
	ATOMIC_SET(&cmsg->flag, flag);
}

uint8_t tfe_cmsg_get_flag(struct tfe_cmsg *cmsg)
{
	if (cmsg == NULL)
		return 0;
    uint8_t flag = 0;
    flag = ATOMIC_READ(&cmsg->flag);
    return flag;
}

int tfe_cmsg_set(struct tfe_cmsg * cmsg, enum tfe_cmsg_tlv_type type, const unsigned char * value, uint16_t size)
{
	if(type >= TFE_CMSG_TLV_NR_MAX)
	{
		return TFE_CMSG_INVALID_TYPE;
	}
	pthread_rwlock_wrlock(&cmsg->rwlock);
	struct tfe_cmsg_tlv *tlv = cmsg->tlvs[type];
	uint16_t length = sizeof(struct tfe_cmsg_tlv) + size;

	// If the current tlv has been set, the previous value will be overwritten
	if (NULL != tlv)
	{
		cmsg->nr_tlvs--;
		cmsg->size -= tlv->length;
		cmsg->tlvs[type] = NULL;
		free(tlv);
		tlv = NULL;
	}
	
	if(tlv == NULL)
	{
		tlv = (struct tfe_cmsg_tlv*)ALLOC(char, length);
		cmsg->nr_tlvs++;
		cmsg->size += length;
	}
	tlv->type = type;
	tlv->length = length;
	memcpy(tlv->value_as_string, value, size);
	cmsg->tlvs[type] = tlv;
	pthread_rwlock_unlock(&cmsg->rwlock);
	return 0;
}

int tfe_cmsg_get_value(struct tfe_cmsg * cmsg, enum tfe_cmsg_tlv_type type, unsigned char * out_value,
    size_t sz_out_value_buf, uint16_t * out_size)
{
	struct tfe_cmsg_tlv *tlv;
	int result = 0;
	size_t value_length = 0;

	if (unlikely(type >= TFE_CMSG_TLV_NR_MAX))
    {
        result = -EINVAL;
        goto errout;
    }

	pthread_rwlock_rdlock(&cmsg->rwlock);
    tlv = cmsg->tlvs[type];
    if (unlikely(tlv == NULL))
    {
        result = -ENOENT;
        goto errout;
    }

    value_length = tlv->length - sizeof(struct tfe_cmsg_tlv);
    if (unlikely(sz_out_value_buf <  value_length))
    {
        result = -ENOBUFS;
        goto errout;
    }

    memcpy(out_value, tlv->value_as_string, value_length);
    *out_size = value_length;
	pthread_rwlock_unlock(&cmsg->rwlock);
    return 0;

errout:
	pthread_rwlock_unlock(&cmsg->rwlock);
	return result;
}

uint16_t tfe_cmsg_serialize_size_get(struct tfe_cmsg *cmsg)
{
	pthread_rwlock_rdlock(&cmsg->rwlock);
	uint16_t size = cmsg->size;
	pthread_rwlock_unlock(&cmsg->rwlock);
	return size;
}

int tfe_cmsg_serialize(struct tfe_cmsg *cmsg, unsigned char *buff, uint16_t bufflen, uint16_t *serialize_len)
{
	//size是serialize之后的实际长度
	pthread_rwlock_rdlock(&cmsg->rwlock);
	uint16_t size = cmsg->size;
	//传入buff是否够长
	if(bufflen < size)
	{
		pthread_rwlock_unlock(&cmsg->rwlock);
		return TFE_CMSG_BUFF_NOT_ENOUGH;
	}
	//size是否正确
	if(size < sizeof(struct tfe_cmsg_serialize_header))
	{
		pthread_rwlock_unlock(&cmsg->rwlock);
		return TFE_CMSG_BUFF_NOT_ENOUGH;
	}
	struct tfe_cmsg_serialize_header *header = (struct tfe_cmsg_serialize_header*)buff;
	header->__magic__[0] = 0x4d;
	header->__magic__[1] = 0x5a;
	header->nr_tlvs = htons(cmsg->nr_tlvs);
	uint16_t offset = sizeof(struct tfe_cmsg_serialize_header);
	//检查nr_tlvs是否合法
	int count = 0;
	for(int i = 0; i < TFE_CMSG_TLV_NR_MAX; i++){
		if(cmsg->tlvs[i] != NULL)
		{
			count++;
		}
	}
	if(count != cmsg->nr_tlvs)
	{
		pthread_rwlock_unlock(&cmsg->rwlock);
		return TFE_CMSG_INVALID_FORMAT;
	}
	//序列化
	for(int i = 0; i < TFE_CMSG_TLV_NR_MAX; i++)
	{
		struct tfe_cmsg_tlv *tlv = cmsg->tlvs[i];
		if(tlv == NULL)
		{
			continue;
		}
		if(i != tlv->type)
		{
			pthread_rwlock_unlock(&cmsg->rwlock);
			return TFE_CMSG_INVALID_FORMAT;
		}
		uint16_t length = tlv->length;
		if(length < sizeof(struct tfe_cmsg_tlv) || offset + length > size)
		{
			pthread_rwlock_unlock(&cmsg->rwlock);
			return TFE_CMSG_INVALID_FORMAT;
		}
		memcpy((char*)header + offset, (void*)tlv, length);
		struct tfe_cmsg_tlv *tlv1 = (struct tfe_cmsg_tlv*)((char*)header + offset);
		tlv1->type = htons(tlv->type);
		tlv1->length = htons(tlv->length);
		offset += length;
	}
	//检查size是否正确
	if(offset != size)
	{
		pthread_rwlock_unlock(&cmsg->rwlock);
		return TFE_CMSG_INVALID_FORMAT;
	}
	*serialize_len = size;
	pthread_rwlock_unlock(&cmsg->rwlock);
	return 0;
}

//反序列化
int tfe_cmsg_deserialize(const unsigned char *data, uint16_t len, struct tfe_cmsg** pcmsg)
{
	struct tfe_cmsg_serialize_header *header = (struct tfe_cmsg_serialize_header*)data;
	struct tfe_cmsg *cmsg = NULL;
	int offset = 0, nr_tlvs = -1;

	if(len < sizeof(struct tfe_cmsg_serialize_header))
	{
		goto error_out;
	}
	if(header->__magic__[0] != 0x4d || header->__magic__[1] != 0x5a)
	{
		goto error_out;
	}

	cmsg = ALLOC(struct tfe_cmsg, 1);
	pthread_rwlock_init(&(cmsg->rwlock), NULL);
	offset = sizeof(struct tfe_cmsg_serialize_header);
	nr_tlvs = ntohs(header->nr_tlvs);
	for(int i = 0; i < nr_tlvs; i++)
	{
		struct tfe_cmsg_tlv *tlv = (struct tfe_cmsg_tlv*)(data + offset);
		if(offset + sizeof(struct tfe_cmsg_tlv) > len)
		{
			goto error_out;
		}

		uint16_t type = ntohs(tlv->type);
		uint16_t length = ntohs(tlv->length);

		if(length < sizeof(struct tfe_cmsg_tlv) || offset + length > len)
		{
			goto error_out;
		}

		int ret = tfe_cmsg_set(cmsg, (enum tfe_cmsg_tlv_type)type,
		    data + offset + sizeof(struct tfe_cmsg_tlv), length - sizeof(struct tfe_cmsg_tlv));

		if(ret < 0)
		{
			goto error_out;
		}
		offset += length;
	}
	cmsg->size = offset;
	*pcmsg = cmsg;
	return 0;

error_out:
	tfe_cmsg_destroy(cmsg);
	return TFE_CMSG_INVALID_FORMAT;
}