File Coverage

blib/lib/IO/Socket/Socks/Wrapper.pm

Criterion	Covered	Total	%
statement	111	223	49.7
branch	43	96	44.7
condition	18	37	48.6
subroutine	19	34	55.8
pod	0	1	0.0
total	191	391	48.8

package IO::Socket::Socks::Wrapper;

use strict;
no warnings 'prototype';
no warnings 'redefine';
use Socket qw(:DEFAULT inet_ntop);
use Errno;
use base 'Exporter';

our $VERSION = '0.15';
our @EXPORT_OK = ('connect', 'wrap_connection');

# cache
# pkg -> ref to pkg::sub || undef(if pkg has no connect)
my %PKGS;

# reference to &IO::Socket::connect
my $io_socket_connect;
sub _io_socket_connect_ref {
	return $io_socket_connect if $io_socket_connect;
	$io_socket_connect = \&IO::Socket::connect;
}

# fake handle to put under event loop while making socks handshake
sub _get_blocking_handles {
	pipe(my $blocking_reader, my $blocking_writer)
		or die 'pipe(): ', $!;

	$blocking_writer->blocking(0);
	$blocking_reader->blocking(0);
	my $garbage = '\0' x 1024;

	my ($writed, $total_writed);
	while ($writed = syswrite($blocking_writer, $garbage)) {
		$total_writed += $writed;

		if ($total_writed > 2097152) {
			# socket with buffer more than 2 mb
			# are u kidding me?
			die "Can't create blocking handle";
		}
	}

	return ($blocking_reader, $blocking_writer);
}

sub _unblock_handles {
	my ($blocking_reader, $blocking_writer) = @_;

	while (sysread($blocking_reader, my $buf, 4096)) {
		vec(my $win, fileno($blocking_writer), 1) = 1;
		last if select(undef, $win, undef, 0);
	}
}

sub import {
	my $mypkg = shift;

	if (@_ == 1 && !ref($_[0]) && $_[0] eq 'wrap_connection') {
		return __PACKAGE__->export_to_level(1, $mypkg, 'wrap_connection');
	}

	while (my ($pkg, $cfg) = splice @_, 0, 2) {
		unless (defined $cfg) {
			$cfg = $pkg;
			$pkg = undef;
		}

		if ($pkg) {
			no strict 'refs';

			my $sub;
			if ($pkg =~ /^(.+)::([^:]+)\(\)$/) {
				$pkg = $1;
				$sub = $2;
			}

			# override in the package

			# try to load if not already available
			# and if not requested to do not load
			unless(delete $cfg->{_norequire} || %{$pkg.'::'}) {
				eval "require $pkg" # make @ISA available
					or die $@;
			}

			if ($sub) {
			# localize IO::Socket::connect overriding
			# in the sub where IO::Socket::connect called
				my $symbol = $pkg.'::'.$sub;
				my $pkg_sub = exists $PKGS{$symbol} ?
				                     $PKGS{$symbol} :
				                     ($PKGS{$symbol} = \&$symbol);

				_io_socket_connect_ref();

				*$symbol = sub {
					local *IO::Socket::IP::connect = local *IO::Socket::connect = sub {
						_connect(@_, $cfg, 1);
					};

					$pkg_sub->(@_);
				};
				next;
			}
			elsif ($pkg->isa('IO::Socket')) {
			# replace IO::Socket::connect
			# if package inherits from IO::Socket
				# save replaceable package version of the connect
				# if it has own
				# will call it from the our new connect
				my $symbol = $pkg.'::connect';
				my $pkg_connect = exists $PKGS{$pkg} ?
				                         $PKGS{$pkg} :
				                         ($PKGS{$pkg} = eval{ *{$symbol}{CODE} } ? \&$symbol : undef);

				*connect = sub {
					_io_socket_connect_ref();

					local *IO::Socket::IP::connect = local *IO::Socket::connect = sub {
						_connect(@_, $cfg, 1);
					};

					my $self = shift;

					if ($pkg_connect) {
					# package has its own connect
						$pkg_connect->($self, @_);
					}
					else {
					# get first parent which has connect sub
					# and call it
						my $ref = ref($self);

						foreach my $parent (@{$pkg.'::ISA'}) {
							if($parent->isa('IO::Socket')) {
								bless $self, $parent;
								my $connected = $self->connect(@_);
								bless $self, $ref;
								return $connected ? $self : undef;
							}
						}
					}
				}
			}
			else {
				# replace package version of connect
				*connect = sub {
					_connect(@_, $cfg);
				}
			}

			$mypkg->export($pkg, 'connect');
		}
		else {
			# override connect() globally
			*connect = sub(*$) {
				my $socket = shift;
				unless (ref $socket) {
					# old-style bareword used
					no strict 'refs';
					my $caller = caller;
					$socket = $caller . '::' . $socket;
					$socket = \*{$socket};
				}

				_connect($socket, @_, $cfg);
			};

			$mypkg->export('CORE::GLOBAL', 'connect');
		}
	}
}

sub wrap_connection {
	require IO::Socket::Socks::Wrapped;
	return  IO::Socket::Socks::Wrapped->new(@_);
}

sub _connect {
	my ($socket, $name, $cfg, $io_socket) = @_;
	my $ref = ref($socket);

	if ($socket->isa('IO::Socket::Socks') || !$cfg) {
		unless ($io_socket and ${*$socket}{'io_socket_timeout'}) {
			return CORE::connect( $socket, $name );
		}

		# use IO::Socket::connect for timeout support
		local *connect = sub { CORE::connect($_[0], $_[1]) };
		return _io_socket_connect_ref->( $socket, $name );
	}

	my ($port, $host);
	if (($port, $host) = eval { unpack_sockaddr_in($name) }) {
		$host = inet_ntoa($host);
	}
	else {
		($port, $host) = unpack_sockaddr_in6($name);
		$host = inet_ntop(AF_INET6, $host);
	}

	# global overriding will not work with `use'
	require IO::Socket::Socks;
	my $io_handler = $cfg->{_io_handler};

	unless ($io_handler || exists $cfg->{Timeout}) {
		$cfg->{Timeout} = $ref && $socket->isa('IO::Socket') && ${*$socket}{'io_socket_timeout'} || 180;
	}

	my $need_nb;

	if ($io_handler) {
		$io_handler = $io_handler->();
		require POSIX;


		my $fd = fileno($socket);
		my $tmp_fd = POSIX::dup($fd) // die 'dup(): ', $!;
		open my $tmp_socket, '+<&=' . $tmp_fd or die 'open(): ', $!;

		my ($blocking_reader, $blocking_writer) = _get_blocking_handles();
		POSIX::dup2(fileno($blocking_writer), $fd) // die 'dup2(): ', $!;

		$io_handler->{blocking_reader} = $blocking_reader;
		$io_handler->{blocking_writer} = $blocking_writer;
		$io_handler->{orig_socket} = $socket;
		Scalar::Util::weaken($io_handler->{orig_socket});
		$socket = $tmp_socket;
	}
	elsif (!$socket->blocking) {
		# without io handler non-blocking connection will not be success
		# so set socket to blocking mode while making socks handshake
		$socket->blocking(1);
		$need_nb = 1;
	}

	my $ok = IO::Socket::Socks->new_from_socket(
		$socket,
		ConnectAddr  => $host,
		ConnectPort  => $port,
		%$cfg
	);

	if ($need_nb) {
		$socket->blocking(0);
	}

	return unless $ok || ($io_handler && $IO::Socket::Socks::SOCKS_ERROR == &IO::Socket::Socks::SOCKS_WANT_WRITE);
	bless $socket, $ref if $ref && !$io_handler; # XXX: should we unbless for GLOB?

	if ($io_handler) {
		my ($r_cb, $w_cb);
		my $done;

		tie *{$io_handler->{orig_socket}}, 'IO::Socket::Socks::Wrapper::Handle', $io_handler->{orig_socket}, sub {
			unless ($done) {
				$io_handler->{unset_read_watcher}->($socket);
				$io_handler->{unset_write_watcher}->($socket);

				if ($io_handler->{destroy_io_watcher}) {
					$io_handler->{destroy_io_watcher}->($socket);
				}

				close $socket;
			}

			# clean circular references
			undef $r_cb;
			undef $w_cb;
		};

		my $on_finish = sub {
			tied(*{$io_handler->{orig_socket}})->handshake_done($done = 1);
			POSIX::dup2(fileno($socket), fileno($io_handler->{orig_socket})) // die 'dup2(): ', $!;
			close $socket;
			_unblock_handles($io_handler->{blocking_reader}, $io_handler->{blocking_writer});
		};

		my $on_error = sub {
			tied(*{$io_handler->{orig_socket}})->handshake_done($done = 1);
			shutdown($socket, 0);
			POSIX::dup2(fileno($socket), fileno($io_handler->{orig_socket}))  // die 'dup2(): ', $!;
			close $socket;
		};

		$r_cb = sub {
			if ($socket->ready) {
				$io_handler->{unset_read_watcher}->($socket);

				if ($io_handler->{destroy_io_watcher}) {
					$io_handler->{destroy_io_watcher}->($socket);
				}

				$on_finish->();
			}
			elsif ($IO::Socket::Socks::SOCKS_ERROR == &IO::Socket::Socks::SOCKS_WANT_WRITE) {
				$io_handler->{unset_read_watcher}->($socket);
				$io_handler->{set_write_watcher}->($socket, $w_cb);
			}
			elsif ($IO::Socket::Socks::SOCKS_ERROR != &IO::Socket::Socks::SOCKS_WANT_READ) {
				$io_handler->{unset_read_watcher}->($socket);

				if ($io_handler->{destroy_io_watcher}) {
					$io_handler->{destroy_io_watcher}->($socket);
				}

				$on_error->();
			}
		};

		$w_cb = sub {
			if ($socket->ready) {
				$io_handler->{unset_write_watcher}->($socket);

				if ($io_handler->{destroy_io_watcher}) {
					$io_handler->{destroy_io_watcher}->($socket);
				}

				$on_finish->();
			}
			elsif ($IO::Socket::Socks::SOCKS_ERROR == &IO::Socket::Socks::SOCKS_WANT_READ) {
				$io_handler->{unset_write_watcher}->($socket);
				$io_handler->{set_read_watcher}->($socket, $r_cb);
			}
			elsif ($IO::Socket::Socks::SOCKS_ERROR != &IO::Socket::Socks::SOCKS_WANT_WRITE) {
				$io_handler->{unset_write_watcher}->($socket);

				if ($io_handler->{destroy_io_watcher}) {
					$io_handler->{destroy_io_watcher}->($socket);
				}

				$on_error->();
			}
		};

		if ($io_handler->{init_io_watcher}) {
			$io_handler->{init_io_watcher}->($socket, $r_cb, $w_cb);
		}

		$io_handler->{set_write_watcher}->($socket, $w_cb);

		$! = Errno::EINPROGRESS;
		return 0;
	}

	return 1;
}

package IO::Socket::Socks::Wrapper::Handle;

use strict;

sub TIEHANDLE {
	my ($class, $orig_handle, $cleanup_cb) = @_;

	open my $self, '+<&=' . fileno($orig_handle)
		or die 'open: ', $!;

	${*$self}{handshake_done} = 0;
	${*$self}{cleanup_cb} = $cleanup_cb;

	bless $self, $class;
}

sub handshake_done {
	my $self = shift;

	if (@_) {
		${*$self}{handshake_done} = $_[0];
	}

	return ${*$self}{handshake_done};
}

sub READ {
	my $self = shift;
	sysread($self, $_[0], $_[1], @_ > 2 ? $_[2] : ());
}

sub WRITE {
	my $self = shift;
	syswrite($self, $_[0], $_[1], @_ > 2 ? $_[2] : ());
}

sub FILENO {
	my $self = shift;
	fileno($self);
}

sub CLOSE {
	my $self = shift;

	unless ($self->handshake_done) {
		$self->handshake_done(1);
		${*$self}{cleanup_cb}->();
	}

	close $self;
}

sub DESTROY {
	my $self = shift;
	${*$self}{cleanup_cb}->();
}

1;

__END__