/* rfbproxy - a record/playback VNC proxy
 * Copyright (C) 2000-3  Tim Waugh <twaugh@redhat.com>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 */

/*
 * This program works by recording everything the client/server says, and
 * then playing it back.  Yes, it _is_ that easy. ;-)
 *
 * I've found that hextile encoding gives the best results.
 *
 * See also: the TODO file.  */

/*
 * The FBS (framebuffer stream) file format is this:
 *
 * <capture file> ::- <version><data>
 * <version>      ::- 'FBS 001.000\n'
 * <data>         ::- <byte-count><raw-data><timestamp><data>
 *                  | <byte-count><raw-data><timestamp>
 * <byte-count>   ::- 32-bit number of bytes, big-endian
 * <raw-data>     ::- data received at timestamp <timestamp>, which is
 *                    <byte-count> bytes in length, padded to multiple
 *                    of 32-bits
 * <timestamp>    ::- 32-bit number of milliseconds since beginning of
 *                    capture, big-endian
 *
 * The RFM (remote framebuffer macro) file format is documented here:
 * <URL:ftp://people.redhat.com/twaugh/rfbplaymacro/script-spec>
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <stdio.h>
#include <ctype.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <getopt.h>
#include <netdb.h>
#include <time.h>

#if HAVE_STDINT_H
# include <stdint.h>
#else
# if HAVE_U_INTXX_T
#  ifndef __FreeBSD__
typedef u_int16_t uint16_t;
typedef u_int32_t uint32_t;
#  endif
# else
typedef unsigned short uint16_t;
typedef unsigned long uint32_t;
# endif
#endif

#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#define VNC_BASE 5900
#define DEFAULT_DISPLAY ":10"
#define DEFAULT_SERVER ":1"
#define BUFSIZE 65535

#ifndef INADDR_LOOPBACK
# define INADDR_LOOPBACK ((in_addr_t) 0x7f000001)
#endif

enum protocol_enum {
	Protocol3_3,
	Protocol3_7
};

static ssize_t read_traffic (int fd, void *buf, size_t len)
{
	static char *buffer;
	static int buffer_length;
	static int buffer_at;
	static int buffer_size = BUFSIZE;

	if (!buffer)
		buffer = malloc (BUFSIZE);

	if (!buffer)
		return 0;

	if (fd < 0) {
		/* Reset the buffer */
		buffer_length = buffer_at = 0;
		return 0;
	}

	for (;;) {
		uint32_t dlen;

		if (len < buffer_length - buffer_at) {
			memcpy (buf, buffer + buffer_at, len);
			buffer_at += len;
			return len;
		}

		if (buffer_length - buffer_at) {
			len = buffer_length - buffer_at;
			memcpy (buf, buffer + buffer_at, len);
			buffer_length = buffer_at = 0;
			return len;
		}

		if (read (fd, &dlen, sizeof (dlen)) <
		    sizeof (dlen))
			return 0;

		buffer_length = ntohl (dlen);
		buffer_at = 0;

		if (buffer_size < buffer_length) {
			while (buffer_size < buffer_length)
				buffer_size <<= 1;
			free (buffer);
			buffer = malloc (buffer_size);
			if (!buffer) {
				fprintf (stderr, "Memory squeeze\n");
				return 0;
			}
		}

		if (read (fd, buffer,
			  4 * ((buffer_length + 3) / 4)) < buffer_length)
			return 0;

		/* Skip timestamp */
		lseek (fd, 4, SEEK_CUR);
	}
}

static void rewind_file (int fd)
{
	/* We have to rewind, and then skip authentication */
	char buf4[4];
	int state = 0;
	enum protocol_enum protocol = Protocol3_3;

	read_traffic (-1, NULL, 0);
	lseek (fd, 12, SEEK_SET); /* Skip version */

	do {
		switch (state) {
		case 0: /* ProtocolVersion */
		{
			char skip[12];
			if (read_traffic (fd, skip, 12) < 12)
				goto too_short;
			if (skip[10] == '7')
				protocol = Protocol3_7;
			state++;
			break;
		}
		case 1: /* Authentication */
		{
			char skip[256];
			uint32_t auth;
			if (protocol == Protocol3_7) {
				if (read_traffic (fd, buf4, 1) < 1)
					goto too_short;
				if (read_traffic (fd, skip, buf4[0]) < buf4[0])
					goto too_short;
				auth = (uint32_t) skip[0];
			} else {
				if (read_traffic (fd, buf4, 4) < 4)
					goto too_short;
				memcpy (&auth, buf4, 4);
				auth = ntohl (auth);
			}
			if (auth == 2)
				if (read_traffic (fd, skip, 20) < 20)
					goto too_short;
			state++;
			if (read_traffic (fd, skip, 20) < 20)
				goto too_short;
			break;
		}
		case 2: /* ServerInitialisation */
		{
			char *skip;
			size_t len;
			if (read_traffic (fd, buf4, 4) < 4)
				goto too_short;
			len = ntohl (*(size_t*)buf4);
			skip = malloc (len);
			if (read_traffic (fd, skip, len) < len)
				goto too_short;
			free (skip);
			state++;
			break;
		}
		}
	} while (state < 3);
	return;

 too_short:
	fprintf (stderr, "Input has finished too early\n");
	exit (1);
}

static ssize_t do_write (int fd, const void *buf, size_t len)
{
	while (len) {
		ssize_t wrote = write (fd, buf, len);
		if (wrote < 0) {
			perror ("write");
			exit (1);
		}
		buf += wrote;
		len -= wrote;
	}
	return len;
}

static ssize_t do_read (int fd, void *buf, size_t len)
{
	while (len) {
		ssize_t got = read (fd, buf, len);
		if (got < 0) {
			perror ("read");
			exit (1);
		}
		if (!got)
			break;
		buf += got;
		len -= got;
	}
	return len;
}

static int write_packet (FILE *f, const void *buf, size_t len,
			 struct timeval *tvp)
{
	uint32_t timestamp = htonl (1000 * tvp->tv_sec + tvp->tv_usec / 1000);
	uint32_t dlen = htonl (len);
	len = 4 * ((len + 3) / 4);
	fwrite (&dlen, 4, 1, f);
	fwrite (buf, 1, len, f);
	fwrite (&timestamp, 4, 1, f);
	return 0;
}

static int do_authentication (int server, int clientr, int clientw, FILE *f,
			      int do_events_instead)
{
	char packet[24];
	size_t packet_size;
	struct timeval start;
	uint32_t auth;
	enum protocol_enum protocol = Protocol3_3;


	start.tv_sec = 0;
	start.tv_usec = 0;

	/* ProtocolVersion */
	if (do_read (server, packet, 12))
		return 1;
	do_write (clientw, packet, 12);
	if (do_read (clientr, packet, 12))
		return 1;
	do_write (server, packet, 12);
	if (!do_events_instead)
		/* Record the protocol in use */
		write_packet (f, packet, 12, &start);
	packet_size = 4;
	if (packet[10] == '7') {
		protocol = Protocol3_7;
		packet_size = 1;
	}

	/* Authentication */
	if (do_read (server, packet, packet_size))
		return 1;

	if (protocol == Protocol3_3) {
		uint32_t noauth = htonl (1);
		do_write (clientw, packet, 4);
		memcpy (&auth, packet, 4);
		auth = ntohl (auth);
		if (!do_events_instead) {
			memcpy (packet, &noauth, 4);
			write_packet (f, packet, 4, &start);
		}
	} else {
		size_t num_types;
		num_types = (size_t) packet[0];
		if (num_types == 0)
			return 1;
		if (do_read (server, packet + 1, num_types))
			return 1;
		do_write (clientw, packet, num_types + 1);
		if (do_read (clientr, packet, 1))
			return 1;
		do_write (server, packet, 1);
		auth = (uint32_t) packet[0];
		if (!do_events_instead) {
			packet[0] = 1;
			packet[1] = 1;
			write_packet (f, packet, 2, &start);
		}
	}

	if (auth == 2) {
		/* Don't record this stuff. */
		if (do_read (server, packet, 16))
			return 1;
		do_write (clientw, packet, 16);
		if (do_read (clientr, packet, 16))
			return 1;
		do_write (server, packet, 16);
		if (do_read (server, packet, 4))
			return 1;
		do_write (clientw, packet, 4);
	}

	/* ClientInitialisation */
	if (do_read (clientr, packet, 1))
		return 1;
	do_write (server, packet, 1);

	/* ServerInitialisation */
	if (do_read (server, packet, 24))
		return 1;
	else {
		uint32_t name_length;
		char *buffer;
		memcpy (&name_length, packet + 20, 4);
		name_length = ntohl (name_length);
		buffer = malloc (name_length);
		if (!buffer)
			return 1;
		if (do_read (server, buffer, name_length)) {
			free (buffer);
			return 1;
		}
		do_write (clientw, packet, 24);
		do_write (clientw, buffer, name_length);
		if (!do_events_instead) {
			write_packet (f, packet, 24, &start);
			write_packet (f, buffer, name_length, &start);
		}
		free (buffer);
	}

	return 0;
}

static size_t variable_part (char *buffer)
{
	int message = (int) *buffer;
	switch (message) {
	case 0: /* SetPixelFormat */
	case 3: /* FramebufferUpdateRequest */
	case 4: /* KeyEvent */
	case 5: /* PointerEvent */
		/* No variable part */
		return 0;
	case 1: /* FixColourMapEntries */
	{
		uint16_t number_of_colours;
		memcpy (&number_of_colours, buffer + 4, 2);
		number_of_colours = ntohs (number_of_colours);
		return number_of_colours * 6;
	}
	case 2: /* SetEncodings */
	{
		uint16_t number_of_encodings;
		memcpy (&number_of_encodings, buffer + 2, 2);
		number_of_encodings = ntohs (number_of_encodings);
		return number_of_encodings * 4;
	}
	case 6: /* ClientCutText */
	{
		uint32_t length;
		memcpy (&length, buffer + 4, 4);
		length = ntohl (length);
		return length;
	}
	} /* switch */

	/* Caught earlier anwyay */
	fprintf (stderr, "Protocol error\n");
	exit (1);
}

/* For recording */
static int process_client_message (char *fixed, char *variable, FILE *f)
{
	static int first = 1;
	static char delayed_output[100];
	static int elapsed;
	static struct timeval last_tv, first_tv;
	struct timeval tv, diff;
	struct timezone tz;
	static unsigned int last_was_key_down;
	static unsigned int current_x, current_y;
	static unsigned char current_buttons;
	int ms;
	int message = (int) *fixed;

	gettimeofday (&tv, &tz);
	if (!last_tv.tv_sec && !last_tv.tv_usec)
		first_tv = last_tv = tv;
	diff.tv_sec = tv.tv_sec - last_tv.tv_sec;
	diff.tv_usec = tv.tv_usec - last_tv.tv_usec;
	ms = diff.tv_sec * 1000 + diff.tv_usec / 1000;

	if (first) {
		first = 0;
		fputs ("RFM 001.000\nshared\n", f);
	} else if (*delayed_output && (!last_was_key_down || message > 4)) {
		/* We need to output a deferred line after calculating
		 * the delay */
		if (ms > 0) {
			char *p = delayed_output + strlen (delayed_output);
			sprintf (p, " delay %dms", ms);
		}
		strcat (delayed_output, "\n");
		fputs (delayed_output, f);
		last_tv = tv;
		*delayed_output = '\0';
	}

	switch (message) {
	case 0: /* SetPixelFormat */
	case 1: /* FixColourMapEntries */
	case 2: /* SetEncodings */
	case 3: /* FramebufferUpdateRequest */
		diff.tv_sec = tv.tv_sec - first_tv.tv_sec;
		diff.tv_usec = tv.tv_usec - first_tv.tv_usec;
		ms = diff.tv_sec * 1000 + diff.tv_usec / 1000;
		if (ms > (1000 * (1 + elapsed))) {
			fprintf (f, "# At %dms from start\n", ms);
			elapsed = ms / 1000;
		}
		return 0;
	case 4: /* KeyEvent */
	{
		char *p = delayed_output;
		const char *down_flag = "up";
		uint32_t key;

		memcpy (&key, fixed + 4, 4);
		key = ntohl (key);

		/* We might be changing key up/down into press */
		if (*delayed_output) {
			/* last_was_key_down is the last key down */
			if (fixed[1] || last_was_key_down != key || ms > 400) {
				/* Can't make a press out of that */
				char *p = delayed_output;
				p += strlen (p);
				if (ms > 0)
					sprintf (p, " delay %dms", ms);
				strcat (delayed_output, "\n");
				fputs (delayed_output, f);
				last_tv = tv;
				*delayed_output = '\0';
				last_was_key_down = 0;
			} else {
				char *p = delayed_output;
				char *end;
				p += strcspn (p, " \t");
				p += strspn (p, " \t");
				end = p + strcspn (p, " \t");
				*end = '\0';
				end = strdup (p);
				sprintf (delayed_output, "press %s", end);
				last_was_key_down = 0;
				break;
			}
		}

		if (fixed[1]) {
			last_was_key_down = key;
			down_flag = "down";
		}
		sprintf (p, "key ");
		p += strlen (p);
		if (key < 256 && isprint ((char) key) && !isspace ((char) key))
			*p++ = (char) key;
		else {
			sprintf (p, "%#x", key);
			p += strlen (p);
		}

		sprintf (p, " %s", down_flag);
		break;
	}
	case 5: /* PointerEvent */
	{
		uint16_t x, y;
		unsigned char buttons = fixed[1];
		memcpy (&x, fixed + 2, 2);
		memcpy (&y, fixed + 4, 2);
		x = ntohs (x);
		y = ntohs (y);

		/* First deal with buttons */
		if (buttons != current_buttons) {
			int i;
			int diff = buttons ^ current_buttons;
			while ((i = ffs (diff))) {
				if (*delayed_output) {
					strcat (delayed_output, "\n");
					fputs (delayed_output, f);
				}
				i--;
				sprintf (delayed_output,
					 "button %d %s", i,
					 (buttons & (1<<i)) ? "down" : "up");
				diff ^= 1<<i;
			}
			current_buttons = buttons;
		}

		/* Now deal with position */
		if (current_x != x || current_y != y) {
			if (*delayed_output) {
				strcat (delayed_output, "\n");
				fputs (delayed_output, f);
			}
			sprintf (delayed_output,
				 "pointer %d %d", x, y);
			current_x = x;
			current_y = y;
		}
		break;
	}
	case 6: /* ClientCutText */
		fputs ("# ClientCutText not supported yet\n", f);
		break;
	default:
		fprintf (stderr, "Protocol error\n");
		exit (1);
	}
	return 0;
}

static int record (const char *file, int clientr, int clientw,
		   struct sockaddr_in server_addr, int do_events_instead, int appenddate)
{
	const char *version = "FBS 001.000\n";
	FILE *f;
	int server = -1;
	struct timeval epoch;
	int first = 1;
	char *buf = malloc (BUFSIZE);
	time_t now;	

	if (!buf) {
		fprintf (stderr, "Couldn't allocate buffer\n");
		exit (1);
	}

	if (appenddate)
	{ /* if we're appending the date, make the new filename in 'buf'.  if we're not, just call
	     fopen directly on 'file' */
		if (strlen(file)+17 > BUFSIZE)
		{	/* Ya, like this is going to happen.  whatever */
			fprintf (stderr, "Filename is bigger than filename buffer size.  Increase BUFSIZE and recompile\n");
			exit (1);
		}
		time(&now);
		strftime(buf+sprintf(buf, "%s-", file), 16, "%Y%m%d-%H%M%S", localtime(&now));
		f = fopen (buf, "wb");
	} else {
		f = fopen (file, "wb");
	}
	if (!f) {
		perror ("fopen");
		exit (1);
	}

	if (!do_events_instead)
		fwrite (version, 1, 12, f);

	server = socket (PF_INET, SOCK_STREAM, IPPROTO_IP);
	if (server == -1) {
		perror ("socket");
		goto out;
	}

	if (connect (server, (struct sockaddr *) &server_addr,
		     sizeof (struct sockaddr_in))) {
		perror ("connect");
		goto out;
	}

	if (do_authentication (server, clientr, clientw,
			       f, do_events_instead)) {
		fprintf (stderr, "Error during authentication\n");
		exit (1);
	}

	for (;;) {
		ssize_t bufs;
		struct timeval tv;
		struct timezone tz;
		fd_set rfds;
		FD_ZERO (&rfds);
		FD_SET (server, &rfds);
		FD_SET (clientr, &rfds);

		if (select(FD_SETSIZE, &rfds, NULL, NULL, NULL) < 0) {
			perror ("select");
			goto out;
		}

		if (FD_ISSET (server, &rfds)) {
			struct timeval diff;

			gettimeofday (&tv, &tz);
			if (first) {
				first = 0;
				epoch = tv;
			}
			bufs = read (server, buf, BUFSIZE);
			if (!bufs) break;
			do_write (clientw, buf, bufs);
			diff.tv_sec = tv.tv_sec - epoch.tv_sec;
			diff.tv_usec = tv.tv_usec - epoch.tv_usec;
			if (!do_events_instead)
				write_packet (f, buf, bufs, &diff);
		}
		else {
			/* We want to actually listen to the
			 * individual client messages.
			 *
			 * The largest non-variable part of a
			 * client->server message is 20 bytes.  */
			static char tmp_buffer[20];
			static char client_buffer[20];
			static char *variable_buffer;
			static size_t variable_bytes_left;
			static size_t variable_bytes_got;
			static size_t client_bytes_got;
			static size_t client_bytes_left;
			static const size_t mlen[] = {
				20, 6, 4, 10, 8, 6, 8
			}; /* message lengths */
			char *at = tmp_buffer;

			/* Read the available data */
			bufs = read (clientr, tmp_buffer, 20);
			if (!bufs) break;
			do_write (server, tmp_buffer, bufs);

			if (!do_events_instead)
				continue;

			while (bufs) {
				size_t length;

				/* Figure out where to put it  */
				if (variable_bytes_left) {
					size_t need = bufs;
					if (variable_bytes_left < need)
						need = variable_bytes_left;
					memcpy (variable_buffer +
						variable_bytes_got,
						at, need);
					variable_bytes_got += need;
					variable_bytes_left -= need;
					at += need;
					bufs -= need;
				} else if (client_bytes_left) {
					size_t need = bufs;
					if (client_bytes_left < need)
						need = client_bytes_left;
					memcpy (client_buffer +
						client_bytes_got,
						at, need);
					client_bytes_got += need;
					client_bytes_left -= need;
					at += need;
					bufs -= need;
				} else {
					/* Clean slate */
					*client_buffer = *at++;
					bufs--;
					client_bytes_got = 1;
				}

				/* Figure out what to do with it */
				if (client_buffer[0] > 6) {
					fprintf (stderr,
						 "Protocol error\n");
					exit (1);
				}
				length = mlen[(int) client_buffer[0]];
				if (client_bytes_got < length) {
					client_bytes_left = (length -
							     client_bytes_got);
					/* Incomplete fixed part */
					continue;
				}

				length = variable_part (client_buffer);
				if (variable_bytes_got < length) {
					int need_alloc = !variable_bytes_left;
					variable_bytes_left = length -
						variable_bytes_got;
					if (need_alloc)
						variable_buffer = malloc
							(variable_bytes_left);
					/* Incomplete variable part */
					continue;
				}

				process_client_message (client_buffer,
							variable_buffer, f);
				if (variable_bytes_got) {
					variable_bytes_got = 0;
					free (variable_buffer);
				}
				client_bytes_got = 0;
			}
		}
	}
 out:
	free (buf);
	if (server != -1)
		close (server);

	if (fclose (f))
		perror ("Error writing file");

	return 0;
}

/* Returns bitmask:
 *
 * bit 0: cycle
 * bit 1: pause
 */
static int handle_keys_during_playback (int clientr, int cycle, int pause)
{
	/* We want to actually listen to the
	 * individual client messages.
	 *
	 * The largest non-variable part of a
	 * client->server message is 20 bytes.  */
	static char tmp_buffer[20];
	static char client_buffer[20];
	static char *variable_buffer;
	static size_t variable_bytes_left;
	static size_t variable_bytes_got;
	static size_t client_bytes_got;
	static size_t client_bytes_left;
	static const size_t mlen[] = {
		20, 6, 4, 10, 8, 6, 8
	}; /* message lengths */
	char *at = tmp_buffer;
	ssize_t bufs;
	int do_pause = 0;
	int do_cycle = 0;

	/* Read the available data */
	bufs = read (clientr, tmp_buffer, 20);
	if (bufs < 1)
		return -1;

	while (bufs) {
		size_t length;

		/* Figure out where to put it  */
		if (variable_bytes_left) {
			size_t need = bufs;
			if (variable_bytes_left < need)
				need = variable_bytes_left;
			memcpy (variable_buffer +
				variable_bytes_got,
				at, need);
			variable_bytes_got += need;
			variable_bytes_left -= need;
			at += need;
			bufs -= need;
		} else if (client_bytes_left) {
			size_t need = bufs;
			if (client_bytes_left < need)
				need = client_bytes_left;
			memcpy (client_buffer +
				client_bytes_got,
				at, need);
			client_bytes_got += need;
			client_bytes_left -= need;
			at += need;
			bufs -= need;
		} else {
			/* Clean slate */
			*client_buffer = *at++;
			bufs--;
			client_bytes_got = 1;
		}

		/* Figure out what to do with it */
		if (client_buffer[0] > 6) {
			fprintf (stderr,
				 "Protocol error (%d)\n", client_buffer[0]);
			exit (1);
		}
		length = mlen[(int) client_buffer[0]];
		if (client_bytes_got < length) {
			client_bytes_left = (length -
					     client_bytes_got);
			/* Incomplete fixed part */
			continue;
		}

		length = variable_part (client_buffer);
		if (variable_bytes_got < length) {
			int need_alloc = !variable_bytes_left;
			variable_bytes_left = length -
				variable_bytes_got;
			if (need_alloc)
				variable_buffer = malloc
					(variable_bytes_left);
			/* Incomplete variable part */
			continue;
		}
		
		if (client_buffer[0] == 4 /* KeyEvent */ &&
		    client_buffer[1] /* Key down */) {
			uint32_t key;
			memcpy (&key, client_buffer + 4, 4);
			key = ntohl (key);
			if (key == pause)
				do_pause = 2 - do_pause;
			if (key == cycle)
				do_cycle = 1;
		}

		if (variable_bytes_got) {
			variable_bytes_got = 0;
			free (variable_buffer);
		}
		client_bytes_got = 0;
	}

	return do_pause | do_cycle;
}

static int playback (const char *file, int clientr, int clientw, int loop,
		     int cycle, int pause)
{
	int fd = open (file, O_RDONLY);
	struct stat st;
	char *map, *buf;
	char *skip;
	uint32_t last_packet = 0;
	int ret = -1;
	int paused = 0;
	int authentication = loop ? 0 : 1;
	int finish = 0;
	enum protocol_enum protocol = Protocol3_3;

	if (fd < 0) {
		fprintf (stderr, "Couldn't open input file\n");
		return -1;
	}

	fstat (fd, &st);
	map = buf = mmap (NULL, st.st_size, PROT_READ | PROT_WRITE,
			  MAP_PRIVATE, fd, 0);
	close (fd);
	if (map == MAP_FAILED) {
		perror ("Couldn't map input file");
		return -1;
	}

	skip = malloc (BUFSIZE);
	if (!skip)
		goto out;

	if (strncmp (buf, "FBS 001.", 8)) {
		fprintf (stderr, "Incorrect version in input file\n");
		return -1;
	}

	buf += 12; /* Skip version */

	if (loop) {
		off_t loop_offset = 0;
		fd = open (file, O_RDONLY);
		rewind_file (fd);
		loop_offset = lseek (fd, 0, SEEK_CUR);
		close (fd);
		buf = map + loop_offset;
	}

	for (;;) {
		uint32_t *bit32;
		size_t len;
		struct timeval tv, deadline;
		struct timezone tz;
		unsigned long ms;
		fd_set rfds;
		FD_ZERO (&rfds);
		FD_SET (clientr, &rfds);

		if (map + st.st_size - buf < (2 * sizeof (uint32_t)))
			break;

		bit32 = (uint32_t *) buf;
		len = ntohl (*bit32);

		if (map + st.st_size - buf < (2 * sizeof (uint32_t)) + len)
			break;

		bit32 = (uint32_t *) (4 + buf + 4 * ((len + 3)/ 4));
		ms = ntohl (*bit32) - last_packet;
		last_packet = ntohl (*bit32);
		tv.tv_sec = ms / 1000;
		tv.tv_usec = 1000 * (ms % 1000);
		gettimeofday (&deadline, &tz);
		deadline.tv_sec += tv.tv_sec;
		deadline.tv_usec += tv.tv_usec;
		if (deadline.tv_usec >= 1000000) {
			deadline.tv_usec -= 1000000;
			deadline.tv_sec++;
		}

		/* Heuristic: if the delay is >= 0.1s, we are at a message
		 * boundary.  THIS WILL BREAK FOR SLOW CONNECTIONS! */
		if (finish && (tv.tv_sec || tv.tv_usec > 100000)) {
			ret = 1;
			goto out;
		}

		while ((paused || tv.tv_usec > 5000) &&
		       select(FD_SETSIZE, &rfds, NULL, NULL, &tv) != 0) {
			if (!authentication && (cycle || pause)) {
				int stuff;
				stuff = handle_keys_during_playback
					(clientr, cycle, pause);
				if (stuff == -1)
					/* Connection closed. */
					goto out;
				if (stuff & 2)
					paused = 1 - paused;
				if (stuff & 1)
					finish = 1;
			} else if (!read (clientr, skip, BUFSIZE))
				/* Connection closed. */
				goto out;

			/* Recalculate delay */
			gettimeofday (&tv, &tz);
			if (tv.tv_sec > deadline.tv_sec ||
			    (tv.tv_sec == deadline.tv_sec &&
			     tv.tv_usec >= deadline.tv_usec))
				/* Deadline already passed */
				break;

			tv.tv_sec = deadline.tv_sec - tv.tv_sec;
			tv.tv_usec = deadline.tv_usec - tv.tv_usec;
			if (tv.tv_usec < 0) {
				tv.tv_usec += 1000000;
				tv.tv_sec--;
			}
		}

		if (paused)
			continue;

		do_write (clientw, buf + sizeof (uint32_t), len);
		buf += 2 * sizeof (uint32_t) + 4 * ((len + 3) / 4);

		if (!cycle && !pause)
			continue;

		/* Assume one packet per message during authentication */
		switch (authentication) {
			static uint32_t auth;
		case 0:
			/* We already authenticated */
			break;

		case 1:
			if (do_read (clientr, skip, 12) != 0)
				goto out;

			if (skip[10] == '7')
				protocol = Protocol3_7;

			/* We just sent ProtocolVersion */
			if (protocol == Protocol3_7) {
				/* Hope that whoever recorded the file
				 * made this easy for us. */
				auth = buf[5];
				auth = 1;
			} else {
				memcpy (&auth, buf + 4, 4);
				auth = ntohl (auth);
				authentication++;
			}
			authentication++;
			break;
		case 2:
			if (do_read (clientr, skip, 1))
				goto out;
			authentication++;
			break;
		case 3:
			/* We just sent Authentication */
			authentication++;
			if (auth != 2)
				authentication += 2;
			break;
		case 4:
			/* We just sent a challenge */
			authentication++;
			break;
		case 5:
			/* First bit of ServerInitialisation */
			authentication++;
			break;
		case 6:
			/* Server name */
			authentication++;
			break;
		case 7:
			/* We just finished authentication */
			authentication = 0;
			/* Read the response "RFB xxx.yyy\n" and the
			 * ClientInitialization byte from our
			 * client. */
			if ((cycle || pause) &&
			    do_read (clientr, skip, 1))
				goto out;
			break;
		}
	}

	ret = 1;

 out:
	free (skip);
	munmap (map, st.st_size);
	return ret;
}

static void version (void)
{
	fprintf (stderr, "rfbproxy version %s\n", VERSION);
}

static void usage (const char *name)
{
	fprintf (stderr,
		 "usage: %s [OPTIONS] ACTION files\n"
		 "where OPTIONS are:\n\n"
		 " -c, --stdout  Use stdin and stdout for communications\n"
		 "               with the client. Useful in conjunction\n"
		 "               with inetd.\n"
		 " -l, --loop    (playback only) When file is finished,\n"
		 "               replay from first FrameBufferUpdate\n"
		 " -d, --date    (record only) Append Date to end of filename\n"
		 " --pause=key   (playback only) When the key is pressed,\n"
		 "               playback will be paused until it is pressed\n"
		 "               again.\n"
		 " --cycle=key   (playback only) When multiple files are specified\n"
		 "               pressing the key will cycle between them.\n"
		 " --type=[screen|events]\n"
		 "               (record only) Capture either screen updates\n"
		 "               (\"screen\") or keyboard/mouse events (\"events\").\n"
		 "               The default is \"screen\".\n"
		 " :n            Occupy VNC display localhost:n (not valid with\n"
		 "               -c option). The default is \""
		 DEFAULT_DISPLAY "\".\n"
		 " --server=[server]:display\n"
		 "               (record only) Use specified VNC server. The\n"
		 "               default is \"" DEFAULT_SERVER "\".\n"
		 "\nACTION is one of:\n"
		 " -r, --record\n"
		 "               Record RFB communications and store them in\n"
		 "               the file.\n"
		 " -p, --playback\n"
		 "               Play back the RFB communications that were\n"
		 "               captured to the file or files.\n",
		 name);
	exit (1);
}

int accept_connection (int port)
{
	int bound;
	int sock;
	struct sockaddr_in sin;
	int on = 1;
	sin.sin_family = AF_INET;
	sin.sin_addr.s_addr = INADDR_ANY;
	sin.sin_port = htons (VNC_BASE + port);
	bound = socket (AF_INET, SOCK_STREAM, IPPROTO_IP);
	if (bound < 0) {
		perror ("socket");
		exit (1);
	}
	setsockopt (bound, SOL_SOCKET, SO_REUSEADDR, &on, sizeof (on));
	if (bind (bound, (struct sockaddr *) &sin,
		  sizeof (struct sockaddr_in))) {
		perror ("bind");
		exit (1);
	}
	listen (bound, 1);
	sock = accept (bound, NULL, 0);
	close (bound);
	return sock;
}

int main (int argc, char *argv[])
{
	int use_stdout = 0;
	int loop = 0;
	int appenddate = 0;
	const char **file = NULL;
	int files = 0;
	char action = '\0';
	char type = '\0';
	char *server = NULL;
	char *display = NULL;
	int clientr, clientw;
	struct sockaddr_in server_addr;
	int orig_optind;
	int cycle = 0, pause = 0;

	/* Options */
	for (;;) {
		static struct option long_options[] = {
			{"playback", 0, 0, 'p'},
			{"record", 0, 0, 'r'},
			{"type", 1, 0, 't'},
			{"date", 0, 0, 'd'},
			{"loop", 0, 0, 'l'},
			{"stdout", 0, 0, 'c'},
			{"server", 1, 0, 's'},
			{"help", 0, 0, 'h'},
			{"version", 0, 0, 'v'},
			{"pause", 1, 0, 'P'},
			{"cycle", 1, 0, 'C'},
			{0, 0, 0, 0}
		};
		int l;
		int c = getopt_long (argc, argv, "prcld",
				     long_options, NULL);
		if (c == -1)
			break;

		switch (c) {
			char *p;
		case 'P':
			pause = strtoul (optarg, &p, 16);
			if (optarg == p)
				pause = *optarg;
			break;
		case 'C':
			cycle = strtoul (optarg, &p, 16);
			if (optarg == p)
				cycle = *optarg;
			break;
		case 'v':
			version ();
			exit (0);
		case 'h':
			version ();
			usage (argv[0]);
		case 'c':
			if (use_stdout)
				usage (argv[0]);
			use_stdout = 1;
			break;
		case 'l':
			if (loop)
				usage (argv[0]);
			loop = 1;
			break;
		case 'd':
			if (appenddate)
				usage (argv[0]);
			appenddate = 1;
			break;
		case 't':
			if (type)
				usage (argv[0]);
			l = strlen (optarg);
			if (!strncmp (optarg, "screen", l))
				type = 's';
			else if (!strncmp (optarg, "events", l))
				type = 'e';
			else usage (argv[0]);
			break;
		case 's':
			if (server)
				usage (argv[0]);
			server = optarg;
			break;
		case 'p':
		case 'r':
			if (action)
				usage (argv[0]);
			action = c;
			break;
		}
	}

	orig_optind = optind;
	for (; optind < argc; optind++) {
		if (!display && argv[optind][0] == ':') {
			display = argv[optind];
			continue;
		}

		files++;
	}
	file = malloc ((files + 1) * sizeof (char *));
	if (!file) {
		fprintf (stderr, "out of memory\n");
		exit (1);
	}
	files = 0;
	for (optind = orig_optind; optind < argc; optind++)
		if (argv[optind][0] != ':')
			file[files++] = argv[optind];
	file[files] = NULL;

	/* Invalid option combinations */
	if (!action ||
	    (loop && action != 'p') ||
	    (type && action != 'r') ||
	    (server && action != 'r') ||
	    (use_stdout && display))
		usage (argv[0]);

	if (files < 1)
		/* No files specified */
		usage (argv[0]);

	/* Defaults */
	if (!server)
		server = strdup (DEFAULT_SERVER);
	if (!display)
		display = strdup (DEFAULT_DISPLAY);
	if (!type)
		type = 's';

	if (action == 'r') {
		int port;
		char *end;
		char *cl = strchr (server, ':');
		if (files > 1)
			/* Can't record to more than one file */
			usage (argv[0]);
		if (!cl)
			usage (argv[0]);
		*cl++ = '\0';
		port = VNC_BASE + strtoul (cl, &end, 10);
		if (cl == end)
			usage (argv[0]);
		server_addr.sin_family = AF_INET;
		if (!server[0])
			server_addr.sin_addr.s_addr = htonl (INADDR_LOOPBACK);
		else {
			struct hostent *hp;
			if ((server_addr.sin_addr.s_addr =
			     inet_addr (server)) == -1) {
				hp = gethostbyname (server);
				if (!hp) {
					perror (server);
					exit (1);
				}
				memcpy (&server_addr.sin_addr.s_addr,
					hp->h_addr_list[0],
					hp->h_length);
			}
		}
		server_addr.sin_port = htons (port);
	}

	/* Get some file descriptors */
	if (use_stdout) {
		clientr = fileno (stdin);
		clientw = fileno (stdout);
	} else {
		unsigned long port;
		char *end;
		display++;
		port = strtoul (display, &end, 10);
		if (display == end)
			usage (argv[0]);
		clientr = clientw = accept_connection (port);
	}

	/* Do it */
	if (action == 'r') {
		record (file[0], clientr, clientw, server_addr, type == 'e', appenddate);
	} else {
		int file_to_play = 0;
		int looping = 0;
		while (playback (file[file_to_play],
				 clientr, clientw, looping,
				 cycle, pause) > 0) {
			file_to_play++;
			if (file_to_play == files) {
				if (!loop)
					break;
				file_to_play = 0;
			}
			looping = 1;
		}
	}

	/* Clean up */
	if (!use_stdout)
		close (clientr);

	return 0;
}


syntax highlighted by Code2HTML, v. 0.9.1