mtproto 0.0.7 → 0.0.9

data/lib/mtproto/client/rpc.rb

@@ -1,140 +1,100 @@
 # frozen_string_literal: true
 
 require_relative '../message_id'
+require_relative 'rpc/response'
 
 module MTProto
   class Client
     class RPC
       CONSTRUCTOR_RPC_RESULT = 0xf35c6d01
-      CONSTRUCTOR_MSGS_ACK = 0x62d6b459
       CONSTRUCTOR_BAD_SERVER_SALT = 0xedab447b
 
+      attr_reader :pending_requests
+
       def initialize(client)
         @client = client
-        @connection_initialized = false
+        @pending_requests = {}
       end
 
-      def call(body, content_related: true)
-        raise 'Auth key not generated' unless @client.auth_key
+      def call(request, response_class)
+        raise 'Mainloop not running' unless @client.mainloop_running?
+        raise 'Auth key not generated' unless @client.auth_key?
         raise 'Session not initialized' unless @client.session
 
-        msg_id = MessageId.generate(time_offset: @client.time_offset)
-        seq_no = @client.session.next_seq_no(content_related: content_related)
+        body = serialize_request(request)
+        msg_id = send_encrypted(body)
+        response = Response.new(msg_id, response_class, body)
+        @pending_requests[msg_id] = response
 
-        encrypted_msg = EncryptedMessage.encrypt(
-          auth_key: @client.auth_key,
-          server_salt: @client.server_salt,
-          session_id: @client.session.session_id,
-          msg_id: msg_id,
-          seq_no: seq_no,
-          body: body
-        )
+        response
+      end
 
-        @client.connection.send(encrypted_msg.serialize)
+      def handle_bad_server_salt(response_body, client)
+        bad_msg_id = response_body[4, 8].unpack1('Q<')
+        new_server_salt = response_body[20, 8].unpack1('Q<')
 
-        response_data = @client.connection.recv(timeout: @client.timeout)
+        client.update_server_salt(new_server_salt)
 
-        decrypted = EncryptedMessage.decrypt(
-          auth_key: @client.auth_key,
-          encrypted_message_data: response_data,
-          sender: :server
-        )
+        response = @pending_requests.delete(bad_msg_id)
+        return unless response
 
-        response_body = decrypted[:body]
-
-        constructor = response_body[0, 4].unpack1('L<')
-
-        # Handle bad_server_salt
-        if constructor == CONSTRUCTOR_BAD_SERVER_SALT
-          # bad_server_salt#edab447b bad_msg_id:long bad_msg_seqno:int error_code:int new_server_salt:long
-          offset = 4
-          bad_msg_id = response_body[offset, 8].unpack1('Q<')
-          offset += 8
-          bad_msg_seqno = response_body[offset, 4].unpack1('L<')
-          offset += 4
-          error_code = response_body[offset, 4].unpack1('L<')
-          offset += 4
-          new_server_salt = response_body[offset, 8].unpack1('Q<')
-
-          # Update server salt and retry
-          update_server_salt(new_server_salt)
-          return call(body, content_related: content_related)
-        end
+        new_msg_id = send_encrypted(response.body_bytes)
+        response.instance_variable_set(:@msg_id, new_msg_id)
+        @pending_requests[new_msg_id] = response
+      end
 
-        if constructor == TL::NewSessionCreated::CONSTRUCTOR
-          session_info = TL::NewSessionCreated.deserialize(response_body)
-          update_server_salt(session_info.server_salt)
-
-          response_data = @client.connection.recv(timeout: @client.timeout)
-          decrypted = EncryptedMessage.decrypt(
-            auth_key: @client.auth_key,
-            encrypted_message_data: response_data,
-            sender: :server
-          )
-          response_body = decrypted[:body]
-          constructor = response_body[0, 4].unpack1('L<')
-        end
+      def handle_rpc_result(response_body)
+        req_msg_id = response_body[4, 8].unpack1('Q<')
+        result = response_body[12..]
 
-        if constructor == TL::MsgContainer::CONSTRUCTOR
-          container = TL::MsgContainer.deserialize(response_body)
-
-          rpc_result = container.messages.find do |msg|
-            msg[:body][0, 4].unpack1('L<') == CONSTRUCTOR_RPC_RESULT
-          end
-
-          return rpc_result[:body][12..] if rpc_result
-
-          new_session = container.messages.find do |msg|
-            msg[:body][0, 4].unpack1('L<') == TL::NewSessionCreated::CONSTRUCTOR
-          end
-          if new_session
-            session_info = TL::NewSessionCreated.deserialize(new_session[:body])
-            update_server_salt(session_info.server_salt)
-
-            other_messages = container.messages.reject do |msg|
-              constructor = msg[:body][0, 4].unpack1('L<')
-              (
-                constructor == TL::NewSessionCreated::CONSTRUCTOR ||
-                constructor == CONSTRUCTOR_MSGS_ACK
-              )
-            end
-
-            return other_messages.first[:body] unless other_messages.empty?
-
-            response_data = @client.connection.recv(timeout: @client.timeout)
-            decrypted = EncryptedMessage.decrypt(
-              auth_key: @client.auth_key,
-              encrypted_message_data: response_data,
-              sender: :server
-            )
-            response_body = decrypted[:body]
-            constructor = response_body[0, 4].unpack1('L<')
-
-            return response_body[12..] if constructor == CONSTRUCTOR_RPC_RESULT
-
-            if constructor == TL::MsgContainer::CONSTRUCTOR
-              container = TL::MsgContainer.deserialize(response_body)
-              rpc_result = container.messages.find do |msg|
-                msg[:body][0, 4].unpack1('L<') == CONSTRUCTOR_RPC_RESULT
-              end
-              return rpc_result[:body][12..] if rpc_result
-            end
-
-            return response_body
-          end
-
-          return container.messages.first[:body]
+        result_constructor = result[0, 4].unpack1('L<')
+        if result_constructor == Type::GzipPacked::CONSTRUCTOR
+          result = Type::GzipPacked.unpack(result)
+          result_constructor = result[0, 4].unpack1('L<')
         end
 
-        return response_body[12..] if constructor == CONSTRUCTOR_RPC_RESULT
+        response = @pending_requests.delete(req_msg_id)
+        return unless response
+
+        if result_constructor == Type::RpcError::CONSTRUCTOR
+          error = Type::RpcError.deserialize(result)
+          response.signal_error(RpcError.new(error.error_code, error.error_message))
+        else
+          response.signal(result)
+        end
+      end
 
-        response_body
+      def signal_all_error(error)
+        @pending_requests.each_value { |response| response.signal_error(error) }
+        @pending_requests.clear
       end
 
       private
 
-      def update_server_salt(server_salt)
-        @client.instance_variable_set(:@server_salt, server_salt)
+      def serialize_request(request)
+        if request.respond_to?(:body)
+          request.body.pack('C*')
+        else
+          request
+        end
+      end
+
+      def send_encrypted(body, content_related: true)
+        msg_id = MessageId.generate(time_offset: @client.time_offset)
+        seq_no = @client.session.next_seq_no(content_related: content_related)
+
+        encrypted_msg = EncryptedMessage.encrypt(
+          auth_key: @client.auth_key,
+          server_salt: @client.server_salt,
+          session_id: @client.session.session_id,
+          msg_id: msg_id,
+          seq_no: seq_no,
+          body: body
+        )
+
+        @client.connection.send(Transport::Packet.new(encrypted_msg.serialize.bytes))
+
+        msg_id
       end
     end
   end

data/lib/mtproto/client.rb

@@ -3,25 +3,59 @@
 require 'securerandom'
 require 'digest'
 require 'base64'
+require 'async'
+require 'async/condition'
 require_relative 'transport/tcp_connection'
+require_relative 'transport/connection'
 require_relative 'transport/abridged_packet_codec'
-require_relative 'tl/message'
+require_relative 'type/message'
 require_relative 'client/rpc'
-require_relative 'tl/method_builder'
+require_relative 'client/api'
+require_relative 'tl/objects/invoke_with_layer'
+require_relative 'tl/objects/init_connection'
+require_relative 'tl/objects/get_config'
+require_relative 'tl/objects/help_config'
 
 module MTProto
   class Client
-    include TL::MethodBuilder
+    extend DelegateMethods
 
-    attr_reader :connection, :server_key, :auth_key, :server_salt, :time_offset, :session, :timeout, :user_id, :access_hash, :dc_number
-    attr_accessor :api_id, :api_hash, :device_model, :system_version, :app_version, :system_lang_code, :lang_pack, :lang_code
+    API_LAYER = 214
 
-    def initialize(api_id: nil, api_hash: nil, host:, port: 443, public_key: nil, dc_number: nil, test_mode: false, timeout: 10)
-      raise ArgumentError, "host is required" if host.nil? || host.empty?
-      raise ArgumentError, "dc_number must be positive. Use test_mode: true for test DCs" if dc_number && dc_number < 0
+    attr_reader :connection, :server_key, :auth_key, :server_salt, :time_offset, :session,
+      :timeout, :user_id, :access_hash, :dc_number
+    attr_accessor :api_id, :api_hash, :device_model, :system_version, :app_version,
+      :system_lang_code, :lang_pack, :lang_code
 
-      codec = Transport::AbridgedPacketCodec.new
-      @connection = Transport::TCPConnection.new(host, port, codec)
+    delegate :connect!, :connected?, to: :connection
+
+    def auth_key?
+      !@auth_key.nil?
+    end
+
+    def mainloop_running?
+      @running
+    end
+
+    def initialize(
+      host:,
+      api_id:,
+      api_hash:,
+      port: 443,
+      public_key: nil,
+      dc_number: nil,
+      test_mode: false,
+      timeout: 10
+    )
+      raise ArgumentError, 'host is required' if host.nil? || host.empty?
+
+      raise ArgumentError, 'dc_number must be positive. Use test_mode: true for test DCs' if dc_number && dc_number < 0
+
+      @api_id = api_id
+      @api_hash = api_hash
+
+      transport = Transport::TCPConnection.new(host, port)
+      @connection = Transport::Connection.new(transport)
       @public_key = public_key
       @dc_number = dc_number
       @test_mode = test_mode
@@ -35,29 +69,76 @@ module MTProto
       @user_id = nil
       @access_hash = nil
 
-      # Client configuration defaults
-      @api_id = api_id || 0
-      @api_hash = api_hash || ''
       @device_model = 'Ruby MTProto'
       @system_version = RUBY_DESCRIPTION
       @app_version = '0.1.0'
       @system_lang_code = 'en'
       @lang_pack = ''
       @lang_code = 'en'
+
+      @receiver_task = nil
+      @running = false
+      @on_update_callbacks = []
     end
 
-    def connect!
-      @connection.connect!
+    def on_update(&block)
+      @on_update_callbacks << block
+    end
+
+    def run_mainloop
+      raise 'Auth key not set' unless auth_key?
+      raise 'Mainloop already running' if @running
+      raise ArgumentError, 'Block is required' unless block_given?
+
+      begin
+        Async do
+          @running = true
+          @receiver_task = Async do
+            @connection.receive do |packet, error|
+              if error
+                warn "[MTProto] Packet read error: #{error.message}"
+                next
+              end
+
+              decrypted = EncryptedMessage.decrypt(
+                auth_key: @auth_key,
+                encrypted_message_data: packet.data.pack('C*'),
+                sender: :server
+              )
+
+              process_message(decrypted[:body])
+            end
+          end
+
+          yield self
+        ensure
+          disconnect!
+        end
+      rescue Interrupt
+        # Ctrl+C
+      end
     end
 
     def disconnect!
-      @connection.close if @connection
+      @running = false
+      @receiver_task&.stop
+      @receiver_task = nil
+
+      rpc.signal_all_error(Transport::ConnectionError.new('Client shutting down'))
+
+      @connection.disconnect! if @connection&.connected?
     end
 
     def exchange_keys!
-      raise ArgumentError, "public_key is required for auth key generation" if @public_key.nil? || @public_key.empty?
-
-      generator = AuthKeyGenerator.new(@connection, @public_key, @dc_number, test_mode: @test_mode, timeout: @timeout)
+      raise ArgumentError, 'public_key is required for auth key generation' if @public_key.nil? || @public_key.empty?
+
+      generator = AuthKeyGenerator.new(
+        @connection,
+        @public_key,
+        @dc_number,
+        test_mode: @test_mode,
+        timeout: @timeout
+      )
       result = generator.generate
 
       @auth_key = generator.auth_key
@@ -72,8 +153,37 @@ module MTProto
       @rpc ||= RPC.new(self)
     end
 
-    def save_session
-      raise 'Cannot save session: auth_key not set' unless @auth_key
+    def api
+      @api ||= API.new(self)
+    end
+
+    def init_connection!
+      return if @connection_initialized
+
+      query = TL::InvokeWithLayer.new(
+        layer: API_LAYER,
+        query: TL::InitConnection.new(
+          api_id: api_id,
+          device_model: device_model,
+          system_version: system_version,
+          app_version: app_version,
+          system_lang_code: system_lang_code,
+          lang_pack: lang_pack,
+          lang_code: lang_code,
+          query: TL::GetConfig.new
+        )
+      )
+
+      response = rpc.call(query, TL::HelpConfig)
+      response.wait!(timeout)
+
+      @connection_initialized = true
+
+      response.body
+    end
+
+    def save_auth_data
+      raise 'Cannot save auth_data: auth_key not set' unless @auth_key
 
       {
         auth_key: Base64.strict_encode64(@auth_key),
@@ -85,18 +195,17 @@ module MTProto
       }
     end
 
-    def load_session(session_data)
-      raise ArgumentError, 'session_data must be a Hash' unless session_data.is_a?(Hash)
-      raise ArgumentError, 'auth_key is required' unless session_data[:auth_key]
-      raise ArgumentError, 'server_salt is required' unless session_data[:server_salt]
+    def load_auth_data(auth_data)
+      raise ArgumentError, 'auth_data must be a Hash' unless auth_data.is_a?(Hash)
+      raise ArgumentError, 'auth_key is required' unless auth_data[:auth_key]
+      raise ArgumentError, 'server_salt is required' unless auth_data[:server_salt]
 
-      @auth_key = Base64.strict_decode64(session_data[:auth_key])
-      @server_salt = session_data[:server_salt]
-      @time_offset = session_data[:time_offset] || 0
-      @user_id = session_data[:user_id]
-      @access_hash = session_data[:access_hash]
+      @auth_key = Base64.strict_decode64(auth_data[:auth_key])
+      @server_salt = auth_data[:server_salt]
+      @time_offset = auth_data[:time_offset] || 0
+      @user_id = auth_data[:user_id]
+      @access_hash = auth_data[:access_hash]
 
-      # Create a fresh session for the new connection
       @session = Session.new
 
       true
@@ -106,5 +215,41 @@ module MTProto
       @user_id = user_id
       @access_hash = access_hash
     end
+
+    def update_server_salt(new_salt)
+      @server_salt = new_salt
+    end
+
+    private
+
+    def process_message(response_body)
+      constructor = response_body[0, 4].unpack1('L<')
+
+      case constructor
+      when RPC::CONSTRUCTOR_BAD_SERVER_SALT
+        rpc.handle_bad_server_salt(response_body, self)
+      when Type::NewSessionCreated::CONSTRUCTOR
+        handle_new_session(response_body)
+      when Type::MsgContainer::CONSTRUCTOR
+        handle_container(response_body)
+      when RPC::CONSTRUCTOR_RPC_RESULT
+        rpc.handle_rpc_result(response_body)
+      else
+        @on_update_callbacks.each { |cb| cb.call(constructor, response_body) }
+      end
+    end
+
+    def handle_new_session(response_body)
+      session_info = Type::NewSessionCreated.deserialize(response_body)
+      @server_salt = session_info.server_salt
+    end
+
+    def handle_container(response_body)
+      container = Type::MsgContainer.deserialize(response_body)
+
+      container.messages.each do |msg|
+        process_message(msg[:body])
+      end
+    end
   end
 end

data/lib/mtproto/crypto/dh_key_exchange.rb

@@ -36,8 +36,7 @@ module MTProto
         random_bytes = SecureRandom.random_bytes(256)
         random_bytes[0] = (random_bytes[0].ord | 0x80).chr
 
-        bn = OpenSSL::BN.new(random_bytes, 2)
-        bn
+        OpenSSL::BN.new(random_bytes, 2)
       end
     end
   end

data/lib/mtproto/crypto/dh_validator.rb

@@ -22,22 +22,22 @@ module MTProto
         raise 'Invalid g: must be 2, 3, 4, 5, 6, or 7' unless [2, 3, 4, 5, 6, 7].include?(g)
 
         p_mod = case g
-        when 2
-          dh_prime % 8 == 7
-        when 3
-          dh_prime % 3 == 2
-        when 4
-          true
-        when 5
-          mod5 = dh_prime % 5
-          mod5 == 1 || mod5 == 4
-        when 6
-          mod24 = dh_prime % 24
-          mod24 == 19 || mod24 == 23
-        when 7
-          mod7 = dh_prime % 7
-          [3, 5, 6].include?(mod7)
-        end
+                when 2
+                  dh_prime % 8 == 7
+                when 3
+                  dh_prime % 3 == 2
+                when 4
+                  true
+                when 5
+                  mod5 = dh_prime % 5
+                  [1, 4].include?(mod5)
+                when 6
+                  mod24 = dh_prime % 24
+                  [19, 23].include?(mod24)
+                when 7
+                  mod7 = dh_prime % 7
+                  [3, 5, 6].include?(mod7)
+                end
 
         raise "g=#{g} is not a valid generator for this prime" unless p_mod
 
@@ -52,9 +52,7 @@ module MTProto
         min_prime = OpenSSL::BN.new(2)**2047
         max_prime = OpenSSL::BN.new(2)**2048
 
-        if dh_prime <= min_prime || dh_prime >= max_prime
-          raise 'dh_prime out of range (must be 2^2047 < p < 2^2048)'
-        end
+        raise 'dh_prime out of range (must be 2^2047 < p < 2^2048)' if dh_prime <= min_prime || dh_prime >= max_prime
 
         true
       end

data/lib/mtproto/crypto/factorization.rb

@@ -29,7 +29,7 @@ module MTProto
 
         while y < x
           x = y
-          y = (x + n / x) / 2
+          y = (x + (n / x)) / 2
         end
 
         x

data/lib/mtproto/crypto/rsa_key.rb

@@ -9,7 +9,7 @@ module MTProto
       attr_reader :key, :fingerprint
 
       def initialize(pem_string)
-        raise ArgumentError, "pem_string is required" if pem_string.nil? || pem_string.empty?
+        raise ArgumentError, 'pem_string is required' if pem_string.nil? || pem_string.empty?
 
         @key = OpenSSL::PKey::RSA.new(pem_string)
         @fingerprint = calculate_fingerprint
@@ -20,7 +20,7 @@ module MTProto
       end
 
       def self.find_by_fingerprint(fingerprints, public_key)
-        raise ArgumentError, "public_key is required" if public_key.nil? || public_key.empty?
+        raise ArgumentError, 'public_key is required' if public_key.nil? || public_key.empty?
 
         key = new(public_key)
         key if fingerprints.include?(key.fingerprint)

data/lib/mtproto/crypto/srp.rb

@@ -0,0 +1,117 @@
+# frozen_string_literal: true
+
+require 'openssl'
+require 'securerandom'
+
+module MTProto
+  module Crypto
+    module SRP
+      SIZE_FOR_HASH = 256
+
+      module_function
+
+      def compute_check(algo:, srp_b:, srp_id:, password:)
+        salt1 = algo[:salt1]
+        salt2 = algo[:salt2]
+        g = algo[:g]
+        p_bytes = algo[:p]
+
+        p_int = bytes_to_int(p_bytes)
+        g_int = g
+        b_int = bytes_to_int(srp_b)
+
+        pw_hash = compute_hash(salt1, salt2, password)
+        x = bytes_to_int(pw_hash)
+
+        p_for_hash = pad_to_256(p_bytes)
+        g_for_hash = int_to_256_bytes(g_int)
+        b_for_hash = pad_to_256(srp_b)
+
+        g_x = mod_pow(g_int, x, p_int)
+        k = bytes_to_int(sha256(p_for_hash, g_for_hash))
+        kg_x = (k * g_x) % p_int
+
+        a_int, a_for_hash, u = generate_and_check_random(g_int, p_int, b_for_hash)
+
+        g_b = (b_int - kg_x) % p_int
+        ux = u * x
+        a_ux = a_int + ux
+        s_int = mod_pow(g_b, a_ux, p_int)
+        k_key = sha256(int_to_256_bytes(s_int))
+
+        m1 = sha256(
+          xor(sha256(p_for_hash), sha256(g_for_hash)),
+          sha256(salt1),
+          sha256(salt2),
+          a_for_hash,
+          b_for_hash,
+          k_key
+        )
+
+        { srp_id: srp_id, a: a_for_hash, m1: m1 }
+      end
+
+      def compute_hash(salt1, salt2, password)
+        hash1 = sha256(salt1, password.encode('utf-8'), salt1)
+        hash2 = sha256(salt2, hash1, salt2)
+        hash3 = pbkdf2_sha512(hash2, salt1, 100_000)
+        sha256(salt2, hash3, salt2)
+      end
+
+      def sha256(*parts)
+        digest = OpenSSL::Digest.new('SHA256')
+        parts.each { |p| digest.update(p) }
+        digest.digest
+      end
+
+      def pbkdf2_sha512(password, salt, iterations)
+        OpenSSL::KDF.pbkdf2_hmac(
+          password,
+          salt: salt,
+          iterations: iterations,
+          length: 64,
+          hash: 'SHA512'
+        )
+      end
+
+      def mod_pow(base, exp, mod)
+        base.to_bn.mod_exp(exp.to_bn, mod.to_bn).to_i
+      end
+
+      def generate_and_check_random(g, p, b_for_hash)
+        loop do
+          random = SecureRandom.random_bytes(256)
+          a_int = bytes_to_int(random)
+          a_big = mod_pow(g, a_int, p)
+          a_for_hash = int_to_256_bytes(a_big)
+          u = bytes_to_int(sha256(a_for_hash, b_for_hash))
+          next if u.zero?
+
+          return [a_int, a_for_hash, u]
+        end
+      end
+
+      def pad_to_256(bytes)
+        bytes = bytes.b
+        return bytes if bytes.bytesize >= SIZE_FOR_HASH
+
+        ("\x00" * (SIZE_FOR_HASH - bytes.bytesize)).b + bytes
+      end
+
+      def int_to_256_bytes(num)
+        hex = num.to_s(16)
+        hex = "0#{hex}" if hex.length.odd?
+        raw = [hex].pack('H*')
+        pad_to_256(raw)
+      end
+
+      def bytes_to_int(bytes)
+        bytes.unpack1('H*').to_i(16)
+      end
+
+      def xor(a, b)
+        a.bytes.zip(b.bytes).map { |x, y| x ^ y }.pack('C*')
+      end
+    end
+  end
+end

data/lib/mtproto/delegate_methods.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module DelegateMethods
+  def delegate(*methods, to:)
+    methods.each do |method|
+      define_method(method) do |*args, **kwargs, &block|
+        instance_variable_get(:"@#{to}").public_send(method, *args, **kwargs, &block)
+      end
+    end
+  end
+end