myvtm/libtoolkit/iobuffer.c

#include "precompile.h"
#include "iobuffer.h"
#include "refcnt.h"
#include "list.h"
#include <assert.h>
#include <stdarg.h>

#include <winpr/string.h>
#include <winpr/wlog.h>

#include "memutil.h"

#define TAG TOOLKIT_TAG(".iobuffer")

struct iobuffer_t
{
	struct list_head sp_entry;
	char *sp_data;
	int sp_capacity;
	int sp_hdr_capacity;
	int sp_wpos;
	int sp_rpos;
	void *sp_user_data;
	DECLARE_REF_COUNT_MEMBER(sp_ref_cnt);
};

#define DEFAULT_CAPACITY			            256
#define DEFAULT_HEADER_CAPACITY		128

static void expand(iobuffer_t *iobuf, int head_capacity, int newcapacity)
{
	if (head_capacity == -1) {
		int t = iobuf->sp_capacity;
		while (newcapacity > t)
			t = t*2;
		if (t > iobuf->sp_capacity) {
			iobuf->sp_data = (char*)realloc(iobuf->sp_data, t);
			iobuf->sp_capacity = t;
		}
	} else if (newcapacity == -1) {
		int t = iobuf->sp_hdr_capacity;
		while (head_capacity > t)
			t = t*2;
		if (t > iobuf->sp_rpos) {
			char *new_data;
			int new_rpos;
			int new_wpos;
			newcapacity = iobuf->sp_capacity + (t - iobuf->sp_hdr_capacity);
			new_data = malloc(newcapacity);
			new_rpos = t;		
			new_wpos = new_rpos + (iobuf->sp_wpos - iobuf->sp_rpos);
			memcpy(new_data+new_rpos, iobuf->sp_data+iobuf->sp_rpos, iobuf->sp_wpos - iobuf->sp_rpos);
			free(iobuf->sp_data);
			iobuf->sp_data = new_data;
			iobuf->sp_rpos = new_rpos;
			iobuf->sp_wpos = new_wpos;
			iobuf->sp_hdr_capacity  = t;
			iobuf->sp_capacity = newcapacity;
		}
	} else {
		int t = iobuf->sp_capacity;
		char *new_data = (char*)malloc(t);
		int new_rpos = head_capacity;
		int new_wpos = new_rpos + (iobuf->sp_wpos - iobuf->sp_rpos);
		while (newcapacity > t)
			t = t*2;
		memcpy(new_data+new_rpos, iobuf->sp_data+iobuf->sp_rpos, iobuf->sp_wpos - iobuf->sp_rpos);
		free(iobuf->sp_data);
		iobuf->sp_data = new_data;
		iobuf->sp_rpos = new_rpos;
		iobuf->sp_wpos = new_wpos;
		iobuf->sp_capacity = t;
	}
}

static __inline int empty_size(iobuffer_t *iobuf)
{
	return iobuf->sp_capacity - iobuf->sp_wpos;
}

static __inline int content_size(iobuffer_t *iobuf)
{
	return iobuf->sp_wpos - iobuf->sp_rpos;
}

static void destroy(iobuffer_t *iobuf)
{
	WLog_DBG(TAG, "iobuffer destroy 0x%0X", iobuf);

	iobuffer_set_capacity(iobuf, 0);
	//memset(iobuf, 0, sizeof(*iobuf));

	//if (iobuf->sp_capacity > 0) {
	//	assert(iobuf->sp_data);
	//	free(iobuf->sp_data);
	//}
	
	free(iobuf);
}

static void iobuffer_write_buf(iobuffer_t *iobuf, const char *buf, int n)
{
	if (n+iobuf->sp_wpos > iobuf->sp_capacity)
		expand(iobuf, -1, n + iobuf->sp_wpos);
	memcpy(iobuf->sp_data+iobuf->sp_wpos, buf, n);
	iobuf->sp_wpos += n;
}

static void iobuffer_write_head_buf(iobuffer_t *iobuf, const char *buf, int n)
{
	if (iobuf->sp_rpos < n)
		expand(iobuf, iobuf->sp_hdr_capacity-iobuf->sp_rpos+n, -1);
	iobuf->sp_rpos -= n;
	memcpy(iobuf->sp_data+iobuf->sp_rpos, buf, n);
}

static int write_7bit_int(unsigned char buf[8], unsigned int v)
{
	int i = 0;
	while (v >= 0x80) {
		buf[i++] = (unsigned char)(v | 0x80);		// set most left bit == 1 
		v >>= 7;
	}
	buf[i++] = (unsigned char)v;
	return i;
}

static void iobuffer_write_7bit_int(iobuffer_t *iobuf, unsigned int v)
{
	char buf[8];
	int n = write_7bit_int(buf, v);
	iobuffer_write_buf(iobuf, buf, n);
}

static void iobuffer_write_head_7bit_int(iobuffer_t *iobuf, unsigned int v)
{
	char buf[8];
	int n = write_7bit_int(buf, v);
	iobuffer_write_head_buf(iobuf, buf, n);
}

static void iobuffer_read_buf(iobuffer_t *iobuf, char *buf, int n)
{
	assert(content_size(iobuf) >= n);
	assert(buf);
	memcpy(buf, iobuf->sp_data+iobuf->sp_rpos, n);
	iobuf->sp_rpos += n;
}

static void iobuffer_read_7bit_int(iobuffer_t *iobuf, int *v)
{
	int count = 0;
	int shift = 0;
	unsigned char b;
	do {
		if (shift == 5*7) // corrupted stream
			assert(0);
		b = iobuf->sp_data[iobuf->sp_rpos ++];
		count |= (b & 0x7F) << shift;
		shift += 7;
	} while (b & 0x80);		// most left bit is flag
	*v = count;
}

#ifdef DEBUD_MEM_TRACE
static iobuffer_t* iobuffer_create_(int head_capacity, int capacity)
#else
TOOLKIT_API iobuffer_t* iobuffer_create(int head_capacity, int capacity)
#endif
{
	iobuffer_t *iobuf;

	if (head_capacity <= 0)
		head_capacity = DEFAULT_HEADER_CAPACITY;
	if (capacity <= 0) {
		capacity = head_capacity + DEFAULT_CAPACITY;
	} else {
		capacity += head_capacity;
	}

	iobuf = MALLOC_T(iobuffer_t);
	memset(iobuf, 0, sizeof(iobuffer_t));
	iobuf->sp_hdr_capacity = head_capacity;
	iobuf->sp_capacity = capacity;
	iobuf->sp_data = malloc(capacity);
	iobuf->sp_user_data = NULL;
	iobuf->sp_entry.next = iobuf->sp_entry.prev = NULL;
	iobuffer_reset(iobuf);
	REF_COUNT_INIT(&iobuf->sp_ref_cnt);


	WLog_DBG(TAG, "iobuffer create 0x%0X", iobuf);

	return iobuf;
}

TOOLKIT_API void iobuffer_destroy(iobuffer_t *iobuf)
{
	iobuffer_dec_ref(iobuf);
}
#ifdef DEBUD_MEM_TRACE
static iobuffer_t* iobuffer_clone_(iobuffer_t* iobuf)
#else
TOOLKIT_API iobuffer_t* iobuffer_clone(iobuffer_t* iobuf)
#endif
{
	iobuffer_t *newiobuf;

	newiobuf = iobuffer_create(iobuf->sp_hdr_capacity, iobuf->sp_capacity);
	memcpy(newiobuf->sp_data+iobuf->sp_rpos, iobuf->sp_data+iobuf->sp_rpos, iobuf->sp_wpos-iobuf->sp_rpos);
	newiobuf->sp_rpos = iobuf->sp_rpos;
	newiobuf->sp_wpos = iobuf->sp_wpos;
	newiobuf->sp_user_data = iobuf->sp_user_data;

	return newiobuf;
}

TOOLKIT_API void iobuffer_reset(iobuffer_t *iobuf)
{
	iobuf->sp_rpos = iobuf->sp_hdr_capacity;
	iobuf->sp_wpos = iobuf->sp_hdr_capacity;
}

TOOLKIT_API int iobuffer_get_capacity(iobuffer_t *iobuf)
{
	return iobuf->sp_capacity;
}

TOOLKIT_API void iobuffer_set_capacity(iobuffer_t *iobuf, int capacity)
{
	iobuf->sp_data = (char*)realloc(iobuf->sp_data, capacity);
	if (iobuf->sp_wpos > capacity)
		iobuf->sp_wpos = capacity;
	if (iobuf->sp_rpos > capacity)
		iobuf->sp_rpos = capacity;
	iobuf->sp_capacity = capacity;
}

TOOLKIT_API char *iobuffer_data(iobuffer_t *iobuf, int offset)
{
	if (offset != -1) {
		return iobuf->sp_data + iobuf->sp_rpos + offset;
	} else {
		return iobuf->sp_data + iobuf->sp_wpos;
	}
}

TOOLKIT_API int iobuffer_get_length(iobuffer_t *iobuf)
{
	return iobuf->sp_wpos - iobuf->sp_rpos;
}

TOOLKIT_API void iobuffer_set_length(iobuffer_t *iobuf, int length)
{
	iobuf->sp_wpos = iobuf->sp_rpos + length;
	if (iobuf->sp_wpos > iobuf->sp_capacity)
		iobuf->sp_wpos = iobuf->sp_capacity;
}

TOOLKIT_API void iobuffer_push_count(iobuffer_t *iobuf, int n)
{
	assert(empty_size(iobuf) >= n);
	iobuf->sp_wpos += n;
}

TOOLKIT_API void iobuffer_pop_count(iobuffer_t *iobuf, int n)
{
	assert(content_size(iobuf) >= n);
	iobuf->sp_rpos += n;
}

TOOLKIT_API void iobuffer_write(iobuffer_t *iobuf, int t, const void *v, int n)
{
	int len = 0;

	switch (t) {
	case IOBUF_T_I8:
		len += 4;
	case IOBUF_T_I4:
		len += 2;
	case IOBUF_T_I2:
		len += 1;
	case IOBUF_T_I1:
		len += 1;
		iobuffer_write_buf(iobuf, v, len);
		break;
	case IOBUF_T_STR:
		if (!v) {
			n = 0;
		} else {
			if (n == -1)
				n = strlen(v);
		}
		iobuffer_write_7bit_int(iobuf, (unsigned int)n);
		if (n > 0)
			iobuffer_write_buf(iobuf, v, n);
		break;
	case IOBUF_T_WSTR:
		if (!v) {
			n = 0;
		} else {
			if (n == -1) {
				// be compatible with winpr  [4/2/2020 16:02 Gifur]
				n = _wcslen(v) << 1;
				//n = wcslen(v)<<1;
			}
		}
		iobuffer_write_7bit_int(iobuf, n);
		if (n > 0)
			iobuffer_write_buf(iobuf, v, n);
		break;
	case IOBUF_T_BUF:
		if (n > 0)
			iobuffer_write_buf(iobuf, v, n);
		break;
	case IOBUF_T_7BIT:
		iobuffer_write_7bit_int(iobuf, *(unsigned int*)v);
		break;
	}
}

TOOLKIT_API void iobuffer_write_head(iobuffer_t *iobuf, int t, const void *v, int n)
{
	int len = 0;

	switch (t) {
	case IOBUF_T_I8:
		len += 4;
	case IOBUF_T_I4:
		len += 2;
	case IOBUF_T_I2:
		len += 1;
	case IOBUF_T_I1:
		len += 1;
		iobuffer_write_head_buf(iobuf, v, len);
		break;
	case IOBUF_T_STR:
		if (!v) {
			n = 0;
		} else {
			if (n == -1) {
				n = strlen(v);
			}
		}
		if (n > 0)
			iobuffer_write_head_buf(iobuf, v, n);
		iobuffer_write_head_7bit_int(iobuf, n);
		break;
	case IOBUF_T_WSTR:
		if (!v) {
			n = 0;
		} else {
			if (n == -1) {
				n = _wcslen(v)<<1;
			}
		}
		if (n > 0)
			iobuffer_write_head_buf(iobuf, v, n);
		iobuffer_write_head_7bit_int(iobuf, n);
		break;
	case IOBUF_T_BUF:
		if (n > 0)
			iobuffer_write_head_buf(iobuf, v, n);
		break;
	case IOBUF_T_7BIT:
		iobuffer_write_head_7bit_int(iobuf, *(unsigned int*)v);
		break;
	default:
		assert(0);
	}
}

TOOLKIT_API void iobuffer_read(iobuffer_t *iobuf, int t, void *v, int *n)
{
	int len = 0;

	switch (t) {
	case IOBUF_T_I8:
		len += 4;
	case IOBUF_T_I4:
		len += 2;
	case IOBUF_T_I2:
		len += 1;
	case IOBUF_T_I1:
		len += 1;
		iobuffer_read_buf(iobuf, v, len);
		break;
	case IOBUF_T_STR:
		if (v) {
			int old_rpos = iobuf->sp_rpos;
			iobuffer_read_7bit_int(iobuf, &len);
			if (!n || len < *n) {
				iobuffer_read_buf(iobuf, v, len);
				*((char*)v + len) = 0;
				if (n)
					*n = len;
			} else {
				iobuf->sp_rpos = old_rpos;
				*n = len;
			}
		} else {
			if (n) {
				int old_rpos = iobuf->sp_rpos;
				iobuffer_read_7bit_int(iobuf, &len);
				iobuf->sp_rpos = old_rpos;
				*n = len;
			} else {
				assert(0);
			}
		}
		break;
	case IOBUF_T_WSTR:
		if (v) {
			int old_rpos = iobuf->sp_rpos;
			iobuffer_read_7bit_int(iobuf, &len);
			if (!n || len < *n) {
				iobuffer_read_buf(iobuf, v, len); /* WCHAR */
				*((char*)v + len) = 0; // wchar termination char is 2 byte
				*((char*)v + len+1) = 0;
				if (n)
					*n = len;
			} else {
				iobuf->sp_rpos = old_rpos;
				*n = len;
			}
		} else {
			if (n) {
				int old_rpos = iobuf->sp_rpos;
				iobuffer_read_7bit_int(iobuf, &len);
				iobuf->sp_rpos = old_rpos;
				*n = len;
			} else {
				assert(0);
			}
		}
		break;
	case IOBUF_T_BUF:
		if (*n > 0)
			iobuffer_read_buf(iobuf, v, *n);
		break;
	case IOBUF_T_7BIT:
		iobuffer_read_7bit_int(iobuf, v);
		break;
	default:
		assert(0);
	}
}

TOOLKIT_API int iobuffer_get_read_state(iobuffer_t *iobuf)
{
	return iobuf->sp_rpos;
}

TOOLKIT_API void iobuffer_restore_read_state(iobuffer_t *iobuf, int state)
{
	iobuf->sp_rpos = state;
}

TOOLKIT_API int iobuffer_get_write_state(iobuffer_t *iobuf)
{
	return iobuf->sp_wpos;
}

TOOLKIT_API void iobuffer_restore_write_state(iobuffer_t *iobuf, int state)
{
	iobuf->sp_wpos = state;
}

TOOLKIT_API void iobuffer_set_user_data(iobuffer_t *iobuf, void *user_data)
{
	iobuf->sp_user_data = user_data;
}

TOOLKIT_API void *iobuffer_get_user_data(iobuffer_t *iobuf)
{
	return iobuf->sp_user_data;
}

IMPLEMENT_REF_COUNT_MT(iobuffer, iobuffer_t, sp_ref_cnt, destroy)

TOOLKIT_API void iobuffer_format_readv(iobuffer_t *iobuf, const char *fmt, va_list arg)
{
	if (fmt) {
		const char *p = fmt;
		while (*p) {
			switch (*p) {
			case '8':
				iobuffer_read(iobuf, IOBUF_T_I8, va_arg(arg, void*), NULL);
				break;
			case '7':  // 7bit int code
				iobuffer_read(iobuf, IOBUF_T_7BIT, va_arg(arg, void*), NULL);
				break;
			case '4':
				iobuffer_read(iobuf, IOBUF_T_I4, va_arg(arg, void*), 0);
				break;
			case '2':
				iobuffer_read(iobuf, IOBUF_T_I2, va_arg(arg, void*), 0);
				break;
			case '1':
				iobuffer_read(iobuf, IOBUF_T_I1, va_arg(arg, void*), 0);
				break;
			case 's':case 'S':
				{
					char **pstr = va_arg(arg, char**);
					int slen;
					iobuffer_read(iobuf, IOBUF_T_STR, NULL, &slen);
					*pstr = (char*)malloc(slen+1);
					iobuffer_read(iobuf, IOBUF_T_STR, *pstr, 0);
				}
				break;
			case 'w': case 'W':
				{
					WCHAR **pstr = va_arg(arg, WCHAR**);
					int slen;
					iobuffer_read(iobuf, IOBUF_T_WSTR, NULL, &slen);
					*pstr = (WCHAR*)malloc(slen+2); // null terminated is two bytes
					iobuffer_read(iobuf, IOBUF_T_WSTR, *pstr, 0);
				}
				break;
			case 'b':case 'B':
				{
					void *buf = va_arg(arg, void*);
					int size = va_arg(arg, int);
					iobuffer_read(iobuf, IOBUF_T_BUF, buf, &size);
				}
				break;
			default:
				assert(0);
				break;
			}
			p++;
		}
	}
}

TOOLKIT_API void iobuffer_format_writev(iobuffer_t *iobuf, const char *fmt, va_list arg)
{
	if (fmt) {
		const char *p = fmt;
		while (*p) {
			switch (*p) {
			case '8':
				iobuffer_write(iobuf, IOBUF_T_I8, va_arg(arg, void*), 0);
				break;
			case '7':
				iobuffer_write(iobuf, IOBUF_T_7BIT, va_arg(arg, void*), 0);
				break;
			case '4':
				iobuffer_write(iobuf, IOBUF_T_I4, va_arg(arg, void*), 0);
				break;
			case '2':
				iobuffer_write(iobuf, IOBUF_T_I2, va_arg(arg, void*), 0);
				break;
			case '1':
				iobuffer_write(iobuf, IOBUF_T_I1, va_arg(arg, void*), 0);
				break;
			case 's':case 'S':
				iobuffer_write(iobuf, IOBUF_T_STR, va_arg(arg, void*), -1);
				break;
			case 'w':case 'W':
				iobuffer_write(iobuf, IOBUF_T_WSTR, va_arg(arg, void*), -1);
				break;
			case 'b':case 'B':
				{
					void *buf = va_arg(arg, void*);
					int size = va_arg(arg, int);
					iobuffer_write(iobuf, IOBUF_T_BUF, buf, size);
				}
				break;
			default:
				assert(0);
				break;
			}
			p++;
		}
	}
}

TOOLKIT_API void iobuffer_format_read(iobuffer_t *iobuf, const char *fmt, ...)
{
	va_list arg;
	va_start(arg, fmt);
	iobuffer_format_readv(iobuf, fmt, arg);
	va_end(arg);
}

TOOLKIT_API void iobuffer_format_write(iobuffer_t *iobuf, const char *fmt, ...)
{
	va_list arg;
	va_start(arg, fmt);
	iobuffer_format_writev(iobuf, fmt, arg);
	va_end(arg);
}

TOOLKIT_API void iobuffer_format_writev_head(iobuffer_t *iobuf, const char *fmt, va_list arg)
{
	if (fmt) {
		const char *p = fmt;
		while (*p) {
			switch (*p) {
			case '8':
				iobuffer_write_head(iobuf, IOBUF_T_I8, va_arg(arg, void*), 0);
				break;
			case '7':
				iobuffer_write_head(iobuf, IOBUF_T_7BIT, va_arg(arg, void*), 0);
				break;
			case '4':
				iobuffer_write_head(iobuf, IOBUF_T_I4, va_arg(arg, void*), 0);
				break;
			case '2':
				iobuffer_write_head(iobuf, IOBUF_T_I2, va_arg(arg, void*), 0);
				break;
			case '1':
				iobuffer_write_head(iobuf, IOBUF_T_I1, va_arg(arg, void*), 0);
				break;
			case 's':case 'S':
				iobuffer_write_head(iobuf, IOBUF_T_STR, va_arg(arg, void*), -1);
				break;
			case 'w':case 'W':
				iobuffer_write_head(iobuf, IOBUF_T_WSTR, va_arg(arg, void*), -1);
				break;
			case 'b':case 'B':
				{
					void *buf = va_arg(arg, void*);
					int size = va_arg(arg, int);
					iobuffer_write_head(iobuf, IOBUF_T_BUF, buf, size);
				}
				break;
			default:
				assert(0);
				break;
			}
			p++;
		}
	}
}

TOOLKIT_API void iobuffer_format_write_head(iobuffer_t *iobuf, const char *fmt, ...)
{
	va_list arg;
	va_start(arg, fmt);
	iobuffer_format_writev_head(iobuf, fmt, arg);
	va_end(arg);
}

struct iobuffer_queue_t
{
	struct list_head iobuf_list;
	int cnt;
};

TOOLKIT_API iobuffer_queue_t *iobuffer_queue_create()
{
	iobuffer_queue_t *queue = MALLOC_T(iobuffer_queue_t);
	queue->cnt = 0;
	INIT_LIST_HEAD(&queue->iobuf_list);
	return queue;
}

TOOLKIT_API void iobuffer_queue_destroy(iobuffer_queue_t *queue)
{
	while (!list_empty(&queue->iobuf_list)) {
		iobuffer_t *iobuf = iobuffer_queue_deque(queue);
		iobuffer_destroy(iobuf);
	}
	free(queue);
}

TOOLKIT_API void iobuffer_queue_enqueue(iobuffer_queue_t *queue, iobuffer_t *iobuf)
{
	list_add_tail(&iobuf->sp_entry, &queue->iobuf_list);
	queue->cnt++;
}

TOOLKIT_API iobuffer_t *iobuffer_queue_deque(iobuffer_queue_t *queue)
{
	iobuffer_t *iob = list_first_entry(&queue->iobuf_list, iobuffer_t, sp_entry);
	list_del(&iob->sp_entry);
	queue->cnt--;
	return iob;
}

TOOLKIT_API int iobuffer_queue_count(iobuffer_queue_t *queue)
{
	return queue->cnt;
}

TOOLKIT_API iobuffer_t *iobuffer_queue_head(iobuffer_queue_t *queue)
{
	return list_first_entry(&queue->iobuf_list, iobuffer_t, sp_entry);
}

TOOLKIT_API iobuffer_t *iobuffer_queue_tail(iobuffer_queue_t *queue)
{
	return list_last_entry(&queue->iobuf_list, iobuffer_t, sp_entry);
}


#ifdef DEBUD_MEM_TRACE
TOOLKIT_API iobuffer_t* iobuffer_create_trace(int head_capacity, int capacity, const char* file, const char* function, int line)
{
	WLog_DBG(TAG, "iobuffer_create_trace: <%s>[%d]{%s}", function, line, file);
	return iobuffer_create_(head_capacity, capacity);
}

TOOLKIT_API iobuffer_t* iobuffer_clone_trace(iobuffer_t* iobuf, const char* file, const char* function, int line)
{
	WLog_DBG(TAG, "%s: <%s>[%d]{%s}", __FUNCTION__, function, line, file);
	return iobuffer_clone_(iobuf);
}

#endif // DEBUD_MEM_TRACE