lib/ruby_smb/dcerpc/ndr.rb in ruby_smb-3.2.3 vs lib/ruby_smb/dcerpc/ndr.rb in ruby_smb-3.2.4

@@ -1313,7 +1313,125 @@
         read(handle.to_s)
       end
     end
   end
 
+  # (IDL/NDR) Pickles as defined in
+  # [(IDL/NDR) # Pickles](https://pubs.opengroup.org/onlinepubs/9668899/chap2.htm#tagcjh_05_01_07)
+  # and
+  # [2.2.6 Type Serialization Version # 1](https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-rpce/9a1d0f97-eac0-49ab-a197-f1a581c2d6a0)
+
+  # [2.2.6.1 Common Type Header for the Serialization Stream](https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-rpce/6d75d40e-e2d2-4420-b9e9-8508a726a9ae)
+  class TypeSerialization1CommonTypeHeader < BinData::Record
+    default_parameter byte_align: 8
+    endian :little
+
+    uint8  :version, initial_value: 1
+    uint8  :endianness, initial_value: 0x10
+    uint16 :common_header_length, initial_value: 8
+    uint32 :filler, initial_value: 0xCCCCCCCC
+  end
+
+  # [2.2.6.2 Private Header for Constructed Type](https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-rpce/63949ba8-bc88-4c0c-9377-23f14b197827)
+  class TypeSerialization1PrivateHeader < BinData::Record
+    default_parameter byte_align: 8
+    endian :little
+
+    uint32 :object_buffer_length, initial_value: -> { buffer_length(@obj) }
+    uint32 :filler, initial_value: 0x00000000
+
+    def buffer_length(obj)
+      parent.respond_to?(:field_length) ? parent.field_length(obj.parent) : 0
+    end
+  end
+
+  # [2.2.6 Type Serialization Version 1](https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-rpce/9a1d0f97-eac0-49ab-a197-f1a581c2d6a0)
+  #
+  # This structure is not meant to be instantiated directly. Instead, the
+  # structure with the fields to be serialized needs to inherit from
+  # TypeSerialization1. This class will take care of adding the
+  # TypeSerialization1PrivateHeader fields in front of any NDR constructed type
+  # structures and setting the buffer length field (:object_buffer_length) to
+  # the correct value.
+  #
+  # Example:
+  #
+  # class TestStruct < RubySMB::Dcerpc::Ndr::NdrStruct
+  #   default_parameters byte_align: 8
+  #   endian :little
+
+  #   rpc_unicode_string :full_name
+  #   ndr_uint32 :user_id
+  # end
+
+  # class TestTypeSerialization1 < RubySMB::Dcerpc::Ndr::TypeSerialization1
+  #   default_parameter byte_align: 8
+  #   endian :little
+  #
+  #   test_struct :data1
+  #   uint32 :value1, initial_value: 5
+  #   uint32 :value2, initial_value: 6
+  #   test_struct :data2
+  #   uint32 :value3, initial_value: 7
+  #   test_struct :data3
+  # end
+  #
+  # This will result in the following structure:
+  # {
+  #   :common_header => {:version=>1, :endianness=>16, :common_header_length=>8, :filler=>3435973836},
+  #   :private_header1 => {:object_buffer_length=>12, :filler=>0},
+  #   :data1 => {:full_name=>{:buffer_length=>0, :maximum_length=>0, :buffer=>:null}, :user_id=>0},
+  #   :value1 => 5,
+  #   :value2 => 6,
+  #   :private_header2 => {:object_buffer_length=>12, :filler=>0},
+  #   :data2 => {:full_name=>{:buffer_length=>0, :maximum_length=>0, :buffer=>:null}, :user_id=>0},
+  #   :value3 => 7,
+  #   :private_header3 => {:object_buffer_length=>12, :filler=>0},
+  #   :data3 => {:full_name=>{:buffer_length=>0, :maximum_length=>0, :buffer=>:null}, :user_id=>0}
+  # }
+  #
+  class TypeSerialization1 < BinData::Record
+    PRIVATE_HEADER_BASE_NAME = 'private_header'.freeze
+
+    default_parameter byte_align: 8
+    endian :little
+    search_prefix :type_serialization1
+
+    common_type_header :common_header
+
+    def field_length(obj)
+      length = 0
+      index = find_index_of(obj)
+      if index
+        each_pair {|n, o| length = o.num_bytes if n == field_names[index + 1]}
+      end
+      length
+    end
+
+    def self.method_missing(symbol, *args, &block)
+      return super if dsl_parser.respond_to?(symbol)
+
+      klass = BinData::RegisteredClasses.lookup(symbol, {endian: dsl_parser.endian, search_prefix: dsl_parser.search_prefix})
+      if klass.new.is_a?(ConstructedTypePlugin)
+        # We cannot have two fields with the same name. So, here we look for
+        # existing TypeSerialization1PrivateHeader fields and just increment
+        # the ending number of the last TypeSerialization1PrivateHeader field
+        # to make the new name unique.
+        names = dsl_parser.fields.find_all do |field|
+          field.prototype.instance_variable_get(:@obj_class) == TypeSerialization1PrivateHeader
+        end.map(&:name).sort
+        if names.empty?
+          new_name = "#{PRIVATE_HEADER_BASE_NAME}1"
+        else
+          num = names.last.match(/#{PRIVATE_HEADER_BASE_NAME}(\d)$/)[1].to_i
+          new_name = "#{PRIVATE_HEADER_BASE_NAME}#{num + 1}"
+        end
+
+        super(:private_header, new_name.to_sym)
+      end
+
+      super
+    end
+  end
+
 end