diff --git a/lib/rex/socket/udp.rb b/lib/rex/socket/udp.rb index 77bcb97..59b232c 100644 --- a/lib/rex/socket/udp.rb +++ b/lib/rex/socket/udp.rb @@ -135,37 +135,60 @@ def send(mesg, flags, host = nil, port = nil) 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 index d4b8416..2aa6bf3 100644 --- a/spec/rex/socket/udp_spec.rb +++ b/spec/rex/socket/udp_spec.rb @@ -1,14 +1,40 @@ -# -*- coding: binary -*- +# -*- coding:binary -*- require 'spec_helper' RSpec.describe Rex::Socket::Udp do - let(:receiver) { UDPSocket.new.tap { |s| s.bind('127.0.0.1', 0) } } + let(:loopback) { '127.0.0.1' } + let(:receiver) { UDPSocket.new.tap { |s| s.bind(loopback, 0) } } let(:recv_port) { receiver.addr[1] } - let(:socket) { described_class.create('LocalHost' => '127.0.0.1') } + let(:socket) { described_class.create('LocalHost' => loopback) } - after do - socket.close rescue nil - receiver.close rescue nil + 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 '#send' do @@ -56,4 +82,29 @@ socket.sendto('data', '127.1.1.1', 1337) 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