/* Copyright 2005 Intel Corporation.  All Rights Reserved.
 * Please distribute under the terms of the GPLv2 or later.
 * Written by Benjamin LaHaise.
 */
#define _GNU_SOURCE	1
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <sys/mman.h>
#include <libaio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <netinet/in.h>

#define NR_EVENTS	2

/* aio helper code */
#define IOCB_CMD_SENDMSG	9
#define IOCB_CMD_RECVMSG	10

int io_sendmsg(io_context_t ctx, struct iocb *iocb, io_callback_t cb, int s, struct msghdr *msg, int flags)
{
	io_prep_pread(iocb, s, msg, sizeof(*msg), flags);
	iocb->aio_lio_opcode = IOCB_CMD_SENDMSG;
	io_set_callback(iocb, cb);
	return io_submit(ctx, 1, &iocb);
}

int io_recvmsg(io_context_t ctx, struct iocb *iocb, io_callback_t cb, int s, struct msghdr *msg, int flags)
{
	io_prep_pread(iocb, s, msg, sizeof(*msg), flags);
	iocb->aio_lio_opcode = IOCB_CMD_RECVMSG;
	io_set_callback(iocb, cb);
	return io_submit(ctx, 1, &iocb);
}

int io_pread(io_context_t ctx, struct iocb *iocb, io_callback_t cb, int fd, void *buf, size_t len, long long pos)
{
	io_prep_pread(iocb, fd, buf, len, pos);
	io_set_callback(iocb, cb);
	return io_submit(ctx, 1, &iocb);
}

int io_pwrite(io_context_t ctx, struct iocb *iocb, io_callback_t cb, int fd, void *buf, size_t len, long long pos)
{
	io_prep_pwrite(iocb, fd, buf, len, pos);
	io_set_callback(iocb, cb);
	return io_submit(ctx, 1, &iocb);
}

/* Our global vars */
long long stdin_pos;
struct iocb stdin_iocb;
char stdin_buf[512];

long long stdout_pos;

struct iovec send_iov;
struct sockaddr_in send_sin;
struct msghdr send_msg = {
	.msg_name = &send_sin,
	.msg_namelen = sizeof(send_sin),
	.msg_iov = &send_iov,
	.msg_iovlen = 1,
	.msg_control = NULL,
	.msg_controllen = 0,
	.msg_flags = 0,
};

int done;
int sockfd = -1;
struct sockaddr_in recv_sin;
struct iocb recv_iocb;
char recv_buf[65536];
struct iovec recv_iov = {
	.iov_base = recv_buf,
	.iov_len = sizeof(recv_buf)
};
struct msghdr recv_msg = {
	.msg_name = &recv_sin,
	.msg_namelen = sizeof(recv_sin),
	.msg_iov = &recv_iov,
	.msg_iovlen = 1,
	.msg_control = NULL,
	.msg_controllen = 0,
	.msg_flags = 0,
};

void recvmsg_pwrite_cb(io_context_t ctx, struct iocb *iocb, long res, long res2);
void recvmsg_cb(io_context_t ctx, struct iocb *iocb, long res, long res2)
{
	long err;
	if (res == 0) {
		fprintf(stderr, "recvmsg_cb: done\n");
		done = 1;
		return;
	}
	if (res < 0) {
		fprintf(stderr, "recvmsg_cb: result err(%ld): %s\n", res, strerror(-res));
		exit(1);
	}
fprintf(stderr, "io_pwrite(%p, %p, %ld)\n", iocb, recv_buf, res);
	err = io_pwrite(ctx, iocb, recvmsg_pwrite_cb, /*stdout*/1, recv_buf, res, stdout_pos);
	if (err != 1) {
		fprintf(stderr, "recvmsg_cb: io_pwrite err(%ld): %s\n", err, strerror(-err));
		exit(1);
	}
}

void recvmsg_pwrite_cb(io_context_t ctx, struct iocb *iocb, long res, long res2)
{
	long err;
	if (res <= 0) {
		fprintf(stderr, "recvmsg_pwrite_cb: result err(%ld): %s\n", res, strerror(-res));
		exit(1);
	}
	stdout_pos += res;
	err = io_recvmsg(ctx, &recv_iocb, recvmsg_cb, sockfd, &recv_msg, 0);
	if (err != 1) {
		fprintf(stderr, "recvmsg_pwrite_cb: io_pwrite err(%ld): %s\n", err, strerror(-err));
		exit(1);
	}
}

void stdin_sendmsg_cb(io_context_t ctx, struct iocb *iocb, long res, long res2);
void stdin_read_cb(io_context_t ctx, struct iocb *iocb, long res, long res2)
{
	long err;
	if (res < 0) {
		fprintf(stderr, "stdin_read_cb: result err(%ld): %s\n", res, strerror(-res));
		exit(1);
	}
	send_iov.iov_base = stdin_buf;
	send_iov.iov_len = res;
	err = io_sendmsg(ctx, iocb, stdin_sendmsg_cb, sockfd, &send_msg, 0);
	if (err != 1) {
		fprintf(stderr, "recvmsg_pwrite_cb: io_pwrite err(%ld): %s\n", err, strerror(-err));
		exit(1);
	}
	stdin_pos += res;
}

void stdin_sendmsg_cb(io_context_t ctx, struct iocb *iocb, long res, long res2)
{
	long err;
	if (res < 0) {
		fprintf(stderr, "stdin_sendmsg_cb: result err(%ld): %s\n",
			res, strerror(-res));
		exit(1);
	}
	err = io_pread(ctx, iocb, stdin_read_cb, /*stdin*/0, stdin_buf,
			sizeof(stdin_buf), stdin_pos);
	if (err != 1) {
		fprintf(stderr, "stdin_sendmsg_cb: io_pread err(%ld): %s\n",
			err, strerror(-err));
		exit(1);
	}
}

int main(int argc, char *argv[])
{
	struct io_event events[NR_EVENTS];
	io_context_t ctx;
	long err;
	long loops = 0;
	int tmp;
	struct sockaddr_in sin;

	fprintf(stderr, "recv_buf: %p, stdin_buf: %p\n", recv_buf, stdin_buf);

	sockfd = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (sockfd < 0) {
		perror("socket");
		return 1;
	}

	tmp = 1;
	if (setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &tmp, sizeof(tmp)) < 0) {
		perror("setsockopt");
		return 1;
	}

	memset(&sin, 0, sizeof(sin));
	sin.sin_family = AF_INET;

	if (bind(sockfd, (struct sockaddr *)&sin, sizeof(sin))) {
		perror("bind");
		return 1;
	}

	tmp = sizeof(sin);
	if (getsockname(sockfd, (struct sockaddr *)&sin, &tmp)) {
		perror("getsockname");
		return 1;
	}

	/* we send messages to ourself */
	send_sin = sin;

	if (tmp != sizeof(sin))
		fprintf(stderr, "warning: getsockname returned an address of length %d when %d was expected\n", tmp, (int)sizeof(sin));

	ctx = 0;

	err = io_setup(NR_EVENTS * 2, &ctx);
	if (err) {
		fprintf(stderr, "main: io_setup: %ld: %s\n", err, strerror(-err));
		exit(1);
	}

	err = io_recvmsg(ctx, &recv_iocb, recvmsg_cb, sockfd, &recv_msg, 0);
	if (err != 1) {
		fprintf(stderr, "main: io_recvmsg failed %ld (%s)\n", err, strerror(-err));
		return 1;
	}

	err = io_pread(ctx, &stdin_iocb, stdin_read_cb, 0, stdin_buf, sizeof stdin_buf, stdin_pos);
	if (err != 1) {
		fprintf(stderr, "main: io_pread failed %ld (%s)\n", err, strerror(-err));
		return 1;
	}

	fprintf(stderr, "ios submitted, waiting for events\n");
	while (!done) {
		long i;
		err = io_getevents(ctx, 1, NR_EVENTS, events, NULL);
		if (err <= 0) {
			fprintf(stderr, "io_getevents: %ld vs %d: %s\n", err,
				NR_EVENTS, strerror(-err));
			exit(1);
		}

		for (i=0; i<err; i++) {
			io_callback_t cb = events[i].data;

			cb(ctx, events[i].obj, events[i].res, events[i].res2);
		}
		loops++;
	}
	fprintf(stderr, "completed %ld loops.  read a total of %Ld bytes, wrote %Ld bytes\n",
		loops, stdin_pos, stdout_pos);
	return 0;
}