/*
 * Copyright © 2020 Red Hat, Inc.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice (including the next
 * paragraph) shall be included in all copies or substantial portions of the
 * Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

#pragma once

#include "config.h"

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

#include "util-macros.h"
#include "util-mem.h"

/**
 * Wrapper to convert an errno-setting syscall into a
 * value-or-negative-errno.
 *
 * Use: int rc = xerrno(foo(bar));
 */
static inline int
xerrno(int value) {
	return value < 0 ? -errno : value;
}

/**
 * Wrapper around close(). It checks for fd != -1 to satisfy coverity and
 * friends and always returns -1.
 */
static inline int
xclose(int fd) {
	if (fd != -1) close(fd);
	return -1;
}

DEFINE_TRIVIAL_CLEANUP_FUNC(int, close);
#define _cleanup_close_ _cleanup_(closep)

DEFINE_TRIVIAL_CLEANUP_FUNC(FILE *, fclose);
#define _cleanup_fclose_ _cleanup_(fclosep)

/**
 * Wrapper around read(). Returns the number of bytes read or a negative
 * errno on failure.
 */
static inline int
xread(int fd, void *buf, size_t count)
{
	return xerrno(read(fd, buf, count));
}

/**
 * Wrapper around read(). Returns the number of bytes read or a negative
 * errno on failure. Any fds passed along with the message
 * are stored in the -1-terminated allocated fds array, to be freed by the
 * caller. Where no fds were passed, the array is NULL.
 */
static inline int
xread_with_fds(int fd, void *buf, size_t count, int **fds)
{
	const size_t MAX_FDS = 32;
	union {
		struct cmsghdr header;
		char control[CMSG_SPACE(MAX_FDS * sizeof(int))];
	} ctrl;

	struct iovec iov = {
		.iov_base = buf,
		.iov_len = count,
	};

	struct msghdr msg = {
		.msg_name = NULL,
		.msg_namelen = 0,
		.msg_iov = &iov,
		.msg_iovlen = 1,
		.msg_control = ctrl.control,
		.msg_controllen = sizeof(ctrl.control),
	};

	int received = xerrno(recvmsg(fd, &msg, 0));
	if (received > 0) {
		*fds = NULL;

		_cleanup_free_ int *fd_return = calloc(MAX_FDS + 1, sizeof(int));
		size_t idx = 0;

		for (struct cmsghdr *hdr = CMSG_FIRSTHDR(&msg); hdr; hdr = CMSG_NXTHDR(&msg, hdr)) {
			if (hdr->cmsg_level != SOL_SOCKET ||
			    hdr->cmsg_type != SCM_RIGHTS)
				continue;

			int *fd = (int *)CMSG_DATA(hdr);
			fd_return[idx++] = *fd;
			if (idx >= MAX_FDS)
				break;
		}
		fd_return[idx] = -1;
		*fds = steal(&fd_return);
	}
	return received;
}

/**
 * Wrapper around write(). Returns the number of bytes written or a negative
 * errno on failure.
 */
static inline int
xwrite(int fd, const void *buf, size_t count)
{
	return xerrno(write(fd, buf, count));
}

/**
 * Wrapper around send() that always sets MSG_NOSIGNAL. Returns the number
 * of bytes written or a negative errno on failure.
 */
static inline int
xsend(int fd, const void *buf, size_t len)
{
	return xerrno(send(fd, buf, len, MSG_NOSIGNAL));
}

/**
 * Wrapper around send() that always sets MSG_NOSIGNAL and allows appending
 * file descriptors to the message.
 *
 * @param fds Array of file descriptors, terminated by -1.
 */
static inline int
xsend_with_fd(int fd, const void *buf, size_t len, int *fds)
{
	size_t nfds = 0;

	for (nfds = 0; fds != NULL && fds[nfds] != -1; nfds++) {
		/* noop */
	}

	if (nfds == 0)
		return xsend(fd, buf, len);

	union {
		struct cmsghdr header;
		char control[CMSG_SPACE(nfds * sizeof(int))];
	} ctrl;

	memset(&ctrl, 0, sizeof(ctrl));

	struct iovec iov = {
		.iov_base = (void*)buf,
		.iov_len = len,
	};

	struct msghdr msg = {
		.msg_name = NULL,
		.msg_namelen = 0,
		.msg_iov = &iov,
		.msg_iovlen = 1,
		.msg_control = ctrl.control,
		.msg_controllen = sizeof(ctrl.control),
	};

	ctrl.header.cmsg_len = CMSG_LEN(nfds * sizeof(int));
	ctrl.header.cmsg_level = SOL_SOCKET;
	ctrl.header.cmsg_type = SCM_RIGHTS;
	memcpy(CMSG_DATA(CMSG_FIRSTHDR(&msg)), fds, nfds * sizeof(int));

	return xerrno(sendmsg(fd, &msg, MSG_NOSIGNAL));
}

/* consider this struct opaque */
struct iobuf {
	size_t sz;
	size_t len;
	char *data;
	int fds[32];
};

/**
 * Create a new iobuf structure with the given size.
 */
static inline struct iobuf *
iobuf_new(size_t size)
{
	struct iobuf *buf = malloc(sizeof(*buf));
	char *data = malloc(size);

	assert(buf);
	assert(data);

	*buf = (struct iobuf) {
		.sz = size,
		.len = 0,
		.data = data,
	};

	int *fd;
	ARRAY_FOR_EACH(buf->fds, fd) {
		*fd = -1;
	}

	return buf;
}

/**
 * Resize the buffer. You should never need this one directly, it's a helper
 * function.
 */
static inline void
iobuf_resize(struct iobuf *buf, size_t to_size)
{
	char *newdata = realloc(buf->data, to_size);
	assert(newdata);

	buf->data = newdata;
	buf->sz = to_size;
}

/**
 * The count of data bytes in this buffer.
 */
static inline size_t
iobuf_len(struct iobuf *buf)
{
	return buf->len;
}

/**
 * Pointer to the data bytes. Note that the buffer is considered binary
 * data. The caller must ensure that any strings stored in the buffer are
 * null-terminated.
 *
 * The returned pointer only valid in the immediate scope, any iobuf
 * function may invalidate the pointer.
 */
static inline const char *
iobuf_data(struct iobuf *buf)
{
	return buf->data;
}

/**
 * Return the next available file descriptor in this buffer or -1.
 * The fd is removed from this buffer and belongs to the caller.
 */
static inline int
iobuf_take_fd(struct iobuf *buf)
{
	int fd = buf->fds[0];
	if (fd != -1)
		memmove(buf->fds, buf->fds + 1, ARRAY_LENGTH(buf->fds) - 1);
	return fd;
}

/**
 * Remove the data bytes from the buffer. The caller must free() the data.
 * The buffer state is the same as iobuf_new() after this call.
 */
static inline char *
iobuf_take_data(struct iobuf *buf)
{
	char *data = buf->data;

	buf->data = NULL;
	buf->len = 0;
	iobuf_resize(buf, buf->sz);

	return data;
}

/**
 * Append len bytes to the buffer. If the data exceeds the current buffer
 * size it is resized automatically.
 */
static inline void
iobuf_append(struct iobuf *buf, const char *data, size_t len)
{
	if (buf->len + len > buf->sz) {
		size_t newsize = buf->len + len;
		iobuf_resize(buf, newsize);
	}

	memcpy(buf->data + buf->len, data, len);
	buf->len += len;
}

/**
 * Append all available data from the file descriptor to the pointer. The
 * file descriptor shold be in O_NONBLOCK or this call will block. If the
 * data exceeds the current buffer size it is resized automatically.
 *
 * @return The number of bytes read or a negative errno on failure. Zero
 * indicates EOF.
 */
static inline int
iobuf_append_from_fd(struct iobuf *buf, int fd)
{
	char data[1024];
	size_t nread = 0;
	ssize_t rc;
	do {
		rc = xread(fd, data, sizeof(data));
		if (rc == 0 || rc == -EAGAIN) {
			break;
		} else if (rc < 0) {
			return rc;
		}

		iobuf_append(buf, data, rc);
		nread += rc;
	} while (rc == sizeof(data));

	return nread == 0 ? rc : (int)nread;
}

/**
 * Append all available data from the file descriptor to the pointer. The
 * file descriptor shold be in O_NONBLOCK or this call will block. If the
 * data exceeds the current buffer size it is resized automatically.
 *
 * Any file descriptors passed through the fd are placed
 *
 * @return The number of bytes read or a negative errno on failure. Zero
 * indicates EOF.
 */
static inline int
iobuf_recv_from_fd(struct iobuf *buf, int fd)
{
	char data[1024];
	size_t nread = 0;
	ssize_t rc;
	do {
		_cleanup_free_ int *fds = NULL;
		rc = xread_with_fds(fd, data, sizeof(data), &fds);
		if (rc == 0 || rc == -EAGAIN) {
			break;
		} else if (rc < 0) {
			return rc;
		}
		iobuf_append(buf, data, rc);

		if (fds) {
			int *fd = fds;
			for (size_t idx = 0; *fd != -1 && idx < ARRAY_LENGTH(buf->fds) - 1; idx++) {
				if (buf->fds[idx] == -1) {
					buf->fds[idx] = *fd;
					fd++;
				}
			}
		}

		nread += rc;
	} while (rc == sizeof(data));

	return nread == 0 ? rc : (int)nread;
}

/**
 * Release the memory associated with this iobuf. Use iobuf_take_data()
 * prevent the data from being free()d.
 */
static inline struct iobuf *
iobuf_free(struct iobuf *buf)
{
	if (buf) {
		free(buf->data);
		buf->sz = 0;
		buf->len = 0;
		buf->data = NULL;

		int fd;
		while ((fd = iobuf_take_fd(buf)) != -1)
			xclose(fd);
		free(buf);
	}
	return NULL;
}

static inline void _iobuf_cleanup(struct iobuf **buf) { iobuf_free(*buf); }
/**
 * Helper macro to auto-free a iobuf struct when the variable goes out of
 * scope. Use like this:
 *    _cleanup_iobuf_ struct iobuf *foo = iobuf_new(64);
 */
#define _cleanup_iobuf_ __attribute__((cleanup(_iobuf_cleanup)))