Re-factor. As per comments on the Pull Request.kwilczynski/gce-metadata

Signed-off-by: Krzysztof Wilczynski <krzysztof.wilczynski@linux.com>

2 files changed, 14 insertions, 10 deletions

diff --git a/lib/ohai/util/socket_helper.rb b/lib/ohai/util/socket_helper.rb
index 0ed7f5c3..8c2fe19d 100644
--- a/lib/ohai/util/socket_helper.rb
+++ b/lib/ohai/util/socket_helper.rb
@@ -20,24 +20,25 @@ module Ohai
   module Util
     module SocketHelper
       def tcp_port_open?(host, port, timeout = 2)
-        saved_lookup = Socket.do_not_reverse_lookup = Socket.do_not_reverse_lookup, true
         Timeout.timeout(timeout) do
           begin
+            Socket.do_not_reverse_lookup = true
             TCPSocket.new(host, port).close
             true
           rescue Errno::ECONNREFUSED, Errno::EHOSTUNREACH
             false
-          rescue SystemCallError, SocketError => e
-            # Check for DNS resolution failure and ignore,
-            # otherwise raise as it might be something serious.
-            raise(e) if e.is_a?(SocketError) && !(e.to_s =~ /getaddrinfo/)
+          rescue SystemCallError
+            false
+          rescue SocketError => e
+            # Check for DNS resolution failure and return false
+            # when we stumble upon such, otherwise raise as it
+            # might be something serious that needs attention.
+            raise(e) unless e.to_s =~ /getaddrinfo/
             false
           end
         end
       rescue Timeout::Error
         false
-      ensure
-        Socket.do_not_reverse_lookup = saved_lookup
       end
     end
   end
diff --git a/spec/unit/util/socket_helper_spec.rb b/spec/unit/util/socket_helper_spec.rb
index 37a178a7..b8015700 100644
--- a/spec/unit/util/socket_helper_spec.rb
+++ b/spec/unit/util/socket_helper_spec.rb
@@ -31,13 +31,11 @@ describe 'Ohai::Util::SocketHelper' do
     allow(@socket).to receive(:close).and_return(nil)
     allow(Socket).to receive(:do_not_reverse_lookup).and_call_original
     allow(TCPSocket).to receive(:new).and_return(@socket)
-
-    expect(Socket).to receive(:do_not_reverse_lookup).once
-    expect(Socket).to receive(:do_not_reverse_lookup=).twice
   end
 
   describe 'when remote node is accessible' do
     it 'should return true when connection is accepted' do
+      expect(Socket).to receive(:do_not_reverse_lookup=).once
       expect(TCPSocket).to receive(:new).with('chef.io', 42)
       expect(@socket).to receive(:close).once
       expect(socket_helper.tcp_port_open?('chef.io', 42)).to be true
@@ -48,6 +46,7 @@ describe 'Ohai::Util::SocketHelper' do
     it 'should return false when connection is refused' do
       allow(TCPSocket).to receive(:new).with('getchef.com', 80).and_raise(Errno::ECONNREFUSED)
 
+      expect(Socket).to receive(:do_not_reverse_lookup=).once
       expect(TCPSocket).to receive(:new).with('getchef.com', 80)
       expect(@socket).not_to receive(:close)
       expect(socket_helper.tcp_port_open?('getchef.com', 80)).to be false
@@ -56,6 +55,7 @@ describe 'Ohai::Util::SocketHelper' do
     it 'should return false when connection cannot be established' do
       allow(TCPSocket).to receive(:new).with('opscode.com', 443).and_raise(Errno::EHOSTUNREACH)
 
+      expect(Socket).to receive(:do_not_reverse_lookup=).once
       expect(TCPSocket).to receive(:new).with('opscode.com', 443)
       expect(@socket).not_to receive(:close)
       expect(socket_helper.tcp_port_open?('opscode.com', 443)).to be false
@@ -69,6 +69,7 @@ describe 'Ohai::Util::SocketHelper' do
       #   SocketError: getaddrinfo: Temporary failure in name resolution
       allow(TCPSocket).to receive(:new).with('acme.com', 8080).and_raise(SocketError, 'getaddrinfo: Name or service not known')
 
+      expect(Socket).to receive(:do_not_reverse_lookup=).once
       expect(TCPSocket).to receive(:new).with('acme.com', 8080)
       expect(@socket).not_to receive(:close)
       expect(socket_helper.tcp_port_open?('acme.com', 8080)).to be false
@@ -77,6 +78,7 @@ describe 'Ohai::Util::SocketHelper' do
     it 'should raise unknown SocketError exception' do
       allow(TCPSocket).to receive(:new).with('NCC-1701-D', 40759).and_raise(SocketError, 'the plasma conduit time-matter field appears to be removed')
 
+      expect(Socket).to receive(:do_not_reverse_lookup=).once
       expect(TCPSocket).to receive(:new).with('NCC-1701-D', 40759)
       expect(@socket).not_to receive(:close)
 
@@ -90,6 +92,7 @@ describe 'Ohai::Util::SocketHelper' do
     it 'should return false when a timeout occurs' do
       allow(Timeout).to receive(:timeout).with(3).and_raise(Timeout::Error)
 
+      expect(Socket).not_to receive(:do_not_reverse_lookup=)
       expect(TCPSocket).not_to receive(:new).with('slow.net', 22)
       expect(@socket).not_to receive(:close)
       expect(socket_helper.tcp_port_open?('slow.net', 22, 3)).to be false