ruby_smb 2.0.2 → 2.0.7

data/spec/lib/ruby_smb/dcerpc/winreg/save_key_response_spec.rb

@@ -0,0 +1,22 @@
+RSpec.describe RubySMB::Dcerpc::Winreg::SaveKeyResponse do
+  subject(:packet) { described_class.new }
+
+  it { is_expected.to respond_to :error_status }
+
+  it 'is little endian' do
+    expect(described_class.fields.instance_variable_get(:@hints)[:endian]).to eq :little
+  end
+
+  describe '#error_status' do
+    it 'is a 32-bit unsigned integer' do
+      expect(packet.error_status).to be_a BinData::Uint32le
+    end
+  end
+
+  describe '#initialize_instance' do
+    it 'sets #opnum to REG_CREATE_KEY constant' do
+      expect(packet.opnum).to eq(RubySMB::Dcerpc::Winreg::REG_CREATE_KEY)
+    end
+  end
+end
+

data/spec/lib/ruby_smb/dcerpc/winreg_spec.rb

@@ -132,8 +132,7 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
   describe '#query_value' do
     let(:handle)                      { double('Handle') }
     let(:value_name)                  { double('Value Name') }
-    let(:query_value_request_packet1) { double('Query Value Request Packet #1') }
-    let(:query_value_request_packet2) { double('Query Value Request Packet #2') }
+    let(:query_value_request_packet) { double('Query Value Request Packet #1') }
     let(:lp_data1)                    { double('LpData #1') }
     let(:lp_data2)                    { double('LpData #2') }
     let(:response1)                   { double('Response #1') }
@@ -142,26 +141,28 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
     let(:query_value_response2)       { double('Query Value Response #2') }
     let(:data)                        { double('Data') }
     let(:lpcb_data)                   { double('LpcbData') }
-    let(:lpcb_data_referent)          { double('Lpcb Data Referent') }
+    let(:lpcb_data_referent)          { double('LpcbData Referent') }
+    let(:lp_data2_referent)           { double('LpData Referent') }
     before :example do
+      allow(described_class::QueryValueRequest).to receive(:new).and_return(query_value_request_packet)
+      allow(query_value_request_packet).to receive_messages(
+        :lp_type=   => nil,
+        :lpcb_data= => nil,
+        :lpcb_len=  => nil,
+        :lp_data=   => nil,
+        :lp_data    => lp_data2,
+      )
+      allow(lp_data2).to receive(:referent).and_return(lp_data2_referent)
+      allow(lp_data2_referent).to receive(:max_count=)
       first_request = true
-      allow(described_class::QueryValueRequest).to receive(:new) do
+      allow(winreg).to receive(:dcerpc_request) do |arg|
         if first_request
           first_request = false
-          query_value_request_packet1
+          response1
         else
-          query_value_request_packet2
+          response2
         end
       end
-      allow(query_value_request_packet1).to receive(:lp_data).and_return(lp_data1)
-      allow(query_value_request_packet2).to receive_messages(
-        :lp_data    => lp_data2,
-        :lpcb_data= => nil
-      )
-      allow(lp_data1).to receive(:referent_identifier=)
-      allow(lp_data2).to receive(:max_count=)
-      allow(winreg).to receive(:dcerpc_request).with(query_value_request_packet1).and_return(response1)
-      allow(winreg).to receive(:dcerpc_request).with(query_value_request_packet2).and_return(response2)
       allow(described_class::QueryValueResponse).to receive(:read).with(response1).and_return(query_value_response1)
       allow(described_class::QueryValueResponse).to receive(:read).with(response2).and_return(query_value_response2)
       allow(query_value_response1).to receive(:error_status).and_return(WindowsError::Win32::ERROR_SUCCESS)
@@ -175,24 +176,22 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
 
     it 'create the expected QueryValueRequest packets' do
       winreg.query_value(handle, value_name)
-      expect(described_class::QueryValueRequest).to have_received(:new).with(hkey: handle, lp_value_name: value_name).twice
-    end
-
-    it 'sets the expected user rights on the first request packet' do
-      winreg.query_value(handle, value_name)
-      expect(lp_data1).to have_received(:referent_identifier=).with(0)
+      expect(described_class::QueryValueRequest).to have_received(:new).with(hkey: handle, lp_value_name: value_name)
     end
 
-    it 'sets the expected user rights on the second request packet' do
+    it 'sets the expected fields on the request packet' do
       winreg.query_value(handle, value_name)
-      expect(lp_data2).to have_received(:max_count=).with(lpcb_data_referent)
-      expect(query_value_request_packet2).to have_received(:lpcb_data=).with(lpcb_data)
+      expect(query_value_request_packet).to have_received(:lp_type=).with(0)
+      expect(query_value_request_packet).to have_received(:lpcb_data=).with(0)
+      expect(query_value_request_packet).to have_received(:lpcb_len=).with(0)
+      expect(query_value_request_packet).to have_received(:lpcb_data=).with(lpcb_data)
+      expect(query_value_request_packet).to have_received(:lp_data=).with([])
+      expect(lp_data2_referent).to have_received(:max_count=).with(lpcb_data_referent)
     end
 
     it 'sends the expected dcerpc requests' do
       winreg.query_value(handle, value_name)
-      expect(winreg).to have_received(:dcerpc_request).with(query_value_request_packet1).once.ordered
-      expect(winreg).to have_received(:dcerpc_request).with(query_value_request_packet2).once.ordered
+      expect(winreg).to have_received(:dcerpc_request).with(query_value_request_packet).twice
     end
 
     context 'when receiving the first response' do
@@ -293,8 +292,20 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
     let(:query_info_key_request_packet) { double('CloseKey Request Packet') }
     let(:response)                      { double('Response') }
     let(:query_info_key_response)       { double('CloseKey Response') }
+    let(:lp_class)                      { double('LpClass') }
+    let(:lp_class_referent)             { double('LpClass referent') }
+    let(:lp_class_buf_ref)              { double('LpClass buffer referent') }
     before :example do
       allow(described_class::QueryInfoKeyRequest).to receive(:new).and_return(query_info_key_request_packet)
+      allow(query_info_key_request_packet).to receive_messages(
+        :lp_class= => nil,
+        :lp_class  => lp_class,
+      )
+      allow(lp_class).to receive(:referent).and_return(lp_class_referent)
+      allow(lp_class_referent).to receive(:actual_count=)
+      allow(lp_class).to receive(:maximum_length=)
+      allow(lp_class).to receive_message_chain(:buffer, :referent => lp_class_buf_ref)
+      allow(lp_class_buf_ref).to receive(:max_count=)
       allow(winreg).to receive(:dcerpc_request).and_return(response)
       allow(described_class::QueryInfoKeyResponse).to receive(:read).and_return(query_info_key_response)
       allow(query_info_key_response).to receive(:error_status).and_return(WindowsError::Win32::ERROR_SUCCESS)
@@ -310,6 +321,14 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
       expect(winreg).to have_received(:dcerpc_request).with(query_info_key_request_packet)
     end
 
+    it 'sets the expected fields on the request packet' do
+      winreg.query_info_key(handle)
+      expect(query_info_key_request_packet).to have_received(:lp_class=).with('')
+      expect(lp_class_referent).to have_received(:actual_count=).with(0)
+      expect(lp_class).to have_received(:maximum_length=).with(1024)
+      expect(lp_class_buf_ref).to have_received(:max_count=).with(1024 / 2)
+    end
+
     it 'creates a QueryInfoKeyResponse structure from the expected dcerpc response' do
       winreg.query_info_key(handle)
       expect(described_class::QueryInfoKeyResponse).to have_received(:read).with(response)
@@ -338,25 +357,27 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
     let(:handle)                   { double('Handle') }
     let(:index)                    { double('Index') }
     let(:enum_key_request_packet)  { double('enum_key Request Packet') }
-    let(:lp_class)                 { double('Lp Class') }
     let(:lp_name)                  { double('Lp Name') }
     let(:buffer)                   { double('Buffer') }
-    let(:lp_class_buffer_referent) { double('Lp Class buffer referent') }
     let(:lp_name_buffer_referent)  { double('Lp Name buffer referent') }
     let(:response)                 { double('Response') }
     let(:enum_key_response)        { double('enum_key Response') }
     let(:result_str)               { double('Result String') }
+    let(:lp_class)                 { double('Lp Class') }
+    let(:lp_class_buffer_referent) { double('Lp Class buffer referent') }
     before :example do
       allow(described_class::EnumKeyRequest).to receive(:new).and_return(enum_key_request_packet)
       allow(enum_key_request_packet).to receive_messages(
         :lpft_last_write_time= => nil,
-        :lp_class              => lp_class,
-        :lp_name               => lp_name
+        :lp_class=             => nil,
+        :lp_name               => lp_name,
+        :lp_class              => lp_class
       )
       allow(lp_class).to receive(:referent).and_return(lp_class_buffer_referent)
+      allow(lp_class_buffer_referent).to receive(:buffer=)
       allow(lp_name).to receive(:buffer).and_return(buffer)
+      allow(lp_name).to receive(:buffer=)
       allow(buffer).to receive(:referent).and_return(lp_name_buffer_referent)
-      allow(lp_class_buffer_referent).to receive(:buffer=)
       allow(lp_name_buffer_referent).to receive(:max_count=)
       allow(winreg).to receive(:dcerpc_request).and_return(response)
       allow(described_class::EnumKeyResponse).to receive(:read).and_return(enum_key_response)
@@ -369,10 +390,12 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
       expect(described_class::EnumKeyRequest).to have_received(:new).with(hkey: handle, dw_index: index)
     end
 
-    it 'sets the expected user rights on the request packet' do
+    it 'sets the expected parameters on the request packet' do
       winreg.enum_key(handle, index)
       expect(enum_key_request_packet).to have_received(:lpft_last_write_time=).with(0)
-      expect(lp_class_buffer_referent).to have_received(:buffer=).with(0)
+      expect(enum_key_request_packet).to have_received(:lp_class=).with('')
+      expect(lp_class_buffer_referent).to have_received(:buffer=).with(:null)
+      expect(lp_name).to have_received(:buffer=).with('')
       expect(lp_name_buffer_referent).to have_received(:max_count=).with(256)
     end
 
@@ -410,7 +433,6 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
     let(:index)                     { double('Index') }
     let(:enum_value_request_packet) { double('EnumValue Request Packet') }
     let(:lp_value_name)             { double('Lp Value Name') }
-    let(:lp_data)                   { double('Lp Data') }
     let(:buffer)                    { double('Buffer') }
     let(:referent)                  { double('Referent') }
     let(:response)                  { double('Response') }
@@ -418,14 +440,11 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
     let(:result_str)                { double('Result String') }
     before :example do
       allow(described_class::EnumValueRequest).to receive(:new).and_return(enum_value_request_packet)
-      allow(enum_value_request_packet).to receive_messages(
-        :lp_value_name => lp_value_name,
-        :lp_data       => lp_data
-      )
+      allow(enum_value_request_packet).to receive(:lp_value_name).and_return(lp_value_name)
       allow(lp_value_name).to receive(:buffer).and_return(buffer)
+      allow(lp_value_name).to receive(:buffer=)
       allow(buffer).to receive(:referent).and_return(referent)
       allow(referent).to receive(:max_count=)
-      allow(lp_data).to receive(:referent_identifier=)
       allow(winreg).to receive(:dcerpc_request).and_return(response)
       allow(described_class::EnumValueResponse).to receive(:read).and_return(enum_value_response)
       allow(enum_value_response).to receive(:error_status).and_return(WindowsError::Win32::ERROR_SUCCESS)
@@ -437,10 +456,10 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
       expect(described_class::EnumValueRequest).to have_received(:new).with(hkey: handle, dw_index: index)
     end
 
-    it 'sets the expected user rights on the request packet' do
+    it 'sets the expected buffer on the request packet' do
       winreg.enum_value(handle, index)
       expect(referent).to have_received(:max_count=).with(256)
-      expect(lp_data).to have_received(:referent_identifier=).with(0)
+      expect(lp_value_name).to have_received(:buffer=).with('')
     end
 
     it 'sends the expected dcerpc request' do
@@ -492,6 +511,11 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
       expect(winreg).to have_received(:bind).with(endpoint: RubySMB::Dcerpc::Winreg)
     end
 
+    it 'does not bind a DCERPC connection if #bind argument is false' do
+      winreg.has_registry_key?(key, bind: false)
+      expect(winreg).to_not have_received(:bind)
+    end
+
     it 'opens the expected root key' do
       winreg.has_registry_key?(key)
       expect(winreg).to have_received(:open_root_key).with(root_key)
@@ -519,6 +543,7 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
     it 'closes the key' do
       winreg.has_registry_key?(key)
       expect(winreg).to have_received(:close_key).with(subkey_handle)
+      expect(winreg).to have_received(:close_key).with(root_key_handle)
     end
 
     it 'returns true when no error occurs' do
@@ -549,6 +574,11 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
       expect(winreg).to have_received(:bind).with(endpoint: RubySMB::Dcerpc::Winreg)
     end
 
+    it 'does not bind a DCERPC connection if #bind argument is false' do
+      winreg.read_registry_key_value(key, value_name, bind: false)
+      expect(winreg).to_not have_received(:bind)
+    end
+
     it 'opens the expected root key' do
       winreg.read_registry_key_value(key, value_name)
       expect(winreg).to have_received(:open_root_key).with(root_key)
@@ -567,6 +597,7 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
     it 'closes the key' do
       winreg.read_registry_key_value(key, value_name)
       expect(winreg).to have_received(:close_key).with(subkey_handle)
+      expect(winreg).to have_received(:close_key).with(root_key_handle)
     end
 
     it 'returns expect registry key value' do
@@ -600,6 +631,11 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
       expect(winreg).to have_received(:bind).with(endpoint: RubySMB::Dcerpc::Winreg)
     end
 
+    it 'does not bind a DCERPC connection if #bind argument is false' do
+      winreg.enum_registry_key(key, bind: false)
+      expect(winreg).to_not have_received(:bind)
+    end
+
     it 'opens the expected root key' do
       winreg.enum_registry_key(key)
       expect(winreg).to have_received(:open_root_key).with(root_key)
@@ -630,6 +666,13 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
     it 'closes the key' do
       winreg.enum_registry_key(key)
       expect(winreg).to have_received(:close_key).with(subkey_handle)
+      expect(winreg).to have_received(:close_key).with(root_key_handle)
+    end
+
+    it 'only closes the key once when there is no subkey' do
+      winreg.enum_registry_key(root_key)
+      expect(winreg).to have_received(:close_key).once
+      expect(winreg).to have_received(:close_key).with(root_key_handle)
     end
 
     it 'returns the expected array of enumerated keys' do
@@ -667,6 +710,11 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
       expect(winreg).to have_received(:bind).with(endpoint: RubySMB::Dcerpc::Winreg)
     end
 
+    it 'does not bind a DCERPC connection if #bind argument is false' do
+      winreg.enum_registry_values(key, bind: false)
+      expect(winreg).to_not have_received(:bind)
+    end
+
     it 'opens the expected root key' do
       winreg.enum_registry_values(key)
       expect(winreg).to have_received(:open_root_key).with(root_key)
@@ -697,6 +745,13 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
     it 'closes the key' do
       winreg.enum_registry_values(key)
       expect(winreg).to have_received(:close_key).with(subkey_handle)
+      expect(winreg).to have_received(:close_key).with(root_key_handle)
+    end
+
+    it 'only closes the key once when there is no subkey' do
+      winreg.enum_registry_values(root_key)
+      expect(winreg).to have_received(:close_key).once
+      expect(winreg).to have_received(:close_key).with(root_key_handle)
     end
 
     it 'returns the expected array of enumerated keys' do
@@ -707,4 +762,135 @@ RSpec.describe RubySMB::Dcerpc::Winreg do
       expect(winreg.enum_registry_values(key)).to eq([value1, value2])
     end
   end
+
+  describe '#create_key' do
+    let(:handle)              { double('Handle') }
+    let(:sub_key)             { double('Sub key') }
+    let(:create_key_request)  { double('CreateKey Request') }
+    let(:response)            { double('Response') }
+    let(:create_key_response) { double('CreateKey Response') }
+    let(:hkey)                { double('hkey') }
+    before :example do
+      allow(described_class::CreateKeyRequest).to receive(:new).and_return(create_key_request)
+      allow(winreg).to receive(:dcerpc_request).and_return(response)
+      allow(described_class::CreateKeyResponse).to receive(:read).and_return(create_key_response)
+      allow(create_key_response).to receive(:error_status).and_return(WindowsError::Win32::ERROR_SUCCESS)
+      allow(create_key_response).to receive(:hkey).and_return(hkey)
+    end
+
+    it 'create the expected CreateKeyRequest packet with the default options' do
+      opts = {
+        hkey:                   handle,
+        lp_sub_key:             sub_key,
+        lp_class:               :null,
+        dw_options:             RubySMB::Dcerpc::Winreg::CreateKeyRequest::REG_KEY_TYPE_VOLATILE,
+        sam_desired:            RubySMB::Dcerpc::Winreg::Regsam.new(maximum: 1),
+        lp_security_attributes: RubySMB::Dcerpc::RpcSecurityAttributes.new,
+        lpdw_disposition:       RubySMB::Dcerpc::Winreg::CreateKeyRequest::REG_CREATED_NEW_KEY
+      }
+      winreg.create_key(handle, sub_key)
+      expect(described_class::CreateKeyRequest).to have_received(:new).with(opts)
+    end
+
+    it 'create the expected CreateKeyRequest packet with custom options' do
+      opts = {
+        hkey:                  handle,
+        lp_sub_key:            sub_key,
+        lp_class:              'MyClass',
+        dw_options:             RubySMB::Dcerpc::Winreg::CreateKeyRequest::REG_KEY_TYPE_SYMLINK,
+        sam_desired:            RubySMB::Dcerpc::Winreg::Regsam.new(key_set_value: 1),
+        lp_security_attributes: RubySMB::Dcerpc::RpcSecurityAttributes.new(n_length: 3),
+        lpdw_disposition:       RubySMB::Dcerpc::Winreg::CreateKeyRequest::REG_OPENED_EXISTING_KEY
+      }
+      winreg.create_key(handle, sub_key, opts)
+      expect(described_class::CreateKeyRequest).to have_received(:new).with(opts)
+    end
+
+    it 'sends the expected dcerpc request' do
+      winreg.create_key(handle, sub_key)
+      expect(winreg).to have_received(:dcerpc_request).with(create_key_request)
+    end
+
+    it 'creates a CreateKeyResponse structure from the expected dcerpc response' do
+      winreg.create_key(handle, sub_key)
+      expect(described_class::CreateKeyResponse).to have_received(:read).with(response)
+    end
+
+    context 'when an IOError occurs while parsing the response' do
+      it 'raises a RubySMB::Dcerpc::Error::InvalidPacket' do
+        allow(described_class::CreateKeyResponse).to receive(:read).and_raise(IOError)
+        expect { winreg.create_key(handle, sub_key) }.to raise_error(RubySMB::Dcerpc::Error::InvalidPacket)
+      end
+    end
+
+    context 'when the response error status is not WindowsError::Win32::ERROR_SUCCESS' do
+      it 'raises a RubySMB::Dcerpc::Error::WinregError' do
+        allow(create_key_response).to receive(:error_status).and_return(WindowsError::Win32::ERROR_INVALID_DATA)
+        expect { winreg.create_key(handle, sub_key) }.to raise_error(RubySMB::Dcerpc::Error::WinregError)
+      end
+    end
+
+    it 'returns the expected key name' do
+      expect(winreg.create_key(handle, sub_key)).to eq(hkey)
+    end
+  end
+
+  describe '#save_key' do
+    let(:handle)              { double('Handle') }
+    let(:filename)            { double('Filename') }
+    let(:save_key_request)    { double('CreateKey Request') }
+    let(:response)            { double('Response') }
+    let(:save_key_response)   { double('CreateKey Response') }
+    let(:hkey)                { double('hkey') }
+    before :example do
+      allow(described_class::SaveKeyRequest).to receive(:new).and_return(save_key_request)
+      allow(winreg).to receive(:dcerpc_request).and_return(response)
+      allow(described_class::SaveKeyResponse).to receive(:read).and_return(save_key_response)
+      allow(save_key_response).to receive(:error_status).and_return(WindowsError::Win32::ERROR_SUCCESS)
+    end
+
+    it 'create the expected SaveKeyRequest packet with the default options' do
+      opts = {
+        hkey:                   handle,
+        lp_file:                filename,
+        lp_security_attributes: :null,
+      }
+      winreg.save_key(handle, filename)
+      expect(described_class::SaveKeyRequest).to have_received(:new).with(opts)
+    end
+
+    it 'create the expected SaveKeyRequest packet with custom options' do
+      opts = {
+        hkey:                   handle,
+        lp_file:                filename,
+        lp_security_attributes: RubySMB::Dcerpc::RpcSecurityAttributes.new,
+      }
+      winreg.save_key(handle, filename, opts)
+      expect(described_class::SaveKeyRequest).to have_received(:new).with(opts)
+    end
+
+    it 'sends the expected dcerpc request' do
+      winreg.save_key(handle, filename)
+      expect(winreg).to have_received(:dcerpc_request).with(save_key_request)
+    end
+
+    it 'creates a SaveKeyResponse structure from the expected dcerpc response' do
+      winreg.save_key(handle, filename)
+      expect(described_class::SaveKeyResponse).to have_received(:read).with(response)
+    end
+
+    context 'when an IOError occurs while parsing the response' do
+      it 'raises a RubySMB::Dcerpc::Error::InvalidPacket' do
+        allow(described_class::SaveKeyResponse).to receive(:read).and_raise(IOError)
+        expect { winreg.save_key(handle, filename) }.to raise_error(RubySMB::Dcerpc::Error::InvalidPacket)
+      end
+    end
+
+    context 'when the response error status is not WindowsError::Win32::ERROR_SUCCESS' do
+      it 'raises a RubySMB::Dcerpc::Error::WinregError' do
+        allow(save_key_response).to receive(:error_status).and_return(WindowsError::Win32::ERROR_INVALID_DATA)
+        expect { winreg.save_key(handle, filename) }.to raise_error(RubySMB::Dcerpc::Error::WinregError)
+      end
+    end
+  end
 end

data/spec/lib/ruby_smb/dispatcher/socket_spec.rb

@@ -102,26 +102,26 @@ RSpec.describe RubySMB::Dispatcher::Socket do
       end
 
       it 'reads the number of bytes specified in the nbss header' do
-        packet_length = 10
-        session_header.packet_length = packet_length
+        stream_protocol_length = 10
+        session_header.stream_protocol_length = stream_protocol_length
         allow(RubySMB::Nbss::SessionHeader).to receive(:read).and_return(session_header)
         expect(response_socket).to receive(:read).with(4).and_return(session_header)
-        expect(response_socket).to receive(:read).with(packet_length).and_return('A' * packet_length)
+        expect(response_socket).to receive(:read).with(stream_protocol_length).and_return('A' * stream_protocol_length)
         smb_socket.recv_packet
       end
 
       context 'when the socket does not return everything the first time' do
         it 'calls #read several times until everything is returned' do
-          packet_length = 67
+          stream_protocol_length = 67
           returned_length = 10
-          session_header.packet_length = packet_length
+          session_header.stream_protocol_length = stream_protocol_length
           allow(RubySMB::Nbss::SessionHeader).to receive(:read).and_return(session_header)
 
           expect(response_socket).to receive(:read).with(4).and_return(session_header)
           loop do
-            expect(response_socket).to receive(:read).with(packet_length).and_return('A' * returned_length).once
-            packet_length -= returned_length
-            break if packet_length <= 0
+            expect(response_socket).to receive(:read).with(stream_protocol_length).and_return('A' * returned_length).once
+            stream_protocol_length -= returned_length
+            break if stream_protocol_length <= 0
           end
           smb_socket.recv_packet
         end
@@ -134,7 +134,7 @@ RSpec.describe RubySMB::Dispatcher::Socket do
 
         context 'when the SMB packet is empty' do
           it 'returns the nbss header only' do
-            session_header.packet_length = 0
+            session_header.stream_protocol_length = 0
             allow(RubySMB::Nbss::SessionHeader).to receive(:read).and_return(session_header)
             expect(smb_socket.recv_packet(full_response: true)).to eq(session_header.to_binary_s)
           end
@@ -148,7 +148,7 @@ RSpec.describe RubySMB::Dispatcher::Socket do
 
         context 'when the SMB packet is empty' do
           it 'returns an empty string' do
-            session_header.packet_length = 0
+            session_header.stream_protocol_length = 0
             allow(RubySMB::Nbss::SessionHeader).to receive(:read).and_return(session_header)
             expect(smb_socket.recv_packet(full_response: false)).to eq('')
           end