feat: connect methods' block behavior

ChangeLog.md

@@ -1,3 +1,11 @@
+### next
+
+* {Ronin::Support::Network::TCP.connect}, {Ronin::Support::Network::UDP.connect}, and
+  {Ronin::Support::Network::HTTP.connect}, when passed a block, return the block's return value.
+* {Ronin::Support::Network::TCP.connect} and {Ronin::Support::Network::UDP.connect} properly close the
+  socket when passed a block that raises an exception.
+
+
 ### 1.0.5 / 2023-12-27
 
 * Fixed a bug in {Ronin::Support::Binary::Stream::Methods#read_string} on Ruby
@@ -704,4 +712,3 @@
   * Require combinatorics ~> 0.3.
   * Require uri-query_params ~> 0.5, >= 0.5.2.
   * Require data_paths ~> 0.2, >= 0.2.1.
-

lib/ronin/support/network/http.rb

@@ -473,8 +473,11 @@ def self.connect(host,port, ssl: port == 443, **kwargs)
           http = new(host,port, ssl: ssl, **kwargs)
 
           if block_given?
-            yield http
-            http.close
+            begin
+              yield http
+            ensure
+              http.close
+            end
           else
             return http
           end

lib/ronin/support/network/tcp.rb

@@ -131,8 +131,11 @@ def self.connect(host,port, bind_host: nil, bind_port: nil)
                    end
 
           if block_given?
-            yield socket
-            socket.close
+            begin
+              yield socket
+            ensure
+              socket.close
+            end
           else
             return socket
           end

lib/ronin/support/network/udp.rb

@@ -137,8 +137,11 @@ def self.connect(host,port, bind_host: nil, bind_port: nil)
           socket.connect(host,port)
 
           if block_given?
-            yield socket
-            socket.close
+            begin
+              yield socket
+            ensure
+              socket.close
+            end
           else
             return socket
           end

spec/network/http_spec.rb

@@ -442,6 +442,14 @@
           described_class.connect(host,port,&b)
         }.to yield_with_args(described_class)
       end
+
+      it "must return the block's return value" do
+        returned_value = described_class.connect(host,port) do
+          :return_value
+        end
+
+        expect(returned_value).to eq(:return_value)
+      end
     end
   end
 

spec/network/tcp_spec.rb

@@ -108,6 +108,30 @@
           expect(socket).to be_closed
         end
 
+        it "must return the block's return value" do
+          returned_value = subject.connect(host,port) do |socket|
+            :return_value
+          end
+
+          expect(returned_value).to eq(:return_value)
+        end
+
+        context "when the block raises an exception" do
+          it "must close the TCPSocket" do
+            socket = nil
+
+            expect do
+              subject.connect(host,port) do |yielded_socket|
+                socket = yielded_socket
+                raise "test exception"
+              end
+            end.to raise_error("test exception")
+
+            expect(socket).to be_kind_of(TCPSocket)
+            expect(socket).to be_closed
+          end
+        end
+
         context "when given the bind_port: keyword argument" do
           let(:bind_port) { 1024 + rand(65535 - 1024) }
 

spec/network/udp_spec.rb

@@ -103,6 +103,30 @@
           expect(socket).to be_closed
         end
 
+        it "must return the block's return value" do
+          returned_value = subject.connect(host,port) do |socket|
+            :return_value
+          end
+
+          expect(returned_value).to eq(:return_value)
+        end
+
+        context "when the block raises an exception" do
+          it "must close the UDPSocket" do
+            socket = nil
+
+            expect do
+              subject.connect(host,port) do |yielded_socket|
+                socket = yielded_socket
+                raise "test exception"
+              end
+            end.to raise_error("test exception")
+
+            expect(socket).to be_kind_of(UDPSocket)
+            expect(socket).to be_closed
+          end
+        end
+
         context "when given the bind_port: keyword argument" do
           it "must bind to the local port" do
             bound_port = nil