omq 0.1.0

data/lib/omq/zmtp/codec/command.rb

@@ -0,0 +1,207 @@
+# frozen_string_literal: true
+
+module OMQ
+  module ZMTP
+    module Codec
+      # ZMTP command encode/decode.
+      #
+      # Command frame body format:
+      #   1 byte:    command name length
+      #   N bytes:   command name
+      #   remaining: command data
+      #
+      # READY command data = property list:
+      #   1 byte:  property name length
+      #   N bytes: property name
+      #   4 bytes: property value length (big-endian)
+      #   N bytes: property value
+      #
+      class Command
+        # @return [String] command name (e.g. "READY", "SUBSCRIBE")
+        #
+        attr_reader :name
+
+        # @return [String] command data (binary)
+        #
+        attr_reader :data
+
+        # @param name [String] command name
+        # @param data [String] command data
+        #
+        def initialize(name, data = "".b)
+          @name = name
+          @data = data.b
+        end
+
+        # Encodes as a command frame body.
+        #
+        # @return [String] binary body (name-length + name + data)
+        #
+        def to_body
+          name_bytes = @name.b
+          buf = IO::Buffer.new(1 + name_bytes.bytesize + @data.bytesize)
+          buf.set_value(:U8, 0, name_bytes.bytesize)
+          buf.set_string(name_bytes, 1)
+          buf.set_string(@data, 1 + name_bytes.bytesize)
+          buf.get_string(0, buf.size, Encoding::BINARY)
+        end
+
+        # Encodes as a complete command Frame.
+        #
+        # @return [Frame]
+        #
+        def to_frame
+          Frame.new(to_body, command: true)
+        end
+
+        # Decodes a command from a frame body.
+        #
+        # @param body [String] binary frame body
+        # @return [Command]
+        # @raise [ProtocolError] on malformed command
+        #
+        def self.from_body(body)
+          body = body.b
+          raise ProtocolError, "command body too short" if body.bytesize < 1
+
+          buf = IO::Buffer.for(body)
+          name_len = buf.get_value(:U8, 0)
+
+          raise ProtocolError, "command name truncated" if body.bytesize < 1 + name_len
+
+          name = buf.get_string(1, name_len, Encoding::BINARY)
+          data = body.byteslice(1 + name_len..)
+          new(name, data)
+        end
+
+        # Builds a READY command with Socket-Type and Identity properties.
+        #
+        # @param socket_type [String] e.g. "REQ", "REP", "PAIR"
+        # @param identity [String] peer identity (can be empty)
+        # @return [Command]
+        #
+        def self.ready(socket_type:, identity: "")
+          props = encode_properties(
+            "Socket-Type" => socket_type,
+            "Identity"    => identity,
+          )
+          new("READY", props)
+        end
+
+        # Builds a SUBSCRIBE command.
+        #
+        # @param prefix [String] subscription prefix
+        # @return [Command]
+        #
+        def self.subscribe(prefix)
+          new("SUBSCRIBE", prefix.b)
+        end
+
+        # Builds a CANCEL command (unsubscribe).
+        #
+        # @param prefix [String] subscription prefix to cancel
+        # @return [Command]
+        #
+        def self.cancel(prefix)
+          new("CANCEL", prefix.b)
+        end
+
+        # Builds a PING command.
+        #
+        # @param ttl [Numeric] time-to-live in seconds (sent as deciseconds)
+        # @param context [String] optional context bytes (up to 16 bytes)
+        # @return [Command]
+        #
+        def self.ping(ttl: 0, context: "".b)
+          # TTL is encoded as 2-byte big-endian value in tenths of a second
+          ttl_ds = (ttl * 10).to_i
+          buf = IO::Buffer.new(2 + context.bytesize)
+          buf.set_value(:U16, 0, ttl_ds)
+          buf.set_string(context.b, 2) if context.bytesize > 0
+          new("PING", buf.get_string(0, buf.size, Encoding::BINARY))
+        end
+
+        # Builds a PONG command.
+        #
+        # @param context [String] context bytes from the PING
+        # @return [Command]
+        #
+        def self.pong(context: "".b)
+          new("PONG", context.b)
+        end
+
+        # Extracts TTL (in seconds) and context from a PING command's data.
+        #
+        # @return [Array(Numeric, String)] [ttl_seconds, context_bytes]
+        #
+        def ping_ttl_and_context
+          buf = IO::Buffer.for(@data)
+          ttl_ds = buf.get_value(:U16, 0)
+          context = @data.bytesize > 2 ? @data.byteslice(2..) : "".b
+          [ttl_ds / 10.0, context]
+        end
+
+        # Parses READY command data as a property list.
+        #
+        # @return [Hash<String, String>] property name => value
+        # @raise [ProtocolError] on malformed properties
+        #
+        def properties
+          self.class.decode_properties(@data)
+        end
+
+        # Encodes a hash of properties into ZMTP property list format.
+        #
+        # @param props [Hash<String, String>]
+        # @return [String] binary property list
+        #
+        def self.encode_properties(props)
+          parts = props.map do |name, value|
+            name_bytes = name.b
+            value_bytes = value.b
+            buf = IO::Buffer.new(1 + name_bytes.bytesize + 4 + value_bytes.bytesize)
+            buf.set_value(:U8, 0, name_bytes.bytesize)
+            buf.set_string(name_bytes, 1)
+            buf.set_value(:U32, 1 + name_bytes.bytesize, value_bytes.bytesize) # big-endian
+            buf.set_string(value_bytes, 1 + name_bytes.bytesize + 4)
+            buf.get_string(0, buf.size, Encoding::BINARY)
+          end
+          parts.join
+        end
+        # Decodes a ZMTP property list from binary data.
+        #
+        # @param data [String] binary property list
+        # @return [Hash<String, String>] property name => value
+        # @raise [ProtocolError] on malformed properties
+        #
+        def self.decode_properties(data)
+          result = {}
+          buf    = IO::Buffer.for(data)
+          offset = 0
+
+          while offset < data.bytesize
+            raise ProtocolError, "property name truncated" if offset + 1 > data.bytesize
+            name_len = buf.get_value(:U8, offset)
+            offset += 1
+
+            raise ProtocolError, "property name truncated" if offset + name_len > data.bytesize
+            name = buf.get_string(offset, name_len, Encoding::BINARY)
+            offset += name_len
+
+            raise ProtocolError, "property value length truncated" if offset + 4 > data.bytesize
+            value_len = buf.get_value(:U32, offset)
+            offset += 4
+
+            raise ProtocolError, "property value truncated" if offset + value_len > data.bytesize
+            value = buf.get_string(offset, value_len, Encoding::BINARY)
+            offset += value_len
+
+            result[name] = value
+          end
+
+          result
+        end
+      end
+    end
+  end
+end

data/lib/omq/zmtp/codec/frame.rb

@@ -0,0 +1,104 @@
+# frozen_string_literal: true
+
+module OMQ
+  module ZMTP
+    module Codec
+      # ZMTP frame encode/decode.
+      #
+      # Wire format:
+      #   Byte 0:   flags (bit 0=MORE, bit 1=LONG, bit 2=COMMAND)
+      #   Next 1-8: size (1-byte if short, 8-byte big-endian if LONG)
+      #   Next N:   body
+      #
+      class Frame
+        FLAGS_MORE    = 0x01
+        FLAGS_LONG    = 0x02
+        FLAGS_COMMAND = 0x04
+
+        # Short frame: 1-byte size, max body 255 bytes.
+        #
+        SHORT_MAX = 255
+
+        # @return [String] frame body (binary)
+        #
+        attr_reader :body
+
+        # @param body [String] frame body
+        # @param more [Boolean] more frames follow
+        # @param command [Boolean] this is a command frame
+        #
+        def initialize(body, more: false, command: false)
+          @body    = body.b
+          @more    = more
+          @command = command
+        end
+
+        # @return [Boolean] true if more frames follow in this message
+        #
+        def more?    = @more
+
+        # @return [Boolean] true if this is a command frame
+        #
+        def command? = @command
+
+        # Encodes to wire bytes.
+        #
+        # @return [String] binary wire representation (flags + size + body)
+        #
+        def to_wire
+          size = @body.bytesize
+          flags = 0
+          flags |= FLAGS_MORE if @more
+          flags |= FLAGS_COMMAND if @command
+
+          if size > SHORT_MAX
+            flags |= FLAGS_LONG
+            buf = IO::Buffer.new(9 + size)
+            buf.set_value(:U8, 0, flags)
+            buf.set_value(:U64, 1, size) # big-endian
+            buf.set_string(@body, 9)
+            buf.get_string(0, 9 + size, Encoding::BINARY)
+          else
+            buf = IO::Buffer.new(2 + size)
+            buf.set_value(:U8, 0, flags)
+            buf.set_value(:U8, 1, size)
+            buf.set_string(@body, 2)
+            buf.get_string(0, 2 + size, Encoding::BINARY)
+          end
+        end
+
+        # Reads one frame from an IO-like object.
+        #
+        # @param io [#read] must support read(n) returning exactly n bytes
+        # @return [Frame]
+        # @raise [ProtocolError] on invalid frame
+        # @raise [EOFError] if the connection is closed
+        #
+        def self.read_from(io)
+          flags_byte = io.read_exactly(1)
+          flags_buf = IO::Buffer.for(flags_byte)
+          flags = flags_buf.get_value(:U8, 0)
+
+          more = (flags & FLAGS_MORE) != 0
+          long = (flags & FLAGS_LONG) != 0
+          command = (flags & FLAGS_COMMAND) != 0
+
+          if long
+            size_bytes = io.read_exactly(8)
+            size_buf = IO::Buffer.for(size_bytes)
+            size = size_buf.get_value(:U64, 0) # big-endian
+          else
+            size_byte = io.read_exactly(1)
+            size_buf = IO::Buffer.for(size_byte)
+            size = size_buf.get_value(:U8, 0)
+          end
+
+          body = size > 0 ? io.read_exactly(size) : "".b
+
+          new(body, more: more, command: command)
+        end
+
+      end
+    end
+  end
+end

data/lib/omq/zmtp/codec/greeting.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+module OMQ
+  module ZMTP
+    module Codec
+      # ZMTP 3.1 greeting encode/decode.
+      #
+      # The greeting is always exactly 64 bytes:
+      #   Offset  Bytes  Field
+      #   0       1      0xFF (signature start)
+      #   1-8     8      0x00 padding
+      #   9       1      0x7F (signature end)
+      #   10      1      major version
+      #   11      1      minor version
+      #   12-31   20     mechanism (null-padded ASCII)
+      #   32      1      as-server flag (0x00 or 0x01)
+      #   33-63   31     filler (0x00)
+      #
+      module Greeting
+        SIZE             = 64
+        SIGNATURE_START  = 0xFF
+        SIGNATURE_END    = 0x7F
+        VERSION_MAJOR    = 3
+        VERSION_MINOR    = 1
+        MECHANISM_OFFSET = 12
+        MECHANISM_LENGTH = 20
+        AS_SERVER_OFFSET = 32
+
+        # Encodes a ZMTP 3.1 greeting.
+        #
+        # @param mechanism [String] security mechanism name (e.g. "NULL")
+        # @param as_server [Boolean] whether this peer is the server
+        # @return [String] 64-byte binary greeting
+        #
+        def self.encode(mechanism: "NULL", as_server: false)
+          buf = IO::Buffer.new(SIZE)
+          buf.clear
+
+          # Signature
+          buf.set_value(:U8, 0, SIGNATURE_START)
+          # bytes 1-8 are already 0x00
+          buf.set_value(:U8, 9, SIGNATURE_END)
+
+          # Version
+          buf.set_value(:U8, 10, VERSION_MAJOR)
+          buf.set_value(:U8, 11, VERSION_MINOR)
+
+          # Mechanism (null-padded)
+          buf.set_string(mechanism.b, MECHANISM_OFFSET)
+
+          # As-server flag
+          buf.set_value(:U8, AS_SERVER_OFFSET, as_server ? 1 : 0)
+
+          # Filler bytes 33-63 are already 0x00
+          buf.get_string(0, SIZE, Encoding::BINARY)
+        end
+
+        # Decodes a ZMTP greeting.
+        #
+        # @param data [String] 64-byte binary greeting
+        # @return [Hash] { major:, minor:, mechanism:, as_server: }
+        # @raise [ProtocolError] on invalid greeting
+        #
+        def self.decode(data)
+          raise ProtocolError, "greeting too short (#{data.bytesize} bytes)" if data.bytesize < SIZE
+
+          buf = IO::Buffer.for(data.b)
+
+          # Validate signature
+          unless buf.get_value(:U8, 0) == SIGNATURE_START &&
+                 buf.get_value(:U8, 9) == SIGNATURE_END
+            raise ProtocolError, "invalid greeting signature"
+          end
+
+          major = buf.get_value(:U8, 10)
+          minor = buf.get_value(:U8, 11)
+
+          unless major >= 3
+            raise ProtocolError, "unsupported ZMTP version #{major}.#{minor} (need >= 3.0)"
+          end
+
+          mechanism = buf.get_string(MECHANISM_OFFSET, MECHANISM_LENGTH, Encoding::BINARY)
+                        .delete("\x00")
+          as_server = buf.get_value(:U8, AS_SERVER_OFFSET) == 1
+
+          {
+            major:     major,
+            minor:     minor,
+            mechanism: mechanism,
+            as_server: as_server,
+          }
+        end
+      end
+    end
+  end
+end

data/lib/omq/zmtp/codec.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module OMQ
+  module ZMTP
+    # ZMTP 3.1 wire protocol codec.
+    #
+    module Codec
+    end
+
+    # Raised on ZMTP protocol violations.
+    #
+    class ProtocolError < RuntimeError; end
+  end
+end
+
+require_relative "codec/greeting"
+require_relative "codec/frame"
+require_relative "codec/command"

data/lib/omq/zmtp/connection.rb

@@ -0,0 +1,233 @@
+# frozen_string_literal: true
+
+module OMQ
+  module ZMTP
+    # Manages one ZMTP peer connection over any transport IO.
+    #
+    # Delegates the security handshake to a Mechanism object (Null, Curve, etc.),
+    # then provides message send/receive and command send/receive on top of the
+    # framing codec.
+    #
+    class Connection
+      # @return [String] peer's socket type (from READY handshake)
+      #
+      attr_reader :peer_socket_type
+
+      # @return [String] peer's identity (from READY handshake)
+      #
+      attr_reader :peer_identity
+
+      # @return [Object] transport IO (#read, #write, #close)
+      #
+      attr_reader :io
+
+      # @param io [#read, #write, #close] transport IO
+      # @param socket_type [String] our socket type name (e.g. "REQ")
+      # @param identity [String] our identity
+      # @param as_server [Boolean] whether we are the server side
+      # @param mechanism [Mechanism::Null, Mechanism::Curve] security mechanism
+      # @param heartbeat_interval [Numeric, nil] heartbeat interval in seconds
+      # @param heartbeat_ttl [Numeric, nil] TTL to send in PING
+      # @param heartbeat_timeout [Numeric, nil] timeout for PONG
+      # @param max_message_size [Integer, nil] max frame size in bytes, nil = unlimited
+      #
+      def initialize(io, socket_type:, identity: "", as_server: false,
+                     mechanism: nil,
+                     heartbeat_interval: nil, heartbeat_ttl: nil, heartbeat_timeout: nil,
+                     max_message_size: nil)
+        @io                  = io
+        @socket_type         = socket_type
+        @identity            = identity
+        @as_server           = as_server
+        @mechanism           = mechanism || Mechanism::Null.new
+        @peer_socket_type    = nil
+        @peer_identity       = nil
+        @mutex               = Mutex.new
+        @heartbeat_interval  = heartbeat_interval
+        @heartbeat_ttl       = heartbeat_ttl || heartbeat_interval
+        @heartbeat_timeout   = heartbeat_timeout || heartbeat_interval
+        @last_received_at    = nil
+        @heartbeat_task      = nil
+        @max_message_size    = max_message_size
+      end
+
+      # Performs the full ZMTP handshake via the configured mechanism.
+      #
+      # @return [void]
+      # @raise [ProtocolError] on handshake failure
+      #
+      def handshake!
+        result = @mechanism.handshake!(
+          @io,
+          as_server:   @as_server,
+          socket_type: @socket_type,
+          identity:    @identity,
+        )
+
+        @peer_socket_type = result[:peer_socket_type]
+        @peer_identity    = result[:peer_identity]
+
+        unless @peer_socket_type
+          raise ProtocolError, "peer READY missing Socket-Type"
+        end
+
+        unless ZMTP::VALID_PEERS[@socket_type.to_sym]&.include?(@peer_socket_type.to_sym)
+          raise ProtocolError,
+                "incompatible socket types: #{@socket_type} cannot connect to #{@peer_socket_type}"
+        end
+      end
+
+      # Sends a multi-frame message.
+      #
+      # @param parts [Array<String>] message frames
+      # @return [void]
+      #
+      def send_message(parts)
+        @mutex.synchronize do
+          parts.each_with_index do |part, i|
+            more = i < parts.size - 1
+            if @mechanism.encrypted?
+              @io.write(@mechanism.encrypt(part.b, more: more))
+            else
+              @io.write(Codec::Frame.new(part, more: more).to_wire)
+            end
+          end
+        end
+      end
+
+      # Receives a multi-frame message.
+      # PING/PONG commands are handled automatically by #read_frame.
+      #
+      # @return [Array<String>] message frames
+      # @raise [EOFError] if connection is closed
+      #
+      def receive_message
+        frames = []
+        loop do
+          frame = read_frame
+          if frame.command?
+            yield frame if block_given?
+            next
+          end
+          frames << frame.body
+          break unless frame.more?
+        end
+        frames
+      end
+
+      # Starts the heartbeat sender task. Call after handshake.
+      #
+      # @return [#stop, nil] the heartbeat task, or nil if disabled
+      #
+      def start_heartbeat
+        return nil unless @heartbeat_interval
+        @last_received_at = monotonic_now
+        @heartbeat_task = Reactor.spawn_pump do
+          loop do
+            sleep @heartbeat_interval
+            # Send PING with TTL
+            send_command(Codec::Command.ping(
+              ttl:     @heartbeat_ttl || 0,
+              context: "".b,
+            ))
+            # Check if peer has gone silent
+            if @heartbeat_timeout && @last_received_at
+              elapsed = monotonic_now - @last_received_at
+              if elapsed > @heartbeat_timeout
+                close
+                break
+              end
+            end
+          end
+        rescue IOError, EOFError
+          # connection closed
+        end
+      end
+
+      # Sends a command.
+      #
+      # @param command [Codec::Command]
+      # @return [void]
+      #
+      def send_command(command)
+        @mutex.synchronize do
+          if @mechanism.encrypted?
+            @io.write(@mechanism.encrypt(command.to_body, command: true))
+          else
+            @io.write(command.to_frame.to_wire)
+          end
+        end
+      end
+
+      # Reads one frame from the wire. Handles PING/PONG automatically.
+      # When using an encrypted mechanism, MESSAGE commands are decrypted
+      # back to ZMTP frames transparently.
+      #
+      # @return [Codec::Frame]
+      # @raise [EOFError] if connection is closed
+      #
+      def read_frame
+        loop do
+          frame = Codec::Frame.read_from(@io)
+          touch_heartbeat
+
+          # When CURVE is active, every wire frame is a MESSAGE envelope
+          # (data frame with "\x07MESSAGE" prefix). Decrypt to recover the
+          # inner ZMTP frame. This check only runs when encrypted? is true,
+          # so user data frames are never misdetected.
+          if @mechanism.encrypted? && frame.body.bytesize > 8 && frame.body.byteslice(0, 8) == "\x07MESSAGE".b
+            frame = @mechanism.decrypt(frame)
+          end
+
+          if @max_message_size && !frame.command? && frame.body.bytesize > @max_message_size
+            close
+            raise ProtocolError, "frame size #{frame.body.bytesize} exceeds max_message_size #{@max_message_size}"
+          end
+          if frame.command?
+            cmd = Codec::Command.from_body(frame.body)
+            case cmd.name
+            when "PING"
+              _, context = cmd.ping_ttl_and_context
+              send_command(Codec::Command.pong(context: context))
+              next
+            when "PONG"
+              next
+            end
+          end
+          return frame
+        end
+      end
+
+      # Closes the connection.
+      #
+      # @return [void]
+      #
+      def close
+        @heartbeat_task&.stop rescue nil
+        @io.close
+      rescue IOError
+        # already closed
+      end
+
+      private
+
+      def touch_heartbeat
+        @last_received_at = monotonic_now if @heartbeat_interval
+      end
+
+      def monotonic_now
+        Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      end
+
+      # Sends one frame to the wire.
+      #
+      # @param frame [Codec::Frame]
+      # @return [void]
+      #
+      def send_frame(frame)
+        @mutex.synchronize { @io.write(frame.to_wire) }
+      end
+
+    end
+  end
+end