diff --git a/lib/rex/socket/udp.rb b/lib/rex/socket/udp.rb index 58ba2c0..9aa81aa 100644 --- a/lib/rex/socket/udp.rb +++ b/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 # diff --git a/spec/rex/socket/udp_spec.rb b/spec/rex/socket/udp_spec.rb new file mode 100644 index 0000000..d56057c --- /dev/null +++ b/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