From 533d4bfc5f6877136d0365df8a6dcb8879428482 Mon Sep 17 00:00:00 2001
From: Jonathan Claudius <jclaudius@mozilla.com>
Date: Thu, 15 Jun 2017 10:24:24 -0400
Subject: [PATCH] Reset error register when falling back to IPv4

---
 lib/ssh_scan/client.rb      | 11 +++++++++++
 lib/ssh_scan/result.rb      |  4 ++++
 lib/ssh_scan/scan_engine.rb |  9 +++++++--
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/lib/ssh_scan/client.rb b/lib/ssh_scan/client.rb
index da1b48ae..9e489e31 100644
--- a/lib/ssh_scan/client.rb
+++ b/lib/ssh_scan/client.rb
@@ -28,6 +28,17 @@ def banner
       @server_banner
     end
 
+    def close()
+      begin
+        unless @sock.nil?
+          @sock.close
+        end
+      rescue
+        @sock = nil
+      end
+      return true
+    end
+
     def connect()
       @error = nil
 
diff --git a/lib/ssh_scan/result.rb b/lib/ssh_scan/result.rb
index 2da79939..5d9ce212 100644
--- a/lib/ssh_scan/result.rb
+++ b/lib/ssh_scan/result.rb
@@ -211,6 +211,10 @@ def error=(error)
       @error = error.to_s
     end
 
+    def unset_error
+      @error = nil
+    end
+
     def error?
       !@error.nil?
     end
diff --git a/lib/ssh_scan/scan_engine.rb b/lib/ssh_scan/scan_engine.rb
index 90b167a8..0945bbc9 100644
--- a/lib/ssh_scan/scan_engine.rb
+++ b/lib/ssh_scan/scan_engine.rb
@@ -36,9 +36,10 @@ def scan_target(socket, opts)
           client = SSHScan::Client.new(
             target.resolve_fqdn_as_ipv4.to_s, port, timeout
           )
-          client.connect()
+          client.connect
           result.set_client_attributes(client)
           kex_result = client.get_kex_result()
+          client.close
           result.set_kex_result(kex_result) unless kex_result.nil?
           result.error = client.error if client.error?
         # If it does resolve as IPv6, we're try IPv6
@@ -46,20 +47,23 @@ def scan_target(socket, opts)
           client = SSHScan::Client.new(
             target.resolve_fqdn_as_ipv6.to_s, port, timeout
           )
-          client.connect()
+          client.connect
           result.set_client_attributes(client)
           kex_result = client.get_kex_result()
+          client.close
           result.set_kex_result(kex_result) unless kex_result.nil?
           result.error = client.error if client.error?
 
           # If resolves as IPv6, but somehow we get an client error, fall-back to IPv4
           if result.error?
+            result.unset_error
             client = SSHScan::Client.new(
               target.resolve_fqdn_as_ipv4.to_s, port, timeout
             )
             client.connect()
             result.set_client_attributes(client)
             kex_result = client.get_kex_result()
+            client.close
             result.set_kex_result(kex_result) unless kex_result.nil?
             result.error = client.error if client.error?
           end
@@ -69,6 +73,7 @@ def scan_target(socket, opts)
         client.connect()
         result.set_client_attributes(client)
         kex_result = client.get_kex_result()
+        client.close
 
         unless kex_result.nil?
           result.set_kex_result(kex_result)