/*
 * vdev_ctrl — eudev-compatible udev control socket
 *
 * Derived from udev-ctrl.c in eudev (Kay Sievers, LGPL-2.1+).
 */

/* struct ucred, SCM_CREDENTIALS, accept4 (vdevd build lacks -D_GNU_SOURCE) */
#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif

#include <sys/socket.h>
#include <sys/un.h>

#include "vdev_ctrl.h"
#include "util.h"

#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

struct vdev_ctrl {
	int refcount;
	int sock;
	char path[sizeof(((struct sockaddr_un *)0)->sun_path)];
	bool bound;
	bool cleanup_socket;
	bool connected;
};

struct vdev_ctrl_connection {
	int refcount;
	struct vdev_ctrl *uctrl;
	int sock;
};

struct vdev_ctrl_msg {
	int refcount;
	struct vdev_ctrl_connection *conn;
	struct vdev_ctrl_msg_wire wire;
};

static int vdev_ctrl_getpeercred(int fd, uid_t *uid_out)
{
	struct ucred cred;
	socklen_t n = sizeof(cred);

	if (getsockopt(fd, SOL_SOCKET, SO_PEERCRED, &cred, &n) < 0)
		return -errno;
	*uid_out = cred.uid;
	return 0;
}

static int vdev_ctrl_accept4(int sock)
{
#if defined(HAVE_ACCEPT4) || defined(__linux__)
	int fd = accept4(sock, NULL, NULL, SOCK_CLOEXEC | SOCK_NONBLOCK);

	if (fd < 0 && (errno == ENOSYS || errno == EOPNOTSUPP)) {
		fd = accept(sock, NULL, NULL);
		if (fd >= 0) {
			fcntl(fd, F_SETFL, fcntl(fd, F_GETFL, 0) | O_NONBLOCK);
			fcntl(fd, F_SETFD, FD_CLOEXEC);
		}
	}
	return fd;
#else
	int fd = accept(sock, NULL, NULL);

	if (fd >= 0) {
		fcntl(fd, F_SETFL, fcntl(fd, F_GETFL, 0) | O_NONBLOCK);
		fcntl(fd, F_SETFD, FD_CLOEXEC);
	}
	return fd;
#endif
}

static struct vdev_ctrl *vdev_ctrl_ref(struct vdev_ctrl *uctrl)
{
	if (uctrl)
		uctrl->refcount++;
	return uctrl;
}

struct vdev_ctrl *vdev_ctrl_unref(struct vdev_ctrl *uctrl)
{
	if (!uctrl || --uctrl->refcount != 0)
		return uctrl;
	if (uctrl->sock >= 0)
		close(uctrl->sock);
	free(uctrl);
	return NULL;
}

struct vdev_ctrl *vdev_ctrl_new(const char *socket_path)
{
	struct vdev_ctrl *uctrl;
	const int on = 1;
	const char *path = socket_path ? socket_path : VDEV_CTRL_DEFAULT_PATH;

	uctrl = calloc(1, sizeof(*uctrl));
	if (!uctrl)
		return NULL;
	uctrl->refcount = 1;
	uctrl->sock = socket(AF_LOCAL, SOCK_SEQPACKET | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
	if (uctrl->sock < 0) {
		free(uctrl);
		return NULL;
	}
	if (setsockopt(uctrl->sock, SOL_SOCKET, SO_PASSCRED, &on, sizeof(on)) < 0)
		vdev_warn("vdev_ctrl: SO_PASSCRED: %s\n", strerror(errno));
	snprintf(uctrl->path, sizeof(uctrl->path), "%s", path);
	return uctrl;
}

int vdev_ctrl_enable_receiving(struct vdev_ctrl *uctrl)
{
	struct sockaddr_un addr;

	if (!uctrl || uctrl->bound)
		return 0;

	memset(&addr, 0, sizeof(addr));
	addr.sun_family = AF_LOCAL;
	snprintf(addr.sun_path, sizeof(addr.sun_path), "%s", uctrl->path);

	if (bind(uctrl->sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
		if (errno == EADDRINUSE) {
			unlink(uctrl->path);
			if (bind(uctrl->sock, (struct sockaddr *)&addr, sizeof(addr)) < 0)
				return -errno;
		} else {
			return -errno;
		}
	}
	if (listen(uctrl->sock, 128) < 0)
		return -errno;

	uctrl->bound = true;
	uctrl->cleanup_socket = true;
	return 0;
}

int vdev_ctrl_cleanup(struct vdev_ctrl *uctrl)
{
	struct stat st_path, st_sock;

	/*
	 * Only unlink if the path still names *our* listen socket.  Another
	 * vdevd may have rebound /run/udev/control after we started; blindly
	 * unlinking would orphan that daemon's clients (ENOENT).
	 */
	if (!uctrl || !uctrl->cleanup_socket)
		return 0;
	if (uctrl->sock >= 0 &&
	    fstat(uctrl->sock, &st_sock) == 0 &&
	    stat(uctrl->path, &st_path) == 0 &&
	    st_path.st_dev == st_sock.st_dev &&
	    st_path.st_ino == st_sock.st_ino)
		unlink(uctrl->path);
	uctrl->cleanup_socket = false;
	return 0;
}

int vdev_ctrl_get_fd(struct vdev_ctrl *uctrl)
{
	return uctrl ? uctrl->sock : -EINVAL;
}

const char *vdev_ctrl_get_path(struct vdev_ctrl *uctrl)
{
	return uctrl ? uctrl->path : NULL;
}

int vdev_ctrl_connect(struct vdev_ctrl *uctrl, int timeout_sec)
{
	struct sockaddr_un addr;
	int err;

	if (!uctrl)
		return -EINVAL;
	if (uctrl->connected)
		return 0;

	memset(&addr, 0, sizeof(addr));
	addr.sun_family = AF_LOCAL;
	snprintf(addr.sun_path, sizeof(addr.sun_path), "%s", uctrl->path);

	if (connect(uctrl->sock, (struct sockaddr *)&addr, sizeof(addr)) == 0) {
		uctrl->connected = true;
		return 0;
	}
	if (errno != EINPROGRESS)
		return -errno;

	for (;;) {
		struct pollfd pfd = { .fd = uctrl->sock, .events = POLLOUT };
		int r = poll(&pfd, 1, timeout_sec > 0 ? timeout_sec * 1000 : -1);

		if (r < 0) {
			if (errno == EINTR)
				continue;
			return -errno;
		}
		if (r == 0)
			return -ETIMEDOUT;
		break;
	}

	err = 0;
	{
		socklen_t n = sizeof(err);

		if (getsockopt(uctrl->sock, SOL_SOCKET, SO_ERROR, &err, &n) < 0)
			return -errno;
	}
	if (err != 0)
		return -err;

	uctrl->connected = true;
	return 0;
}

int vdev_ctrl_peer_pid(struct vdev_ctrl *uctrl, pid_t *pid_out, int timeout_sec)
{
	struct ucred cred;
	socklen_t n = sizeof(cred);
	int rc;

	if (!uctrl || !pid_out)
		return -EINVAL;

	rc = vdev_ctrl_connect(uctrl, timeout_sec);
	if (rc != 0)
		return rc;

	if (getsockopt(uctrl->sock, SOL_SOCKET, SO_PEERCRED, &cred, &n) < 0)
		return -errno;
	if (cred.pid <= 0)
		return -ESRCH;

	*pid_out = cred.pid;
	return 0;
}

struct vdev_ctrl_connection *vdev_ctrl_get_connection(struct vdev_ctrl *uctrl)
{
	struct vdev_ctrl_connection *conn;
	uid_t uid = (uid_t)-1;
	const int on = 1;
	int fd;

	if (!uctrl)
		return NULL;

	fd = vdev_ctrl_accept4(uctrl->sock);
	if (fd < 0) {
		if (errno != EINTR && errno != EAGAIN)
			vdev_error("vdev_ctrl: accept: %s\n", strerror(errno));
		return NULL;
	}

	if (vdev_ctrl_getpeercred(fd, &uid) < 0 || uid != 0) {
		vdev_error("vdev_ctrl: rejected connection uid=%u\n", (unsigned)uid);
		close(fd);
		return NULL;
	}

	conn = calloc(1, sizeof(*conn));
	if (!conn) {
		close(fd);
		return NULL;
	}
	conn->refcount = 1;
	conn->uctrl = vdev_ctrl_ref(uctrl);
	conn->sock = fd;
	if (setsockopt(conn->sock, SOL_SOCKET, SO_PASSCRED, &on, sizeof(on)) < 0)
		vdev_warn("vdev_ctrl: conn SO_PASSCRED: %s\n", strerror(errno));
	return conn;
}

struct vdev_ctrl_connection *vdev_ctrl_connection_unref(
	struct vdev_ctrl_connection *conn)
{
	if (!conn || --conn->refcount != 0)
		return conn;
	if (conn->sock >= 0)
		close(conn->sock);
	vdev_ctrl_unref(conn->uctrl);
	free(conn);
	return NULL;
}

struct vdev_ctrl_msg *vdev_ctrl_receive_msg(struct vdev_ctrl_connection *conn)
{
	struct vdev_ctrl_msg *msg;
	struct cmsghdr *cmsg;
	struct iovec iov;
	char cred_msg[CMSG_SPACE(sizeof(struct ucred))];
	struct msghdr smsg = {
		.msg_iov = &iov,
		.msg_iovlen = 1,
		.msg_control = cred_msg,
		.msg_controllen = sizeof(cred_msg),
	};
	struct ucred *cred;
	ssize_t size;

	if (!conn)
		return NULL;

	msg = calloc(1, sizeof(*msg));
	if (!msg)
		return NULL;
	msg->refcount = 1;
	msg->conn = conn;
	conn->refcount++;

	for (;;) {
		struct pollfd pfd = { .fd = conn->sock, .events = POLLIN };
		int r = poll(&pfd, 1, 10000);

		if (r < 0) {
			if (errno == EINTR)
				continue;
			goto err;
		}
		if (r == 0) {
			vdev_error("vdev_ctrl: receive timeout\n");
			goto err;
		}
		if (!(pfd.revents & POLLIN))
			goto err;
		break;
	}

	iov.iov_base = &msg->wire;
	iov.iov_len = sizeof(msg->wire);
	size = recvmsg(conn->sock, &smsg, 0);
	if (size < 0) {
		vdev_error("vdev_ctrl: recvmsg: %s\n", strerror(errno));
		goto err;
	}

	cmsg = CMSG_FIRSTHDR(&smsg);
	if (!cmsg || cmsg->cmsg_type != SCM_CREDENTIALS) {
		vdev_error("vdev_ctrl: missing SCM_CREDENTIALS\n");
		goto err;
	}
	cred = (struct ucred *)CMSG_DATA(cmsg);
	if (cred->uid != 0) {
		vdev_error("vdev_ctrl: message uid=%u ignored\n", (unsigned)cred->uid);
		goto err;
	}
	if (msg->wire.magic != VDEV_CTRL_MAGIC) {
		vdev_error("vdev_ctrl: bad magic 0x%08x\n", msg->wire.magic);
		goto err;
	}
	return msg;

err:
	vdev_ctrl_msg_unref(msg);
	return NULL;
}

struct vdev_ctrl_msg *vdev_ctrl_msg_unref(struct vdev_ctrl_msg *msg)
{
	if (!msg || --msg->refcount != 0)
		return msg;
	vdev_ctrl_connection_unref(msg->conn);
	free(msg);
	return NULL;
}

enum vdev_ctrl_msg_type vdev_ctrl_msg_type(struct vdev_ctrl_msg *msg)
{
	return msg ? msg->wire.type : VDEV_CTRL_UNKNOWN;
}

int vdev_ctrl_msg_intval(struct vdev_ctrl_msg *msg)
{
	return msg ? msg->wire.intval : -1;
}

const char *vdev_ctrl_msg_buf(struct vdev_ctrl_msg *msg)
{
	return msg ? msg->wire.buf : NULL;
}

static int vdev_ctrl_send(struct vdev_ctrl *uctrl, enum vdev_ctrl_msg_type type,
						  int intval, const char *buf, int timeout_sec)
{
	struct vdev_ctrl_msg_wire wire;
	int err = 0;
	int rc;

	if (!uctrl)
		return -EINVAL;

	memset(&wire, 0, sizeof(wire));
	snprintf(wire.version, sizeof(wire.version), "%s", VDEV_CTRL_VERSION_STR);
	wire.magic = VDEV_CTRL_MAGIC;
	wire.type = type;
	if (buf)
		snprintf(wire.buf, sizeof(wire.buf), "%s", buf);
	else
		wire.intval = intval;

	rc = vdev_ctrl_connect(uctrl, timeout_sec);
	if (rc != 0)
		return rc;

	if (send(uctrl->sock, &wire, sizeof(wire), 0) < 0)
		return -errno;

	for (;;) {
		struct pollfd pfd = { .fd = uctrl->sock, .events = POLLIN };
		int r = poll(&pfd, 1, timeout_sec > 0 ? timeout_sec * 1000 : -1);

		if (r < 0) {
			if (errno == EINTR)
				continue;
			return -errno;
		}
		if (r > 0 && (pfd.revents & POLLERR))
			break;
		if (r == 0)
			return -ETIMEDOUT;
		break;
	}
	return err;
}

int vdev_ctrl_send_set_log_level(struct vdev_ctrl *uctrl, int priority,
								 int timeout_sec)
{
	return vdev_ctrl_send(uctrl, VDEV_CTRL_SET_LOG_LEVEL, priority, NULL,
						  timeout_sec);
}

int vdev_ctrl_send_stop_exec_queue(struct vdev_ctrl *uctrl, int timeout_sec)
{
	return vdev_ctrl_send(uctrl, VDEV_CTRL_STOP_EXEC_QUEUE, 0, NULL,
						  timeout_sec);
}

int vdev_ctrl_send_start_exec_queue(struct vdev_ctrl *uctrl, int timeout_sec)
{
	return vdev_ctrl_send(uctrl, VDEV_CTRL_START_EXEC_QUEUE, 0, NULL,
						  timeout_sec);
}

int vdev_ctrl_send_reload(struct vdev_ctrl *uctrl, int timeout_sec)
{
	return vdev_ctrl_send(uctrl, VDEV_CTRL_RELOAD, 0, NULL, timeout_sec);
}

int vdev_ctrl_send_set_env(struct vdev_ctrl *uctrl, const char *key,
						   int timeout_sec)
{
	return vdev_ctrl_send(uctrl, VDEV_CTRL_SET_ENV, 0, key, timeout_sec);
}

int vdev_ctrl_send_set_children_max(struct vdev_ctrl *uctrl, int count,
									int timeout_sec)
{
	return vdev_ctrl_send(uctrl, VDEV_CTRL_SET_CHILDREN_MAX, count, NULL,
						  timeout_sec);
}

int vdev_ctrl_send_ping(struct vdev_ctrl *uctrl, int timeout_sec)
{
	return vdev_ctrl_send(uctrl, VDEV_CTRL_PING, 0, NULL, timeout_sec);
}

int vdev_ctrl_send_exit(struct vdev_ctrl *uctrl, int timeout_sec)
{
	return vdev_ctrl_send(uctrl, VDEV_CTRL_EXIT, 0, NULL, timeout_sec);
}

int vdev_ctrl_send_promote(struct vdev_ctrl *uctrl, const char *config_path,
						   int timeout_sec)
{
	return vdev_ctrl_send(uctrl, VDEV_CTRL_PROMOTE, 0,
						  (config_path && config_path[0]) ? config_path : NULL,
						  timeout_sec);
}

void vdev_ctrl_queue_update_busy(int busy)
{
	struct stat st;

	/*
	 * /run/udev is often a symlink to /dev/metadata/udev.  If "queue"
	 * already exists as a directory there, do not use it as the eudev
	 * busy flag (unlink would fail and udevadm settle would hang).
	 */
	if (stat(VDEV_CTRL_QUEUE_PATH, &st) == 0 && S_ISDIR(st.st_mode))
		return;

	if (busy) {
		int fd;

		vdev_mkdirs("/run/udev", 0, 0755);
		fd = open(VDEV_CTRL_QUEUE_PATH, O_CREAT | O_WRONLY | O_CLOEXEC, 0644);
		if (fd >= 0)
			close(fd);
	} else {
		unlink(VDEV_CTRL_QUEUE_PATH);
	}
}
