Skip to content

Commit

Permalink
feat: TCP.connect and UDP.connect block behavior
Browse files Browse the repository at this point in the history
- return the block's return value
- close the socket even if the block raises an exception
  • Loading branch information
flavorjones committed May 4, 2024
1 parent 0b7efc4 commit 0fc8bdb
Show file tree
Hide file tree
Showing 5 changed files with 66 additions and 5 deletions.
9 changes: 8 additions & 1 deletion ChangeLog.md
Original file line number Diff line number Diff line change
@@ -1,3 +1,11 @@
### next

* {Ronin::Support::Network::TCP.close} and {Ronin::Support::Network::UDP.close}, when passed a
block, return the block's return value.
* {Ronin::Support::Network::TCP.close} and {Ronin::Support::Network::UDP.close} 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
Expand Down Expand Up @@ -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.

7 changes: 5 additions & 2 deletions lib/ronin/support/network/tcp.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 5 additions & 2 deletions lib/ronin/support/network/udp.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
24 changes: 24 additions & 0 deletions spec/network/tcp_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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) }

Expand Down
24 changes: 24 additions & 0 deletions spec/network/udp_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 0fc8bdb

Please sign in to comment.