mtproto 0.0.8 → 0.0.10

data/lib/mtproto/tl/objects/updates_difference.rb

@@ -0,0 +1,152 @@
+# frozen_string_literal: true
+
+module MTProto
+  module TL
+    class UpdatesDifference
+      attr_reader :type, :date, :seq, :new_messages, :users, :pts, :state
+
+      def initialize(type:, date: nil, seq: nil, new_messages: [], users: [], pts: nil, state: nil)
+        @type = type
+        @date = date
+        @seq = seq
+        @new_messages = new_messages
+        @users = users
+        @pts = pts
+        @state = state
+      end
+
+      def self.parse(data)
+        constructor = data[0, 4].unpack1('L<')
+
+        case constructor
+        when Constructors::UPDATES_DIFFERENCE_EMPTY
+          parse_empty(data)
+        when Constructors::UPDATES_DIFFERENCE, Constructors::UPDATES_DIFFERENCE_SLICE
+          parse_difference(data, constructor)
+        when Constructors::UPDATES_DIFFERENCE_TOO_LONG
+          new(type: :too_long, pts: data[4, 4].unpack1('L<'))
+        else
+          warn "Unknown Difference constructor: 0x#{constructor.to_s(16)}, treating as empty"
+          new(type: :empty, date: Time.now.to_i, seq: 0)
+        end
+      end
+
+      class << self
+        private
+
+        def parse_empty(data)
+          new(
+            type: :empty,
+            date: data[4, 4].unpack1('L<'),
+            seq: data[8, 4].unpack1('L<')
+          )
+        end
+
+        def parse_difference(data, constructor)
+          offset = 4
+
+          messages, offset = parse_messages_vector(data, offset)
+          offset = skip_vector(data, offset) # encrypted_messages
+          offset = skip_vector(data, offset) # other_updates
+          offset = skip_vector(data, offset) # chats
+          users, = parse_users_vector(data, offset)
+
+          new(
+            type: constructor == Constructors::UPDATES_DIFFERENCE ? :difference : :slice,
+            new_messages: messages,
+            users: users
+          )
+        end
+
+        def parse_messages_vector(data, offset)
+          offset += 4 # vector constructor
+          count = data[offset, 4].unpack1('L<')
+          offset += 4
+
+          messages = []
+          count.times do
+            msg, offset = parse_message(data, offset)
+            messages << msg if msg
+          end
+          [messages, offset]
+        end
+
+        def parse_message(data, offset)
+          msg_constructor = data[offset, 4].unpack1('L<')
+          offset += 4
+
+          return [nil, offset] unless msg_constructor == Constructors::MESSAGE
+
+          flags = data[offset, 4].unpack1('L<')
+          offset += 4
+          id = data[offset, 4].unpack1('L<')
+          offset += 4
+          offset += 12 if flags.anybits?(1 << 8) # from_id
+          offset += 4 if flags.anybits?(1 << 29) # from_boosts_applied
+          offset += 12 # peer_id
+          offset += 8 if flags.anybits?(1 << 11) # via_bot_id
+          date = data[offset, 4].unpack1('L<')
+          offset += 4
+          message_text, offset = read_tl_string(data, offset)
+
+          [{ id: id, date: date, message: message_text, flags: flags }, offset]
+        rescue StandardError
+          [nil, offset]
+        end
+
+        def parse_users_vector(data, offset)
+          offset += 4 # vector constructor
+          count = data[offset, 4].unpack1('L<')
+          offset += 4
+
+          users = []
+          count.times do
+            user_constructor = data[offset, 4].unpack1('L<')
+            offset += 4
+            next unless user_constructor == Constructors::USER
+
+            flags = data[offset, 4].unpack1('L<')
+            offset += 4
+            offset += 4 # flags2
+            user_id = data[offset, 8].unpack1('Q<')
+            offset += 8
+
+            access_hash = nil
+            if flags.anybits?(1 << 0)
+              access_hash = data[offset, 8].unpack1('Q<')
+              offset += 8
+            end
+
+            users << { id: user_id, access_hash: access_hash }
+          rescue StandardError
+            break
+          end
+          [users, offset]
+        end
+
+        def skip_vector(data, offset)
+          offset += 4 # vector constructor
+          count = data[offset, 4].unpack1('L<')
+          offset += 4
+          count.times { offset += 4 }
+          offset
+        end
+
+        def read_tl_string(data, offset)
+          first_byte = data.getbyte(offset)
+          if first_byte == 254
+            len = data.getbyte(offset + 1) |
+                  (data.getbyte(offset + 2) << 8) |
+                  (data.getbyte(offset + 3) << 16)
+            total = 4 + len
+          else
+            len = first_byte
+            total = 1 + len
+          end
+          padding = (4 - (total % 4)) % 4
+          [data[offset + (total - len), len], offset + total + padding]
+        end
+      end
+    end
+  end
+end

data/lib/mtproto/tl/objects/updates_state.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module MTProto
+  module TL
+    class UpdatesState
+      attr_accessor :pts, :qts, :date, :seq, :unread_count
+
+      def initialize(pts:, qts:, date:, seq:, unread_count:)
+        @pts = pts
+        @qts = qts
+        @date = date
+        @seq = seq
+        @unread_count = unread_count
+      end
+
+      def self.parse(data)
+        constructor = data[0, 4].unpack1('L<')
+        raise UnexpectedConstructorError, constructor unless constructor == Constructors::UPDATES_STATE
+
+        offset = 4
+        pts = data[offset, 4].unpack1('L<')
+        offset += 4
+        qts = data[offset, 4].unpack1('L<')
+        offset += 4
+        date = data[offset, 4].unpack1('L<')
+        offset += 4
+        seq = data[offset, 4].unpack1('L<')
+        offset += 4
+        unread_count = data[offset, 4].unpack1('L<')
+
+        new(pts: pts, qts: qts, date: date, seq: seq, unread_count: unread_count)
+      end
+    end
+  end
+end

data/lib/mtproto/tl/objects/users.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+module MTProto
+  module TL
+    class Users
+      User = Struct.new(:id, :access_hash, :first_name, :last_name, :username, :phone, :flags, :flags2,
+        keyword_init: true)
+
+      def self.parse(data)
+        constructor = data[0, 4].unpack1('L<')
+        raise "Expected Vector constructor, got 0x#{constructor.to_s(16)}" unless constructor == Constructors::VECTOR
+
+        offset = 4
+        count = data[offset, 4].unpack1('L<')
+        offset += 4
+
+        users = []
+        count.times do
+          user_constructor = data[offset, 4].unpack1('L<')
+          offset += 4
+          next unless user_constructor == Constructors::USER
+
+          flags = data[offset, 4].unpack1('L<')
+          offset += 4
+
+          flags2 = data[offset, 4].unpack1('L<')
+          offset += 4
+
+          id = data[offset, 8].unpack1('Q<')
+          offset += 8
+
+          access_hash = nil
+          if flags.anybits?(1 << 0)
+            access_hash = data[offset, 8].unpack1('Q<')
+            offset += 8
+          end
+
+          first_name = nil
+          first_name, offset = read_tl_string(data, offset) if flags.anybits?(1 << 1)
+
+          last_name = nil
+          last_name, offset = read_tl_string(data, offset) if flags.anybits?(1 << 2)
+
+          username = nil
+          username, offset = read_tl_string(data, offset) if flags.anybits?(1 << 3)
+
+          phone = nil
+          phone, offset = read_tl_string(data, offset) if flags.anybits?(1 << 4)
+
+          users << User.new(
+            id: id, access_hash: access_hash,
+            first_name: first_name, last_name: last_name,
+            username: username, phone: phone,
+            flags: flags, flags2: flags2
+          )
+
+          break
+        end
+
+        users
+      end
+
+      class << self
+        private
+
+        def read_tl_string(data, offset)
+          first_byte = data.getbyte(offset)
+          if first_byte == 254
+            len = data.getbyte(offset + 1) |
+                  (data.getbyte(offset + 2) << 8) |
+                  (data.getbyte(offset + 3) << 16)
+            total = 4 + len
+          else
+            len = first_byte
+            total = 1 + len
+          end
+          padding = (4 - (total % 4)) % 4
+          [data[offset + (total - len), len], offset + total + padding]
+        end
+      end
+    end
+  end
+end

data/lib/mtproto/tl/schema.rb

@@ -0,0 +1,102 @@
+# frozen_string_literal: true
+
+require 'json'
+
+module MTProto
+  module TL
+    class Schema
+      PRIMITIVES = {
+        'int' => 4, 'long' => 8, 'double' => 8,
+        'int128' => 16, 'int256' => 32
+      }.freeze
+
+      def initialize(path)
+        raw = JSON.parse(File.read(path))
+        @constructors = {}
+        raw['constructors'].each do |c|
+          uid = c['id'].to_i & 0xFFFFFFFF
+          @constructors[uid] = c['params']
+        end
+      end
+
+      def skip(data, offset)
+        constructor = data[offset, 4].unpack1('L<')
+        offset += 4
+        skip_params(data, offset, constructor)
+      end
+
+      def skip_vector(data, offset, &block)
+        offset += 4 # vector constructor 0x1cb5c415
+        count = data[offset, 4].unpack1('L<')
+        offset += 4
+        count.times { offset = block ? yield(data, offset) : skip(data, offset) }
+        offset
+      end
+
+      def skip_params(data, offset, constructor)
+        params = @constructors[constructor]
+        raise "Unknown constructor: 0x#{constructor.to_s(16).rjust(8, '0')}" unless params
+
+        flag_values = {}
+
+        params.each do |param|
+          type = param['type']
+
+          if type == '#'
+            flag_values[param['name']] = data[offset, 4].unpack1('L<')
+            offset += 4
+            next
+          end
+
+          if (m = type.match(/\A(\w+)\.(\d+)\?(.+)\z/))
+            flag_name = m[1]
+            bit = m[2].to_i
+            inner_type = m[3]
+            next if inner_type == 'true'
+            next unless flag_values[flag_name]&.anybits?(1 << bit)
+
+            offset = skip_type(data, offset, inner_type)
+          else
+            offset = skip_type(data, offset, type)
+          end
+        end
+        offset
+      end
+
+      private
+
+      def skip_type(data, offset, type)
+        if (size = PRIMITIVES[type])
+          return offset + size
+        end
+
+        case type
+        when 'string', 'bytes'
+          skip_tl_string(data, offset)
+        when 'Bool'
+          offset + 4
+        when /\AVector<(.+)>\z/
+          inner = Regexp.last_match(1)
+          skip_vector(data, offset) { |d, o| skip_type(d, o, inner) }
+        else
+          skip(data, offset)
+        end
+      end
+
+      def skip_tl_string(data, offset)
+        first_byte = data.getbyte(offset)
+        if first_byte == 254
+          len = data.getbyte(offset + 1) |
+                (data.getbyte(offset + 2) << 8) |
+                (data.getbyte(offset + 3) << 16)
+          total = 4 + len
+        else
+          len = first_byte
+          total = 1 + len
+        end
+        padding = (4 - (total % 4)) % 4
+        offset + total + padding
+      end
+    end
+  end
+end

data/lib/mtproto/transport/abridged_packet_codec.rb

@@ -6,29 +6,52 @@ module MTProto
       TAG = "\xef".b
       OBFUSCATE_TAG = "\xef\xef\xef\xef".b
 
-      def encode_packet(data)
-        length = data.bytesize >> 2
+      def initialize(stream)
+        @stream = stream
+      end
+
+      def send(packet)
+        length = packet.size >> 2
 
-        length_prefix = if length < 127
-          [length].pack('C')
+        if length < 127
+          @stream.write [length].pack('C')
         else
-          "\x7f".b + [length].pack('L<')[0, 3]
+          @stream.write "\x7f".b + [length].pack('L<')[0, 3]
         end
 
-        length_prefix + data
+        @stream.write packet.data.pack('C*')
       end
 
-      def decode_packet(data)
-        length = data.unpack1('C')
-        offset = 1
+      MAX_PACKET_SIZE = 1_048_576
+
+      def scan
+        first_byte = read_exactly(1)
+        length = first_byte.unpack1('C')
 
         if length >= 127
-          length = (data[1, 3] + "\x00").unpack1('L<')
-          offset = 4
+          length_bytes = read_exactly(3)
+          length = "#{length_bytes}\x00".unpack1('L<')
         end
 
         actual_length = length << 2
-        data[offset, actual_length]
+
+        raise PacketReadError, "Packet too large: #{actual_length} bytes" if actual_length > MAX_PACKET_SIZE
+
+        data = read_exactly(actual_length)
+        Packet.new(data.bytes)
+      end
+
+      private
+
+      def read_exactly(bytes_needed)
+        result = ''.b
+        while result.bytesize < bytes_needed
+          chunk = @stream.read(bytes_needed - result.bytesize)
+          raise IOError, 'EOF while reading' if chunk.nil? || chunk.empty?
+
+          result += chunk
+        end
+        result
       end
     end
   end

data/lib/mtproto/transport/connection.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module MTProto
+  module Transport
+    class Connection
+      extend DelegateMethods
+
+      attr_reader :transport
+
+      def initialize(transport, wait_timeout: 60, read_timeout: 30)
+        @transport = transport
+        @transport.wait_timeout = wait_timeout
+        @transport.read_timeout = read_timeout
+      end
+
+      delegate :connect!, :disconnect!, :connected?, :not_connected?,
+        :send, :receive,
+        :wait_timeout, :wait_timeout=,
+        :read_timeout, :read_timeout=,
+        to: :transport
+    end
+  end
+end

data/lib/mtproto/transport/errors.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module MTProto
+  module Transport
+    class ConnectionError < StandardError; end
+    class NotConnectedError < ConnectionError; end
+    class ConnectionClosedError < ConnectionError; end
+    class ReceiveTimeoutError < ConnectionError; end
+    class PacketReadError < ConnectionError; end
+  end
+end

data/lib/mtproto/transport/packet.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module MTProto
+  module Transport
+    class Packet
+      attr_reader :data
+
+      def initialize(data)
+        raise ArgumentError, 'data must be an Array' unless data.is_a?(Array)
+
+        @data = data
+      end
+
+      def size
+        data.length
+      end
+    end
+  end
+end

data/lib/mtproto/transport/tcp_connection.rb

@@ -5,86 +5,97 @@ require 'timeout'
 
 module MTProto
   module Transport
-    class ConnectionError < StandardError; end
-
     class TCPConnection
-      attr_reader :host, :port, :codec
+      attr_reader :host, :port
+      attr_accessor :wait_timeout, :read_timeout
 
-      def initialize(host, port, codec)
+      def initialize(host, port, codec_class: AbridgedPacketCodec, wait_timeout: 60, read_timeout: 3)
         @host = host
         @port = port
-        @codec = codec
+        @codec_class = codec_class
+        @wait_timeout = wait_timeout
+        @read_timeout = read_timeout
         @socket = nil
+        @codec = nil
       end
 
       def connect!
         return if connected?
 
         @socket = TCPSocket.new(@host, @port)
+        @codec = @codec_class.new(@socket)
 
-        send_init_tag if @codec.class.const_defined?(:TAG)
+        write_init_tag
       end
 
-      def connected?
-        !@socket.nil? && !@socket.closed?
+      def disconnect!
+        return unless @socket
+
+        @socket.close
+      rescue StandardError
+        nil
+      ensure
+        @socket = nil
+        @codec = nil
       end
 
-      def send(data)
-        raise ConnectionError, 'Not connected' unless connected?
+      def not_connected?
+        @socket.nil? || @socket.closed?
+      end
 
-        encoded = @codec.encode_packet(data)
-        @socket.write(encoded)
+      def connected?
+        not not_connected?
       end
 
-      def recv(timeout: 60)
-        raise ConnectionError, 'Not connected' unless connected?
+      def send(packet)
+        raise NotConnectedError, 'Not connected' unless connected?
 
-        Timeout.timeout(timeout) do
-          read_packet
-        end
-      rescue Timeout::Error
-        raise ConnectionError, 'Receive timeout'
+        @codec.send(packet)
       end
 
-      def close
-        return unless @socket
+      def receive(&)
+        raise NotConnectedError, 'Not connected' unless connected?
 
-        @socket.close
-      rescue StandardError
-        nil
-      ensure
-        @socket = nil
+        if block_given?
+          receive_loop(&)
+        else
+          receive_once
+        end
       end
 
       private
 
-      def send_init_tag
-        tag = @codec.class.const_get(:TAG)
-        @socket.write(tag) if tag
-      end
+      def receive_once
+        return nil unless @socket.wait_readable(@wait_timeout)
 
-      def read_packet
-        first_byte = read_exactly(1)
-        length = first_byte.unpack1('C')
+        read_packet
+      end
 
-        if length >= 127
-          length_bytes = read_exactly(3)
-          length = (length_bytes + "\x00").unpack1('L<')
+      def receive_loop
+        loop do
+          @socket.wait_readable
+          packet = read_packet
+          yield packet, nil
+        rescue ConnectionClosedError, NotConnectedError
+          raise
+        rescue PacketReadError => e
+          yield nil, e
         end
+      end
 
-        actual_length = length << 2
-        read_exactly(actual_length)
+      def read_packet
+        Timeout.timeout(@read_timeout) { @codec.scan }
+      rescue Timeout::Error
+        raise PacketReadError, 'Packet read timeout'
+      rescue IOError, Errno::ECONNRESET => e
+        raise ConnectionClosedError, e.message
       end
 
-      def read_exactly(bytes_needed)
-        result = ''.b
-        while result.bytesize < bytes_needed
-          chunk = @socket.read(bytes_needed - result.bytesize)
-          raise ConnectionError, 'EOF while reading' if chunk.nil? || chunk.empty?
+      def write_init_tag
+        return unless @codec_class.const_defined?(:TAG)
 
-          result += chunk
-        end
-        result
+        tag = @codec_class.const_get(:TAG)
+        @socket.write tag if tag
       end
     end
   end