messagepack 1.0.0

data/lib/messagepack/extensions/registry.rb

@@ -0,0 +1,154 @@
+# frozen_string_literal: true
+
+module Messagepack
+  # Extension registries for packer and unpacker.
+  #
+  # The packer registry maps Ruby classes to extension type IDs and packing procs.
+  # The unpacker registry maps extension type IDs to unpacking procs.
+  #
+  module ExtensionRegistry
+    # Packer registry for serializing custom types.
+    #
+    # Uses hash-based lookup with ancestor chain caching.
+    #
+    class Packer
+      def initialize
+        @registry = {}  # klass => [type_id, proc, flags]
+        @cache = {}     # klass => [type_id, proc, flags] (ancestor cache)
+      end
+
+      def dup
+        copy = self.class.new
+        copy.instance_variable_set(:@registry, @registry.dup)
+        copy.instance_variable_set(:@cache, {})
+        copy
+      end
+
+      alias clone dup
+
+      def register(type_id, klass, proc, flags: 0)
+        @registry[klass] = [type_id, proc, flags]
+        @cache.clear
+      end
+
+      def lookup(value)
+        klass = value.class
+        return @registry[klass] if @registry.key?(klass)
+
+        # Check cache
+        return @cache[klass] if @cache.key?(klass)
+
+        # Search ancestors and modules
+        @registry.each do |registered_class, data|
+          # Check for inheritance (klass <= registered_class)
+          if klass <= registered_class
+            @cache[klass] = data
+            return data
+          end
+
+          # Check if registered_class is a module (but not a class) and klass includes it
+          # In Ruby, Class is not a subclass of Module, so we need to check the class
+          if registered_class.is_a?(Class)
+            # registered_class is a Class, already handled by inheritance check above
+            next
+          end
+
+          # registered_class is a Module
+          # Check if klass includes the module
+          if klass.include?(registered_class)
+            @cache[klass] = data
+            return data
+          end
+
+          # Check if the object's singleton class includes the module
+          # This handles cases like: obj.extend(Mod)
+          # Some values don't support singleton_class:
+          # - Immediate values: nil, true, false, integers, floats, symbols
+          # - Frozen values (frozen_string_literal: true makes string literals frozen)
+          unless value.nil? || value == true || value == false ||
+                 value.is_a?(Integer) || value.is_a?(Float) || value.is_a?(Symbol) ||
+                 value.frozen?
+            if value.singleton_class.include?(registered_class)
+              @cache[klass] = data
+              return data
+            end
+          end
+        end
+
+        nil
+      end
+
+      def registered_types
+        @registry.map { |klass, (type_id, proc, flags)|
+          { type: type_id, class: klass, packer: proc }
+        }
+      end
+
+      def type_registered?(klass_or_type)
+        if klass_or_type.is_a?(Integer)
+          @registry.any? { |_, (type_id, _, _)| type_id == klass_or_type }
+        else
+          @registry.key?(klass_or_type)
+        end
+      end
+
+      def clear
+        @registry.clear
+        @cache.clear
+      end
+    end
+
+    # Unpacker registry for deserializing custom types.
+    #
+    # Uses direct array for O(1) type ID lookup.
+    #
+    class Unpacker
+      ARRAY_SIZE = 256  # Signed byte range: -128 to 127
+
+      def initialize
+        @array = ::Array.new(ARRAY_SIZE)  # [klass, proc, flags]
+      end
+
+      def dup
+        copy = self.class.new
+        copy.instance_variable_set(:@array, @array.dup)
+        copy
+      end
+
+      alias clone dup
+
+      def register(type_id, klass, proc, flags: 0)
+        index = type_id + 128  # Offset for signed byte
+        raise IndexError, "type_id out of range: #{type_id}" unless (0...ARRAY_SIZE).cover?(index)
+        @array[index] = [klass, proc, flags]
+      end
+
+      def lookup(type_id)
+        index = type_id + 128
+        return nil unless (0...ARRAY_SIZE).cover?(index)
+        @array[index]
+      end
+
+      def registered_types
+        @array.each_with_index.map { |(klass, proc, flags), index|
+          next unless proc  # Include entries that have a proc, even if klass is nil
+          { type: index - 128, class: klass, unpacker: proc }
+        }.compact
+      end
+
+      def type_registered?(klass_or_type)
+        if klass_or_type.is_a?(Integer)
+          index = klass_or_type + 128
+          # Check if any data (klass, proc, or flags) exists at this type ID index
+          (0...ARRAY_SIZE).cover?(index) && @array[index]&.at(1)
+        else
+          @array.any? { |(klass, _, _)| klass == klass_or_type }
+        end
+      end
+
+      def clear
+        @array.fill(nil)
+      end
+    end
+  end
+end

data/lib/messagepack/extensions/symbol.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+# Symbol extension for MessagePack.
+#
+# This adds support for serializing Symbol objects directly.
+# When registered, symbols are serialized as their string representation.
+#
+class Symbol
+  # to_msgpack_ext is supposed to return a binary string.
+  # The canonical way to do it for symbols would be:
+  #  [to_s].pack('A*')
+  # However in this instance we can take a shortcut
+  if method_defined?(:name)
+    alias_method :to_msgpack_ext, :name
+  else
+    alias_method :to_msgpack_ext, :to_s
+  end
+
+  # Reconstruct symbol from binary payload.
+  #
+  # @param data [String] Binary payload (symbol name as string)
+  # @return [Symbol] The reconstructed symbol
+  #
+  def self.from_msgpack_ext(data)
+    # from_msgpack_ext is supposed to parse a binary string.
+    # The canonical way to do it for symbols would be:
+    #  data.unpack1('A*').to_sym
+    # However in this instance we can take a shortcut
+
+    # We assume the string encoding is UTF-8, and let Ruby create either
+    # an ASCII symbol or UTF-8 symbol.
+    data.force_encoding(Encoding::UTF_8).to_sym
+  rescue EncodingError
+    # If somehow the string wasn't valid UTF-8 not valid ASCII, we fallback
+    # to what has been the historical behavior of creating a binary symbol
+    data.force_encoding(Encoding::BINARY).to_sym
+  end
+end

data/lib/messagepack/extensions/timestamp.rb

@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+require_relative 'value'
+
+module Messagepack
+  # Timestamp extension type (type -1).
+  #
+  # Represents a high-resolution timestamp with second and nanosecond precision.
+  # Implements the MessagePack timestamp specification:
+  # https://github.com/msgpack/msgpack/blob/master/spec.md#timestamp-extension-type
+  #
+  class Timestamp
+    TYPE = -1
+
+    TIMESTAMP32_MAX_SEC = (1 << 32) - 1
+    TIMESTAMP64_MAX_SEC = (1 << 34) - 1
+
+    attr_reader :sec, :nsec
+
+    def initialize(sec, nsec)
+      @sec = sec
+      @nsec = nsec
+    end
+
+    # Deserialize from binary payload.
+    #
+    # @param data [String] Binary payload from MessagePack
+    # @return [Timestamp] Deserialized timestamp
+    #
+    def self.from_msgpack_ext(data)
+      case data.bytesize
+      when 4
+        # timestamp32 (sec: uint32be)
+        sec = data.unpack1('L>')
+        new(sec, 0)
+      when 8
+        # timestamp64 (nsec: uint30be, sec: uint34be)
+        n, s = data.unpack('L>2')
+        sec = ((n & 0b11) << 32) | s
+        nsec = n >> 2
+        new(sec, nsec)
+      when 12
+        # timestamp96 (nsec: uint32be, sec: int64be)
+        nsec, sec = data.unpack('L>q>')
+        new(sec, nsec)
+      else
+        raise MalformedFormatError, "Invalid timestamp data size: #{data.bytesize}"
+      end
+    end
+
+    # Serialize to binary payload.
+    #
+    # @return [String] Binary payload for MessagePack extension
+    #
+    def to_msgpack_ext
+      self.class.to_msgpack_ext(@sec, @nsec)
+    end
+
+    # Class helper for serialization.
+    #
+    # @param sec [Integer] Seconds since epoch
+    # @param nsec [Integer] Nanoseconds within second (0-999999999)
+    # @return [String] Binary payload
+    #
+    def self.to_msgpack_ext(sec, nsec)
+      if sec >= 0 && nsec >= 0 && sec <= TIMESTAMP64_MAX_SEC
+        if nsec == 0 && sec <= TIMESTAMP32_MAX_SEC
+          # timestamp32: 4 bytes
+          [sec].pack('L>')
+        else
+          # timestamp64: 8 bytes
+          nsec30 = nsec << 2
+          sec_high2 = sec >> 32
+          sec_low32 = sec & 0xffffffff
+          [nsec30 | sec_high2, sec_low32].pack('L>2')
+        end
+      else
+        # timestamp96: 12 bytes
+        [nsec, sec].pack('L>q>')
+      end
+    end
+
+    # Convert to Time object
+    def to_time
+      Time.at(@sec, @nsec, :nanosecond)
+    end
+
+    # Create from Time object
+    def self.from_time(time)
+      new(time.tv_sec, time.tv_nsec)
+    end
+
+    def ==(other)
+      other.class == self.class && @sec == other.sec && @nsec == other.nsec
+    end
+
+    alias eql? ==
+
+    def hash
+      [@sec, @nsec].hash
+    end
+
+    # String representation
+    def to_s
+      "Timestamp(#{@sec}, #{@nsec})"
+    end
+
+    alias inspect to_s
+  end
+end

data/lib/messagepack/extensions/value.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module Messagepack
+  # ExtensionValue represents a raw MessagePack extension type.
+  #
+  # This is used when unpacking extension types that don't have a
+  # registered handler, or for manual extension construction.
+  #
+  class ExtensionValue
+    attr_accessor :type, :payload
+
+    def initialize(type, payload)
+      @type = type
+      @payload = payload
+    end
+
+    def ==(other)
+      return false unless other.is_a?(ExtensionValue)
+      @type == other.type && @payload == other.payload
+    end
+
+    alias eql? ==
+
+    def hash
+      [@type, @payload].hash
+    end
+
+    # Convert back to MessagePack format
+    def to_msgpack(packer = nil)
+      if packer.is_a?(Messagepack::Packer)
+        packer.write_extension(@type, @payload)
+        packer
+      else
+        Messagepack::Packer.new.write_extension(@type, @payload).full_pack
+      end
+    end
+  end
+end

data/lib/messagepack/factory.rb

@@ -0,0 +1,349 @@
+# frozen_string_literal: true
+
+require_relative 'packer'
+require_relative 'unpacker'
+require_relative 'extensions/registry'
+
+module Messagepack
+  # Factory for creating packer/unpacker instances with custom type registration.
+  #
+  # Usage:
+  #   factory = Messagepack::Factory.new
+  #   factory.register_type(0x01, MyClass, packer: :to_msgpack, unpacker: :from_msgpack)
+  #   packer = factory.packer
+  #   unpacker = factory.unpacker
+  #
+  class Factory
+    attr_reader :packer_registry, :unpacker_registry
+
+    def initialize
+      @packer_registry = ExtensionRegistry::Packer.new
+      @unpacker_registry = ExtensionRegistry::Unpacker.new
+      @frozen = false
+    end
+
+    def freeze
+      @frozen = true
+      @packer_registry.freeze
+      @unpacker_registry.freeze
+      super
+    end
+
+    # Ensure that when the factory is duplicated (for pool), the registries are also duplicated
+    def initialize_copy(other)
+      super
+      @packer_registry = @packer_registry.dup
+      @unpacker_registry = @unpacker_registry.dup
+      @frozen = false
+    end
+
+    # Register a custom type for packing and unpacking.
+    #
+    # @param type_id [Integer] Extension type ID (-128 to 127)
+    # @param klass [Class] Ruby class to register
+    # @param options [Hash] Registration options
+    # @option options [Proc, Method, Symbol, String] :packer Serialization proc/method
+    # @option options [Proc, Method, Symbol, String] :unpacker Deserialization proc/method
+    # @option options [Boolean] :recursive Whether packer/unpacker is passed to proc
+    #
+    def register_type(type_id, klass, options = {})
+      raise FrozenError, "can't modify frozen Messagepack::Factory" if frozen?
+
+      packer = normalize_packer(options[:packer], klass)
+      unpacker = normalize_unpacker(options[:unpacker], klass)
+
+      register_type_internal(type_id, klass, packer, unpacker, options)
+    end
+
+    # Get list of registered types.
+    #
+    # @param selector [Symbol] :packer, :unpacker, or :both
+    # @return [Array<Hash>] List of registered types
+    #
+    def registered_types(selector = :both)
+      case selector
+      when :packer
+        @packer_registry.registered_types
+      when :unpacker
+        @unpacker_registry.registered_types
+      when :both
+        packer_types = @packer_registry.registered_types
+        unpacker_types = @unpacker_registry.registered_types
+
+        # Merge by type_id
+        all_types = {}
+        packer_types.each { |t| all_types[t[:type]] = t.merge(unpacker: nil) }
+        unpacker_types.each do |t|
+          if all_types.key?(t[:type])
+            all_types[t[:type]][:unpacker] = t[:unpacker]
+          else
+            all_types[t[:type]] = t.merge(packer: nil)
+          end
+        end
+
+        all_types.values
+      else
+        raise ArgumentError, "Invalid selector: #{selector}"
+      end
+    end
+
+    # Check if a type is registered.
+    #
+    # @param klass_or_type [Class, Integer] Class or type ID to check
+    # @param selector [Symbol] :packer, :unpacker, or :both
+    # @return [Boolean]
+    #
+    def type_registered?(klass_or_type, selector = :both)
+      case selector
+      when :packer
+        @packer_registry.type_registered?(klass_or_type)
+      when :unpacker
+        @unpacker_registry.type_registered?(klass_or_type)
+      when :both
+        @packer_registry.type_registered?(klass_or_type) ||
+          @unpacker_registry.type_registered?(klass_or_type)
+      else
+        raise ArgumentError, "Invalid selector: #{selector}"
+      end
+    end
+
+    # Create a new packer instance.
+    #
+    # @param io [IO] Optional IO object for streaming output
+    # @param options [Hash] Options to pass to Packer
+    # @return [Packer] New packer with registered types
+    #
+    def packer(io = nil, options = nil)
+      Packer.new(io, options).tap do |pk|
+        pk.extension_registry = @packer_registry.dup
+      end
+    end
+
+    # Create a new unpacker instance.
+    #
+    # @param io [IO] Optional IO object for streaming input
+    # @param options [Hash] Options to pass to Unpacker
+    # @return [Unpacker] New unpacker with registered types
+    #
+    def unpacker(io = nil, **options)
+      Unpacker.new(io, **options).tap do |uk|
+        uk.extension_registry = @unpacker_registry.dup
+      end
+    end
+
+    # Serialize an object to MessagePack binary.
+    #
+    # @param object Object to serialize
+    # @param io [IO] Optional IO to write to
+    # @param options [Hash] Options to pass to Packer
+    # @return [String, nil] Binary string if io is nil
+    #
+    def dump(object, *args)
+      io = args.first if args.first.respond_to?(:write)
+      options = args.last if args.last.is_a?(Hash)
+      packer(io, **(options || {})).tap { |pk| pk.write(object); pk.flush }.full_pack
+    end
+
+    # Deserialize MessagePack binary to Ruby object.
+    #
+    # @param data [String, IO] Binary data or IO to read from
+    # @param options [Hash] Options to pass to Unpacker
+    # @return [Object] Deserialized object
+    #
+    def load(data, options = nil)
+      # Check if data is an IO-like object
+      if data.respond_to?(:read)
+        unpacker(data, **(options || {})).full_unpack
+      else
+        unpacker(nil, **(options || {})).tap { |uk| uk.feed(data) }.full_unpack
+      end
+    end
+
+    alias :pack :dump
+    alias :unpack :load
+
+    # Create a pool of packers/unpackers for thread-safe reuse.
+    #
+    # @param size [Integer] Number of packers/unpackers in the pool
+    # @param options [Hash] Options for packer/unpacker creation
+    # @return [Pool] New pool instance
+    #
+    def pool(size = 1, **options)
+      Pool.new(frozen? ? self : dup.freeze, size, options)
+    end
+
+    private
+
+    def normalize_packer(packer, klass)
+      case packer
+      when nil
+        packer
+      when Proc
+        packer
+      when String, Symbol
+        # Create a proc that calls the method on the object being packed
+        ->(obj) { obj.send(packer) }
+      when Method
+        packer.to_proc
+      else
+        if packer.respond_to?(:call)
+          packer.method(:call).to_proc
+        else
+          raise ::TypeError, "invalid packer: #{packer.inspect}"
+        end
+      end
+    end
+
+    def normalize_unpacker(unpacker, klass)
+      case unpacker
+      when nil, Proc
+        unpacker
+      when String, Symbol
+        klass.method(unpacker).to_proc
+      when Method
+        unpacker.to_proc
+      else
+        if unpacker.respond_to?(:call)
+          unpacker.method(:call).to_proc
+        else
+          raise ::TypeError, "invalid unpacker: #{unpacker.inspect}"
+        end
+      end
+    end
+
+    def register_type_internal(type_id, klass, packer, unpacker, options)
+      # Validate oversized_integer_extension option
+      if options[:oversized_integer_extension]
+        unless klass == Integer
+          raise ArgumentError, "oversized_integer_extension can only be used with Integer class"
+        end
+      end
+
+      # Only use default methods when packer/unpacker are not specified at all
+      # (i.e., the keys don't exist in options)
+      has_packer = options.key?(:packer)
+      has_unpacker = options.key?(:unpacker)
+
+      if !has_packer && !has_unpacker && packer.nil?
+        packer = :to_msgpack_ext if klass.method_defined?(:to_msgpack_ext)
+      end
+
+      if !has_packer && !has_unpacker && unpacker.nil?
+        unpacker = :from_msgpack_ext if klass.respond_to?(:from_msgpack_ext)
+      end
+
+      # Normalize packer/unpacker (convert symbols to procs)
+      if packer.is_a?(Symbol)
+        method_name = packer
+        packer = ->(obj) { obj.send(method_name) }
+      end
+
+      if unpacker.is_a?(Symbol)
+        unpacker = klass.method(unpacker).to_proc
+      end
+
+      flags = 0
+      flags |= 0x01 if options[:recursive]
+      flags |= 0x02 if options[:oversized_integer_extension]
+
+      @packer_registry.register(type_id, klass, packer, flags: flags) if packer
+      @unpacker_registry.register(type_id, klass, unpacker, flags: flags) if unpacker
+    end
+
+    # Pool for thread-safe packer/unpacker reuse.
+    #
+    class Pool
+      def initialize(factory, size, options)
+        @factory = factory
+        @size = size
+        @options = options.empty? ? nil : options
+        @packers = []
+        @unpackers = []
+        @mutex = Mutex.new
+      end
+
+      # Deserialize data.
+      #
+      # @param data [String] Binary data
+      # @return [Object] Deserialized object
+      #
+      def load(data)
+        with_unpacker { |uk| uk.feed(data); uk.full_unpack }
+      end
+
+      # Serialize object.
+      #
+      # @param object Object to serialize
+      # @return [String] Binary data
+      #
+      def dump(object)
+        with_packer { |pk| pk.write(object); pk.full_pack }
+      end
+
+      alias :pack :dump
+      alias :unpack :load
+
+      # Execute block with a packer from the pool.
+      #
+      def with_packer
+        packer = nil
+        @mutex.synchronize do
+          packer = @packers.pop || @factory.packer(**(@options || {}))
+        end
+
+        yield packer
+      ensure
+        @mutex.synchronize { @packers << packer.reset } if packer
+      end
+
+      # Execute block with an unpacker from the pool.
+      #
+      def with_unpacker
+        unpacker = nil
+        @mutex.synchronize do
+          unpacker = @unpackers.pop || @factory.unpacker(**(@options || {}))
+        end
+
+        yield unpacker
+      ensure
+        @mutex.synchronize { @unpackers << unpacker.reset } if unpacker
+      end
+
+      # Get a packer from the pool and yield it to the block.
+      #
+      # @return [Object] Result of the block
+      #
+      def packer
+        packer = nil
+        @mutex.synchronize do
+          packer = @packers.pop || @factory.packer(**(@options || {}))
+        end
+
+        # Set frozen flag to prevent type registration
+        packer.freeze_for_pool
+
+        yield packer
+      ensure
+        @mutex.synchronize { @packers << packer.reset } if packer
+      end
+
+      # Get an unpacker from the pool and yield it to the block.
+      #
+      # @return [Object] Result of the block
+      #
+      def unpacker
+        unpacker = nil
+        @mutex.synchronize do
+          unpacker = @unpackers.pop || @factory.unpacker(**(@options || {}))
+        end
+
+        # Set frozen flag to prevent type registration
+        unpacker.freeze_for_pool
+
+        yield unpacker
+      ensure
+        @mutex.synchronize { @unpackers << unpacker.reset } if unpacker
+      end
+    end
+  end
+end