path: root/utils.c

/*
 * utils.c: utilities
 *
 * Copyright (c) 2014, Přemysl Janouch <p.janouch@gmail.com>
 * All rights reserved.
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
 * SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
 * OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
 * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 */

#include <getopt.h>

#if defined __GNUC__
#define ATTRIBUTE_PRINTF(x, y) __attribute__ ((format (printf, x, y)))
#else // ! __GNUC__
#define ATTRIBUTE_PRINTF(x, y)
#endif // ! __GNUC__

#if defined __GNUC__ && __GNUC__ >= 4
#define ATTRIBUTE_SENTINEL __attribute__ ((sentinel))
#else // ! __GNUC__ || __GNUC__ < 4
#define ATTRIBUTE_SENTINEL
#endif // ! __GNUC__ || __GNUC__ < 4

#define N_ELEMENTS(a) (sizeof (a) / sizeof ((a)[0]))

#define BLOCK_START  do {
#define BLOCK_END    } while (0)

// --- Safe memory management --------------------------------------------------

// When a memory allocation fails and we need the memory, we're usually pretty
// much fucked.  Use the non-prefixed versions when there's a legitimate
// worry that an unrealistic amount of memory may be requested for allocation.

static void *
xmalloc (size_t n)
{
	void *p = malloc (n);
	if (!p)
	{
		perror ("malloc");
		exit (EXIT_FAILURE);
	}
	return p;
}

static void *
xcalloc (size_t n, size_t m)
{
	void *p = calloc (n, m);
	if (!p && n && m)
	{
		perror ("calloc");
		exit (EXIT_FAILURE);
	}
	return p;
}

static void *
xrealloc (void *o, size_t n)
{
	void *p = realloc (o, n);
	if (!p && n)
	{
		perror ("realloc");
		exit (EXIT_FAILURE);
	}
	return p;
}

// --- Double-linked list helpers ----------------------------------------------

#define LIST_HEADER(type)                                                      \
	type *next;                                                                \
	type *prev;

#define LIST_PREPEND(head, link)                                               \
	BLOCK_START                                                                \
		(link)->prev = NULL;                                                   \
		(link)->next = (head);                                                 \
		if ((link)->next)                                                      \
			(link)->next->prev = (link);                                       \
		(head) = (link);                                                       \
	BLOCK_END

#define LIST_UNLINK(head, link)                                                \
	BLOCK_START                                                                \
		if ((link)->prev)                                                      \
			(link)->prev->next = (link)->next;                                 \
		else                                                                   \
			(head) = (link)->next;                                             \
		if ((link)->next)                                                      \
			(link)->next->prev = (link)->prev;                                 \
	BLOCK_END

#define LIST_APPEND_WITH_TAIL(head, tail, link)                                \
	BLOCK_START                                                                \
		(link)->prev = (tail);                                                 \
		(link)->next = NULL;                                                   \
		if ((link)->prev)                                                      \
			(link)->prev->next = (link);                                       \
		else                                                                   \
			(head) = (link);                                                   \
		(tail) = (link);                                                       \
	BLOCK_END

#define LIST_UNLINK_WITH_TAIL(head, tail, link)                                \
	BLOCK_START                                                                \
		if ((tail) == (link))                                                  \
			(tail) = (link)->prev;                                             \
		LIST_UNLINK ((head), (link));                                          \
	BLOCK_END

// --- Dynamically allocated strings -------------------------------------------

// Basically a string builder to abstract away manual memory management.

struct str
{
	char *str;                          ///< String data, null terminated
	size_t alloc;                       ///< How many bytes are allocated
	size_t len;                         ///< How long the string actually is
};

/// We don't care about allocations that are way too large for the content, as
/// long as the allocation is below the given threshold.  (Trivial heuristics.)
#define STR_SHRINK_THRESHOLD (1 << 20)

static void
str_init (struct str *self)
{
	self->alloc = 16;
	self->len = 0;
	self->str = strcpy (xmalloc (self->alloc), "");
}

static void
str_free (struct str *self)
{
	free (self->str);
	self->str = NULL;
	self->alloc = 0;
	self->len = 0;
}

static char *
str_steal (struct str *self)
{
	char *str = self->str;
	self->str = NULL;
	str_free (self);
	return str;
}

static void
str_ensure_space (struct str *self, size_t n)
{
	// We allocate at least one more byte for the terminating null character
	size_t new_alloc = self->alloc;
	while (new_alloc <= self->len + n)
		new_alloc <<= 1;
	if (new_alloc != self->alloc)
		self->str = xrealloc (self->str, (self->alloc = new_alloc));
}

static void
str_append_data (struct str *self, const void *data, size_t n)
{
	str_ensure_space (self, n);
	memcpy (self->str + self->len, data, n);
	self->len += n;
	self->str[self->len] = '\0';
}

static void
str_append_c (struct str *self, char c)
{
	str_append_data (self, &c, 1);
}

static void
str_append (struct str *self, const char *s)
{
	str_append_data (self, s, strlen (s));
}

static int
str_append_vprintf (struct str *self, const char *fmt, va_list va)
{
	va_list ap;
	int size;

	va_copy (ap, va);
	size = vsnprintf (NULL, 0, fmt, ap);
	va_end (ap);

	if (size < 0)
		return -1;

	va_copy (ap, va);
	str_ensure_space (self, size);
	size = vsnprintf (self->str + self->len, self->alloc - self->len, fmt, ap);
	va_end (ap);

	if (size > 0)
		self->len += size;

	return size;
}

static int
str_append_printf (struct str *self, const char *fmt, ...)
	ATTRIBUTE_PRINTF (2, 3);

static int
str_append_printf (struct str *self, const char *fmt, ...)
{
	va_list ap;

	va_start (ap, fmt);
	int size = str_append_vprintf (self, fmt, ap);
	va_end (ap);
	return size;
}

static void
str_remove_slice (struct str *self, size_t start, size_t length)
{
	size_t end = start + length;
	if (end > self->len)
		end = self->len;
	memmove (self->str + start, self->str + end, self->len - end);
	self->str[self->len -= length] = '\0';

	// Shrink the string if the allocation becomes way too large
	if (self->alloc >= STR_SHRINK_THRESHOLD && self->len < (self->alloc >> 2))
		self->str = xrealloc (self->str, self->alloc >>= 2);
}

// --- Utilities ---------------------------------------------------------------

static bool
set_blocking (int fd, bool blocking)
{
	int flags = fcntl (fd, F_GETFL);

	bool prev = !(flags & O_NONBLOCK);
	if (blocking)
		flags &= ~O_NONBLOCK;
	else
		flags |=  O_NONBLOCK;

	(void) fcntl (fd, F_SETFL, flags);
	return prev;
}

static void
xclose (int fd)
{
	while (close (fd) == -1)
		if (errno != EINTR)
			break;
}

static char *xstrdup_printf (const char *, ...) ATTRIBUTE_PRINTF (1, 2);

static char *
xstrdup_printf (const char *format, ...)
{
	va_list ap;
	struct str tmp;
	str_init (&tmp);
	va_start (ap, format);
	str_append_vprintf (&tmp, format, ap);
	va_end (ap);
	return str_steal (&tmp);
}

static char *
format_host_port_pair (const char *host, const char *port)
{
	// IPv6 addresses mess with the "colon notation"; let's go with RFC 2732
	if (strchr (host, ':'))
		return xstrdup_printf ("[%s]:%s", host, port);
	return xstrdup_printf ("%s:%s", host, port);
}

static bool
xstrtoul (unsigned long *out, const char *s, int base)
{
	char *end;
	errno = 0;
	*out = strtoul (s, &end, base);
	return errno == 0 && !*end && end != s;
}

// --- libuv-style write adaptor -----------------------------------------------

// Makes it possible to use iovec to write multiple data chunks at once.

typedef struct write_req write_req_t;
struct write_req
{
	LIST_HEADER (write_req_t)
	struct iovec data;                  ///< Data to be written
};

typedef struct write_queue write_queue_t;
struct write_queue
{
	write_req_t *head;                  ///< The head of the queue
	write_req_t *tail;                  ///< The tail of the queue
	size_t head_offset;                 ///< Offset into the head
	size_t len;
};

static void
write_queue_init (struct write_queue *self)
{
	self->head = self->tail = NULL;
	self->head_offset = 0;
	self->len = 0;
}

static void
write_queue_free (struct write_queue *self)
{
	for (write_req_t *iter = self->head, *next; iter; iter = next)
	{
		next = iter->next;
		free (iter->data.iov_base);
		free (iter);
	}
}

static void
write_queue_add (struct write_queue *self, write_req_t *req)
{
	LIST_APPEND_WITH_TAIL (self->head, self->tail, req);
	self->len++;
}

static void
write_queue_processed (struct write_queue *self, size_t len)
{
	while (self->head
		&& self->head_offset + len >= self->head->data.iov_len)
	{
		write_req_t *head = self->head;
		len -= (head->data.iov_len - self->head_offset);
		self->head_offset = 0;

		LIST_UNLINK_WITH_TAIL (self->head, self->tail, head);
		self->len--;
		free (head->data.iov_base);
		free (head);
	}
	self->head_offset += len;
}

static bool
write_queue_is_empty (struct write_queue *self)
{
	return self->head == NULL;
}

// --- Message reader ----------------------------------------------------------

struct msg_reader
{
	struct str buf;                     ///< Input buffer
	uint64_t offset;                    ///< Current offset in the buffer
};

static void
msg_reader_init (struct msg_reader *self)
{
	str_init (&self->buf);
	self->offset = 0;
}

static void
msg_reader_free (struct msg_reader *self)
{
	str_free (&self->buf);
}

static void
msg_reader_compact (struct msg_reader *self)
{
	str_remove_slice (&self->buf, 0, self->offset);
	self->offset = 0;
}

static void
msg_reader_feed (struct msg_reader *self, const void *data, size_t len)
{
	// TODO: have some mechanism to prevent flooding
	msg_reader_compact (self);
	str_append_data (&self->buf, data, len);
}

static void *
msg_reader_get (struct msg_reader *self, size_t *len)
{
	// Try to read in the length of the message
	if (self->offset + sizeof (uint64_t) > self->buf.len)
		return NULL;

	uint8_t *x = (uint8_t *) self->buf.str + self->offset;
	uint64_t msg_len
		= (uint64_t) x[0] << 56 | (uint64_t) x[1] << 48
		| (uint64_t) x[2] << 40 | (uint64_t) x[3] << 32
		| (uint64_t) x[4] << 24 | (uint64_t) x[5] << 16
		| (uint64_t) x[6] << 8  | (uint64_t) x[7];

	if (msg_len < sizeof msg_len)
	{
		// The message is shorter than its header
		// TODO: have some mechanism to report errors
		return NULL;
	}

	if (self->offset + msg_len < self->offset)
	{
		// Trying to read an insane amount of data but whatever
		msg_reader_compact (self);
		return NULL;
	}

	// Check if we've got the full message in the buffer and return it
	if (self->offset + msg_len > self->buf.len)
		return NULL;

	// We have to subtract the header from the reported length
	void *data = self->buf.str + self->offset + sizeof msg_len;
	self->offset += msg_len;
	*len = msg_len - sizeof msg_len;
	return data;
}

// --- Message unpacker --------------------------------------------------------

struct msg_unpacker
{
	const char *data;
	size_t offset;
	size_t len;
};

static void
msg_unpacker_init (struct msg_unpacker *self, const void *data, size_t len)
{
	self->data = data;
	self->len = len;
	self->offset = 0;
}

static size_t
msg_unpacker_get_available (struct msg_unpacker *self)
{
	return self->len - self->offset;
}

#define UNPACKER_INT_BEGIN                                                     \
	if (self->len - self->offset < sizeof *value)                              \
		return false;                                                          \
	uint8_t *x = (uint8_t *) self->data + self->offset;                        \
	self->offset += sizeof *value;

static bool
msg_unpacker_u8 (struct msg_unpacker *self, uint8_t *value)
{
	UNPACKER_INT_BEGIN
	*value = x[0];
	return true;
}

static bool
msg_unpacker_i32 (struct msg_unpacker *self, int32_t *value)
{
	UNPACKER_INT_BEGIN
	*value
		= (uint32_t) x[0] << 24 | (uint32_t) x[1] << 16
		| (uint32_t) x[2] << 8  | (uint32_t) x[3];
	return true;
}

static bool
msg_unpacker_u64 (struct msg_unpacker *self, uint64_t *value)
{
	UNPACKER_INT_BEGIN
	*value
		= (uint64_t) x[0] << 56 | (uint64_t) x[1] << 48
		| (uint64_t) x[2] << 40 | (uint64_t) x[3] << 32
		| (uint64_t) x[4] << 24 | (uint64_t) x[5] << 16
		| (uint64_t) x[6] << 8  | (uint64_t) x[7];
	return true;
}

#undef UNPACKER_INT_BEGIN

// --- Message packer and writer -----------------------------------------------

struct msg_writer
{
	struct str buf;                     ///< Holds the message data
};

static void
msg_writer_init (struct msg_writer *self)
{
	str_init (&self->buf);
	// Placeholder for message length
	str_append_data (&self->buf, "\x00\x00\x00\x00" "\x00\x00\x00\x00", 8);
}

static void
msg_writer_u8 (struct msg_writer *self, uint8_t x)
{
	str_append_data (&self->buf, &x, 1);
}

static void
msg_writer_i32 (struct msg_writer *self, int32_t x)
{
	uint32_t u = x;
	uint8_t tmp[4] = { u >> 24, u >> 16, u >> 8, u };
	str_append_data (&self->buf, tmp, sizeof tmp);
}

static void
msg_writer_u64 (struct msg_writer *self, uint64_t x)
{
	uint8_t tmp[8] =
		{ x >> 56, x >> 48, x >> 40, x >> 32, x >> 24, x >> 16, x >> 8, x };
	str_append_data (&self->buf, tmp, sizeof tmp);
}

static void *
msg_writer_flush (struct msg_writer *self, size_t *len)
{
	// Update the message length
	uint64_t x = self->buf.len;
	uint8_t tmp[8] =
		{ x >> 56, x >> 48, x >> 40, x >> 32, x >> 24, x >> 16, x >> 8, x };
	memcpy (self->buf.str, tmp, sizeof tmp);

	*len = x;
	return str_steal (&self->buf);
}

// --- Option handler ----------------------------------------------------------

// Simple wrapper for the getopt_long API to make it easier to use and maintain.

#define OPT_USAGE_ALIGNMENT_COLUMN 30   ///< Alignment for option descriptions

enum
{
	OPT_OPTIONAL_ARG  = (1 << 0),       ///< The argument is optional
	OPT_LONG_ONLY     = (1 << 1)        ///< Ignore the short name in opt_string
};

// All options need to have both a short name, and a long name.  The short name
// is what is returned from opt_handler_get().  It is possible to define a value
// completely out of the character range combined with the OPT_LONG_ONLY flag.
//
// When `arg_hint' is defined, the option is assumed to have an argument.

struct opt
{
	int short_name;                     ///< The single-letter name
	const char *long_name;              ///< The long name
	const char *arg_hint;               ///< Option argument hint
	int flags;                          ///< Option flags
	const char *description;            ///< Option description
};

struct opt_handler
{
	int argc;                           ///< The number of program arguments
	char **argv;                        ///< Program arguments

	const char *arg_hint;               ///< Program arguments hint
	const char *description;            ///< Description of the program

	const struct opt *opts;             ///< The list of options
	size_t opts_len;                    ///< The length of the option array

	struct option *options;             ///< The list of options for getopt
	char *opt_string;                   ///< The `optstring' for getopt
};

static void
opt_handler_free (struct opt_handler *self)
{
	free (self->options);
	free (self->opt_string);
}

static void
opt_handler_init (struct opt_handler *self, int argc, char **argv,
	const struct opt *opts, const char *arg_hint, const char *description)
{
	memset (self, 0, sizeof *self);
	self->argc = argc;
	self->argv = argv;
	self->arg_hint = arg_hint;
	self->description = description;

	size_t len = 0;
	for (const struct opt *iter = opts; iter->long_name; iter++)
		len++;

	self->opts = opts;
	self->opts_len = len;
	self->options = xcalloc (len + 1, sizeof *self->options);

	struct str opt_string;
	str_init (&opt_string);

	for (size_t i = 0; i < len; i++)
	{
		const struct opt *opt = opts + i;
		struct option *mapped = self->options + i;

		mapped->name = opt->long_name;
		if (!opt->arg_hint)
			mapped->has_arg = no_argument;
		else if (opt->flags & OPT_OPTIONAL_ARG)
			mapped->has_arg = optional_argument;
		else
			mapped->has_arg = required_argument;
		mapped->val = opt->short_name;

		if (opt->flags & OPT_LONG_ONLY)
			continue;

		str_append_c (&opt_string, opt->short_name);
		if (opt->arg_hint)
		{
			str_append_c (&opt_string, ':');
			if (opt->flags & OPT_OPTIONAL_ARG)
				str_append_c (&opt_string, ':');
		}
	}

	self->opt_string = str_steal (&opt_string);
}

static void
opt_handler_usage (struct opt_handler *self, FILE *stream)
{
	struct str usage;
	str_init (&usage);

	str_append_printf (&usage, "Usage: %s [OPTION]... %s\n",
		self->argv[0], self->arg_hint ? self->arg_hint : "");
	str_append_printf (&usage, "%s\n\n", self->description);

	for (size_t i = 0; i < self->opts_len; i++)
	{
		struct str row;
		str_init (&row);

		const struct opt *opt = self->opts + i;
		if (!(opt->flags & OPT_LONG_ONLY))
			str_append_printf (&row, "  -%c, ", opt->short_name);
		else
			str_append (&row, "      ");
		str_append_printf (&row, "--%s", opt->long_name);
		if (opt->arg_hint)
			str_append_printf (&row, (opt->flags & OPT_OPTIONAL_ARG)
				? " [%s]" : " %s", opt->arg_hint);

		// TODO: keep the indent if there are multiple lines
		if (row.len + 2 <= OPT_USAGE_ALIGNMENT_COLUMN)
		{
			str_append (&row, "  ");
			str_append_printf (&usage, "%-*s%s\n",
				OPT_USAGE_ALIGNMENT_COLUMN, row.str, opt->description);
		}
		else
			str_append_printf (&usage, "%s\n%-*s%s\n", row.str,
				OPT_USAGE_ALIGNMENT_COLUMN, "", opt->description);

		str_free (&row);
	}

	fputs (usage.str, stream);
	str_free (&usage);
}

static int
opt_handler_get (struct opt_handler *self)
{
	return getopt_long (self->argc, self->argv,
		self->opt_string, self->options, NULL);
}