POE::Component::Server::DNS - A non-blocking, concurrent DNS server POE component

  use strict;
  use Net::DNS::RR;
  use POE qw(Component::Server::DNS);

  my $dns_server = POE::Component::Server::DNS->spawn( alias => 'dns_server' );

  POE::Session->create(
        package_states => [ 'main' => [ qw(_start handler log) ], ],
  );

  $poe_kernel->run();
  exit 0;

  sub _start {
    my ($kernel,$heap) = @_[KERNEL,HEAP];

    # Tell the component that we want log events to go to 'log'
    $kernel->post( 'dns_server', 'log_event', 'log' );

    # register a handler for any foobar.com suffixed domains
    $kernel->post( 'dns_server', add_handler, 
	{ 
	  event => 'handler', 
	  label => 'foobar', 
	  match => 'foobar\.com$',
        } 
    );
    undef;
  }

  sub handler {
    my ($qname,$qclass,$qtype,$callback) = @_[ARG0..ARG3];
    my ($rcode, @ans, @auth, @add);

    if ($qtype eq "A") {
      my ($ttl, $rdata) = (3600, "10.1.2.3");
      push @ans, Net::DNS::RR->new("$qname $ttl $qclass $qtype $rdata");
      $rcode = "NOERROR";
    } else { 
      $rcode = "NXDOMAIN";
    }

    $callback->($rcode, \@ans, \@auth, \@add, { aa => 1 });
    undef;
  }

  sub log {
    my ($ip_port,$net_dns_packet) = @_[ARG0..ARG1];
    $net_dns_packet->print();
    undef;
  }

  ARG0, query name
  ARG1, query class
  ARG2, query type
  ARG3, a callback coderef
  ARG4, the IP address and port of the requestor, 'IPaddr:port'

  $callback->( $rcode, \@ans, \@auth, \@add, { aa => 1 } );

  ARG0, the IP address and port of the requestor, 'IPaddr:port';
  ARG1, the Net::DNS::Packet object;

package POE::Component::Server::DNS;

use strict;
use warnings;
use POE qw(Component::Client::DNS Wheel::ReadWrite Component::Client::DNS::Recursive Wheel::SocketFactory);
use Socket;
use Net::DNS::RR;
use IO::Socket::INET;
use vars qw($VERSION);

$VERSION = '0.20';

sub spawn {
  my $package = shift;
  my %args = @_;
  $args{lc $_} = delete $args{$_} for keys %args;

  my $options = delete $args{options};

  my $self = bless \%args, $package;

  $self->{_handlers} = [ ];

  unless ( $self->{no_clients} ) {
      $self->{_localhost} = Net::DNS::RR->new('localhost. 0 A 127.0.0.1');
  }

  $self->{session_id} = POE::Session->create(
	object_states => [ 
		$self => { shutdown => '_shutdown', },
		$self => [ qw(_start _dns_incoming _dns_err _dns_response _dns_recursive add_handler del_handler _handled_req _sock_up _sock_err log_event) ],
	],
	heap => $self,
	options => ( $options && ref $options eq 'HASH' ? $options : { } ),
  )->ID();

  return $self;
}

sub _start {
  my ($kernel,$self,$session) = @_[KERNEL,OBJECT,SESSION];

  $self->{session_id} = $session->ID();

  if ( $self->{alias} ) {
     $kernel->alias_set($self->{alias});
  } 
  else {
     $kernel->alias_set("$self");
     $self->{alias} = "$self";
  }

  unless ( $self->{no_clients} ) {
      $self->{resolver_opts} = { } unless $self->{resolver_opts} and ref $self->{resolver_opts} eq 'HASH';
      delete $self->{resolver_opts}->{Alias};
      $self->{resolver} = POE::Component::Client::DNS->spawn( Alias => "resolver" . $self->session_id(), %{ $self->{resolver_opts} } );
#      $self->{recursive}->hints( { event => '_hints' } );
  }

  $self->{factory} = POE::Wheel::SocketFactory->new(
	SocketProtocol => 'udp',
	BindAddress => $self->{address} || INADDR_ANY,
	BindPort => ( defined $self->{port} ? $self->{port} : 53 ),
	SuccessEvent   => '_sock_up',
	FailureEvent   => '_sock_err',
  );

  undef;
}

sub _sock_up {
  my ($kernel,$self,$dns_socket) = @_[KERNEL,OBJECT,ARG0];
  $self->{_sockport} = ( sockaddr_in( getsockname($dns_socket) ) )[0];
  delete $self->{factory};
  $self->{dnsrw} = POE::Wheel::ReadWrite->new(
	    Handle => $dns_socket,
	    Driver => DNS::Driver::SendRecv->new(),
	    Filter => DNS::Filter::UDPDNS->new(),
	    InputEvent => '_dns_incoming',
	    ErrorEvent => '_dns_err',
	);
  undef;
}

sub _sock_err {
  my ($operation, $errnum, $errstr, $wheel_id) = @_[ARG0..ARG3];
  delete $_[OBJECT]->{factory};
  die "Wheel $wheel_id generated $operation error $errnum: $errstr\n";
  undef;
}

sub session_id {
  return $_[0]->{session_id};
}

sub resolver {
  return $_[0]->{resolver};
}

sub sockport {
  return $_[0]->{_sockport};
}

sub shutdown {
  my $self = shift;
  $poe_kernel->post( $self->session_id() => 'shutdown' );
}

sub _shutdown {
  my ($kernel,$self) = @_[KERNEL,OBJECT];
  $kernel->alarm_remove_all();
  $kernel->alias_remove( $_ ) for $kernel->alias_list( $_[SESSION] );
  delete $self->{dnsrw};
  unless ( $self->{no_clients} ) {
      $self->{resolver}->shutdown();
      #$self->{recursive}->shutdown();
  }
  $kernel->refcount_decrement( $_->{session}, __PACKAGE__ ) for @{ $self->{_handlers} };
  $kernel->refcount_decrement( $_, __PACKAGE__ ) for keys %{ $self->{_sessions} };
  delete $self->{_handlers};
  undef;
}

sub log_event {
  my ($kernel,$self,$sender,$event) = @_[KERNEL,OBJECT,SENDER,ARG0];
  $sender = $sender->ID();

  if ( exists $self->{_sessions}->{ $sender } and !$event ) {
	delete $self->{_sessions}->{ $sender };
	$kernel->refcount_decrement( $sender => __PACKAGE__ );
	return;
  }
  
  if ( exists $self->{_sessions}->{ $sender } ) {
	$self->{_sessions}->{ $sender } = $event;
	return;
  }

  $self->{_sessions}->{ $sender } = $event;
  $kernel->refcount_increment( $sender => __PACKAGE__ );
  return;
}

sub add_handler {
  my ($kernel,$self,$sender) = @_[KERNEL,OBJECT,SENDER];
  $sender = $sender->ID();

  # Get the arguments
  my $args;
  if (ref($_[ARG0]) eq 'HASH') {
	$args = { %{ $_[ARG0] } };
  } 
  else {
	warn "first parameter must be a ref hash, trying to adjust. "
		."(fix this to get rid of this message)";
	$args = { @_[ARG0 .. $#_ ] };
  }

  $args->{ lc $_ } = delete $args->{$_} for keys %{ $args };

  unless ( $args->{label} ) {
	warn "you must supply a label argument, to make it unique\n";
	return;
  }

  if ( grep { $_->{label} eq $args->{label} } @{ $self->{_handlers} } ) {
	warn "you must supply a unique label argument, I already have that one\n";
	return;
  }

  unless ( $args->{event} ) {
	warn "you must supply an event argument, otherwise where do I send the replies to\n";
	return;
  }

  unless ( $args->{match} ) {
	warn "you must supply a match argument, otherwise what's the point\n";
	return;
  }

  my $regex;
  eval { $regex = qr/$args->{match}/ };

  if ( $@ ) { 
	warn "The match argument you supplied was fubar, please try harder\n";
	return;
  } 
  else {
	$args->{match} = $regex;
  }

  $args->{session} = $sender unless $args->{session};
  if ( my $ref = $kernel->alias_resolve( $args->{session} ) ) {
	$args->{session} = $ref->ID();
  } 
  else {
	$args->{session} = $sender->ID();
  }

  $kernel->refcount_increment( $args->{session}, __PACKAGE__ );

  push @{ $self->{_handlers} }, $args;

  undef;
}

sub del_handler {
  my ($kernel,$self,$label) = @_[KERNEL,OBJECT,ARG0];
  return unless $label;
  
  my $i = 0; my $rec;
  for ( @{ $self->{_handlers} } ) {
    if ( $_->{label} eq $label ) {
      splice( @{ $self->{_handlers} }, $i, 1 );
      $rec = $_;
      last;
    }
  }

  $kernel->refcount_decrement( $rec->{session}, __PACKAGE__ );
  undef;
}

sub _dns_incoming {
  my($kernel,$self,$session,$dnsq) = @_[KERNEL,OBJECT,SESSION,ARG0];

  my($q) = $dnsq->question();
  return unless $q;

  foreach my $handler ( @{ $self->{_handlers} } ) {
	next unless $q->qname =~ $handler->{match};
	my $callback = $session->callback( '_handled_req', $dnsq );
	$kernel->post(
			$handler->{session},
			$handler->{event},
			$q->qname,
			$q->qclass,
			$q->qtype,
			$callback,
			$dnsq->answerfrom );
	return;
  }
  
  if ( $self->{no_clients} ) {
    # Refuse unhandled requests, like an authoritative-only
    #  BIND server would.
    $dnsq->header->rcode('REFUSED');
    $dnsq->header->aa(0);
    $dnsq->header->ra(0);
    $dnsq->header->ad(0);
    $self->_dispatch_log( $dnsq );
    $self->{dnsrw}->put( $dnsq ) if $self->{dnsrw};
    return;
  }

  if ( $q->qname =~ /^localhost\.*$/i ) {
	$dnsq->push( answer => $self->{_localhost} );
	$self->_dispatch_log( $dnsq );
  	$self->{dnsrw}->put( $dnsq ) if $self->{dnsrw};
	return;
  }

  if ( $self->{forward_only} ) {
    my %query = (
      class   => $q->qclass,
      type    => $q->qtype,
      host    => $q->qname,
      context => [ $dnsq->answerfrom, $dnsq->header->id ],   	   
      event   => '_dns_response', 
    );
  
    my $response = $self->{resolver}->resolve( %query );
    $kernel->yield( '_dns_response', $response ) if $response;
  
  } 
  else {
#    $self->{recursive}->query_dorecursion( { event => '_dns_recursive', data => [ $dnsq, $dnsq->answerfrom, $dnsq->header->id ], }, $q->qname, $q->qtype, $q->qclass );
    POE::Component::Client::DNS::Recursive->resolve(
	event   => '_dns_recursive',
	context => [ $dnsq, $dnsq->answerfrom, $dnsq->header->id ],
	host    => $q->qname,
	type    => $q->qtype,
	class   => $q->qclass,
    );
  }

  undef;
}

sub _handled_req {
  my ($kernel,$self,$passthru,$passback) = @_[KERNEL,OBJECT,ARG0,ARG1];
  my $reply = $passthru->[0];
  my ($rcode, $ans, $auth, $add, $headermask) = @{ $passback };
  $reply->header->rcode($rcode);
  $reply->push("answer",     @$ans)  if $ans;
  $reply->push("authority",  @$auth) if $auth;
  $reply->push("additional", @$add)  if $add;
  if (!defined ($headermask)) {
	$reply->header->ra($self->{no_clients} ? 0 : 1);
	$reply->header->ad(0);
  } 
  else {
	$reply->header->aa(1) if $headermask->{'aa'};
	$reply->header->ra(1) if $headermask->{'ra'};
	$reply->header->ad(1) if $headermask->{'ad'};
  }
	
  $reply->header->qr(1);
  $self->_dispatch_log( $reply );
  $self->{dnsrw}->put( $reply ) if $self->{dnsrw};
  undef;
}

sub _dns_err {
  my($kernel,$self, $op, $errnum, $errstr) = @_[KERNEL,OBJECT, ARG0..ARG2];
  warn "DNS readwrite: $op generated error $errnum: $errstr\n";
  $kernel->yield('shutdown');
  undef;
}

sub _dns_recursive {
  my ($kernel,$self,$data) = @_[KERNEL,OBJECT,ARG0];
  return if $data->{error};
  my ($dnsq,$answerfrom,$id) = @{ $data->{context} };
  my $response = $data->{response};
  if ( $response ) {
    $response->header->id( $id );
    $response->answerfrom( $answerfrom );
    $self->_dispatch_log( $response );
    $self->{dnsrw}->put( $response ) if $self->{dnsrw};
    return;
  }
  $dnsq->header->rcode('NXDOMAIN');
  $self->_dispatch_log( $dnsq );
  $self->{dnsrw}->put( $dnsq ) if $self->{dnsrw};
  undef;
}

sub _dns_response {
  my ($kernel,$self,$reply) = @_[KERNEL,OBJECT,ARG0];

  my ($answerfrom,$id) = @{ $reply->{context} };
  my $response = delete $reply->{response};
  $response->header->id( $id );
  $response->answerfrom( $answerfrom );
  $self->_dispatch_log( $response );
  $self->{dnsrw}->put( $response ) if $self->{dnsrw};
  undef;
}

sub _dispatch_log {
  my $self = shift;
  my $packet = shift || return;
  my $af = $packet->answerfrom;
  $poe_kernel->post( $_, $self->{_sessions}->{$_}, $af, $packet ) for keys %{ $self->{_sessions} };
  return 1;
}

package DNS::Driver::SendRecv;

use strict;
use POE::Driver;
use Socket;

sub new {
    my $class = shift;
    my $self = []; # the output queue
    bless $self, $class;
}

sub get {
    my $self = shift;
    my $fh = shift;

    my @ret;
    while (1) {
        my $from = recv($fh, my $buffer = '', 4096, 0 );
        last if !$from;
        push @ret, [ $from, $buffer ];
    }
    return if !@ret;
    return \@ret;
}

sub put {
    my $self = shift;
    my $data = shift;

    push @$self, @$data;
    my $sum = 0;
    $sum += length( $_->[1] ) for @$self;
    return $sum;
}

sub flush {
    my $self = shift;
    my $fh = shift;

    while ( @$self ) {
        my $n = send($fh, $self->[0][1], 0, $self->[0][0])
            or return;
        $n == length($self->[0][1])
            or die "Couldn't write complete message to socket: $!\n";
        shift @$self;
    }
}

package DNS::Filter::UDPDNS;

use strict;
use POE::Filter;
use Socket;
use Net::DNS::Packet;

sub new {
    my $class = shift;
    bless {}, $class;
}

sub get {
    my $self = shift;
    my $data = shift;

    my @ret;
    for my $d ( @$data ) {
        ref($d) eq "ARRAY"
            or die "UDPDNS filter expected arrayrefs for input\n";
        my($port, $inet) = sockaddr_in($d->[0]);
        my $inetstr = inet_ntoa($inet);
        my($p, $err) = Net::DNS::Packet->new(\$d->[1]);
        if ( !$p ) {
            warn "Cannot create DNS question for packet received from " .
                "$inetstr: $err\n";
        } else {
            $p->answerfrom("$inetstr:$port");
            push @ret, $p;
        }
    }
    return \@ret;
}

sub put {
    my $self = shift;
    my $data = shift;

    my @ret;
    for my $d ( @$data ) {
        my($inetstr, $port) = split /:/, $d->answerfrom();
        if ( !defined $port ) {
            warn "answerfrom not set in DNS packet, no destination known\n";
        } else {
            push @ret,
                [ pack_sockaddr_in($port, inet_aton($inetstr)), $d->data ];
        }
    }
    return \@ret;
}

1;
__END__