File Coverage

blib/lib/POE/Component/Client/Ping.pm

Criterion	Covered	Total	%
statement	206	233	88.4
branch	65	98	66.3
condition	22	33	66.6
subroutine	23	25	92.0
pod	0	7	0.0
total	316	396	79.8

# License and documentation are after __END__.
# vim: ts=2 sw=2 expandtab

package POE::Component::Client::Ping;

use warnings;
use strict;

use Exporter;
use vars qw(@ISA @EXPORT_OK %EXPORT_TAGS);

@ISA = qw(Exporter);
@EXPORT_OK = qw(
  REQ_ADDRESS REQ_TIMEOUT REQ_TIME REQ_USER_ARGS
  RES_ADDRESS RES_ROUNDTRIP RES_TIME RES_TTL
);
%EXPORT_TAGS = (
  const => [
    qw(
      REQ_ADDRESS REQ_TIMEOUT REQ_TIME REQ_USER_ARGS
      RES_ADDRESS RES_ROUNDTRIP RES_TIME RES_TTL
    )
  ]
);

use vars qw($VERSION $PKTSIZE);
$VERSION = '1.176';
$PKTSIZE = $^O eq 'linux' ? 3_000 : 100;

use Carp qw(croak);
use Symbol qw(gensym);
use Socket;
use Time::HiRes qw(time);

use POE::Session;

sub DEBUG        () { 0 } # Enable more information.
sub DEBUG_SOCKET () { 0 } # Watch the socket open and close.
sub DEBUG_PBS    () { 0 } # Watch ping_by_seq management.

# ping_by_seq structure offsets.

sub PBS_POSTBACK     () { 0 };
sub PBS_SESSION      () { 1 };
sub PBS_ADDRESS      () { 2 };
sub PBS_REQUEST_TIME () { 3 };

# request_packet offsets
sub REQ_ADDRESS       () { 0 };
sub REQ_TIMEOUT       () { 1 };
sub REQ_TIME          () { 2 };
sub REQ_USER_ARGS     () { 3 };

# response_packet offsets
sub RES_ADDRESS       () { 0 };
sub RES_ROUNDTRIP     () { 1 };
sub RES_TIME          () { 2 };
sub RES_TTL           () { 3 };

# ICMP echo constants. Types, structures, and fields.  Cribbed
# mercilessly from Net::Ping.

sub ICMP_ECHOREPLY () { 0 }
sub ICMP_ECHO      () { 8 }
sub ICMP_STRUCT    () { 'C2 S3 A' }
sub ICMP_SUBCODE   () { 0 }
sub ICMP_FLAGS     () { 0 }
sub ICMP_PORT      () { 0 }

# "Static" variables which will be shared across multiple instances.

my $master_seq = 0;

# Spawn a new PoCo::Client::Ping session.  This basically is a
# constructor, but it isn't named "new" because it doesn't create a
# usable object.  Instead, it spawns the object off as a session.
# Randal Schwartz gave me heck about calling spawny things "new", so I
# blame him for this naming convention.

sub spawn {
  my $type = shift;

  croak "$type requires an even number of parameters" if @_ % 2;
  my %params = @_;

  croak "$type requires root privilege" unless can_open_socket();

  my $alias         = delete $params{Alias};
  $alias            = "pinger" unless defined $alias and length $alias;

  my $timeout       = delete $params{Timeout};
  $timeout          = 1 unless defined $timeout and $timeout >= 0;

  my $onereply      = delete $params{OneReply};
  my $socket        = delete $params{Socket};
  my $parallelism   = delete $params{Parallelism} || -1;
  my $rcvbuf        = delete $params{BufferSize};
  my $always_decode = delete $params{AlwaysDecodeAddress};
  my $retry         = delete $params{Retry};
  my $payload       = delete $params{Payload};

  # 56 data bytes :)
  $payload = 'Use POE!' x 7 unless defined $payload;

  croak(
    "$type doesn't know these parameters: ", join(', ', sort keys %params)
  ) if scalar keys %params;

  POE::Session->create(
    inline_states => {
      _start   => \&poco_ping_start,
      ping     => \&poco_ping_ping,
      clear    => \&poco_ping_clear,
      got_pong => \&poco_ping_pong,
      _default => \&poco_ping_default,
    },
    heap => {
      alias         => $alias,
      always_decode => $always_decode,
      data          => $payload,
      data_size     => length($payload),
      keep_socket   => (defined $socket) || 0,
      onereply      => $onereply,
      rcvbuf        => $rcvbuf,
      retry         => $retry || 0,
      socket_handle => $socket,
      timeout       => $timeout,

      # Active query tracking.
      ping_by_seq   => { },  # keyed on sequence number
      addr_to_seq   => { },  # keyed on request address, then sender

      # Queue to manage throttling.
      parallelism   => $parallelism, # how many pings can we send at once
      queue         => [ ], # ordered list of throttled pings
      pending       => { }, # data for the sequence ids of queued pings
      outstanding   => 0,   # How many pings are we awaiting replies for
    },
  );

  undef;
}


# Start the pinger session.

sub poco_ping_start {
  $_[KERNEL]->alias_set( $_[HEAP]->{alias} );
}


# (NOT A POE EVENT HANDLER)
# Test whether this process can open raw sockets.
sub can_open_socket {
  my $socket = eval { _create_handle() };
  return 0 if $@ or not $socket;
  return 1;
}


# (NOT A POE EVENT HANDLER)
# Create a raw socket to send ICMP packets down.
# (optionally) mess with the size of the buffers on the socket.

sub _create_handle {
  DEBUG_SOCKET and warn "opening a raw socket for icmp";

  my $protocol = Socket::IPPROTO_ICMP;

  my $socket = gensym();
  socket($socket, PF_INET, SOCK_RAW, $protocol)
    or die "can't create icmp socket: $!";

  return $socket;
}

### NOT A POE EVENT HANDLER

sub _setup_handle {
  my ($kernel, $heap) = @_;

  if ($heap->{rcvbuf}) {
    unless (
      setsockopt(
        $heap->{socket_handle}, SOL_SOCKET,
        SO_RCVBUF, pack("I", $heap->{rcvbuf})
      )
    ) {
        warn("setsockopt rcvbuf size ($heap->{rcvbuf}) failed: $!");
    }
  }

  if ($heap->{parallelism} && $heap->{parallelism} == -1) {
    my $rcvbuf = getsockopt($heap->{socket_handle}, SOL_SOCKET, SO_RCVBUF);
    if ($rcvbuf) {
      my $size = unpack("I", $rcvbuf);
      my $max_parallel = int($size / $PKTSIZE);
      if ($max_parallel > 8) {
        $max_parallel -= 8;
      }
      elsif ($max_parallel < 1) {
        $max_parallel = 1;
      }
      $heap->{parallelism} = $max_parallel;
    }
  }

  $kernel->select_read($heap->{socket_handle}, 'got_pong');
}

# Request a ping.  This code borrows heavily from Net::Ping.

sub poco_ping_ping {
  my (
    $kernel, $heap, $sender,
    $event, $address, $timeout, $tries_left
  ) = @_[
    KERNEL, HEAP, SENDER,
    ARG0, ARG1, ARG2, ARG3
  ];

  $tries_left ||= $heap->{retry};

  DEBUG and warn "ping requested for $address ($tries_left try/tries left)\n";

  _do_ping($kernel, $heap, $sender, $event, $address, $timeout, $tries_left);
}


sub _do_ping {
  my ($kernel, $heap, $sender, $event, $address, $timeout, $tries_left) = @_;

  # No current pings.  Open a socket, or setup the existing one.
  unless (scalar(keys %{$heap->{ping_by_seq}})) {
    unless (defined $heap->{socket_handle}) {
      $heap->{socket_handle} = _create_handle();
    }
    _setup_handle($kernel, $heap);
  }

  # Get the timeout, or default to the one set for the component.
  $timeout = $heap->{timeout} unless defined $timeout and $timeout > 0;
  $tries_left = $heap->{retry} unless defined $tries_left;

  # Find an unused sequence number.
  while (1) {
    $master_seq = ($master_seq + 1) & 0xFFFF;
    last unless exists $heap->{ping_by_seq}->{$master_seq};
  }

  my $checksum = 0;

  # Build the message without a checksum.
  my $msg = pack(
    ICMP_STRUCT . $heap->{data_size},
    ICMP_ECHO, ICMP_SUBCODE,
    $checksum, ($$ & 0xFFFF), $master_seq, $heap->{data}
  );

  ### Begin checksum calculation section.

  # Sum up short integers in the packet.
  my $shorts = int(length($msg) / 2);
  foreach my $short (unpack "S$shorts", $msg) {
    $checksum += $short;
  }

  # If there's an odd byte, add that in as well.
  $checksum += ord(substr($msg, -1)) if length($msg) % 2;

  # Fold the high short into the low one twice, and then complement.
  $checksum = ($checksum >> 16) + ($checksum & 0xFFFF);
  $checksum = ~( ($checksum >> 16) + $checksum) & 0xFFFF;

  ### Cease checksum calculation section.

  # Rebuild the message with the checksum this time.
  $msg = pack(
    ICMP_STRUCT . $heap->{data_size},
    ICMP_ECHO, ICMP_SUBCODE, $checksum, ($$ & 0xFFFF), $master_seq,
    $heap->{data}
  );

  # Record information about the ping request.

  my ($event_name, @user_args);
  if (ref($event) eq "ARRAY") {
    ($event_name, @user_args) = @$event;
  }
  else {
    $event_name = $event;
  }

  # Build an address to send the ping at.
  # TODO - This blocks, so resolve them first.
  # TODO - This assumes four-octet addresses are IPv4.

  my $usable_address = $address;
  if ($heap->{always_decode} || length($address) != 4) {
    $usable_address = inet_aton($address);
  }

  # Return failure if an address was not resolvable.  This simulates
  # the postback behavior.

  unless (defined $usable_address) {
    $kernel->post(
      $sender, $event_name,
      [
        $address,   # REQ_ADDRESS
        $timeout,   # REQ_TIMEOUT
        time(),     # REQ_TIME
        @user_args, # REQ_USER_ARGS
        ],
        [
        undef,      # RES_ADDRESS
        undef,      # RES_ROUNDTRIP
        time(),     # RES_TIME
        undef,      # RES_TTL
      ],
    );
    _check_for_close($kernel, $heap);
    return;
  }

  my $socket_address = pack_sockaddr_in(ICMP_PORT, $usable_address);

  push(@{$heap->{queue}}, $master_seq);
  $heap->{pending}->{$master_seq} = [
    $msg,               # PEND_MSG
    $socket_address,    # PEND_ADDR
    $sender,            # PEND_SENDER
    $event,             # PEND_EVENT
    $address,           # PEND_ADDR ???
    $timeout,           # PEND_TIMEOUT
  ];

  if ($tries_left and $tries_left > 1) {
    $heap->{retrydata}->{$master_seq} = [
      $sender,     # RD_SENDER
      $event,      # RD_EVENT
      $address,    # RD_ADDRESS
      $timeout,    # RD_TIMEOUT
      $tries_left, # RD_RETRY
    ];
  }

  _send_next_packet($kernel, $heap);
}


sub _send_next_packet {
  my ($kernel, $heap) = @_;
  return unless (scalar @{$heap->{queue}});

  if ($heap->{parallelism} && $heap->{outstanding} >= $heap->{parallelism}) {
    # We want to throttle back since we're still waiting for pings
    # so, let's just leave this till later
    DEBUG and warn(
      "throttled since there are $heap->{outstanding} pings outstanding. " .
      "queue size=" . (scalar @{$heap->{queue}}) . "\n"
    );
    return;
  }

  my $seq = shift(@{$heap->{queue}});

  # May have been cleared by caller
  return unless (exists $heap->{pending}->{$seq});

  my $ping_info = delete $heap->{pending}->{$seq};
  my (
    $msg,               # PEND_MSG
    $socket_address,    # PEND_ADDR
    $sender,            # PEND_SENDER
    $event,             # PEND_EVENT
    $address,           # PEND_ADDR ???
    $timeout,           # PEND_TIMEOUT
  ) = @$ping_info;

  # Send the packet.  If send() fails, then we bail with an error.
  my @user_args = ();
  ($event, @user_args) = @$event if ref($event) eq "ARRAY";

  DEBUG and warn "sending packet sequence number $seq\n";
  unless (send($heap->{socket_handle}, $msg, ICMP_FLAGS, $socket_address)) {
    $kernel->post(
      $sender, $event,
      [ $address,    # REQ_ADDRESS
        $timeout,    # REQ_TIMEOUT
        time(),      # REQ_TIME
        @user_args,  # REQ_USER_ARGS
      ],
      [ undef,   # RES_ADDRESS
        undef,   # RES_ROUNDTRIP
        time(),  # RES_TIME
        undef,   # RES_TTL
      ],
    );
    _check_for_close($kernel, $heap);
    return;
  }

  # Record the message's length.  This is constant, but we do it here
  # anyway.  It's also used to flag when we start requesting replies.
  $heap->{message_length} = length($msg);
  $heap->{outstanding}++;

  # Set a timeout based on the sequence number.
  $kernel->delay( $seq => $timeout );

  DEBUG_PBS and warn "recording ping_by_seq($seq)";
  $heap->{ping_by_seq}->{$seq} = [
    # PBS_POSTBACK
    $sender->postback(
      $event,
      $address,    # REQ_ADDRESS
      $timeout,    # REQ_TIMEOUT
      time(),      # REQ_TIME
      @user_args,  # REQ_USER_ARGS
    ),
    "$sender",   # PBS_SESSION (stringified to weaken reference)
    $address,    # PBS_ADDRESS
    time()       # PBS_REQUEST_TIME
  ];

  # Duplicate pings?  Forcibly time out the previous one.
  if (exists $heap->{addr_to_seq}->{$sender}->{$address}) {
    DEBUG and warn "Duplicate ping. Canceling $address";

    my $now = time();
    my $ping_info = _end_ping_by_requester_and_address(
      $kernel, $heap, $sender, $address
    );

    $ping_info->[PBS_POSTBACK]->( undef, undef, $now, undef );
  }

  $heap->{addr_to_seq}->{$sender}->{$address} = $seq;
}

# Clear a ping postback by address.  The sender+address pair are a
# unique ID into the pinger's data.

sub poco_ping_clear {
  my ($kernel, $heap, $sender, $address) = @_[KERNEL, HEAP, SENDER, ARG0];

  # Is the sender still waiting for anything?
  return unless exists $heap->{addr_to_seq}->{$sender};

  # Try to clear a single ping if an address was specified.
  if (defined $address) {
    _end_ping_by_requester_and_address($kernel, $heap, $sender, $address);
  }

  # No address was specified.  Clear all the pings for this session.
  else {
    _end_pings_by_requester($kernel, $heap, $sender);
  }

  _check_for_close($kernel, $heap);
}

# (NOT A POE EVENT HANDLER)
# Check to see if no more pings are waiting.  Close the socket if so.

sub _check_for_close {
  my ($kernel, $heap) = @_;

  return unless exists $heap->{socket_handle};

  return if scalar keys %{$heap->{ping_by_seq}};

  DEBUG_SOCKET and warn "stopping raw socket watcher";
  $kernel->select_read( $heap->{socket_handle} );

  return if $heap->{keep_socket};

  DEBUG_SOCKET and warn "closing raw socket";
  delete $heap->{socket_handle};
}

# (NOT A POE EVENT HANDLER)
# Clean up after we're done with a ping.
# Remove it from all tracking hashes.
# Determine if the socket should be unthrottled or shut down.

sub _end_ping_by_sequence {
  my ($kernel, $heap, $seq) = @_;

  # Delete the ping information.  Cache a copy for other cleanup.
  DEBUG_PBS and warn "removing ping by sequence ($seq)";
  my $ping_info = delete $heap->{ping_by_seq}->{$seq};
  return unless $ping_info;

  # Stop its associated timeout.
  $kernel->delay($seq);

  # Stop mapping the session+address to this sequence number.
  my $pbs_session = $ping_info->[PBS_SESSION];
  delete $heap->{addr_to_seq}->{$pbs_session}->{$ping_info->[PBS_ADDRESS]};

  # Stop tracking the session if that was its last address.
  delete $heap->{addr_to_seq}->{$pbs_session} unless (
    scalar(keys %{$heap->{addr_to_seq}->{$pbs_session}})
  );

  $heap->{outstanding}--;

  return $ping_info;
}


sub _end_ping_by_requester_and_address {
  my ($kernel, $heap, $sender, $address) = @_;

  return unless exists $heap->{addr_to_seq}->{$sender};
  my $addr_to_seq_rec = $heap->{addr_to_seq}->{$sender};

  my $seq = delete $addr_to_seq_rec->{$address};
  unless ($seq) {
    # TODO - Why?
    delete $heap->{pending}->{$sender}->{$address};
    return;
  }

  # Stop tracking the sender if that was the last address.
  delete $heap->{addr_to_seq}->{$sender} unless scalar(
    keys %{$heap->{addr_to_seq}->{$sender}}
  );

  # Discard the postback for the discarded sequence number.
  DEBUG_PBS and warn "removing ping_by_seq($seq)";
  my $ping_info = delete $heap->{ping_by_seq}->{$seq};
  $kernel->delay($seq);

  $heap->{outstanding}--;

  return $ping_info;
}


sub _end_pings_by_requester {
  my ($kernel, $heap, $sender) = @_;

  return unless exists $heap->{addr_to_seq}->{$sender};
  my $addr_to_seq_rec = delete $heap->{addr_to_seq}->{$sender};

  # Discard cross references.

  foreach my $seq (values %$addr_to_seq_rec) {
    DEBUG_PBS and warn "removing ping_by_seq($seq)";
    delete $heap->{ping_by_seq}->{$seq};
    $kernel->delay($seq);

    $heap->{outstanding}--;
  }

  return;
}



# Something has arrived.  Try to match it against something being
# waited for.

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

  # Record the receive time for possible use later.
  my $now = time();

  # Receive a message on the ICMP port.
  my $recv_message = '';
  my $from_saddr = recv($socket, $recv_message, 1500, ICMP_FLAGS);
  return unless $from_saddr;

  # We haven't yet sent a message, so don't bother with whatever we've
  # received.
  return unless defined $heap->{message_length};

  # Unpack the packet's sender address.
  my ($from_port, $from_ip) = unpack_sockaddr_in($from_saddr);

  # Get the response packet's time to live.
  my ($ihl, $from_ttl) = unpack('C1@7C1', $recv_message);
  $ihl &= 0x0F;

  # Unpack the packet itself.
  my (
    $from_type, $from_subcode,
    $from_checksum, $from_pid, $from_seq, $from_message
  )  = unpack(
    '@'.$ihl*4 . ICMP_STRUCT.$heap->{data_size}, $recv_message
  );

  DEBUG and do {
    warn ",----- packet from ", inet_ntoa($from_ip), ", port $from_port\n";
    warn "| type = $from_type / subtype = $from_subcode\n";
    warn "| checksum = $from_checksum, pid = $from_pid, seq = $from_seq\n";
    warn "| message: ", unpack("H*", $from_message), "\n";
    warn "`------------------------------------------------------------\n";
  };

  # Not an ICMP echo reply.  Move along.
  return unless $from_type == ICMP_ECHOREPLY;

  DEBUG and warn "it's an ICMP echo reply";

  # Not from this process.  Move along.
  return unless $from_pid == ($$ & 0xFFFF);

  DEBUG and warn "it's from the current process";

  # Not waiting for a response with that sequence number.  Move along.
  return unless exists $heap->{ping_by_seq}->{$from_seq};

  DEBUG and warn "it's one we're waiting for ($from_seq)";

  # This is the response we're looking for.  Calculate the round trip
  # time, and map it to a postback.
  my $trip_time = $now - $heap->{ping_by_seq}->{$from_seq}->[PBS_REQUEST_TIME];
  $heap->{ping_by_seq}->{$from_seq}->[PBS_POSTBACK]->(
    inet_ntoa($from_ip), $trip_time, $now, $from_ttl
  );

  # It's a single-reply ping.  Clean up after it.
  if ($heap->{onereply}) {
    _end_ping_by_sequence($kernel, $heap, $from_seq);
    _send_next_packet($kernel, $heap);
    _check_for_close($kernel, $heap);
  }
}

# Default's used to catch ping timeouts, which are named after the
# packed socket addresses being pinged.  We always send the timeout so
# the other session knows that a ping period has ended.

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

  # Record the receive time for possible use later.
  my $now = time();

  # Are we waiting for this sequence number?  We should be!
  unless (exists $heap->{ping_by_seq}->{$seq}) {
    warn "this shouldn't technically be displayed ($seq)" if (
      DEBUG and $seq =~ /^\d+$/
    );
    return;
  }

  my $ping_info = _end_ping_by_sequence($kernel, $heap, $seq);

  my $retryinfo = delete $heap->{retrydata}->{$seq};
  if ($retryinfo) {
    my ($sender, $event, $address, $timeout, $remaining) = @{$retryinfo};
    DEBUG and warn("retrying ping for $address (", $remaining - 1, ")\n");
    _do_ping(
      $kernel, $heap, $sender, $event, $address, $timeout,
      $remaining - 1
    );
    return;
  }

  # Post a timer tick back to the session.  This marks the end of
  # the request/response transaction.
  $ping_info->[PBS_POSTBACK]->( undef, undef, $now, undef );
  _send_next_packet($kernel, $heap);
  _check_for_close($kernel, $heap);

  return;
}

1;

__END__