Fix #recvfrom to match stdlib

lib/rex/socket/udp.rb

@@ -117,37 +117,60 @@ def sendto(gram, peerhost, peerport, flags = 0)
   end
 
   #
-  # Receives a datagram and returns the data and host:port of the requestor
-  # as [ data, host, port ].
+  # Receives a datagram and returns the data and sender address information as
+  # [ data, [address_family, port, host, host] ], matching stdlib
+  # UDPSocket#recvfrom. Like the stdlib method, this blocks until a datagram is
+  # available and has no timeout of its own (see #timed_recvfrom for a variant
+  # that does). The host appears in both the hostname and numeric address
+  # positions; no reverse-DNS lookup is performed.
+  #
+  # @param maxlen [Integer] maximum number of bytes to receive
+  # @param flags [Integer] flags passed to the underlying recvfrom(2) call (default: 0)
+  # @return [Array(String, Array)] the datagram and the sender address information
+  #
+  def recvfrom(maxlen, flags = 0)
+    # Block until the socket is readable to mirror the stdlib's blocking
+    # UDPSocket#recvfrom; a nil timeout waits indefinitely.
+    ::IO.select([ fd ], nil, nil, nil)
+    data, saddr = recvfrom_nonblock(maxlen, flags)
+    [ data, sender_addr_info(saddr) ]
+  end
+
+  #
+  # Receives a datagram like #recvfrom but waits at most +timeout+ seconds for
+  # one to arrive, returning nil if the timeout elapses first. The return value
+  # otherwise matches #recvfrom: [ data, [address_family, port, host, host] ].
+  #
+  # @param maxlen [Integer] maximum number of bytes to receive
+  # @param timeout [Numeric] seconds to wait for a datagram before giving up
+  # @return [Array(String, Array), nil] the datagram and sender address
+  #   information, or nil if no datagram arrived within +timeout+ seconds
   #
-  def recvfrom(length = 65535, timeout=def_read_timeout)
+  def timed_recvfrom(maxlen = 65535, timeout = def_read_timeout)
+    return nil unless ::IO.select([ fd ], nil, nil, timeout)
 
-    begin
-      if ((rv = ::IO.select([ fd ], nil, nil, timeout)) and
-          (rv[0]) and (rv[0][0] == fd)
-         )
-          data, saddr    = recvfrom_nonblock(length)
-          af, host, port = Rex::Socket.from_sockaddr(saddr)
+    data, saddr = recvfrom_nonblock(maxlen)
+    [ data, sender_addr_info(saddr) ]
+  rescue ::Timeout::Error, ::Errno::ECONNREFUSED
+    nil
+  end
 
-          return [ data, host, port ]
-      else
-        return [ '', nil, nil ]
-      end
-    rescue ::Timeout::Error
-      return [ '', nil, nil ]
-    rescue ::Interrupt
-      raise $!
-    rescue ::Exception
-      return [ '', nil, nil ]
-    end
+  #
+  # Converts a packed sockaddr into the stdlib UDPSocket#recvfrom-style sender
+  # address tuple [ address_family, port, host, host ].
+  #
+  def sender_addr_info(saddr)
+    af, host, port = Rex::Socket.from_sockaddr(saddr)
+    af_name = ::Socket.constants.grep(/^AF_/).find { |c| ::Socket.const_get(c) == af }.to_s
+    [ af_name, port, host, host ]
   end
+  private :sender_addr_info
 
   #
-  # Calls recvfrom and only returns the data
+  # Calls #timed_read and returns the data
   #
   def get(timeout=nil)
-    data, saddr, sport = recvfrom(65535, timeout)
-    return data
+    timed_read(65535, timeout)
   end
 
   #

spec/rex/socket/udp_spec.rb

@@ -0,0 +1,61 @@
+# -*- coding:binary -*-
+require 'spec_helper'
+
+RSpec.describe Rex::Socket::Udp do
+  let(:loopback) { '127.0.0.1' }
+
+  def make_server
+    described_class.create('LocalHost' => loopback, 'LocalPort' => 0)
+  end
+
+  def make_client(port)
+    described_class.create('PeerHost' => loopback, 'PeerPort' => port)
+  end
+
+  describe '#recvfrom' do
+    it 'returns the data and a stdlib-style sender address tuple' do
+      server = make_server
+      client = make_client(server.local_address.ip_port)
+      client.write('hello')
+
+      data, addr = server.recvfrom(65535)
+      expect(data).to eq('hello')
+      expect(addr).to be_an(Array)
+      expect(addr.length).to eq(4)
+
+      af_name, port, host, numeric = addr
+      expect(af_name).to eq('AF_INET')
+      expect(port).to be_a(Integer)
+      expect(host).to eq(loopback)
+      expect(numeric).to eq(loopback)
+    ensure
+      client&.close
+      server&.close
+    end
+  end
+
+  describe '#timed_recvfrom' do
+    it 'returns the datagram and sender address when one arrives in time' do
+      server = make_server
+      client = make_client(server.local_address.ip_port)
+      client.write('ping')
+
+      result = server.timed_recvfrom(65535, 5)
+      expect(result).to_not be_nil
+
+      data, addr = result
+      expect(data).to eq('ping')
+      expect(addr).to eq(['AF_INET', addr[1], loopback, loopback])
+    ensure
+      client&.close
+      server&.close
+    end
+
+    it 'returns nil when no datagram arrives before the timeout' do
+      server = make_server
+      expect(server.timed_recvfrom(65535, 0.1)).to be_nil
+    ensure
+      server&.close
+    end
+  end
+end