#!/usr/bin/perl -w
# test_relay

use strict;
use IO::Socket;
use IO::Select;
use IO::File;
use constant SHUT_WR => 1;
use Errno;

my $max_buffer_size = 1024;
my $n_connections = 0;

$| = 1;

use POSIX ":sys_wait_h";

$SIG{CHLD} = sub { while (waitpid(-1,WNOHANG) > 0) {} };

my $listeners = 20;

if (@ARGV != 2) {
	print STDERR "usage: $0 (listen socket) (relay to socket)\n";
	exit 1;
}

my ($listen_sockname, $relay_sockname) = @ARGV;

my $listen_handle = listen_to($listen_sockname);

sub listen_to {
	my $sockname = shift;
	my $handle;
	if ($sockname =~ /^\d{1,5}$/) {
		my $port = $sockname;
		# serve tcp on the specified host and port
		$handle = IO::Socket::INET->new( Proto     => "tcp",
					         LocalPort => $port,
					         Listen    => $listeners,
					         Reuse     => 1 )
			or die "can't serve on TCP port $port: $!";
		print STDERR "[Serving on TCP port $port]\n";
	} else {
		# serve unix domain on the specified named socket
		unlink $sockname;
		$handle = IO::Socket::UNIX->new( Local  => $sockname,
						 Listen => $listeners )
			or die "cannot serve on UNIX socket $sockname: $!";
		chmod 0777, $sockname;
		print STDERR "[Serving on UNIX socket $sockname]\n";
	}
	return $handle;
}

sub nb_connect_to {
	my $sockname = shift;
	my $handle;
	$sockname = "127.0.0.1:$sockname" if $sockname !~ /^\d{1,5}$/;
	if ($sockname =~ /:/) {
		my ($host, $port) = split /:/, $sockname;
		# create a tcp connection to the specified host and port
		$handle = IO::Socket::INET->new;
		socket($handle, PF_INET, SOCK_STREAM, getprotobyname("tcp"));
		non_block($handle);
		connect($handle, sockaddr_in($port, inet_aton($host)))
			or $!{EINPROGRESS} or die "can't connect to port $port on $host: $!";
		print STDERR "[Connecting to $host:$port]\n";
	} else {
		# create a UNIX sockets connection to the specified named socket
		$handle = IO::Socket::UNIX->new;
		socket($handle, PF_UNIX, SOCK_STREAM, 0);
		non_block($handle);
		connect($handle, sockaddr_un($sockname))
			or die "can't connect to socket $sockname: $!";
		print STDERR "[Connecting to $sockname]\n";
	}
	return $handle;
}

sub connect_to {
	my $sockname = shift;
	my $handle;
	$sockname = "127.0.0.1:$sockname" if $sockname =~ /^\d{1,5}$/;
	if ($sockname =~ /:/) {
		my ($host, $port) = split /:/, $sockname;
		# create a tcp connection to the specified host and port
		$handle = IO::Socket::INET->new( Proto     => "tcp",
						 PeerAddr  => $host,
						 PeerPort  => $port )
			or die "can't connect to port $port on $host: $!";
		print STDERR "[Connected to $host:$port]\n";
	} else {
		# create a UNIX sockets connection to the specified named socket
		$handle = IO::Socket::UNIX->new(Peer => $sockname)
			or die "can't connect to socket $sockname: $!";
		print STDERR "[Connected to $sockname]\n";
	}
	return $handle;
}

my $select_read = IO::Select->new;
my $select_write = IO::Select->new;

$select_read->add($listen_handle);
non_block($listen_handle);

my %record_by_handle = (
	$listen_handle => {
		name => 'listen',
		type => 'listen'
	}
);

# each record is of this form:
# {
#	name => ...,	    # socket `name' - some identifier for human to read
# 	type => 'client',   # or 'listen' or 'server'
#	write_buffer => '', # not for type 'listen'
#	peer => ...	    # the socket to relay to, not for type 'listen'
# }

print STDERR "[waiting for connections]\n";

while (1) {
#	print "select read: "; printfds($select_read);
#	print "select write: "; printfds($select_write);
	sub printfds {
		my $sel = shift;
		print join ",", map {$_->fileno} $sel->handles;
		print "\n";
	}
	
	my ($can_read, $can_write, $has_errors) = IO::Select->select($select_read, $select_write, $select_read);

	for my $handle (@$has_errors) {
		my $record = $record_by_handle{$handle};
		my $name = $record->{name};
		print STDERR "$name: socket error\n";
	}

	for my $handle (@$can_write) {
#		print "RTW: ", $handle->fileno, "\n";
		my $record = $record_by_handle{$handle};
		my $name = $record->{name};
		my $peer = $record->{peer};
		my $peer_record = $record_by_handle{$peer};
		my $buffer_length = length($record->{write_buffer});
		if ($buffer_length == 0) {
			# we need to close the connection
			shutdown($handle, SHUT_WR) or warn "shutdown: $!";
			$select_write->remove($handle);
			$record->{closed} = 1; # this means `no more will be written'
#			print STDERR "$name: closed for writing\n";
			if ($peer_record->{closed}) {
				delete $record_by_handle{$handle};
				delete $record_by_handle{$peer};
			}
		} else {
			my $bytes_written;
			if ($bytes_written = $handle->syswrite($record->{write_buffer}, length $record->{write_buffer}, 0)) {
				substr $record->{write_buffer}, 0, $bytes_written, '';
				if ($buffer_length == $max_buffer_size) {
					# the buffer was full, now is not, so can read again
					$select_read->add($peer);
				}
				if ($buffer_length == $bytes_written) {
					# the buffer is now empty, so stop selecting to write
					$select_write->remove($handle);
				}
			} elsif (! defined $bytes_written) {
				print STDERR "$name: write error: $!\n";
			}
		}
	}

	for my $handle (@$can_read) {
		my $record = $record_by_handle{$handle};
		my $name = $record->{name};
		if ($record->{type} eq 'listen') {
			my $new_client_handle = $handle->accept;
			non_block($new_client_handle);
			#my $new_server_handle = nb_connect_to($relay_sockname);
			my $new_server_handle = connect_to($relay_sockname);

			++$n_connections;
			my $log_handle = IO::File->new("log.$n_connections", "w")
				or die "cannot create log for connection $n_connections";
			my $new_client_record = {
				type => 'client',
				name => 'client '.$n_connections,
				write_buffer => '',
				peer => $new_server_handle,
				log => $log_handle
			};
			my $new_server_record = {
				type => 'server',
				name => 'server '.$n_connections,
				write_buffer => '',
				peer => $new_client_handle
			};
			$record_by_handle{$new_client_handle} = $new_client_record;
			$record_by_handle{$new_server_handle} = $new_server_record;
			$select_read->add($new_client_handle);
			$select_read->add($new_server_handle);
		} else {
#			print "RTR: ", $handle->fileno, "\n";
			my $peer = $record->{peer};
			my $peer_record = $record_by_handle{$peer};
			my $buffer_length = length($peer_record->{write_buffer});
			my $bytes_to_read = $max_buffer_size - $buffer_length;
			$bytes_to_read <= 0 and die "internal error: selecting to read on a blocked connection\n";

			my $bytes_read;
			if ($bytes_read = $handle->sysread($peer_record->{write_buffer}, $bytes_to_read, $buffer_length)) {
				my $log = substr $peer_record->{write_buffer}, $buffer_length, $bytes_read;
				$log .= "\\n" unless $log =~ /\n$/;
				if ($record->{type} eq 'client') {
					$log =~ s/^/C>  /mg;
					print $log;
					$record->{log}->print($log) or die "cannot write to log";
				} else {
					$log =~ s/^/ S> /mg;
					print $log;
					$peer_record->{log}->print($log) or die "cannot write to log";
				}
				$buffer_length == 0 and $select_write->add($peer);
				if ($bytes_read == $bytes_to_read) {
					# the buffer is full
					$select_read->remove($handle);
				}
			} elsif (! defined $bytes_read) {
				print STDERR "$name: read error: $!\n";
			} else {
				# 0 -> connection closed / EOF
				print STDERR "$name: closed\n"; # for reading\n";
				$select_read->remove($handle);
				$buffer_length == 0 and $select_write->add($peer); # `write' the EOF
			}
		}
	}
}

use Fcntl;

sub non_block {
	my $handle = shift;
	my $flags;
	$flags = fcntl($handle, F_GETFL, 0)
		and fcntl($handle, F_SETFL, $flags | O_NONBLOCK)
		or die "non_block: $!";
}