Net::Proxy - Framework for proxying network connections in many ways

received

package Net::Proxy;
use strict;
use warnings;
use Carp;
use Scalar::Util qw( refaddr reftype );
use IO::Select;
use POSIX 'strftime';

our $VERSION = '0.12';

# interal socket information table
my %SOCK_INFO;
my %LISTENER;
my %CLOSING;
my $READERS;
my $WRITERS;
my %PROXY;
my %STATS;

# Net::Proxy attributes
my %CONNECTOR = (
    in  => {},
    out => {},
);
my $VERBOSITY = 0; # be silent by default
my $BUFFSIZE  = 16384;

#
# some logging-related methods
#
sub set_verbosity { $VERBOSITY = $_[1]; }
{
    my $i;
    for my $meth (qw( error notice info debug )) {
        no strict 'refs';
        my $level = $i++;
        *$meth = sub {
            return if $VERBOSITY < $level;
            print STDERR strftime "%Y-%m-%d %H:%M:%S $_[1]\n", localtime;
        };
    }
}

#
# constructor
#
sub new {
    my ( $class, $args ) = @_;
    my $self = bless \do { my $anon }, $class;

    croak "Argument to new() must be a HASHREF" if ref $args ne 'HASH';

    for my $conn (qw( in out )) {

        # check arguments
        croak "'$conn' connector required" if !exists $args->{$conn};

        croak "'$conn' connector must be a HASHREF"
            if ref $args->{$conn} ne 'HASH';

        croak "'type' key required for '$conn' connector"
            if !exists $args->{$conn}{type};

        croak "'hook' key is not a CODE reference for '$conn' connector"
            if $args->{$conn}{hook}
            && reftype( $args->{$conn}{hook} ) ne 'CODE';
 
        # load the class
        my $class = 'Net::Proxy::Connector::' . $args->{$conn}{type};
        eval "require $class";
        croak "Couldn't load $class for '$conn' connector: $@" if $@;

        # create and store the Connector object
        $args->{$conn}{_proxy_} = $self;
        $CONNECTOR{$conn}{ refaddr $self} = $class->new( $args->{$conn} );
        $CONNECTOR{$conn}{ refaddr $self}->set_proxy($self);
    }

    return $self;
}

sub register { $PROXY{ refaddr $_[0] } = $_[0]; }
sub unregister { delete $PROXY{ refaddr $_[0] }; }

#
# The Net::Proxy attributes
#
sub in_connector  { return $CONNECTOR{in}{ refaddr $_[0] }; }
sub out_connector { return $CONNECTOR{out}{ refaddr $_[0] }; }

#
# create the socket setter/getter methods
# these are actually Net::Proxy clas methods
#
BEGIN {
    my $n = 0;
    my $buffer_id;
    for my $attr (qw( peer connector state nick buffer callback )) {
        no strict 'refs';
        my $i = $n;
        *{"get_$attr"} = sub { $SOCK_INFO{ refaddr $_[1] }[$i]; };
        *{"set_$attr"} = sub { $SOCK_INFO{ refaddr $_[1] }[$i] = $_[2]; };
        $buffer_id = $n if $attr eq 'buffer';
        $n++;
    }
    # special shortcut
    sub add_to_buffer { $SOCK_INFO{ refaddr $_[1] }[$buffer_id] .= $_[2]; }
}

#
# create statistical methods
#
for my $info (qw( opened closed )) {
    no strict 'refs';
    *{"stat_inc_$info"} = sub {
        $STATS{ refaddr $_[0]}{$info}++;
        $STATS{total}{$info}++;
    };
    *{"stat_$info"}       = sub { $STATS{ refaddr $_[0]}{$info} || 0; };
    *{"stat_total_$info"} = sub { $STATS{total}{$info} || 0; };
}

#
# socket-related methods
#
sub add_listeners {
    my ( $class, @socks ) = @_;
    for my $sock (@socks) {
        Net::Proxy->notice( 'Add ' . Net::Proxy->get_nick($sock) );
        $LISTENER{ refaddr $sock} = $sock;
    }
    return;
}

sub close_sockets {
    my ( $class, @socks ) = @_;

  SOCKET:
    for my $sock (@socks) {
        if( my $data = Net::Proxy->get_buffer( $sock ) ) {
            ## Net::Proxy->debug( length($data) . ' bytes left to write on ' . Net::Proxy->get_nick( $sock ) );
            $CLOSING{ refaddr $sock} = $sock;
            next SOCKET;
        }

        Net::Proxy->notice( 'Closing ' . Net::Proxy->get_nick( $sock ) );

        # clean up connector
        if ( my $conn = Net::Proxy->get_connector($sock) ) {
            $conn->close($sock) if $conn->can('close');

            # count connections to the proxy "in connectors" only
            my $proxy = $conn->get_proxy();
            if ( refaddr $conn == refaddr $proxy->in_connector()
                && !_is_listener($sock) )
            {
                $proxy->stat_inc_closed();
            }
        }

        # clean up internal structures
        delete $SOCK_INFO{ refaddr $sock};
        delete $LISTENER{ refaddr $sock};
        delete $CLOSING{ refaddr $sock};

        # clean up sockets
        $READERS->remove($sock);
        $WRITERS->remove($sock);
        $sock->close();
    }

    return;
}

#
# select() stuff
#
sub watch_reader_sockets {
    my ( $class, @socks ) = @_;
    $READERS->add(@socks);
    return;
}

sub watch_writer_sockets {
    my ( $class, @socks ) = @_;
    $WRITERS->add(@socks);
    return;
}

sub remove_writer_sockets {
    my ( $class, @socks ) = @_;
    $WRITERS->remove(@socks);
    return;
}

#
# destructor
#
sub DESTROY {
    my ($self) = @_;
    delete $CONNECTOR{in}{ refaddr $self};
    delete $CONNECTOR{out}{ refaddr $self};
}

#
# the mainloop itself
#
sub mainloop {
    my ( $class, $max_connections ) = @_;
    $max_connections ||= 0;

    # initialise the loop
    $READERS = IO::Select->new();
    $WRITERS = IO::Select->new();

    # initialise all proxies
    for my $proxy ( values %PROXY ) {
        my $in    = $proxy->in_connector();
        my @socks = $in->listen();
        Net::Proxy->add_listeners(@socks);
        Net::Proxy->watch_reader_sockets(@socks);
        Net::Proxy->set_connector( $_, $in ) for @socks;
    }

    my $continue = 1;
    for my $signal (qw( INT HUP )) {
        $SIG{$signal} = sub {
            Net::Proxy->notice("Caught $signal signal");
            $continue = 0;
        };
    }

    # loop indefinitely
    while ( $continue and my @ready = IO::Select->select( $READERS, $WRITERS ) ) {

        ## Net::Proxy->debug( 0+@{$ready[0]} . " sockets ready for reading" );
        ## Net::Proxy->debug( join "\n  ", "Readers:", map { Net::Proxy->get_nick($_) } $READERS->handles() );
        ## Net::Proxy->debug( 0+@{$ready[1]} . " sockets ready for writing" );
        ## Net::Proxy->debug( join "\n  ", "Writers:", map { Net::Proxy->get_nick($_) } $WRITERS->handles() );

        # first read
    READER:
        for my $sock (@{$ready[0]}) {
            if ( _is_listener($sock) ) {

                # accept the new connection and connect to the destination
                Net::Proxy->get_connector($sock)->new_connection_on($sock);
            }
            else {

                # have we read too much?
                my $peer = Net::Proxy->get_peer($sock);
                next READER
                    if !$peer
                    || length( Net::Proxy->get_buffer($peer) ) >= $BUFFSIZE;

                # read the data
                if ( my $conn = Net::Proxy->get_connector($sock) ) {
                    my $data = $conn->read_from($sock);
                    next READER if !defined $data;

                    if ($peer) {

                        # run the hook on incoming data
                        my $callback = Net::Proxy->get_callback( $sock );
                        $callback->( \$data, $sock, $conn )
                            if $callback && defined $data;

                        Net::Proxy->add_to_buffer( $peer, $data );
                        Net::Proxy->watch_writer_sockets($peer);

                        ## Net::Proxy->debug( "Will write " . length( Net::Proxy->get_buffer($peer)). " bytes to " .  Net::Proxy->get_nick( $peer ));
                    }
                }
            }
        }

        # then write
        for my $sock (@{$ready[1]}) {
            my $conn = Net::Proxy->get_connector($sock);
            $conn->write_to($sock);
        }

    }
    continue {
        if( %CLOSING ) {
            Net::Proxy->close_sockets( values %CLOSING );
        }
        if( $max_connections ) {

            # stop after that many connections
            last if Net::Proxy->stat_total_closed() == $max_connections;

            # prevent new connections
            if ( %LISTENER
                && Net::Proxy->stat_total_opened() == $max_connections )
            {
                Net::Proxy->close_sockets( values %LISTENER );
            }
        }
    }

    # close all remaining sockets
    Net::Proxy->close_sockets( $READERS->handles(), $WRITERS->handles() );
}

#
# helper private FUNCTIONS
#
sub _is_listener { return exists $LISTENER{ refaddr $_[0] }; }

1;

__END__

Net::Proxy - Framework for proxying network connections in many ways

Index

Code Index:

NAME

SYNOPSIS

DESCRIPTION

METHODS

Class methods

Instance methods

Statistical methods

CONNECTORS

Connector hooks

Available connectors

Summary

AUTHOR

BUGS

TODO

SUPPORT

COPYRIGHT

LICENSE