#include <sys/socket.h>
#include <sys/poll.h>
#include <sys/time.h>
#include <sys/un.h>
#include <sys/uio.h>
#include <string.h>
#include <stdlib.h>
#include <signal.h>
#include <stdio.h>
#include <stdarg.h>
#include <time.h>
#include <unistd.h>
#include <errno.h>
#include <netdb.h>
#include <fcntl.h>
#include <stddef.h>

#define SOCKET "test-socket"

static int do_recv(int sock, char *buffer, size_t count, struct cmsgcred **cred, int dontwait)
{
	struct cmsgcred *cred_buf;
	unsigned char control[CMSG_SPACE(sizeof(*cred_buf)) + 1024];
	struct cmsghdr  *cmsg;
	struct msghdr   msg;
	struct iovec    iov;
	int             len;
	char name_buf[256];

	if (dontwait)
		dontwait = MSG_DONTWAIT;

	iov.iov_base = buffer;
	iov.iov_len = count;

	memset(&msg, 0, sizeof(msg));
#ifdef MSG_NAME
	msg.msg_name = name_buf;
	msg.msg_namelen = sizeof(name_buf);
#endif
	msg.msg_iov = &iov;
	msg.msg_iovlen = 1;
	msg.msg_control = control;
	msg.msg_controllen = sizeof(control);
	msg.msg_flags = dontwait;

	len = recvmsg(sock, &msg, dontwait);

	if (len < 0)
		return len;

	cmsg = CMSG_FIRSTHDR(&msg);
	while (cmsg) {
		if (cmsg->cmsg_level == SOL_SOCKET && cmsg->cmsg_type == SCM_CREDS) {
			cred_buf = malloc(sizeof(*cred_buf));
			memcpy(cred_buf, CMSG_DATA(cmsg), sizeof(*cred_buf));
			*cred = cred_buf;
			break;
		}
		cmsg = CMSG_NXTHDR(&msg, cmsg);
	}

	return len;
}

static int do_send(int sock, const char *buffer, size_t count)
{
	unsigned char   control[CMSG_SPACE(sizeof(struct cmsgcred))];
	struct cmsghdr  *cmsg;
	struct msghdr   msg;
	struct iovec    iov;

	iov.iov_base = (void *)buffer;
	iov.iov_len = count;

	memset(&msg, 0, sizeof(msg));
	msg.msg_iov = &iov;
	msg.msg_iovlen = 1;
	
	memset(&control, 0, sizeof(control));
	msg.msg_control = control;
	msg.msg_controllen = sizeof(control);

	cmsg = CMSG_FIRSTHDR(&msg);
	cmsg->cmsg_len = CMSG_LEN(sizeof(struct cmsgcred));
	cmsg->cmsg_level = SOL_SOCKET;
	cmsg->cmsg_type = SCM_CREDS;

	return sendmsg(sock, &msg, 0);
}

#define DO(_a) \
	do { \
		int r = (_a); \
		if (r < 0) { \
			perror(#_a); \
			exit(1); \
		} \
	} while(0)

int main(int argc, char **argv)
{
	int sock;
	int server = 0;
	struct sockaddr_storage addr_buf;
	struct sockaddr_un *addr_un = (struct sockaddr_un *)&addr_buf;
	struct sockaddr *addr = (struct sockaddr *)&addr_buf;
	socklen_t addr_len;
	int v;
	struct pollfd f;
	
	if (argc > 1 && !strcmp(argv[1], "server"))
		server = 1;
	
	addr_un->sun_family = AF_UNIX;
	strcpy(addr_un->sun_path, SOCKET); /* kids, don't do this in real code, this is a buffer overflow in waiting */
	addr_len = offsetof(struct sockaddr_un, sun_path) + strlen(SOCKET);
	
	DO(sock = socket(AF_LOCAL, SOCK_STREAM, 0));
	
	if (server) {
		int csock;
		struct cmsgcred *creds, *creds2;
		int v2, r;

		unlink(SOCKET);
		DO(bind(sock, addr, addr_len));
		DO(listen(sock, 1));
		DO(csock = accept(sock, NULL, NULL));
		close(sock);
		DO(do_recv(csock, (char *)&v, sizeof(int), &creds, 0));
#ifdef SECOND_RECV
		r = do_recv(csock, (char *)&v2, sizeof(int), &creds2, 1); /* set MSG_DONTWAIT on second message */
		if (r == 0 || errno != EAGAIN) {
			printf("Warning: second receive didn't give us -EAGAIN, got rv=%d, errno=%d (%m).\n", r, errno);
		}
#endif
		if (!creds) {
			printf("no creds, v = %d\n", v);
			return 0;
		}
		printf("creds->uid = %d, v = %d\n", (int)creds->cmcred_uid, v);
		return 0;
	}
	
	DO(connect(sock, addr, addr_len));
	
	v = 42;
	DO(do_send(sock, (const char *)&v, sizeof(int)));
	f.fd = sock;
	f.events = POLLIN;
	f.revents = 0;
	/* wait for server to exit and close socket */
	poll(&f, 1, -1);
	printf("Send successful.\n");
	return 0;
}