resolver.rb in rubydns-0.5.3

- old
+ new

@@ -1,6 +1,6 @@
-# Copyright (c) 2012 Samuel G. D. Williams. <http://www.oriontransfer.co.nz>
+# Copyright, 2012, by Samuel G. D. Williams. <http://www.codeotaku.com>
 # 
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
 # in the Software without restriction, including without limitation the rights
 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
@@ -109,16 +109,26 @@
 			def run!
 				try_next_server!
 			end
 			
 			def process_response!(response)
-				if response.tc != 0
+				if Exception === response
+					@logger.warn "[#{@message.id}] Failure while processing response #{exception}!" 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
-					# We hardcode this behaviour for now.
+					
 					try_next_server!
+				elsif response.id != @message.id
+					@logger.warn "[#{@message.id}] Received response with incorrect message id: #{response.id}" if @request.logger
+					
+					try_next_server!
 				else
 					@logger.warn "[#{@message.id}] Received valid response #{response.inspect}" if @logger
+					
 					succeed response
 				end
 			end
 			
 			private
@@ -173,17 +183,15 @@
 				end
 				
 				def receive_data(data)
 					# Receiving response from remote DNS server...
 					message = RubyDNS::decode_message(data)
-					
+						
 					# The message id must match, and it can't be truncated:
-					if message.id == @request.message.id
-						@request.process_response!(message)
-					else
-						@request.logger.warn "[#{@request.message.id}] Received response with incorrect message id: #{message.id}" if @request.logger
-					end
+					@request.process_response!(message)
+				rescue Resolv::DNS::DecodeError => error
+					@request.process_response!(error)
 				end
 			end
 			
 			module TCPRequestHandler
 				def self.open(host, port, request)
@@ -212,22 +220,19 @@
 						if @buffer.size > 2
 							@length = @buffer.string.byteslice(0, 2).unpack('n')[0]
 						end
 					end
 
-					if @buffer.size == (@length + 2)
+					# If we have received more data than expected, should this be an error?
+					if @buffer.size >= (@length + 2)
 						data = @buffer.string.byteslice(2, @length)
 						
 						message = RubyDNS::decode_message(data)
 						
-						if message.id == @request.message.id
-							@request.process_response!(message)
-						else
-							@request.logger.warn "[#{@request.message.id}] Received response with incorrect message id: #{message.id}" if @request.logger
-						end
-					elsif @buffer.size > (@length + 2)
-						@request.try_next_server!
+						@request.process_response!(message)
 					end
+				rescue Resolv::DNS::DecodeError => error
+					@request.process_response!(error)
 				end
 			end
 		end
 	end
 end
\ No newline at end of file