parse redis diskstore files (which contain hashes)

diskstore_inspector.rb

# encoding: UTF-8

# dependencies:
# - wget http://dist.schmorp.de/liblzf/liblzf-3.6.tar.gz
# - gem install lzfruby

require 'lzfruby'
require 'stringio'

class DiskstoreInspector

  REDIS_RDB_6BITLEN = 0
  REDIS_RDB_14BITLEN = 1
  REDIS_RDB_32BITLEN = 2
  REDIS_RDB_ENCVAL = 3
  REDIS_RDB_LENERR = (2 ** 32) - 1 # UINT_MAX

  REDIS_RDB_ENC_INT8 = 0
  REDIS_RDB_ENC_INT16 = 1
  REDIS_RDB_ENC_INT32 = 2
  REDIS_RDB_ENC_LZF = 3

  REDIS_STRING = 0
  REDIS_LIST = 1
  REDIS_SET = 2
  REDIS_ZSET = 3
  REDIS_HASH = 4

  def self.parse(bytes)
    io = StringIO.new(bytes)
    read_key_value(io)
  end

  def self.read_key_value(io)
    type, _ = read_len(io)
    raise "not a hash" unless type == REDIS_HASH
    key = read_string(io)
    value = read_hash(io)
    [key, value]
  end

  def self.read_hash(io)
    result = {}
    len, _ = read_len(io)
    len.times do
      key, val = read_string(io), read_string(io)
      result[key] = val
    end
    result
  end

  def self.read_len(io)
    buf = io.read(1).bytes.to_a
    type = (buf[0] & 0xc0) >> 6
    if (type == REDIS_RDB_ENCVAL)
      return buf[0] & 0x3F, true
    elsif (type == REDIS_RDB_6BITLEN)
      return buf[0] & 0x3F, false
    elsif (type == REDIS_RDB_14BITLEN)
      b = io.read(1).bytes.to_a
      return ((buf[0] & 0x3F) << 8) | b[0], false
    else
      return io.read(4).unpack('N'), false
    end
  end

  def self.read_integer(io, enctype)
    if (enctype == REDIS_RDB_ENC_INT8)
      return io.read(1).bytes.to_a[0]
    elsif (enctype == REDIS_RDB_ENC_INT16)
      buf = io.read(2).bytes.to_a
      return buf[0] | (buf[1] << 8)
    elsif (enctype == REDIS_RDB_ENC_INT32)
      buf = io.read(4).bytes.to_a
      return buf[0] | (buf[1] << 8) | (buf[2] << 16) | (buf[3] << 24)
    else
      raise "Unknown RDB integer encoding type"
    end
  end

  def self.read_string(io)
    len, encoded = read_len(io)
    if encoded
      case len
      when REDIS_RDB_ENC_INT8, REDIS_RDB_ENC_INT16, REDIS_RDB_ENC_INT32
        return read_integer(io, len).to_s
      when REDIS_RDB_ENC_LZF
        return read_lzf_string(io)
      else
        raise "Unknown RDB encoding type"
      end
    elsif len == REDIS_RDB_LENERR
      return nil
    else
      return io.read(len)
    end
  end

  def self.read_lzf_string(io)
    clen, _ = read_len(io)
    len, _ = read_len(io)
    header = "ZV\x01" + (clen >> 8).chr + (clen & 0xff).chr + (len >> 8).chr + (len & 0xff).chr
    compressed = StringIO.new(header + io.read(clen))
    uncompressed = StringIO.new
    LZF.decompress(compressed, uncompressed)
    uncompressed.string
  end

end