From 8d961c19384aee16dc088b734b9144adb40fa209 Mon Sep 17 00:00:00 2001
From: soh_cah_toa <soh_cah_toa@cybercrud.net>
Date: Fri, 19 Jul 2013 22:15:25 -0400
Subject: [PATCH] Added support for rules that fail to resolve (e.g. NXDOMAIN).
 Included unit tests.

---
 extensions/dns/ruby/rubydns.rb |  5 ++-
 test/unit/extensions/tc_dns.rb | 77 ++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/extensions/dns/ruby/rubydns.rb b/extensions/dns/ruby/rubydns.rb
index 5f54f8fcb..37ef067e5 100644
--- a/extensions/dns/ruby/rubydns.rb
+++ b/extensions/dns/ruby/rubydns.rb
@@ -189,7 +189,8 @@ module RubyDNS
     # New method that parses response callback and returns RDATA as an array
     def parse_response(block)
       # Extract response arguments into an array
-      args = /(?<=respond!\().*(?=\))/.match(block).to_s.split(/,\s*/)
+      methods = '(respond|failure)'
+      args = /(?<=\.#{methods}!\().*(?=\))/.match(block).to_s.split(/,\s*/)
 
       result = []
 
@@ -199,6 +200,8 @@ module RubyDNS
 
         if /Name\.create\((.*)\)/.match(elem)
           arg = $1
+        elsif /:(NoError|FormErr|ServFail|NXDomain|NotImp|Refused|NotAuth)/.match(elem)
+          arg = $1.upcase
         else
           int_test = elem.to_i
           arg = (int_test != 0 ? int_test : elem)
diff --git a/test/unit/extensions/tc_dns.rb b/test/unit/extensions/tc_dns.rb
index 5ad8a8d50..a1792e2aa 100644
--- a/test/unit/extensions/tc_dns.rb
+++ b/test/unit/extensions/tc_dns.rb
@@ -183,6 +183,65 @@ class TC_Dns < Test::Unit::TestCase
     assert_equal(0, ruleset.length)
   end
 
+  # Tests each supported type of query failure
+  def test_14_failure_types
+    begin
+      id = @@dns.add_rule('noerror.beef.com', IN::A) do |transaction|
+        transaction.failure!(:NoError)
+      end
+
+      check_failure_status(id, :NoError)
+    end
+
+    begin
+      id = @@dns.add_rule('formerr.beef.com', IN::A) do |transaction|
+        transaction.failure!(:FormErr)
+      end
+
+      check_failure_status(id, :FormErr)
+    end
+
+    begin
+      id = @@dns.add_rule('servfail.beef.com', IN::A) do |transaction|
+        transaction.failure!(:ServFail)
+      end
+
+      check_failure_status(id, :ServFail)
+    end
+
+    begin
+      id = @@dns.add_rule('nxdomain.beef.com', IN::A) do |transaction|
+        transaction.failure!(:NXDomain)
+      end
+
+      check_failure_status(id, :NXDomain)
+    end
+
+    begin
+      id = @@dns.add_rule('notimp.beef.com', IN::A) do |transaction|
+        transaction.failure!(:NotImp)
+      end
+
+      check_failure_status(id, :NotImp)
+    end
+
+    begin
+      id = @@dns.add_rule('refused.beef.com', IN::A) do |transaction|
+        transaction.failure!(:Refused)
+      end
+
+      check_failure_status(id, :Refused)
+    end
+
+    begin
+      id = @@dns.add_rule('notauth.beef.com', IN::A) do |transaction|
+        transaction.failure!(:NotAuth)
+      end
+
+      check_failure_status(id, :NotAuth)
+    end
+  end
+
   private
 
   # Compares each key in hash 'rule' with the respective key in hash 'expected'
@@ -192,6 +251,24 @@ class TC_Dns < Test::Unit::TestCase
     assert_equal(expected[:response], rule[:response][0])
   end
 
+  # Compares output of dig command against regex
+  def check_dns_response(regex, type, pattern)
+    address = @@dns_config['address']
+    port    = @@dns_config['port']
+
+    dig_output = `dig @#{address} -p #{port} -t #{type} #{pattern}`
+    assert_match(regex, dig_output)
+  end
+
+  # Confirms that a query for the rule given in 'id' returns a 'type' failure status
+  def check_failure_status(id, type)
+    rule   = @@dns.get_rule(id)
+    status = type.to_s.force_encoding('UTF-8').upcase
+
+    assert_equal(status, rule[:response][0])
+    check_dns_response(/status: #{status}/, rule[:type], rule[:pattern])
+  end
+
 end
 
 # Suppresses unnecessary output from RubyDNS