Improve connect methods' block behavior (#505)

* Fixed `TCP.connect`, `UDP.connect`, `POP3.connect`, `SMTP.connect`, and `HTTP.connect` to close the connection even if the given block raises an exception. * Ensure that `TCP.connect`, `UDP.connect`, `POP3.connect`, `SMTP.connect`, and `HTTP.connect` returns the return value of the block. --------- Co-authored-by: Postmodern <[email protected]>
ronin-rb · May 6, 2024 · b3a883e · b3a883e
1 parent 0b7efc4
commit b3a883e
Show file tree

Hide file tree

Showing 11 changed files with 149 additions and 11 deletions.
diff --git a/ChangeLog.md b/ChangeLog.md
@@ -1,3 +1,11 @@
+### 1.0.6 / 2024-XX-XX
+
+* {Ronin::Support::Network::TCP.connect}, {Ronin::Support::Network::UDP.connect}, and
+  {Ronin::Support::Network::HTTP.connect}, when given a block, now returns the block's return value.
+* {Ronin::Support::Network::TCP.connect} and {Ronin::Support::Network::UDP.connect} properly closes 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.
-
diff --git a/lib/ronin/support/network/http.rb b/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

diff --git a/lib/ronin/support/network/pop3/mixin.rb b/lib/ronin/support/network/pop3/mixin.rb
@@ -100,8 +100,11 @@ def pop3_connect(host, port: DEFAULT_PORT,
             pop3.start(user,password)
 
             if block_given?
-              yield pop3
-              pop3.finish
+              begin
+                yield pop3
+              ensure
+                pop3.finish
+              end
             else
               return pop3
             end

diff --git a/lib/ronin/support/network/smtp/mixin.rb b/lib/ronin/support/network/smtp/mixin.rb
@@ -164,8 +164,11 @@ def smtp_connect(host, user: ,
             smtp.start(helo,user,password,auth)
 
             if block_given?
-              yield smtp
-              smtp.finish
+              begin
+                yield smtp
+              ensure
+                smtp.finish
+              end
             else
               return smtp
             end

diff --git a/lib/ronin/support/network/tcp.rb b/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

diff --git a/lib/ronin/support/network/udp.rb b/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

diff --git a/spec/network/http_spec.rb b/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
 

diff --git a/spec/network/pop3/mixin_spec.rb b/spec/network/pop3/mixin_spec.rb
@@ -58,6 +58,16 @@
           expect(yielded_pop3).to be_kind_of(Net::POP3)
         end
 
+        it "must return the block's return value" do
+          pending "need valid POP3 credentials"
+
+          returned_value = subject.pop3_connect(host, port: port, ssl: true) do |pop3|
+            :return_value
+          end
+
+          expect(returned_value).to be(:return_value)
+        end
+
         it "must finish the POP3 session after yielding it" do
           pending "need valid POP3 credentials"
 
@@ -72,6 +82,26 @@
           expect(was_started).to be(true)
           expect(pop3).to_not be_started
         end
+
+        context "when the block raises an exception" do
+          it "must finish the POP3 session after yielding it" do
+            pending "need valid POP3 credentials"
+
+            pop3        = nil
+            was_started = nil
+
+            expect do
+              subject.pop3_connect(host, port: port, ssl: true) do |yielded_pop3|
+                pop3        = yielded_pop3
+                was_started = pop3.started?
+                raise "test exception"
+              end
+            end.to raise_error("test exception")
+
+            expect(was_started).to be(true)
+            expect(pop3).to_not be_started
+          end
+        end
       end
     end
   end

diff --git a/spec/network/smtp/mixin_spec.rb b/spec/network/smtp/mixin_spec.rb
@@ -56,6 +56,16 @@
           expect(yielded_smtp).to be_kind_of(Net::SMTP)
         end
 
+        it "must return the block's return value" do
+          pending "need valid SMTP credentials"
+
+          returned_value = subject.smtp_connect(host) do |smtp|
+            :return_value
+          end
+
+          expect(returned_value).to be(:return_value)
+        end
+
         it "must finish the SMTP session after yielding it" do
           pending "need valid SMTP credentials"
 
@@ -70,6 +80,26 @@
           expect(was_started).to be(true)
           expect(smtp).to_not be_started
         end
+
+        context "when the block raises an exception" do
+          it "must finish the SMTP session after yielding it" do
+            pending "need valid SMTP credentials"
+
+            smtp        = nil
+            was_started = nil
+
+            expect do
+              subject.smtp_connect(host) do |yielded_smtp|
+                smtp        = yielded_smtp
+                was_started = smtp.started?
+                raise "test exception"
+              end
+            end.to raise_error("test exception")
+
+            expect(was_started).to be(true)
+            expect(smtp).to_not be_started
+          end
+        end
       end
     end
   end

diff --git a/spec/network/tcp_spec.rb b/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) }
 

diff --git a/spec/network/udp_spec.rb b/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