riak-client 1.4.5 → 2.0.0.rc1

data/lib/riak/client/beefcake/object_methods.rb

@@ -31,7 +31,7 @@ module Riak
           return robject if pbuf.respond_to?(:unchanged) && pbuf.unchanged # Reloading
           robject.vclock = Base64.encode64(pbuf.vclock).chomp if pbuf.vclock
           robject.key = maybe_unescape(pbuf.key) if pbuf.respond_to?(:key) && pbuf.key # Put w/o key
-          robject.siblings = pbuf.content.map do |c|
+          robject.siblings = (pbuf.content || []).map do |c|
             RContent.new(robject) do |sibling|
               load_content(c, sibling)
             end

data/lib/riak/client/beefcake/protocol.rb

@@ -0,0 +1,105 @@
+require 'riak/client/beefcake/messages'
+require 'riak/client/beefcake/message_codes'
+require 'riak/errors/failed_request'
+require 'riak/errors/protobuffs_error'
+
+module Riak
+  class Client
+    class BeefcakeProtobuffsBackend < ProtobuffsBackend
+      class Protocol
+        include Riak::Util::Translation
+        attr_reader :socket
+
+        def initialize(socket)
+          @socket = socket
+        end
+
+        # Encodes and writes a Riak-formatted message, including protocol buffer
+        # payload if given.
+        # 
+        # @param [Symbol, Integer] code the symbolic or numeric code for the 
+        #   message
+        # @param [Beefcake::Message, nil] message the protocol buffer message
+        #   payload, or nil if the message carries no payload
+        def write(code, message=nil)
+          if code.is_a? Symbol
+            code = BeefcakeMessageCodes.index code
+          end
+
+          serialized = serialize message
+
+          header = [serialized.length + 1, code].pack 'NC'
+
+          payload = header + serialized
+
+          socket.write payload
+          socket.flush
+        end
+
+        # Receives a Riak-formatted message, and returns the symbolic name of
+        # the message along with the string payload from the network.
+        #
+        # @return [Array<Symbol, String>]
+        def receive
+          header = socket.read 5
+          
+          raise ProtobuffsFailedHeader.new if header.nil?
+          message_length, code = header.unpack 'NC'
+          body_length = message_length - 1
+          body = nil
+          body = socket.read body_length if body_length > 0
+
+          name = BeefcakeMessageCodes[code]
+
+          return name, body
+        end
+
+        # Receives a Riak-formatted message, checks the symbolic name against
+        # the given code, decodes it if it matches, and can optionally return
+        # success if the payload is empty.
+        #
+        # @param [Symbol] code the code for the message
+        # @param [Class, nil] decoder_class the class to attempt to decode 
+        #   the payload with
+        # @param [Hash] options
+        # @option options [Boolean] :empty_body_acceptable Whether to accept
+        #   an empty body and not attempt decoding. In this case, this method
+        #   will return the symbol `:empty` instead of a `Beefcake::Message`
+        #   instance
+        # @return [Beefcake::Message, :empty]
+        # @raise {ProtobuffsErrorResponse} if the message from Riak was a
+        #   255-ErrorResp
+        # @raise {ProtobuffsUnexpectedResponse} if the message from riak did
+        #   not match `code`
+        def expect(code, decoder_class=nil, options={ })
+          code = BeefcakeMessageCodes[code] unless code.is_a? Symbol
+          name, body = receive
+          
+          if name == :ErrorResp
+            raise ProtobuffsErrorResponse.new RpbErrorResp.decode(body)
+          end
+
+          if name != code
+            raise ProtobuffsUnexpectedResponse.new name, code
+          end
+
+          return true if decoder_class.nil?
+
+          return :empty if body.nil? && options[:empty_body_acceptable]
+
+          return decoder_class.decode body
+        end
+
+        private
+
+        def serialize(message)
+          return '' if message.nil?
+          return message if message.is_a? String
+          return message.encode.to_s if message.is_a? Beefcake::Message
+
+          raise ArgumentError.new t('pbc.unknown_serialize', message: message)
+        end
+      end
+    end
+  end
+end

data/lib/riak/client/beefcake/socket.rb

@@ -0,0 +1,243 @@
+require 'openssl'
+require 'r509/cert/validator'
+require 'riak/client/beefcake/messages'
+require 'riak/errors/connection_error'
+
+module Riak
+  class Client
+    class BeefcakeProtobuffsBackend
+      # A factory class for making sockets, whether secure or not
+      # @api private
+      class BeefcakeSocket
+        include Client::BeefcakeMessageCodes
+        # Only create class methods, don't initialize
+        class << self
+          def new(host, port, options={})
+            return start_tcp_socket(host, port) if options[:authentication].blank?
+            return start_tls_socket(host, port, options[:authentication])
+          end
+
+          private
+          def start_tcp_socket(host, port)
+            TCPSocket.new(host, port).tap do |sock|
+              sock.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, true)
+            end
+          end
+
+          def start_tls_socket(host, port, authentication)
+            raise Riak::UserConfigurationError.new if authentication[:username]
+
+            tcp = start_tcp_socket(host, port)
+            TlsInitiator.new(tcp, host, authentication).tls_socket
+          end
+          
+          # Wrap up the logic to turn a TCP socket into a TLS socket.
+          # Depends on Beefcake, which should be relatively safe.
+          class TlsInitiator
+            BC = ::Riak::Client::BeefcakeProtobuffsBackend
+            include Util::Translation
+
+            # Create a TLS Initiator
+            #
+            # @param tcp_socket [TCPSocket] the {TCPSocket} to start TLS on
+            # @param authentication [Hash] a hash of authentication details
+            def initialize(tcp_socket, host, authentication)
+              @sock = @tcp = tcp_socket
+              @host = host
+              @auth = authentication
+            end
+
+            # Return the SSLSocket that has a TLS session running. (TLS is a
+            # better and safer SSL).
+            #
+            # @return [OpenSSL::SSL::SSLSocket]
+            def tls_socket
+              configure_context
+              start_tls
+              validate_session
+              send_authentication
+              validate_connection
+              return @tls
+            end
+
+            private
+            def riak_cert
+              @riak_cert ||= R509::Cert.new cert: @tls.peer_cert
+            end
+
+            # Set up an SSL context with appropriate defaults for Riak TLS
+            def configure_context
+              @context = OpenSSL::SSL::SSLContext.new
+
+              # Replace insecure defaults
+              @context.ssl_version = @auth[:ssl_version] || :TLSv1_2_client
+              @context.verify_mode = @auth[:verify_mode] || OpenSSL::SSL::VERIFY_PEER
+
+              cert_ify
+              key_ify
+
+              # Defer to defaults
+              %w{ cert key client_ca ca_file ca_path timeout }.each do |k|
+                @context.send(:"#{k}=", @auth[k.to_sym]) if @auth[k.to_sym]
+              end
+            end
+
+            # Convert cert and client_ca fields to X509 Certs
+            def cert_ify
+              %w{ cert client_ca }.each do |k|
+                candidate = @auth[k.to_sym]
+                next if candidate.nil?
+                next if candidate.is_a? OpenSSL::X509::Certificate
+                
+                @auth[k.to_sym] = OpenSSL::X509::Certificate.new try_load candidate
+              end
+            end
+
+            def key_ify
+              candidate = @auth[:key]
+              return if candidate.nil?
+              return if candidate.is_a? OpenSSL::PKey::PKey
+
+              candidate = try_load candidate
+
+              pkey_class_names = OpenSSL::PKey.
+                constants.
+                reject{|s| s.to_s =~ /Error$/}
+
+              pkey_classes = pkey_class_names.map{ |n| OpenSSL::PKey.const_get n }
+
+              pkey_classes.each do |klass|
+                begin
+                  successfully_initialized = klass.new candidate
+                  @auth[:key] = successfully_initialized
+                  return
+                rescue
+                  next
+                end
+              end
+
+              # Don't try and guess what the key is
+              raise TlsError::UnknownKeyTypeError.new
+            end
+
+            # Figure out if the given string is the data itself or a path to the data
+            def try_load(data_or_path)
+              begin
+                data_or_path = File.read data_or_path
+              rescue Errno::ENOENT
+                # couldn't read the file, it might be a string containing
+                # a key
+              rescue Errno::ENAMETOOLONG
+                # the filename is too long, it's almost certainly a string 
+                # containing a key
+              rescue => e
+                raise TlsError::ReadDataError.new e, data_or_path
+              end
+
+              return data_or_path
+            end
+
+            # Attempt to exchange the TCP socket for a TLS socket.
+            def start_tls
+              write_message :StartTls
+              expect_message :StartTls
+              # Swap the tls socket in for the tcp socket, so write_message and
+              # read_message continue working
+              @sock = @tls = OpenSSL::SSL::SSLSocket.new @tcp, @context
+              @tls.connect
+            end
+
+            # Validate the TLS session
+            def validate_session
+              if @auth[:verify_hostname] &&
+                  !OpenSSL::SSL::verify_certificate_identity(riak_cert.cert, @host)
+                raise TlsError::CertHostMismatchError.new
+              end
+
+              unless riak_cert.valid?
+                raise TlsError::CertNotValidError.new
+              end
+
+              validator = R509::Cert::Validator.new riak_cert
+
+              validator_options = {}
+              validator_options[:ocsp] = !!@auth[:ocsp]
+              validator_options[:crl] = !!@auth[:crl]
+              validator_options[:crl_file] = @auth[:crl_file]
+
+              unless validator.validate(validator_options)
+                raise TlsError::CertRevokedError.new
+              end
+            end
+
+            def validator_options
+              o = {
+                ocsp: !!@auth[:ocsp],
+                crl: !!@auth[:crl]
+              }
+              
+              if @auth[:crl_file]
+                o[:crl_file] = @auth[:crl_file]
+                o[:crl] = true
+              end
+
+              return o
+            end
+
+            # Send an AuthReq with the authentication data. Rely on beefcake
+            # discarding message parts it doesn't understand.
+            def send_authentication
+              req = BC::RpbAuthReq.new @auth
+              write_message :AuthReq, req.encode
+              expect_message :AuthResp
+            end
+
+            # Ping the Riak node and make sure it actually works.
+            def validate_connection
+              write_message :PingReq
+              expect_message :PingResp
+            end
+
+            # Write a protocol buffers message to whatever the current
+            # socket is.
+            def write_message(code, message='')
+              if code.is_a? Symbol
+                code = BeefcakeMessageCodes.index code
+              end
+
+              header = [message.length+1, code].pack 'NC'
+              @sock.write header + message
+            end
+
+            def read_message
+              header = @sock.read 5
+              raise TlsError.new(t('ssl.eof_during_init')) if header.nil?
+              len, code = header.unpack 'NC'
+              decode = BeefcakeMessageCodes[code]
+              return decode, '' if len == 1
+              
+              message = @sock.read(len - 1)
+              return decode, message
+            end
+
+            def expect_message(expected_code)
+              if expected_code.is_a? Numeric
+                expected_code = BeefcakeMessageCodes[code]
+              end
+
+              candidate_code, message = read_message
+              return message if expected_code == candidate_code
+
+              raise TlsError.new(t('ssl.unexpected_during_init',
+                                   expected: expected_code.inspect,
+                                   actual: candidate_code.inspect,
+                                   body: message.inspect
+                                   ))
+              
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/riak/client/beefcake_protobuffs_backend.rb

@@ -1,7 +1,7 @@
 require 'base64'
 require 'riak/json'
 require 'riak/client'
-require 'riak/failed_request'
+require 'riak/errors/failed_request'
 require 'riak/client/protobuffs_backend'
 
 module Riak
@@ -11,13 +11,50 @@ module Riak
         begin
           require 'beefcake'
           require 'riak/client/beefcake/messages'
+          require 'riak/client/beefcake/message_overlay'
           require 'riak/client/beefcake/object_methods'
+          require 'riak/client/beefcake/crdt_operator'
+          require 'riak/client/beefcake/crdt_loader'
+          require 'riak/client/beefcake/protocol'
+          require 'riak/client/beefcake/socket'
           true
         rescue LoadError, NameError
           false
         end
       end
 
+      def protocol
+        p = Protocol.new socket
+        yield p
+      end
+
+      def new_socket
+        BeefcakeSocket.new @node.host, @node.pb_port, authentication: client.authentication
+      end
+
+      def ping
+        protocol do |p|
+          p.write :PingReq
+          p.expect :PingResp
+        end
+      end
+
+      def get_client_id
+        protocol do |p|
+          p.write :GetClientIdReq
+          p.expect(:GetClientIdResp, RpbGetClientIdResp).client_id
+        end
+      end
+
+      def server_info
+        resp = protocol do |p|
+          p.write :GetServerInfoReq
+          p.expect(:GetServerInfoResp, RpbGetServerInfoResp)
+        end
+
+        { node: resp.node, server_version: resp.server_version }
+      end
+
       def set_client_id(id)
         value = case id
                 when Integer
@@ -26,16 +63,29 @@ module Riak
                   id.to_s
                 end
         req = RpbSetClientIdReq.new(:client_id => value)
-        write_protobuff(:SetClientIdReq, req)
-        decode_response
+        protocol do |p|
+          p.write :SetClientIdReq, req
+          p.expect :SetClientIdResp
+        end
+        return true
       end
 
       def fetch_object(bucket, key, options={})
         options = prune_unsupported_options(:GetReq, normalize_quorums(options))
         bucket = Bucket === bucket ? bucket.name : bucket
         req = RpbGetReq.new(options.merge(:bucket => maybe_encode(bucket), :key => maybe_encode(key)))
-        write_protobuff(:GetReq, req)
-        decode_response(RObject.new(client.bucket(bucket), key))
+
+        resp = protocol do |p|
+          p.write :GetReq, req
+          p.expect :GetResp, RpbGetResp, empty_body_acceptable: true
+        end
+
+        if :empty == resp
+          raise Riak::ProtobuffsFailedRequest.new(:not_found, t('not_found'))
+        end
+
+        template = RObject.new(client.bucket(bucket), key)
+        load_object(resp, template)
       end
 
       def reload_object(robject, options={})
@@ -44,11 +94,21 @@ module Riak
         options[:key] = maybe_encode(robject.key)
         options[:if_modified] = maybe_encode Base64.decode64(robject.vclock) if robject.vclock
         req = RpbGetReq.new(prune_unsupported_options(:GetReq, options))
-        write_protobuff(:GetReq, req)
-        decode_response(robject)
+
+        resp = protocol do |p|
+          p.write :GetReq, req
+          p.expect :GetResp, RpbGetResp, empty_body_acceptable: true
+        end
+
+        if :empty == resp
+          raise Riak::ProtobuffsFailedRequest.new(:not_found, t('not_found'))
+        end
+
+        load_object(resp, robject)
       end
 
       def store_object(robject, options={})
+        options[:return_body] ||= options[:returnbody]
         options = normalize_quorums(options)
         if robject.prevent_stale_writes
           unless pb_conditionals?
@@ -62,8 +122,15 @@ module Riak
           end
         end
         req = dump_object(robject, prune_unsupported_options(:PutReq, options))
-        write_protobuff(:PutReq, req)
-        decode_response(robject)
+
+        resp = protocol do |p|
+          p.write(:PutReq, req)
+          p.expect :PutResp, RpbPutResp, empty_body_acceptable: true
+        end
+
+        return true if :empty == resp
+
+        load_object resp, robject
       end
 
       def delete_object(bucket, key, options={})
@@ -73,8 +140,13 @@ module Riak
         options[:key] = maybe_encode(key)
         options[:vclock] = Base64.decode64(options[:vclock]) if options[:vclock]
         req = RpbDelReq.new(prune_unsupported_options(:DelReq, options))
-        write_protobuff(:DelReq, req)
-        decode_response
+
+        protocol do |p|
+          p.write :DelReq, req
+          p.expect :DelResp
+        end
+        
+        return true
       end
 
       def get_counter(bucket, key, options={})
@@ -85,9 +157,17 @@ module Riak
         options[:key] = key
         
         request = RpbCounterGetReq.new options
-        write_protobuff :CounterGetReq, request
         
-        decode_response
+        resp = protocol do |p|
+          p.write :CounterGetReq, request
+          p.expect :CounterGetResp, RpbCounterGetResp, empty_body_acceptable: true
+        end
+        
+        if :empty == resp
+          return 0
+        end
+
+        return resp.value || 0
       end
 
       def post_counter(bucket, key, amount, options={})
@@ -96,53 +176,83 @@ module Riak
         options = normalize_quorums(options)
         options[:bucket] = bucket
         options[:key] = key
-        # TODO: raise if ammount doesn't fit in sint64
+        # TODO: raise if amount doesn't fit in sint64
         options[:amount] = amount
-
+        options[:returnvalue] = options[:returnvalue] || options[:return_value]
+        
         request = RpbCounterUpdateReq.new options
-        write_protobuff :CounterUpdateReq, request
 
-        decode_response
+        resp = protocol do |p|
+          p.write :CounterUpdateReq, request
+          p.expect :CounterUpdateResp, RpbCounterUpdateResp, empty_body_acceptable: true
+        end
+
+        return nil if :empty == resp
+        
+        return resp.value
       end
 
-      def get_bucket_props(bucket)
+      def get_bucket_props(bucket, options = {  })
         bucket = bucket.name if Bucket === bucket
+
         req = RpbGetBucketReq.new(:bucket => maybe_encode(bucket))
-        write_protobuff(:GetBucketReq, req)
-        resp = normalize_quorums decode_response
+        req.type = options[:type] if options[:type]
+
+        resp_message = protocol do |p|
+          p.write :GetBucketReq, req
+          p.expect :GetBucketResp, RpbGetBucketResp
+        end
+
+        resp = normalize_quorums resp_message.props.to_hash.stringify_keys
         normalized = normalize_hooks resp
         normalized.stringify_keys
       end
 
-      def set_bucket_props(bucket, props)
+      def set_bucket_props(bucket, props, type=nil)
         bucket = bucket.name if Bucket === bucket
-        props = props.slice('n_val', 'allow_mult')
-        req = RpbSetBucketReq.new(:bucket => maybe_encode(bucket), :props => RpbBucketProps.new(props))
-        write_protobuff(:SetBucketReq, req)
-        decode_response
+        req = RpbSetBucketReq.new(
+                                  bucket: maybe_encode(bucket),
+                                  props: RpbBucketProps.new(props.symbolize_keys),
+                                  type: type)
+
+        protocol do |p|
+          p.write :SetBucketReq, req
+          p.expect :SetBucketResp
+        end
       end
 
       def reset_bucket_props(bucket)
         bucket = bucket.name if Bucket === bucket
         req = RpbResetBucketReq.new(:bucket => maybe_encode(bucket))
-        write_protobuff(:ResetBucketReq)
-        decode_response
+
+        protocol do |p|
+          p.write :ResetBucketReq, req
+          p.expect :ResetBucketResp
+        end
       end
 
       def list_keys(bucket, options={}, &block)
         bucket = bucket.name if Bucket === bucket
         req = RpbListKeysReq.new(options.merge(:bucket => maybe_encode(bucket)))
-        write_protobuff(:ListKeysReq, req)
+
         keys = []
-        while msg = decode_response
-          break if msg.done
-          if block_given?
-            yield msg.keys
-          else
-            keys += msg.keys
+
+        protocol do |p|
+          p.write :ListKeysReq, req
+
+          while msg = p.expect(:ListKeysResp, RpbListKeysResp)
+            break if msg.done
+            if block_given?
+              yield msg.keys
+            else
+              keys += msg.keys
+            end
           end
         end
-        block_given? || keys
+
+        return keys unless block_given?
+
+        return true
       end
 
       # override the simple list_buckets
@@ -155,24 +265,35 @@ module Riak
         
         request = RpbListBucketsReq.new options
 
-        write_protobuff :ListBucketsReq, request
+        resp = protocol do |p|
+          p.write :ListBucketsReq, request
+
+          p.expect :ListBucketsResp, RpbListBucketsResp, empty_body_acceptable: true
+        end
+
+        return [] if :empty == resp
 
-        decode_response
+        resp.buckets
       end
 
       def mapred(mr, &block)
         raise MapReduceError.new(t("empty_map_reduce_query")) if mr.query.empty? && !mapred_phaseless?
         req = RpbMapRedReq.new(:request => mr.to_json, :content_type => "application/json")
-        write_protobuff(:MapRedReq, req)
+        
         results = MapReduce::Results.new(mr)
-        while msg = decode_response
-          break if msg.done
-          if block_given?
-            yield msg.phase, JSON.parse(msg.response)
-          else
-            results.add msg.phase, JSON.parse(msg.response)
+        
+        protocol do |p|
+          p.write :MapRedReq, req
+          while msg = p.expect(:MapRedResp, RpbMapRedResp)
+            break if msg.done
+            if block_given?
+              yield msg.phase, JSON.parse(msg.response)
+            else
+              results.add msg.phase, JSON.parse(msg.response)
+            end
           end
         end
+        
         block_given? || results.report
       end
 
@@ -197,8 +318,11 @@ module Riak
         options[:stream] = block_given?
 
         req = RpbIndexReq.new(options)
-        write_protobuff(:IndexReq, req)
-        decode_index_response(&block)
+
+        protocol do |p|
+          p.write :IndexReq, req
+          decode_index_response(p, &block)
+        end
       end
 
       def search(index, query, options={})
@@ -206,35 +330,97 @@ module Riak
         options = options.symbolize_keys
         options[:op] = options.delete(:'q.op') if options[:'q.op']
         req = RpbSearchQueryReq.new(options.merge(:index => index || 'search', :q => query))
-        write_protobuff(:SearchQueryReq, req)
-        decode_response
+
+        resp = protocol do |p|
+          p.write :SearchQueryReq, req
+          p.expect :SearchQueryResp, RpbSearchQueryResp
+        end
+
+        resp.docs = [] if resp.docs.nil?
+
+        ret = { 'max_score' => resp.max_score, 'num_found' => resp.num_found }
+        ret['docs'] = resp.docs.map { |d| decode_doc d }
+
+        return ret
+      end
+
+      def create_search_index(name, schema=nil, n_val=nil)
+        index = RpbYokozunaIndex.new(:name => name, :schema => schema, :n_val => n_val)
+        req = RpbYokozunaIndexPutReq.new(:index => index)
+
+        protocol do |p|
+          p.write :YokozunaIndexPutReq, req
+          p.expect :PutResp
+        end
+      end
+
+      def get_search_index(name)
+        req = RpbYokozunaIndexGetReq.new(:name => name)
+        resp = protocol do |p|
+          p.write :YokozunaIndexGetReq, req
+          p.expect :YokozunaIndexGetResp, RpbYokozunaIndexGetResp, empty_body_acceptable: true
+        end
+        
+        if :empty == resp
+          raise Riak::ProtobuffsFailedRequest.new(:not_found, t('not_found'))
+        end
+
+        if resp.index && Array === resp
+          resp.index.map{|index| {:name => index.name, :schema => index.schema, :n_val => index.n_val} }
+        else
+          resp
+        end
+      end
+
+      def delete_search_index(name)
+        req = RpbYokozunaIndexDeleteReq.new(:name => name)
+        protocol do |p|
+          p.write :YokozunaIndexDeleteReq, req
+          p.expect :DelResp
+        end
+        true
+      end
+
+      def create_search_schema(name, content)
+        schema = RpbYokozunaSchema.new(:name => name, :content => content)
+        req = RpbYokozunaSchemaPutReq.new(:schema => schema)
+
+        protocol do |p|
+          p.write :YokozunaSchemaPutReq, req
+          p.expect :PutResp
+        end
+        true
+      end
+
+      def get_search_schema(name)
+        req = RpbYokozunaSchemaGetReq.new(:name => name)
+
+        resp = protocol do |p|
+          p.write :YokozunaSchemaGetReq, req
+          p.expect :YokozunaSchemaGetResp, RpbYokozunaSchemaGetResp
+        end
+
+        resp.schema ? resp.schema : resp
       end
 
-      private
       def write_protobuff(code, message)
         encoded = message.encode
         header = [encoded.length+1, MESSAGE_CODES.index(code)].pack("NC")
         socket.write(header + encoded)
       end
 
+      private
       def decode_response(*args)
         header = socket.read(5)
-        raise SocketError, "Unexpected EOF on PBC socket" if header.nil?
+        raise ProtobuffsFailedHeader.new if header.nil?
         msglen, msgcode = header.unpack("NC")
         if msglen == 1
           case MESSAGE_CODES[msgcode]
-          when :PingResp, 
-               :SetClientIdResp, 
-               :PutResp, 
-               :DelResp, 
-               :SetBucketResp, 
-               :ResetBucketResp
-            true
-          when :ListBucketsResp, 
-               :ListKeysResp, 
+          when :ListBucketsResp,  
                :IndexResp
             []
-          when :GetResp
+          when :GetResp,
+               :YokozunaSchemaGetResp
             raise Riak::ProtobuffsFailedRequest.new(:not_found, t('not_found'))
           when :CounterGetResp,
                :CounterUpdateResp
@@ -248,44 +434,6 @@ module Riak
           when :ErrorResp
             res = RpbErrorResp.decode(message)
             raise Riak::ProtobuffsFailedRequest.new(res.errcode, res.errmsg)
-          when :GetClientIdResp
-            res = RpbGetClientIdResp.decode(message)
-            res.client_id
-          when :GetServerInfoResp
-            res = RpbGetServerInfoResp.decode(message)
-            {:node => res.node, :server_version => res.server_version}
-          when :GetResp
-            res = RpbGetResp.decode(message)
-            load_object(res, args.first)
-          when :PutResp
-            res = RpbPutResp.decode(message)
-            load_object(res, args.first)
-          when :ListBucketsResp
-            res = RpbListBucketsResp.decode(message)
-            res.buckets
-          when :ListKeysResp
-            RpbListKeysResp.decode(message)
-          when :GetBucketResp
-            res = RpbGetBucketResp.decode(message)
-            res.props.to_hash.stringify_keys
-          when :MapRedResp
-            RpbMapRedResp.decode(message)
-          when :IndexResp
-            res = RpbIndexResp.decode(message)
-            IndexCollection.new_from_protobuf res
-          when :SearchQueryResp
-            res = RpbSearchQueryResp.decode(message)
-            { 'docs' => res.docs.map {|d| decode_doc(d) },
-              'max_score' => res.max_score,
-              'num_found' => res.num_found }
-          when :CSBucketResp
-            res = RpbCSBucketResp.decode message
-          when :CounterUpdateResp
-            res = RpbCounterUpdateResp.decode message
-            res.value || nil
-          when :CounterGetResp
-            res = RpbCounterGetResp.decode message
-            res.value || 0
           end
         end
       rescue SystemCallError, SocketError => e
@@ -312,45 +460,37 @@ module Riak
         end
       end
 
-      def decode_index_response
+      def decode_index_response(p)
         loop do
-          header = socket.read(5)
-          raise SocketError, "Unexpected EOF on PBC socket" if header.nil?
-          msglen, msgcode = header.unpack("NC")
-          code = MESSAGE_CODES[msgcode]
-          if code == :ErrorResp
-            resp = RpbErrorResp.decode socket.read msglen - 1
-            message = resp.errmsg
-            if match = message.match(/indexes_not_supported,(\w+)/)
-              message = t('index.wrong_backend', backend: match[1])
-            end
-            raise ProtobuffsFailedRequest.new resp.errcode, message
-          elsif code != :IndexResp
-            teardown # close socket, we don't know what's going on anymore
-            inner = ProtobuffsFailedRequest.new code, t('protobuffs.unexpected_message')
-            raise Innertube::Pool::BadResource, inner
-          end
+          resp = p.expect :IndexResp, RpbIndexResp, empty_body_acceptable: true
 
-          if msglen == 1
+          if :empty == resp
             return if block_given?
             return IndexCollection.new_from_protobuf(RpbIndexResp.decode(''))
           end
 
-          message = RpbIndexResp.decode socket.read msglen - 1
-
           if !block_given?
-            return IndexCollection.new_from_protobuf(message)
+            return IndexCollection.new_from_protobuf(resp)
           end
           
-          content = message.keys || message.results || []
+          content = resp.keys || resp.results || []
           yield content
           
-          return if message.done
+          return if resp.done
+        end
+      rescue ProtobuffsErrorResponse => err
+        if match = err.message.match(/indexes_not_supported,(\w+)/)
+          old_err = err
+          err = ProtobuffsFailedRequest.new(:indexes_not_supported, 
+                                            t('index.wrong_backend', backend: match[1])
+                                            )
         end
+
+        raise err
       end
 
       def decode_doc(doc)
-        Hash[doc.properties.map {|p| [ force_utf8(p.key), force_utf8(p.value) ] }]
+        Hash[doc.fields.map {|p| [ force_utf8(p.key), force_utf8(p.value) ] }]
       end
 
       def force_utf8(str)