trace_tcp.c

/*
	trace tcp connection handling

	Copyright Jeroen Vreeken (jeroen@vreeken.net), 2013

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.

*/

#include <trace/trace_def.h>
#include <trace/trace.h>
#include <log/log.h>
#include <utils/tcp_connect.h>

#include <unistd.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <sys/ioctl.h>
#include <stdio.h>

#define TRACE_BUF_BLOCK	4096

struct trace *trace_alloc(void)
{
	return calloc(sizeof(struct trace), 1);
}

struct trace *trace_open(char *host, int port)
{
	struct trace *trace;
	
	trace = calloc(sizeof(struct trace), 1);
	if (!trace)
		return NULL;

	trace->host = strdup(host);
	trace->port = port;
	
	trace->fd = tcp_connect(host, port);
	if (trace->fd >= 0) {
		ioctl(trace->fd, FIONBIO, &(int){1});
	}
	
	return trace;
}

void trace_initialize_fd(struct trace *trace, int fd)
{
	trace->fd = fd;
	trace->rx_len = 0;
	
	if (trace->fd >= 0) {
		ioctl(trace->fd, FIONBIO, &(int){1});
	}
}

void trace_close(struct trace *trace)
{
	log_send(LOG_T_DEBUG, "Close trace on request");
	if (trace->fd >= 0)
		close(trace->fd);
	trace->fd = -1;	
	trace->recover = false;
}

void trace_free(struct trace *trace)
{
	trace_close(trace);

	free(trace->name);
	free(trace->rx_buffer);
	free(trace->host);
	free(trace);
}

int trace_fd_get(struct trace *trace)
{
	return trace->fd;
}

bool trace_fd_read(struct trace *trace)
{
	ssize_t ret;
	int i;
	bool rx_ready = false;

	if (trace->rx_buf_size - trace->rx_len < TRACE_BUF_BLOCK) {
		unsigned char *newbuf;
		
		newbuf = realloc(trace->rx_buffer, trace->rx_buf_size +	TRACE_BUF_BLOCK);
		if (newbuf) {
			trace->rx_buffer = newbuf;
			trace->rx_buf_size += TRACE_BUF_BLOCK;
		} else {
			log_send(LOG_T_WARNING, 
			    "Could not enlarge buffer for receiving traces");
			return false;
		}
	}
	
	ret = read(trace->fd, trace->rx_buffer + trace->rx_len, TRACE_BUF_BLOCK);

	if (ret > 0) {
		trace->rx_len += ret;
	}

	for (i = 0; i < trace->rx_len; i++) {
		if (trace->rx_buffer[i] == TRACE_END)
			rx_ready = true;
	}
	
	if (((ret < 0 && errno != EAGAIN) || ret == 0) && !rx_ready) {
		log_send(LOG_T_DEBUG, "error during read, closing");
		close(trace->fd);
		trace->fd = -1;
		
		if (trace->handler_close)
			trace->handler_close(trace);
	}
	
	return rx_ready;
}

struct trace_pkt *trace_packet_get(struct trace *trace)
{
	int i;
	size_t len = 0, buflen;
	struct trace_pkt *pkt;
	
	while(trace->rx_len && trace->rx_buffer[0] == TRACE_END) {
		memmove(trace->rx_buffer, trace->rx_buffer+1, trace->rx_len-1);
		trace->rx_len--;
	}

	for (i = 0; i < trace->rx_len; i++) {
		if (trace->rx_buffer[i] == TRACE_END) {
			len = i;
			break;
		}
	}
	
	if (!len) {
		return NULL;
	}
	
	buflen = len;
	
	pkt = trace_packet_new();
	if (!pkt)
		goto err_pkt;
	if (trace_packet_resize(pkt, len))
		goto err_pkt_data;

	memcpy(pkt->data, trace->rx_buffer, len);
	pkt->type = pkt->data[0];

	for (i = 0; i < len - 1; i++) {
		if (pkt->data[i] == TRACE_ESC) {
			if (pkt->data[i+1] == TRACE_ESC_END)
				pkt->data[i] = TRACE_END;
			memmove(&pkt->data[i+1], &pkt->data[i+2], len - i - 2);
			len--;
		}
	}
	pkt->len = len;
	
	trace->rx_len -= buflen + 1;
	memmove(trace->rx_buffer, trace->rx_buffer + buflen + 1, trace->rx_len);
	
	return pkt;

err_pkt_data:
	free(pkt);
err_pkt:
	return NULL;
}

int trace_packet_write(struct trace *trace, struct trace_pkt *pkt)
{
	int i;
	int start = 0;
	ssize_t r;
	
	if (!pkt->len)
		return 0;
	
	while (start < pkt->len) for (i = start; i < pkt->len; i++) {
		if (pkt->data[i] == TRACE_ESC || pkt->data[i] == TRACE_END) {
			unsigned char *seq;
			
			r = write(trace->fd, pkt->data + start, i - start);
			if (r != i - start)
				goto err_write;
			
			start = i + 1;
			if (pkt->data[i] == TRACE_ESC)
				seq = (unsigned char[2]){ TRACE_ESC, TRACE_ESC_ESC };
			else
				seq = (unsigned char[2]){ TRACE_ESC, TRACE_ESC_END };
			r = write(trace->fd, seq, 2);
			if (r != 2)
				goto err_write;
		} else if (i == pkt->len -1) {
			r = write(trace->fd, pkt->data + start, i - start + 1);
			if (r != i - start + 1)
				goto err_write;
			start = i + 1;
		}
	}
	r = write(trace->fd, &(char[1]){ TRACE_END }, 1);
	if (r < 0)
		goto err_write;
	
	return 0;

err_write:
	log_send(LOG_T_DEBUG, "Error during write, closing %d", trace->fd);
	if (trace->fd >= 0)
		close(trace->fd);
	trace->fd = -1;
		
	if (trace->handler_close)
		trace->handler_close(trace);
	return -1;
}