#!/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;
my $proxy = $ENV{https_proxy} or
	die "https_proxy is not defined int the environment\n";
$proxy =~ s,^http://,,;
$proxy =~ s,/$,,;
my ($proxy_host, $proxy_port) = split /:/, $proxy;
$proxy_host and $proxy_port or
	die "https_proxy environment variable has bad syntax\n";

$| = 1;

use POSIX ":sys_wait_h";

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

my $listeners = 20;

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

my %record_by_handle;

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

my ($listen_sockname, $relay_sockname);
while (my ($listen_sockname, $relay_sockname) = splice @ARGV, 0, 2) {
	my $listen_handle = listen_to($listen_sockname);
	$select_read->add($listen_handle);
	non_block($listen_handle);

	$record_by_handle{$listen_handle} = {
		name => 'listen '.$listen_sockname,
		type => 'listen',
		relay_sockname => $relay_sockname,
	};
}

# 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) {
#		print STDERR "handle: $handle\n";
		my $record = $record_by_handle{$handle};
#		print STDERR "record: $record\n";
		my $name = $record->{name};
#		print STDERR "name: $name\n";
		if ($record->{type} eq 'listen') {
			my $new_client_handle = $handle->accept;
			non_block($new_client_handle);
			my $new_server_handle = nb_connect_to_proxy();

			++$n_connections;
			my $new_client_record = {
				type => 'client',
				name => 'client '.$n_connections,
				write_buffer => '',
				peer => $new_server_handle,
				eat_headers => 0,
			};
			my $new_server_record = {
				type => 'server',
				name => 'server '.$n_connections,
				write_buffer => "CONNECT ".$record->{relay_sockname}." HTTP/1.0\r\n\r\n",
				peer => $new_client_handle,
				eat_headers => 1,
				eat_headers_buffer => '',
			};
			$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);
			$select_write->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 ($record->{eat_headers}) {
#				print STDERR "eating headers\n";
				if ($bytes_read = $handle->sysread($record->{eat_headers_buffer}, $bytes_to_read, length($record->{eat_headers_buffer}))) {
					if ($record->{eat_headers_buffer} =~ s/.*?(\r\n\r\n|\n\n)//s) {
						if (length ($record->{eat_headers_buffer}) > 0) {
							$peer_record->{write_buffer} .= $record->{eat_headers_buffer};
							$buffer_length == 0 and $select_write->add($peer);
							if (length($peer_record->{write_buffer}) == $max_buffer_size) {
								# the buffer is full
								$select_read->remove($handle);
							}
						}
						$record->{eat_headers_buffer} = "";
						$record->{eat_headers} = 0;
					}
				} 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
				}
			} else {
				if ($bytes_read = $handle->sysread($peer_record->{write_buffer}, $bytes_to_read, $buffer_length)) {
					$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 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_proxy {
	my $handle;
	# 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($proxy_port, inet_aton($proxy_host)))
		or $!{EINPROGRESS} or die "can't connect to port $proxy_port on $proxy_host: $!";
#		print STDERR "[Connecting to $host:$port]\n";
	return $handle;
}

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: $!";
}