dalli 0.9.0 → 0.9.1

data/lib/dalli/sasl/base64.rb

@@ -0,0 +1,14 @@
+begin
+  require 'base64'
+rescue LoadError
+  # Ruby 1.9 compat
+  module Base64
+    def self.encode64(data)
+      [data].pack('m')
+    end
+
+    def self.decode64(data64)
+      data64.unpack('m')[0]
+    end
+  end
+end

data/lib/dalli/sasl/digest_md5.rb

@@ -0,0 +1,175 @@
+require 'digest/md5'
+
+module SASL
+  ##
+  # RFC 2831:
+  # http://tools.ietf.org/html/rfc2831
+  class DigestMD5 < Mechanism
+    attr_writer :cnonce
+
+    def initialize(*a)
+      super
+      @nonce_count = 0
+    end
+
+    def start
+      @state = nil
+      unless defined? @nonce
+        ['auth', nil]
+      else
+        # reauthentication
+        receive('challenge', '')
+      end
+    end
+
+    def receive(message_name, content)
+      if message_name == 'challenge'
+        c = decode_challenge(content)
+
+        unless c['rspauth']
+          response = {}
+          if defined?(@nonce) && response['nonce'].nil?
+            # Could be reauth
+          else
+            # No reauth:
+            @nonce_count = 0
+          end
+          @nonce ||= c['nonce']
+          response['nonce'] = @nonce
+          response['charset'] = 'utf-8'
+          response['username'] = preferences.username
+          response['realm'] = c['realm'] || preferences.realm
+          @cnonce = generate_nonce unless defined? @cnonce
+          response['cnonce'] = @cnonce
+          @nc = next_nc
+          response['nc'] = @nc
+          @qop = c['qop'] || 'auth'
+          response['qop'] = @qop
+          response['digest-uri'] = preferences.digest_uri #"memcached/#{self.hostname}"
+          response['response'] = response_value(response['nonce'], response['nc'], response['cnonce'], response['qop'])
+          ['response', encode_response(response)]
+        else
+          rspauth_expected = response_value(@nonce, @nc, @cnonce, @qop, '')
+          p :rspauth_received=>c['rspauth'], :rspauth_expected=>rspauth_expected
+          if c['rspauth'] == rspauth_expected
+            ['response', nil]
+          else
+            # Bogus server?
+            @state = :failure
+            ['failure', nil]
+          end
+        end
+      else
+        # No challenge? Might be success or failure
+        super
+      end
+    end
+
+    private
+
+    def decode_challenge(text)
+      challenge = {}
+      
+      state = :key
+      key = ''
+      value = ''
+
+      text.scan(/./) do |ch|
+        if state == :key
+          if ch == '='
+            state = :value
+          elsif ch =~ /\S/
+            key += ch
+          end
+          
+        elsif state == :value
+          if ch == ','
+            challenge[key] = value
+            key = ''
+            value = ''
+            state = :key
+          elsif ch == '"' and value == ''
+            state = :quote
+          else
+            value += ch
+          end
+
+        elsif state == :quote
+          if ch == '"'
+            state = :value
+          else
+            value += ch
+          end
+        end
+      end
+      challenge[key] = value unless key == ''
+      
+      p :decode_challenge => challenge
+      challenge
+    end
+
+    def encode_response(response)
+      p :encode_response => response
+      response.collect do |k,v|
+        if v.include?('"')
+          v.sub!('\\', '\\\\')
+          v.sub!('"', '\\"')
+          "#{k}=\"#{v}\""
+        else
+          "#{k}=#{v}"
+        end
+      end.join(',')
+    end
+
+    def generate_nonce
+      nonce = ''
+      while nonce.length < 16
+        c = rand(128).chr
+        nonce += c if c =~ /^[a-zA-Z0-9]$/
+      end
+      nonce
+    end
+
+    ##
+    # Function from RFC2831
+    def h(s); Digest::MD5.digest(s); end
+    ##
+    # Function from RFC2831
+    def hh(s); Digest::MD5.hexdigest(s); end
+    
+    ##
+    # Calculate the value for the response field
+    def response_value(nonce, nc, cnonce, qop, a2_prefix='AUTHENTICATE')
+      p :response_value => {:nonce=>nonce,
+        :cnonce=>cnonce,
+        :qop=>qop,
+        :username=>preferences.username,
+        :realm=>preferences.realm,
+        :password=>preferences.password,
+        :authzid=>preferences.authzid}
+      a1_h = h("#{preferences.username}:#{preferences.realm}:#{preferences.password}")
+      a1 = "#{a1_h}:#{nonce}:#{cnonce}"
+      if preferences.authzid
+        a1 += ":#{preferences.authzid}"
+      end
+      if qop && (qop.downcase == 'auth-int' || qop.downcase == 'auth-conf')
+        a2 = "#{a2_prefix}:#{preferences.digest_uri}:00000000000000000000000000000000"
+      else
+        a2 = "#{a2_prefix}:#{preferences.digest_uri}"
+      end
+      hh("#{hh(a1)}:#{nonce}:#{nc}:#{cnonce}:#{qop}:#{hh(a2)}")
+    end
+
+    def next_nc
+      @nonce_count += 1
+      s = @nonce_count.to_s
+      s = "0#{s}" while s.length < 8
+      s
+    end
+  end
+  
+  # TODO: need to test
+  #MECHANISMS['DIGEST-MD5'] = SASL::DigestMD5
+  
+end
+

data/lib/dalli/sasl/plain.rb

@@ -0,0 +1,18 @@
+module SASL
+  ##
+  # RFC 4616:
+  # http://tools.ietf.org/html/rfc4616
+  class Plain < Mechanism
+    
+    def start
+      @state = nil
+      message = [preferences.authzid.to_s,
+                 preferences.username,
+                 preferences.password].join("\000")
+      ['auth', message]
+    end
+  end
+  
+  MECHANISMS['PLAIN'] = SASL::Plain
+  
+end

data/lib/dalli/server.rb

@@ -16,6 +16,7 @@ module Dalli
       Dalli.logger.debug { "#{@hostname}:#{@port} running memcached v#{request(:version)}" }
     end
     
+    # Chokepoint method for instrumentation
     def request(op, *args)
       begin
         send(op, *args)
@@ -24,8 +25,8 @@ module Dalli
       rescue Dalli::DalliError
         raise
       rescue Exception => ex
-        puts "Unexpected exception: #{ex.class.name}: #{ex.message}"
-        puts ex.backtrace.join("\n")
+        Dalli.logger.error "Unexpected exception in Dalli: #{ex.class.name}: #{ex.message}"
+        Dalli.logger.error ex.backtrace.join("\n\t")
         down!
       end
     end
@@ -44,6 +45,8 @@ module Dalli
     def unlock!
     end
 
+    # NOTE: Additional public methods should be overridden in Dalli::Threadsafe
+
     private
 
     def down!
@@ -66,7 +69,7 @@ module Dalli
       write(req)
     end
 
-    def set(key, value, ttl=0)
+    def set(key, value, ttl)
       raise Dalli::DalliError, "Value too large, memcached can only store 1MB of data per key" if value.size > ONE_MB
 
       req = [REQUEST, OPCODES[:set], key.size, 8, 0, 0, value.size + key.size + 8, 0, 0, 0, ttl, key, value].pack(FORMAT[:set])
@@ -80,10 +83,10 @@ module Dalli
       generic_response
     end
 
-    def add(key, value, ttl)
+    def add(key, value, ttl, cas)
       raise Dalli::DalliError, "Value too large, memcached can only store 1MB of data per key" if value.size > ONE_MB
 
-      req = [REQUEST, OPCODES[:add], key.size, 8, 0, 0, value.size + key.size + 8, 0, 0, 0, ttl, key, value].pack(FORMAT[:add])
+      req = [REQUEST, OPCODES[:add], key.size, 8, 0, 0, value.size + key.size + 8, 0, cas, 0, ttl, key, value].pack(FORMAT[:add])
       write(req)
       generic_response
     end
@@ -99,15 +102,29 @@ module Dalli
       write(req)
       generic_response
     end
-    
-    def decr(key, count)
-      raise NotImplementedError
+
+    def decr(key, count, ttl, default)
+      expiry = default ? ttl : 0xFFFFFFFF
+      default ||= 0
+      (h, l) = split(count)
+      (dh, dl) = split(default)
+      req = [REQUEST, OPCODES[:decr], key.size, 20, 0, 0, key.size + 20, 0, 0, h, l, dh, dl, expiry, key].pack(FORMAT[:decr])
+      write(req)
+      body = generic_response
+      body ? longlong(*body.unpack('NN')) : body
     end
     
-    def incr(key, count)
-      raise NotImplementedError
+    def incr(key, count, ttl, default)
+      expiry = default ? ttl : 0xFFFFFFFF
+      default ||= 0
+      (h, l) = split(count)
+      (dh, dl) = split(default)
+      req = [REQUEST, OPCODES[:incr], key.size, 20, 0, 0, key.size + 20, 0, 0, h, l, dh, dl, expiry, key].pack(FORMAT[:incr])
+      write(req)
+      body = generic_response
+      body ? longlong(*body.unpack('NN')) : body
     end
-
+    
     # Noop is a keepalive operation but also used to demarcate the end of a set of pipelined commands.
     # We need to read all the responses at once.
     def noop
@@ -140,6 +157,27 @@ module Dalli
       keyvalue_response
     end
 
+    def cas(key)
+      req = [REQUEST, OPCODES[:get], key.size, 0, 0, 0, key.size, 0, 0, key].pack(FORMAT[:get])
+      write(req)
+      cas_response
+    end
+
+    def cas_response
+      header = read(24)
+      raise Dalli::NetworkError, 'No response' if !header
+      (extras, status, count, _, cas) = header.unpack(CAS_HEADER)
+      data = read(count) if count > 0
+      if status == 1
+        nil
+      elsif status != 0
+        raise Dalli::DalliError, "Response error #{status}: #{RESPONSE_CODES[status]}"
+      elsif data
+        data = data[extras..-1] if extras != 0
+      end
+      [data, cas]
+    end
+
     def generic_response
       header = read(24)
       raise Dalli::NetworkError, 'No response' if !header
@@ -147,6 +185,8 @@ module Dalli
       data = read(count) if count > 0
       if status == 1
         nil
+      elsif status == 2
+        false # Not stored, normal status for add operation
       elsif status != 0
         raise Dalli::DalliError, "Response error #{status}: #{RESPONSE_CODES[status]}"
       elsif data
@@ -209,6 +249,7 @@ module Dalli
         # end ugly code
 
         sock.setsockopt Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1
+        sasl_authentication(sock) if Dalli::Server.need_auth?
         sock
       end
     end
@@ -237,7 +278,6 @@ module Dalli
             raise Timeout::Error, "IO timeout"
           end
         end
-        raise Errno::EINVAL, "Not enough data to fulfill read request: #{value.inspect}" if value.size != count
         value
       rescue Errno::ECONNRESET, Errno::EPIPE, Errno::ECONNABORTED, Errno::EBADF, Errno::EINVAL, Timeout::Error, EOFError
         down!
@@ -245,6 +285,15 @@ module Dalli
       end
     end
 
+    def split(n)
+      [0xFFFFFFFF & n, n >> 32]
+    end
+
+    def longlong(a, b)
+      a | (b << 32)
+    end
+
+    CAS_HEADER = '@4vnNNQ'
     NORMAL_HEADER = '@4vnN'
     KV_HEADER = '@2n@6nN'
 
@@ -259,6 +308,7 @@ module Dalli
       4 => 'Invalid arguments',
       5 => 'Item not stored',
       6 => 'Incr/decr on a non-numeric value',
+      0x20 => 'Authentication required',
       0x81 => 'Unknown command',
       0x82 => 'Out of memory',
     }
@@ -278,6 +328,9 @@ module Dalli
       :append => 0x0E,
       :prepend => 0x0F,
       :stat => 0x10,
+      :auth_negotiation => 0x20,
+      :auth_request => 0x21,
+      :auth_continue => 0x22,
     }
     
     HEADER = "CCnCCnNNQ"
@@ -296,8 +349,67 @@ module Dalli
       :stat => 'a*',
       :append => 'a*a*',
       :prepend => 'a*a*',
+      :auth_request => 'a*a*',
+      :auth_continue => 'a*a*',
     }
     FORMAT = OP_FORMAT.inject({}) { |memo, (k, v)| memo[k] = HEADER + v; memo }
+
+
+    #######
+    # SASL authentication support for NorthScale
+    #######
+
+    def self.need_auth?
+      ENV['MEMCACHE_USERNAME']
+    end
     
+    def init_sasl
+      require 'dalli/sasl/base'
+      require 'dalli/sasl/base64'
+      require 'dalli/sasl/digest_md5'
+      require 'dalli/sasl/plain'
+    end
+
+    def username
+      ENV['MEMCACHE_USERNAME']
+    end
+
+    def password
+      ENV['MEMCACHE_PASSWORD']
+    end
+
+    def sasl_authentication(socket)
+      init_sasl if !defined?(::SASL)
+      
+      # negotiate
+      req = [REQUEST, OPCODES[:auth_negotiation], 0, 0, 0, 0, 0, 0, 0].pack(FORMAT[:noop])
+      socket.write(req)
+      header = socket.read(24)
+      raise Dalli::NetworkError, 'No response' if !header
+      (extras, status, count) = header.unpack(NORMAL_HEADER)
+      raise Dalli::NetworkError, "Unexpected message format: #{extras} #{count}" unless extras == 0 && count > 0
+      content = socket.read(count)
+      return (Dalli.logger.debug("Authentication not required/supported by server")) if status == 0x81
+      mechanisms = content.split(' ')
+
+      # request
+      sasl = ::SASL.new(mechanisms)
+      msg = sasl.start[1]
+      mechanism = sasl.name
+      #p [mechanism, msg]
+      req = [REQUEST, OPCODES[:auth_request], mechanism.size, 0, 0, 0, mechanism.size + msg.size, 0, 0, mechanism, msg].pack(FORMAT[:auth_request])
+      socket.write(req)
+
+      header = socket.read(24)
+      raise Dalli::NetworkError, 'No response' if !header
+      (extras, status, count) = header.unpack(NORMAL_HEADER)
+      raise Dalli::NetworkError, "Unexpected message format: #{extras} #{count}" unless extras == 0 && count > 0
+      content = socket.read(count)
+      raise Dalli::NetworkError, "Error authenticating: #{status}" unless status == 0x21
+      (step, msg) = sasl.receive('challenge', content)
+      raise Dalli::NetworkError, "Authentication failed" if sasl.failed? || step != 'response'
+
+      
+    end
   end
 end

data/lib/dalli/version.rb

@@ -1,3 +1,3 @@
 module Dalli
-  VERSION = '0.9.0'
+  VERSION = '0.9.1'
 end