rubydns 0.8.5 → 0.9.0

data/lib/rubydns/resolver.rb

@@ -18,19 +18,24 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-require_relative 'message'
-require_relative 'binary_string'
+require_relative 'handler'
 
 require 'securerandom'
+require 'celluloid/io'
 
 module RubyDNS
 	class InvalidProtocolError < StandardError
 	end
 	
+	class InvalidResponseError < StandardError
+	end
+	
 	class ResolutionFailure < StandardError
 	end
 	
 	class Resolver
+		include Celluloid::IO
+		
 		# Servers are specified in the same manor as options[:listen], e.g.
 		#   [:tcp/:udp, address, port]
 		# In the case of multiple servers, they will be checked in sequence.
@@ -38,6 +43,8 @@ module RubyDNS
 			@servers = servers
 			
 			@options = options
+			
+			@logger = options[:logger] || Celluloid.logger
 		end
 
 		# Provides the next sequence identification number which is used to keep track of DNS messages.
@@ -47,54 +54,138 @@ module RubyDNS
 		end
 
 		# Look up a named resource of the given resource_class.
-		def query(name, resource_class = Resolv::DNS::Resource::IN::A, &block)
+		def query(name, resource_class = Resolv::DNS::Resource::IN::A)
 			message = Resolv::DNS::Message.new(next_id!)
 			message.rd = 1
 			message.add_question name, resource_class
 			
-			send_message(message, &block)
+			dispatch_request(message)
 		end
-
-		def send_message(message, &block)
-			Request.fetch(message, @servers, @options, &block)
-		end
-
-		# Yields a list of `Resolv::IPv4` and `Resolv::IPv6` addresses for the given `name` and `resource_class`.
-		def addresses_for(name, resource_class = Resolv::DNS::Resource::IN::A, &block)
-			query(name, resource_class) do |response|
-				# Resolv::DNS::Name doesn't retain the trailing dot.
-				name = name.sub(/\.$/, '')
-				
-				case response
-				when Message
-					yield response.answer.select{|record| record[0].to_s == name}.collect{|record| record[2].address}
-				else
-					yield []
+		
+		# Yields a list of `Resolv::IPv4` and `Resolv::IPv6` addresses for the given `name` and `resource_class`. Raises a ResolutionFailure if no severs respond.
+		def addresses_for(name, resource_class = Resolv::DNS::Resource::IN::A, options = {})
+			(options[:retries] || 5).times do
+				response = query(name, resource_class)
+				
+				if response
+					# Resolv::DNS::Name doesn't retain the trailing dot.
+					name = name.sub(/\.$/, '')
+					
+					return response.answer.select{|record| record[0].to_s == name}.collect{|record| record[2].address}
 				end
+				
+				# Wait 10ms before trying again:
+				sleep 0.01
 			end
+			
+			abort ResolutionFailure.new("No server replied.")
 		end
-
-		# Manages a single DNS question message across one or more servers.
-		class Request
-			include EventMachine::Deferrable
+		
+		def request_timeout
+			@options[:timeout] || 1
+		end
+		
+		# Send the message to available servers. If no servers respond correctly, nil is returned. This result indicates a failure of the resolver to correctly contact any server and get a valid response.
+		def dispatch_request(message)
+			request = Request.new(message, @servers)
 			
-			def self.fetch(*args)
-				request = self.new(*args)
-				
-				request.callback do |message|
-					yield message
-				end
+			request.each do |server|
+				@logger.debug "[#{message.id}] Sending request to server #{server.inspect}" if @logger
 				
-				request.errback do |error|
-					# In the case of a timeout, error will be nil, so we make one up.
+				begin
+					response = nil
+					
+					timeout(request_timeout) do
+						response = try_server(request, server)
+					end
 					
-					yield error
+					if valid_response(message, response)
+						return response
+					end
+				rescue Task::TimeoutError
+					@logger.debug "[#{message.id}] Request timed out!" if @logger
+				rescue InvalidResponseError
+					@logger.warn "[#{message.id}] Invalid response from network: #{$!}!" if @logger
+				rescue DecodeError
+					@logger.warn "[#{message.id}] Error while decoding data from network: #{$!}!" if @logger
+				rescue IOError
+					@logger.warn "[#{message.id}] Error while reading from network: #{$!}!" if @logger
 				end
-				
-				request.run!
 			end
 			
-			def initialize(message, servers, options = {}, &block)
+			return nil
+		end
+		
+		private
+		
+		def try_server(request, server)
+			case server[0]
+			when :udp
+				try_udp_server(request, server[1], server[2])
+			when :tcp
+				try_tcp_server(request, server[1], server[2])
+			else
+				raise InvalidProtocolError.new(server)
+			end
+		end
+		
+		def valid_response(message, response)
+			if response.tc != 0
+				@logger.warn "[#{message.id}] Received truncated response!" if @logger
+			elsif response.id != message.id
+				@logger.warn "[#{message.id}] Received response with incorrect message id: #{response.id}!" if @logger
+			else
+				@logger.debug "[#{message.id}] Received valid response with #{response.answer.count} answer(s)." if @logger
+		
+				return true
+			end
+			
+			return false
+		end
+		
+		def try_udp_server(request, host, port)
+			socket = UDPSocket.new
+			
+			socket.send(request.packet, 0, host, port)
+			
+			data, (_, remote_port) = socket.recvfrom(UDP_TRUNCATION_SIZE)
+			# Need to check host, otherwise security issue.
+			
+			# May indicate some kind of spoofing attack:
+			if port != remote_port
+				raise InvalidResponseError.new("Data was not received from correct remote port (#{port} != #{remote_port})")
+			end
+			
+			message = RubyDNS::decode_message(data)
+		ensure
+			socket.close if socket
+		end
+		
+		def try_tcp_server(request, host, port)
+			begin
+				socket = TCPSocket.new(host, port)
+			rescue Errno::EALREADY
+				raise IOError.new("Could not connect to remote host!")
+			end
+			
+			StreamTransport.write_chunk(socket, request.packet)
+			
+			input_data = StreamTransport.read_chunk(socket)
+			
+			message = RubyDNS::decode_message(input_data)
+		rescue Errno::ECONNREFUSED => error
+			raise IOError.new(error.message)
+		rescue Errno::EPIPE => error
+			raise IOError.new(error.message)
+		rescue Errno::ECONNRESET => error
+			raise IOError.new(error.message)
+		ensure
+			socket.close if socket
+		end
+		
+		# Manages a single DNS question message across one or more servers.
+		class Request
+			def initialize(message, servers)
 				@message = message
 				@packet = message.encode
 				
@@ -104,167 +195,24 @@ module RubyDNS
 				if @packet.bytesize > UDP_TRUNCATION_SIZE
 					@servers.delete_if{|server| server[0] == :udp}
 				end
-				
-				# Measured in seconds:
-				@timeout = options[:timeout] || 1
-				
-				@logger = options[:logger]
 			end
 			
 			attr :message
 			attr :packet
 			attr :logger
 			
-			def run!
-				try_next_server!
-			end
-			
-			# Once either an exception or message is received, we update the status of this request.
-			def process_response!(response)
-				finish_request!
-				
-				if Exception === response
-					@logger.warn "[#{@message.id}] Failure while processing response #{response}!" if @logger
-					RubyDNS.log_exception(@logger, response) if @logger
-					
-					try_next_server!
-				elsif response.tc != 0
-					@logger.warn "[#{@message.id}] Received truncated response!" if @logger
+			def each(&block)
+				@servers.each do |server|
+					next if @packet.bytesize > UDP_TRUNCATION_SIZE
 					
-					try_next_server!
-				elsif response.id != @message.id
-					@logger.warn "[#{@message.id}] Received response with incorrect message id: #{response.id}" if @logger
-					
-					try_next_server!
-				else
-					@logger.debug "[#{@message.id}] Received valid response #{response.inspect}" if @logger
-					
-					succeed response
-				end
-			end
-			
-			private
-			
-			# Closes any connections and cancels any timeout.
-			def finish_request!
-				cancel_timeout
-				
-				# Cancel an existing request if it is in flight:
-				if @request
-					@request.close_connection
-					@request = nil
-				end
-			end
-			
-			def try_next_server!
-				if @servers.size > 0
-					@server = @servers.shift
-					
-					@logger.debug "[#{@message.id}] Sending request to server #{@server.inspect}" if @logger
-					
-					# We make requests one at a time to the given server, naturally the servers are ordered in terms of priority.
-					case @server[0]
-					when :udp
-						@request = UDPRequestHandler.open(@server[1], @server[2], self)
-					when :tcp
-						@request = TCPRequestHandler.open(@server[1], @server[2], self)
-					else
-						raise InvalidProtocolError.new(@server)
-					end
-					
-					# Setting up the timeout...
-					timeout(@timeout)
-				else
-					fail ResolutionFailure.new("No available servers responded to the request.")
-				end
-			end
-			
-			def timeout seconds
-				cancel_timeout
-				
-				@deferred_timeout = EventMachine::Timer.new(seconds) do
-					@logger.debug "[#{@message.id}] Request timed out!" if @logger
-					
-					finish_request!
-					
-					try_next_server!
-				end
-			end
-			
-			module UDPRequestHandler
-				def self.open(host, port, request)
-					# Open a datagram socket... a random socket chosen by the OS by specifying 0 for the port:
-					EventMachine::open_datagram_socket('', 0, self, request, host, port)
-				end
-				
-				def initialize(request, host, port)
-					@request = request
-					@host = host
-					@port = port
-				end
-				
-				def post_init
-					# Sending question to remote DNS server...
-					send_datagram(@request.packet, @host, @port)
-				end
-				
-				def receive_data(data)
-					# local_port, local_ip = Socket.unpack_sockaddr_in(get_sockname)
-					# puts "Socket name: #{local_ip}:#{local_port}"
-					
-					# Receiving response from remote DNS server...
-					message = RubyDNS::decode_message(data)
-					
-					# The message id must match, and it can't be truncated:
-					@request.process_response!(message)
-				rescue Resolv::DNS::DecodeError => error
-					@request.process_response!(error)
+					yield server
 				end
 			end
-			
-			module TCPRequestHandler
-				def self.open(host, port, request)
-					EventMachine::connect(host, port, TCPRequestHandler, request)
-				end
-				
-				def initialize(request)
-					@request = request
-					@buffer = nil
-					@length = nil
-				end
-				
-				def post_init
-					data = @request.packet
-					
-					send_data([data.bytesize].pack('n'))
-					send_data data
-				end
-				
-				def receive_data(data)
-					# We buffer data until we've received the entire packet:
-					@buffer ||= BinaryStringIO.new
-					@buffer.write(data)
 
-					# If we've received enough data and we haven't figured out the length yet...
-					if @length == nil and @buffer.size > 2
-						# Extract the length from the buffer:
-						@length = @buffer.string.byteslice(0, 2).unpack('n')[0]
-					end
-
-					# If we know what the length is, and we've got that much data, we can decode the message:
-					if @length != nil and @buffer.size >= (@length + 2)
-						data = @buffer.string.byteslice(2, @length)
-						
-						message = RubyDNS::decode_message(data)
-						
-						@request.process_response!(message)
-					end
-					
-					# If we have received more data than expected, should this be an error?
-				rescue Resolv::DNS::DecodeError => error
-					@request.process_response!(error)
-				end
+			def update_id!(id)
+				@message.id = id
+				@packet = @message.encode
 			end
 		end
 	end
-end
+end

data/lib/rubydns/server.rb

@@ -18,14 +18,29 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-require 'fiber'
+require 'celluloid/io'
 
 require_relative 'transaction'
 require_relative 'logger'
 
 module RubyDNS
+	class UDPSocketWrapper < Celluloid::IO::UDPSocket
+		def initialize(socket)
+			@socket = socket
+		end
+	end
+	
+	class TCPServerWrapper < Celluloid::IO::TCPServer
+		def initialize(server)
+			@server = server
+		end
+	end
 	
 	class Server
+		include Celluloid::IO
+		
+		finalizer :shutdown
+		
 		# The default server interfaces
 		DEFAULT_INTERFACES = [[:udp, "0.0.0.0", 53], [:tcp, "0.0.0.0", 53]]
 		
@@ -37,8 +52,11 @@ module RubyDNS
 		#		end
 		#	end
 		#
-		def initialize(options)
-			@logger = options[:logger] || Logger.new($stderr)
+		def initialize(options = {})
+			@handlers = []
+			
+			@logger = options[:logger] || Celluloid.logger
+			@interfaces = options[:listen] || DEFAULT_INTERFACES
 		end
 
 		attr_accessor :logger
@@ -47,56 +65,51 @@ module RubyDNS
 		def fire(event_name)
 		end
 		
+		def shutdown
+			fire(:stop)
+		end
+		
 		# Give a name and a record type, try to match a rule and use it for processing the given arguments.
 		def process(name, resource_class, transaction)
 			raise NotImplementedError.new
 		end
 		
-		# Process a block with the current fiber. To resume processing from the block, call `fiber.resume`. You shouldn't call `fiber.resume` until after the top level block has returned.
-		def defer(&block)
-			fiber = Fiber.current
-			
-			yield(fiber)
-			
-			Fiber.yield
-		end
-		
 		# Process an incoming DNS message. Returns a serialized message to be sent back to the client.
 		def process_query(query, options = {}, &block)
 			start_time = Time.now
 			
-			# Setup answer
-			answer = Resolv::DNS::Message::new(query.id)
-			answer.qr = 1                 # 0 = Query, 1 = Response
-			answer.opcode = query.opcode  # Type of Query; copy from query
-			answer.aa = 1                 # Is this an authoritative response: 0 = No, 1 = Yes
-			answer.rd = query.rd          # Is Recursion Desired, copied from query
-			answer.ra = 0                 # Does name server support recursion: 0 = No, 1 = Yes
-			answer.rcode = 0              # Response code: 0 = No errors
+			# Setup response
+			response = Resolv::DNS::Message::new(query.id)
+			response.qr = 1                 # 0 = Query, 1 = Response
+			response.opcode = query.opcode  # Type of Query; copy from query
+			response.aa = 1                 # Is this an authoritative response: 0 = No, 1 = Yes
+			response.rd = query.rd          # Is Recursion Desired, copied from query
+			response.ra = 0                 # Does name server support recursion: 0 = No, 1 = Yes
+			response.rcode = 0              # Response code: 0 = No errors
 			
-			Fiber.new do
-				transaction = nil
-				
-				begin
-					query.question.each do |question, resource_class|
-						@logger.debug {"[#{query.id}] Processing question #{question} #{resource_class}..."}
-				
-						transaction = Transaction.new(self, query, question, resource_class, answer, options)
-						
-						transaction.process
-					end
-				rescue
-					@logger.error {"[#{query.id}] Exception thrown while processing #{transaction}!"}
-					RubyDNS.log_exception(@logger, $!)
-				
-					answer.rcode = Resolv::DNS::RCode::ServFail
+			transaction = nil
+			
+			begin
+				query.question.each do |question, resource_class|
+					@logger.debug {"<#{query.id}> Processing question #{question} #{resource_class}..."}
+			
+					transaction = Transaction.new(self, query, question, resource_class, response, options)
+					
+					transaction.process
 				end
-				
-				yield answer
-				
-				end_time = Time.now
-				@logger.debug {"[#{query.id}] Time to process request: #{end_time - start_time}s"}
-			end.resume
+			rescue Celluloid::ResumableError
+				raise
+			rescue StandardError => error
+				@logger.error "<#{query.id}> Exception thrown while processing #{transaction}!"
+				RubyDNS.log_exception(@logger, error)
+			
+				response.rcode = Resolv::DNS::RCode::ServFail
+			end
+			
+			end_time = Time.now
+			@logger.debug {"<#{query.id}> Time to process request: #{end_time - start_time}s"}
+			
+			return response
 		end
 		
 		#
@@ -113,41 +126,40 @@ module RubyDNS
 		#   Process::Sys.setuid(server_uid)
 		#   INTERFACES = [socket]
 		#
-		def run(options = {})
+		def run
 			@logger.info "Starting RubyDNS server (v#{RubyDNS::VERSION})..."
-		
-			interfaces = options[:listen] || DEFAULT_INTERFACES
-		
+			
 			fire(:setup)
-		
+			
 			# Setup server sockets
-			interfaces.each do |spec|
+			@interfaces.each do |spec|
 				if spec.is_a?(BasicSocket)
 					spec.do_not_reverse_lookup
-					optval = spec.getsockopt(Socket::SOL_SOCKET, Socket::SO_TYPE)
-					protocol = optval.unpack("i")[0]
+					protocol = spec.getsockopt(Socket::SOL_SOCKET, Socket::SO_TYPE).unpack("i")[0]
 					ip = spec.local_address.ip_address
 					port = spec.local_address.ip_port
+					
 					case protocol
 					when Socket::SOCK_DGRAM
-						@logger.info "Attaching to pre-existing UDP socket #{ip}:#{port}"
-						EventMachine.attach(spec, UDPHandler, self)
+						@logger.info "<> Attaching to pre-existing UDP socket #{ip}:#{port}"
+						link UDPSocketHandler.new(self, UDPSocketWrapper.new(spec))
 					when Socket::SOCK_STREAM
-						@logger.info "Attaching to pre-existing TCP socket #{ip}:#{port}"
-						EventMachine.attach(spec, TCPHandler, self)
+						@logger.info "<> Attaching to pre-existing TCP socket #{ip}:#{port}"
+						link TCPSocketHandler.new(self, TCPServerWrapper.new(spec))
 					else
-						@logger.error "Ignoring unknown socket protocol: #{protocol}"
+						raise ArgumentError.new("Unknown socket protocol: #{protocol}")
 					end
+				elsif spec[0] == :udp
+					@logger.info "<> Listening on #{spec.join(':')}"
+					link UDPHandler.new(self, spec[1], spec[2])
+				elsif spec[0] == :tcp
+					@logger.info "<> Listening on #{spec.join(':')}"
+					link TCPHandler.new(self, spec[1], spec[2])
 				else
-					@logger.info "Listening on #{spec.join(':')}"
-					if spec[0] == :udp
-						EventMachine.open_datagram_socket(spec[1], spec[2], UDPHandler, self)
-					elsif spec[0] == :tcp
-						EventMachine.start_server(spec[1], spec[2], TCPHandler, self)
-					end
+					raise ArgumentError.new("Invalid connection specification: #{spec.inspect}")
 				end
 			end
-		
+			
 			fire(:start)
 		end
 	end
@@ -175,9 +187,9 @@ module RubyDNS
 			end
 			
 			# Invoke the rule, if it matches the incoming request, it is evaluated and returns `true`, otherwise returns `false`.
-			def call(server, name, resource_class, *args)
+			def call(server, name, resource_class, transaction)
 				unless match(name, resource_class)
-					server.logger.debug "Resource class #{resource_class} failed to match #{@pattern[1].inspect}!"
+					server.logger.debug "<#{transaction.query.id}> Resource class #{resource_class} failed to match #{@pattern[1].inspect}!"
 					
 					return false
 				end
@@ -188,31 +200,31 @@ module RubyDNS
 					match_data = @pattern[0].match(name)
 					
 					if match_data
-						server.logger.debug "Regexp pattern matched with #{match_data.inspect}."
+						server.logger.debug "<#{transaction.query.id}> Regexp pattern matched with #{match_data.inspect}."
 						
-						@callback[*args, match_data]
+						@callback[transaction, match_data]
 						
 						return true
 					end
 				when String
 					if @pattern[0] == name
-						server.logger.debug "String pattern matched."
+						server.logger.debug "<#{transaction.query.id}> String pattern matched."
 						
-						@callback[*args]
+						@callback[transaction]
 						
 						return true
 					end
 				else
 					if (@pattern[0].call(name, resource_class) rescue false)
-						server.logger.debug "Callable pattern matched."
+						server.logger.debug "<#{transaction.query.id}> Callable pattern matched."
 						
-						@callback[*args]
+						@callback[transaction]
 						
 						return true
 					end
 				end
 				
-				server.logger.debug "No pattern matched."
+				server.logger.debug "<#{transaction.query.id}> No pattern matched."
 				
 				# We failed to match the pattern.
 				return false
@@ -223,6 +235,9 @@ module RubyDNS
 			end
 		end
 		
+		# Don't wrap the block going into initialize.
+		execute_block_on_receiver :initialize
+		
 		# Instantiate a server with a block
 		#
 		#	server = Server.new do
@@ -289,32 +304,23 @@ module RubyDNS
 		end
 		
 		# Give a name and a record type, try to match a rule and use it for processing the given arguments.
-		def process(name, resource_class, *args)
-			@logger.debug {"Searching for #{name} #{resource_class.name}"}
+		def process(name, resource_class, transaction)
+			@logger.debug {"<#{transaction.query.id}> Searching for #{name} #{resource_class.name}"}
 			
 			@rules.each do |rule|
-				@logger.debug {"Checking rule #{rule}..."}
+				@logger.debug {"<#{transaction.query.id}> Checking rule #{rule}..."}
 				
 				catch (:next) do
 					# If the rule returns true, we assume that it was successful and no further rules need to be evaluated.
-					return if rule.call(self, name, resource_class, *args)
+					return if rule.call(self, name, resource_class, transaction)
 				end
 			end
 			
 			if @otherwise
-				@otherwise.call(*args)
+				@otherwise.call(transaction)
 			else
-				@logger.warn "Failed to handle #{name} #{resource_class.name}!"
+				@logger.warn "<#{transaction.query.id}> Failed to handle #{name} #{resource_class.name}!"
 			end
 		end
-		
-		# Process a block with the current fiber. To resume processing from the block, call `fiber.resume`. You shouldn't call `fiber.resume` until after the top level block has returned.
-		def defer(&block)
-			fiber = Fiber.current
-			
-			yield(fiber)
-			
-			Fiber.yield
-		end
 	end
 end