rubydb 0.1.0

data/lib/rubydb/protocol/decoder.rb

@@ -0,0 +1,145 @@
+# frozen_string_literal: true
+
+require "zlib"
+require "base64"
+
+module RubyDB
+  module Protocol
+    class Decoder
+      attr_reader :stats
+
+      def initialize(options = {})
+        @options = options
+        @encryption_key = options[:encryption_key]
+        @stats = {
+          messages_decoded: 0,
+          bytes_decoded: 0,
+          decompressions: 0,
+          decryption_ops: 0,
+          errors: 0
+        }
+        @lock = Mutex.new
+      end
+
+      def decode(data, format = Encoder::FORMAT_JSON)
+        @lock.synchronize do
+          @stats[:bytes_decoded] += data.bytesize
+
+          begin
+            if @encryption_key && encrypted?(data)
+              data = decrypt(data)
+              @stats[:decryption_ops] += 1
+            end
+
+            if compressed?(data)
+              data = decompress(data)
+              @stats[:decompressions] += 1
+            end
+
+            message = case format
+            when Encoder::FORMAT_JSON
+              decode_json(data)
+            when Encoder::FORMAT_MSGPACK
+              decode_msgpack(data)
+            when Encoder::FORMAT_BINARY
+              decode_binary(data)
+            else
+              decode_json(data)
+            end
+
+            @stats[:messages_decoded] += 1
+            message
+
+          rescue => e
+            @stats[:errors] += 1
+            raise ProtocolError, "Decoding failed: #{e.message}"
+          end
+        end
+      end
+
+      def decode_json(data)
+        Message.from_json(data.strip)
+      end
+
+      def decode_msgpack(data)
+        begin
+          require "msgpack"
+          hash = MessagePack.unpack(data)
+          msg = Message.new(hash[:type], hash[:payload] || {})
+          msg.instance_variable_set(:@id, hash[:id])
+          msg.instance_variable_set(:@created_at, Time.parse(hash[:created_at]))
+          msg
+        rescue LoadError
+          decode_json(data)
+        end
+      end
+
+      def decode_binary(data)
+        offset = 0
+
+        type_len = data[offset].unpack("C").first
+        offset += 1
+        type = data[offset, type_len].to_sym
+        offset += type_len
+
+        id_len = data[offset].unpack("C").first
+        offset += 1
+        id = data[offset, id_len]
+        offset += id_len
+
+        timestamp = data[offset, 8].unpack("Q>").first
+        offset += 8
+
+        payload_len = data[offset, 4].unpack("L>").first
+        offset += 4
+        payload_data = data[offset, payload_len]
+        payload = JSON.parse(payload_data, symbolize_names: true)
+
+        msg = Message.new(type, payload)
+        msg.instance_variable_set(:@id, id)
+        msg.instance_variable_set(:@created_at, Time.at(timestamp))
+        msg
+      end
+
+      def compressed?(data)
+        data.bytesize >= 2 && data[0].ord == 0x78 && (data[1].ord == 0x01 || data[1].ord == 0x9C || data[1].ord == 0xDA)
+      end
+
+      def decompress(data)
+        Zlib::Inflate.inflate(data)
+      rescue => e
+        raise ProtocolError, "Decompression failed: #{e.message}"
+      end
+
+      def encrypted?(data)
+        false
+      end
+
+      def decrypt(data)
+        return data unless @encryption_key
+
+        begin
+          require "openssl"
+          iv = data[0, 16]
+          encrypted_data = data[16..-1]
+
+          cipher = OpenSSL::Cipher.new("aes-256-cbc")
+          cipher.decrypt
+          cipher.key = @encryption_key
+          cipher.iv = iv
+          cipher.update(encrypted_data) + cipher.final
+        rescue => e
+          raise ProtocolError, "Decryption failed: #{e.message}"
+        end
+      end
+
+      def stats
+        @lock.synchronize do
+          @stats.merge({
+            bytes_per_message: @stats[:messages_decoded] > 0 ? @stats[:bytes_decoded] / @stats[:messages_decoded] : 0
+          })
+        end
+      end
+    end
+  end
+end

data/lib/rubydb/protocol/encoder.rb

@@ -0,0 +1,136 @@
+# frozen_string_literal: true
+
+require "zlib"
+require "base64"
+
+module RubyDB
+  module Protocol
+    class Encoder
+      FORMAT_JSON = :json
+      FORMAT_MSGPACK = :msgpack
+      FORMAT_BINARY = :binary
+
+      attr_reader :stats
+
+      def initialize(format = FORMAT_JSON, options = {})
+        @format = format
+        @options = options
+        @compression = options[:compression] || false
+        @encryption = options[:encryption] || false
+        @encryption_key = options[:encryption_key]
+        @stats = {
+          messages_encoded: 0,
+          bytes_encoded: 0,
+          compressions: 0,
+          encryption_ops: 0,
+          errors: 0
+        }
+        @lock = Mutex.new
+      end
+
+      def encode(message)
+        @lock.synchronize do
+          @stats[:messages_encoded] += 1
+
+          begin
+          data = case @format
+            when FORMAT_JSON
+              encode_json(message)
+            when FORMAT_MSGPACK
+              encode_msgpack(message)
+            when FORMAT_BINARY
+              encode_binary(message)
+            else
+              encode_json(message)
+            end
+
+            # JSON messages are transported as one line per frame so a
+            # persistent TCP connection can separate responses reliably.
+            data << "\n" if @format == FORMAT_JSON && !data.end_with?("\n")
+
+            if @compression
+              data = compress(data)
+              @stats[:compressions] += 1
+            end
+
+            if @encryption && @encryption_key
+              data = encrypt(data)
+              @stats[:encryption_ops] += 1
+            end
+
+            @stats[:bytes_encoded] += data.bytesize
+            data
+
+          rescue => e
+            @stats[:errors] += 1
+            raise ProtocolError, "Encoding failed: #{e.message}"
+          end
+        end
+      end
+
+      def encode_json(message)
+        message.to_json
+      end
+
+      def encode_msgpack(message)
+        begin
+          require "msgpack"
+          MessagePack.pack(message.to_hash)
+        rescue LoadError
+          message.to_json
+        end
+      end
+
+      def encode_binary(message)
+        data = ""
+        data << [message.type.to_s.bytesize].pack("C")
+        data << message.type.to_s
+        data << [message.id.bytesize].pack("C")
+        data << message.id
+        data << [message.created_at.to_i].pack("Q>")
+        payload = JSON.generate(message.payload)
+        data << [payload.bytesize].pack("L>")
+        data << payload
+        data
+      end
+
+      def compress(data)
+        Zlib::Deflate.deflate(data)
+      rescue => e
+        raise ProtocolError, "Compression failed: #{e.message}"
+      end
+
+      def encrypt(data)
+        return data unless @encryption_key
+
+        begin
+          require "openssl"
+          cipher = OpenSSL::Cipher.new("aes-256-cbc")
+          cipher.encrypt
+          cipher.key = @encryption_key
+          cipher.iv = cipher.random_iv
+          encrypted = cipher.update(data) + cipher.final
+          iv = cipher.iv
+          iv + encrypted
+        rescue => e
+          raise ProtocolError, "Encryption failed: #{e.message}"
+        end
+      end
+
+      def format
+        @format
+      end
+
+      def stats
+        @lock.synchronize do
+          @stats.merge({
+            format: @format,
+            compression: @compression,
+            encryption: @encryption,
+            bytes_per_message: @stats[:messages_encoded] > 0 ? @stats[:bytes_encoded] / @stats[:messages_encoded] : 0
+          })
+        end
+      end
+    end
+  end
+end

data/lib/rubydb/protocol/handshake.rb

@@ -0,0 +1,305 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require "digest"
+require "openssl"
+require "base64"
+
+module RubyDB
+  module Protocol
+    class Handshake
+      STATE_INIT = :init
+      STATE_AUTHENTICATING = :authenticating
+      STATE_AUTHENTICATED = :authenticated
+      STATE_NEGOTIATING = :negotiating
+      STATE_COMPLETE = :complete
+      STATE_FAILED = :failed
+
+      attr_reader :state, :client_info, :server_info
+
+      def initialize(server_info = {})
+        @state = STATE_INIT
+        @client_info = {}
+        server_info = (server_info || {}).dup
+        @credentials = server_info.delete(:authentication_credentials) || {}
+        @server_info = {
+          version: ProtocolVersion.current_string,
+          protocol_version: ProtocolVersion.current,
+          name: "RubyDB",
+          pid: Process.pid,
+          started_at: Time.now.iso8601,
+          auth_methods: ["none", "password", "md5", "scram-sha-256"],
+          default_auth: "none"
+        }.merge(server_info || {})
+        @challenge = nil
+        @scram_salt = (@credentials[:scram_salt] || @credentials["scram_salt"])
+        @scram_salt = SecureRandom.random_bytes(16) unless @scram_salt
+        @scram_iterations = Integer(@credentials[:scram_iterations] || @credentials["scram_iterations"] || 120_000)
+        @scram_server_signature = nil
+        @authenticated = false
+        @capabilities = nil
+        @start_time = Time.now
+        @timeout = 30
+        @max_attempts = 3
+        @attempts = 0
+        @lock = Mutex.new
+      end
+
+      def start(client_info = {})
+        @lock.synchronize do
+          @state = STATE_INIT
+          @client_info = client_info
+          @attempts += 1
+
+          client_version = client_info[:protocol_version]
+          unless client_version
+            return error_response("Protocol version not specified")
+          end
+
+          unless ProtocolVersion.supported?(client_version)
+            return error_response("Unsupported protocol version: #{ProtocolVersion.to_string(client_version)}")
+          end
+
+          unless ProtocolVersion.compatible?(client_version, @server_info[:protocol_version])
+            return error_response("Incompatible protocol version")
+          end
+
+          @state = STATE_AUTHENTICATING
+          authentication_response(client_info)
+        end
+      end
+
+      def authenticate(credentials)
+        @lock.synchronize do
+          return error_response("Handshake not initialized") if @state == STATE_INIT
+
+          @state = STATE_AUTHENTICATING
+
+          # The server selects the required mechanism; clients cannot lower
+          # the server's authentication policy by requesting "none".
+          auth_method = @server_info[:default_auth]
+
+          case auth_method
+          when "none"
+            @authenticated = true
+            @state = STATE_AUTHENTICATED
+            return success_response
+          when "password"
+            @authenticated = authenticate_password(credentials)
+          when "md5"
+            @authenticated = authenticate_md5(credentials)
+          when "scram-sha-256"
+            @authenticated = authenticate_scram(credentials)
+          else
+            return error_response("Unsupported authentication method: #{auth_method}")
+          end
+
+          if @authenticated
+            @state = STATE_AUTHENTICATED
+            success_response
+          else
+            @state = STATE_FAILED
+            error_response("Authentication failed")
+          end
+        end
+      end
+
+      def negotiate(capabilities)
+        @lock.synchronize do
+          return error_response("Not authenticated") unless @authenticated
+
+          @state = STATE_NEGOTIATING
+          @capabilities = capabilities
+
+          negotiated = negotiate_capabilities(capabilities)
+
+          @state = STATE_COMPLETE
+          negotiated_response(negotiated)
+        end
+      end
+
+      def complete?
+        @state == STATE_COMPLETE
+      end
+
+      def authenticated?
+        @authenticated
+      end
+
+      def failed?
+        @state == STATE_FAILED
+      end
+
+      def timed_out?
+        Time.now - @start_time > @timeout
+      end
+
+      def max_attempts_reached?
+        @attempts >= @max_attempts
+      end
+
+      def reset
+        @lock.synchronize do
+          @state = STATE_INIT
+          @authenticated = false
+          @challenge = nil
+          @capabilities = nil
+          @attempts = 0
+          @start_time = Time.now
+        end
+      end
+
+      def to_hash
+        {
+          state: @state,
+          authenticated: @authenticated,
+          client_info: @client_info,
+          server_info: @server_info,
+          capabilities: @capabilities,
+          attempts: @attempts,
+          elapsed: (Time.now - @start_time).round(2)
+        }
+      end
+
+      private
+
+      def success_response
+        response = {
+          success: true,
+          state: @state,
+          server_info: @server_info,
+          message: "Handshake successful"
+        }
+        response[:server_signature] = @scram_server_signature if @scram_server_signature
+        response
+      end
+
+      def error_response(message)
+        {
+          success: false,
+          state: @state,
+          error: message,
+          attempts_remaining: @max_attempts - @attempts
+        }
+      end
+
+      def authentication_response(client_info)
+        response = {
+          success: true,
+          auth_methods: @server_info[:auth_methods],
+          default_auth: @server_info[:default_auth],
+          challenge: generate_challenge
+        }
+        if @server_info[:default_auth] == "scram-sha-256"
+          response[:scram_salt] = Base64.strict_encode64(@scram_salt)
+          response[:scram_iterations] = @scram_iterations
+        end
+        response
+      end
+
+      def negotiated_response(negotiated)
+        {
+          success: true,
+          state: @state,
+          negotiated: negotiated,
+          server_info: @server_info
+        }
+      end
+
+      def generate_challenge
+        @challenge = SecureRandom.hex(16)
+      end
+
+      def authenticate_password(credentials)
+        username = credentials[:username]
+        password = credentials[:password]
+        return false if username.nil? || password.nil?
+
+        users = @credentials[:users] || @credentials["users"]
+        if users
+          expected = users[username] || users[username.to_s] || users[username.to_sym]
+          return secure_equal?(password.to_s, expected.to_s) unless expected.nil?
+          return false
+        end
+
+        expected_username = @credentials[:username] || @credentials["username"]
+        expected_password = @credentials[:password] || @credentials["password"]
+        return false if expected_username.nil? || expected_password.nil?
+
+        secure_equal?(username.to_s, expected_username.to_s) &&
+          secure_equal?(password.to_s, expected_password.to_s)
+      end
+
+      def secure_equal?(left, right)
+        return false unless left.bytesize == right.bytesize
+
+        result = 0
+        left.bytes.each_with_index { |byte, index| result |= byte ^ right.getbyte(index) }
+        result.zero?
+      end
+
+      def authenticate_md5(credentials)
+        username = credentials[:username]
+        password = credentials[:password]
+
+        return false if username.nil? || password.nil?
+
+        expected = Digest::MD5.hexdigest(password + @challenge)
+        actual = credentials[:md5_hash]
+
+        actual == expected
+      end
+
+      def authenticate_scram(credentials)
+        data = credentials[:scram_data]
+        return false unless data.is_a?(Hash)
+
+        username = credentials[:username].to_s
+        client_first_bare = data[:client_first_bare] || data["client_first_bare"]
+        client_final = data[:client_final_without_proof] || data["client_final_without_proof"]
+        proof = data[:proof] || data["proof"]
+        client_nonce = data[:client_nonce] || data["client_nonce"]
+        return false if [client_first_bare, client_final, proof, client_nonce].any?(&:nil?)
+        return false unless client_first_bare.include?("n=#{username},")
+        return false unless client_first_bare.include?("r=#{client_nonce}")
+
+        server_nonce = @challenge.to_s
+        expected_nonce = client_nonce.to_s + server_nonce
+        return false unless client_final == "c=biws,r=#{expected_nonce}"
+
+        password = @credentials[:password] || @credentials["password"]
+        stored_key = @credentials[:scram_stored_key] || @credentials["scram_stored_key"]
+        server_key = @credentials[:scram_server_key] || @credentials["scram_server_key"]
+        return false unless password || (stored_key && server_key)
+
+        salted = OpenSSL::PKCS5.pbkdf2_hmac(password.to_s, @scram_salt, @scram_iterations, 32, OpenSSL::Digest::SHA256.new) if password
+        stored_key ||= Digest::SHA256.digest(OpenSSL::HMAC.digest("SHA256", salted, "Client Key"))
+        server_key ||= OpenSSL::HMAC.digest("SHA256", salted, "Server Key")
+        stored_key = Base64.decode64(stored_key) if stored_key.is_a?(String) && stored_key.match?(/\A[A-Za-z0-9+\/=]+\z/) && stored_key.bytesize != 32
+        server_key = Base64.decode64(server_key) if server_key.is_a?(String) && server_key.match?(/\A[A-Za-z0-9+\/=]+\z/) && server_key.bytesize != 32
+
+        auth_message = [client_first_bare, "r=#{expected_nonce},s=#{Base64.strict_encode64(@scram_salt)},i=#{@scram_iterations}", client_final].join(",")
+        proof_bytes = Base64.decode64(proof.to_s)
+        client_signature = OpenSSL::HMAC.digest("SHA256", stored_key, auth_message)
+        client_key = proof_bytes.bytes.each_with_index.map { |byte, index| byte ^ client_signature.getbyte(index) }.pack("C*")
+        valid = secure_equal?(Digest::SHA256.digest(client_key), stored_key)
+        if valid
+          @scram_server_signature = Base64.strict_encode64(OpenSSL::HMAC.digest("SHA256", server_key, auth_message))
+        end
+        valid
+      end
+
+      def negotiate_capabilities(capabilities)
+        {
+          protocol_version: ProtocolVersion.current,
+          compression: capabilities[:compression] && @server_info[:compression_supported],
+          encryption: capabilities[:encryption] && @server_info[:encryption_supported],
+          pipeline: capabilities[:pipeline] && @server_info[:pipeline_supported],
+          batch_size: [capabilities[:batch_size] || 100, @server_info[:max_batch_size] || 1000].min,
+          timeout: [capabilities[:timeout] || 30, @server_info[:max_timeout] || 300].min,
+          max_rows: [capabilities[:max_rows] || 10000, @server_info[:max_rows] || 100000].min
+        }
+      end
+    end
+  end
+end

data/lib/rubydb/protocol/message.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+require "json"
+require "digest"
+require "time"
+
+module RubyDB
+  module Protocol
+    class Message
+      TYPE_HANDSHAKE = :handshake
+      TYPE_HANDSHAKE_RESPONSE = :handshake_response
+      TYPE_QUERY = :query
+      TYPE_QUERY_RESPONSE = :query_response
+      TYPE_PREPARE = :prepare
+      TYPE_PREPARE_RESPONSE = :prepare_response
+      TYPE_EXECUTE = :execute
+      TYPE_EXECUTE_RESPONSE = :execute_response
+      TYPE_CLOSE = :close
+      TYPE_CLOSE_RESPONSE = :close_response
+      TYPE_PING = :ping
+      TYPE_PONG = :pong
+      TYPE_ERROR = :error
+      TYPE_NOTICE = :notice
+      TYPE_PARAMETER_DESCRIPTION = :parameter_description
+      TYPE_ROW_DESCRIPTION = :row_description
+      TYPE_DATA_ROW = :data_row
+      TYPE_COMMAND_COMPLETE = :command_complete
+      TYPE_READY_FOR_QUERY = :ready_for_query
+      TYPE_BEGIN = :begin
+      TYPE_COMMIT = :commit
+      TYPE_ROLLBACK = :rollback
+      TYPE_SYNCHRONIZE = :synchronize
+      TYPE_CANCEL = :cancel
+      TYPE_TERMINATE = :terminate
+      TYPE_AUTHENTICATION = :authentication
+      TYPE_AUTHENTICATION_RESPONSE = :authentication_response
+
+      attr_reader :type, :id, :created_at, :payload
+
+      def initialize(type, payload = {})
+        @type = type
+        @id = "msg_#{Time.now.to_i}_#{rand(10000)}"
+        @created_at = Time.now
+        @payload = payload
+        @compressed = false
+        @encrypted = false
+        @checksum = nil
+      end
+
+      def compress
+        return self if @compressed
+        @compressed = true
+        self
+      end
+
+      def decompress
+        return self unless @compressed
+        @compressed = false
+        self
+      end
+
+      def encrypt(key)
+        return self if @encrypted
+        @encrypted = true
+        self
+      end
+
+      def decrypt(key)
+        return self unless @encrypted
+        @encrypted = false
+        self
+      end
+
+      def calculate_checksum
+        data = "#{@type}#{@id}#{JSON.generate(@payload)}"
+        Digest::SHA256.hexdigest(data)[0...16]
+      end
+
+      def verify_checksum
+        return true if @checksum.nil?
+        @checksum == calculate_checksum
+      end
+
+      def to_hash
+        {
+          type: @type,
+          id: @id,
+          created_at: @created_at.iso8601,
+          payload: @payload,
+          compressed: @compressed,
+          encrypted: @encrypted,
+          checksum: @checksum
+        }
+      end
+
+      def to_json
+        JSON.generate(to_hash)
+      end
+
+      def self.from_json(json_data)
+        data = JSON.parse(json_data, symbolize_names: true)
+        type = data[:type].respond_to?(:to_sym) ? data[:type].to_sym : data[:type]
+        msg = new(type, data[:payload] || {})
+        msg.instance_variable_set(:@id, data[:id])
+        msg.instance_variable_set(:@created_at, Time.parse(data[:created_at]))
+        msg.instance_variable_set(:@compressed, data[:compressed] || false)
+        msg.instance_variable_set(:@encrypted, data[:encrypted] || false)
+        msg.instance_variable_set(:@checksum, data[:checksum])
+        msg
+      end
+
+      def inspect
+        "#<Message type=#{@type} id=#{@id}>"
+      end
+
+      def to_s
+        inspect
+      end
+    end
+  end
+end