mongo-proxy 1.0.0

data/lib/mongo-proxy/wire.rb

@@ -0,0 +1,526 @@
+require 'bson'
+
+# This is a set of functions for dealing with the Mongo wire protocol. It has
+# methods for moving between Ruby hashes and Mongo wire segments, including
+# their embedded BSON. The MongoDB wire protocol is documented at:
+#
+module WireMongo
+
+  HEADER_SIZE           = 16
+
+  OP_REPLY              = :reply
+  OP_MSG                = :msg
+  OP_UPDATE             = :update
+  OP_INSERT             = :insert
+  OP_QUERY              = :query
+  OP_GET_MORE           = :get_more
+  OP_DELETE             = :delete
+  OP_KILL_CURSORS       = :kill_cursors
+
+  OPS = {
+    1    => OP_REPLY,
+    1000 => OP_MSG,
+    2001 => OP_UPDATE,
+    2002 => OP_INSERT,
+    2004 => OP_QUERY,
+    2005 => OP_GET_MORE,
+    2006 => OP_DELETE,
+    2007 => OP_KILL_CURSORS
+  }
+  OPS_INVERTED = OPS.invert
+
+  FLAG_UPDATE_UPSERT        = 1
+  FLAG_UPDATE_MULTIUPDATE   = (1 << 1)
+  FLAG_DELETE_MULTI         = 1
+
+  # Parse out an arbitrary binary mongo message, returning a hash
+  # representation for easy manipulation.
+  def self.receive socket
+    if socket.is_a?(String)
+      socket = StringIO.new(socket)
+      socket.set_encoding('UTF-8', 'UTF-8')
+    end
+
+    chunk1, x = receive_header(socket)
+    return nil, nil unless x && chunk1
+
+    parsed = {}
+
+    chunk2 = socket.read(x[:messageLength] - HEADER_SIZE)
+
+    case x[:opCode]
+    when OP_REPLY
+      parsed = receive_reply(chunk2)
+    when OP_UPDATE
+      parsed = receive_update(chunk2)
+    when OP_INSERT
+      parsed = receive_insert(chunk2)
+    when OP_QUERY
+      parsed = receive_query(chunk2)
+    when OP_DELETE
+      parsed = receive_delete(chunk2)
+    when OP_GET_MORE
+      parsed = receive_get_more(chunk2)
+    when OP_KILL_CURSORS
+      parsed = receive_kill_cursors(chunk2)
+    else
+      puts "could not parse message type :#{x[:opCode]}:"
+    end
+
+    parsed[:header] = x
+    full = chunk1 + chunk2
+    full = full.force_encoding('UTF-8')
+    return full, parsed
+
+  rescue Exception => e
+    @@log.warn "failed to read from socket #{socket.to_s}"
+    return nil
+  end
+
+  # Write a hash document representation into its corresponding binary form.
+  # This method can be used with documents in the format that receive returns,
+  # making it easy to parse a message, change it, and re-encode it.
+  def self.write doc
+    body = nil
+
+    case doc[:header][:opCode]
+    when OP_REPLY
+      body = write_reply(doc)
+    when OP_UPDATE
+      body = write_update(doc)
+    when OP_INSERT
+      body = write_insert(doc)
+    when OP_QUERY
+      body = write_query(doc)
+    when OP_DELETE
+      body = write_delete(doc)
+    when OP_GET_MORE
+      body = write_get_more(doc)
+    when OP_KILL_CURSORS
+      body = write_kill_cursors(doc)
+    else
+      puts "could not write message type :#{doc[:header][:opCode]}:"
+      return nil
+    end
+
+    body = body.force_encoding('UTF-8')
+
+    return write_header(doc[:header], body)
+  end
+
+  # Receive the Mongo Wire message header from a stream.
+  #
+  # int32 :messageLength - Length in bytes of subsequent message.
+  # int32 :requestID - Identifier of this message.
+  # int32 :responseTo - RequestID from the original request.
+  # int32 :opCode - Message type.
+  def self.receive_header(stream)
+    chunk = stream.read(HEADER_SIZE)
+    return nil unless chunk != nil && chunk.bytesize == HEADER_SIZE
+
+    x = {}
+    x[:messageLength], x[:requestID], x[:responseTo], x[:opCode] = chunk.unpack('VVVV')
+    x[:opCode] = OPS[x[:opCode]]
+    return chunk, x
+  end
+
+  def self.write_header doc, body
+    raise 'no requestID' unless doc[:requestID]
+    raise 'no opCode' unless doc[:opCode]
+    response_to = (doc[:responseTo] or 0)
+    length = body.bytesize + HEADER_SIZE
+
+    header = [length, doc[:requestID], response_to, OPS_INVERTED[doc[:opCode]]].pack('VVVV')
+    header = header.force_encoding('UTF-8')
+    return header + body
+  end
+
+  def self.receive_bson(chunk, start, max = 10000)
+    docs = []
+
+    while start < chunk.bytesize and docs.size < max
+      bsonLength = chunk[start..(start + 4)].unpack('V')[0]
+      doc = nil
+
+      begin
+        doc = BSON.deserialize(chunk[start..(start + bsonLength - 1)])
+      rescue Exception => e
+        puts 'could not deserialize BSON:'
+        pp chunk[start..(start + bsonLength)]
+        return nil, nil
+      end
+
+      docs << doc
+      start += bsonLength
+    end
+
+    return docs, start
+  end
+
+  def self.write_bson(docs)
+    docs = [docs] if docs.is_a? Hash
+    
+    x = ''
+    docs.each do |doc|
+      x << BSON.serialize(doc).to_s
+    end
+
+    return x
+  end
+
+  def self.min(a, b)
+    (a > b ? b : a)
+  end
+
+  def self.parse_full_collection(full_collection)
+    x = full_collection.split('.')
+    return x[0], x[1..-1].join('.')
+  end
+
+  def self.build_full_collection(database, collection)
+    return "#{database}.#{collection}"
+  end
+
+  # OP_REPLY: 1
+  # A reply to a client request.
+  #
+  # header :header - Message header.
+  # int32 :responseFlags - A bit vector of response flags.
+  # int64 :cursorID - ID of open cursor, if there is one. 0 otherwise.
+  # int32 :startingFrom - Offset in cursor of this reply message.
+  # int64 :numberReturned - Number of documents in the reply.
+  def self.receive_reply(chunk)
+    x = {}
+    x[:responseFlags], x[:cursorID], x[:startingFrom], x[:numberReturned] = chunk.unpack('VQ<VV')
+    x[:documents], _ = receive_bson(chunk, 20, x[:numberReturned])
+    return x
+  end
+
+  def self.build_reply(documents, request_id, response_to,
+      response_flags = 0, cursor_id = 0, starting_from = 0)
+    documents = [documents] if documents.is_a?(Hash)
+
+    return {
+      :responseFlags => response_flags,
+      :startingFrom => starting_from,
+      :numberReturned => documents.size,
+      :cursorID => cursor_id,
+      :documents => documents,
+      :header => {
+        :requestID => request_id,
+        :responseTo => response_to,
+        :opCode => OP_REPLY
+      }
+    }
+  end
+
+  def self.write_reply(doc)
+    raise 'no responseTo' unless doc[:header][:responseTo]
+    raise 'no documents' unless doc[:documents]
+    responseFlags = (doc[:responseFlags] || 0)
+    cursorId = (doc[:cursorID] || 0)
+    startingFrom = (doc[:startingFrom] || 0)
+    numberReturned = doc[:numberReturned]
+
+    msg = [responseFlags, cursorId, startingFrom, numberReturned].pack('VQ<VV')
+    msg << write_bson(doc[:documents])
+
+    return msg
+  end
+
+  # OP_UPDATE: 2001
+  # A MongoDB update query message.
+  #
+  # header :header - Message header.
+  # int32 - An empty value.
+  # string :database.:collection - Database and collection name for update.
+  # int32 :flags - Bit vector of update flags.
+  # document :selector - BSON document representing update target.
+  # document :update - BSON document representing the update to perform.
+  def self.receive_update(chunk)
+    x = {}
+    _, full, x[:flags] = chunk.unpack('VZ*V')
+    x[:database], x[:collection] = parse_full_collection(full)
+# TODO break out flags
+    docs, _ = receive_bson(chunk, full.bytesize + 9, 2)
+    x[:selector] = docs[0]
+    x[:update] = docs[1]
+
+    return x
+  end
+
+  def self.build_update(request_id, database_name, collection_name, selector, update, flags = [])
+    flag = 0
+    flag = (flag | FLAG_UPDATE_UPSERT) if flags.include?(:upsert)
+    flag = (flag | FLAG_UPDATE_MULTIUPDATE) if flags.include?(:multi)
+
+    return {
+      :header => {
+        :opCode => OP_UPDATE,
+        :requestID => request_id,
+        :responseTo => 0
+      },
+      :database => database_name,
+      :collection => collection_name,
+      :selector => selector,
+      :update => update,
+      :flags => flag
+    }
+  end
+
+  def self.write_update doc
+    raise 'missing collection info' unless doc[:database] && doc[:collection]
+    raise 'missing selector' unless doc[:selector]
+    raise 'missing update' unless doc[:update]
+    flags = (doc[:flags] or 0)
+
+    msg = [0, build_full_collection(doc[:database], doc[:collection]), flags].pack('VZ*V')
+    msg << write_bson(doc[:selector])
+    msg << write_bson(doc[:update])
+
+    return msg
+  end
+
+  # OP_INSERT: 2002
+  #
+  # header :header - Message header.
+  # int32 :flags - Bit vector flags.
+  # string :database.:collection - Database + collection name.
+  # document[] :documents - An array of BSON documents.
+  def self.receive_insert(chunk)
+    x = {}
+    x[:flags], full = chunk.unpack('VZ*')
+    x[:database], x[:collection] = parse_full_collection(full)
+    x[:documents], _ = receive_bson(chunk, full.bytesize + 5)
+
+    return x
+  end
+
+  def self.build_insert(request_id, database_name, collection_name, documents, flags = 0)
+    documents = [documents] if documents.is_a?(Hash)
+    return {
+      :flags => flags,
+      :database => database_name,
+      :collection => collection_name,
+      :documents => documents,
+      :header => {
+        :requestID => request_id,
+        :responseTo => 0,
+        :opCode => OP_INSERT
+      }
+    }
+  end
+
+  def self.write_insert(doc)
+    raise 'missing full collection' unless doc[:database] && doc[:collection]
+    raise 'missing documents' unless doc[:documents]
+    flags = (doc[:flags] or 0)
+    docs = doc[:documents]
+    docs = [docs] if docs.is_a? Hash
+
+    msg = [flags, build_full_collection(doc[:database], doc[:collection])].pack('VZ*')
+    msg << write_bson(docs)
+
+    return msg
+  end
+
+  # OP_QUERY: 2004
+  #
+  # header :header - Message header.
+  # int32 :flags - A bit vector of query flags.
+  # string :database.:collection - Database + collection name.
+  # int32 :numberToSkip - Offset for results.
+  # int32 :numberToReturn - Limit for results.
+  # document :query - BSON document of query.
+  # document :returnFieldsSelector - Optional BSON document to select fields in response.
+  def self.receive_query(chunk)
+    x = {}
+    x[:flags], full, x[:numberToSkip], x[:numberToReturn] = chunk.unpack('VZ*VV')
+    start = 3 * 4 + full.bytesize + 1
+    x[:database], x[:collection] = parse_full_collection(full)
+    docs, start = receive_bson(chunk, start, 2)
+    x[:query] = docs[0]
+    x[:returnFieldSelector] = (docs.size > 1 ? docs[1] : nil)
+
+    return x
+  end
+
+  def self.build_query(request_id, database_name, collection_name,
+      query = {}, fields = nil, num_to_return = 4294967295, number_to_skip = 0, flags = 0)
+    {
+      :header => {
+        :opCode => OP_QUERY,
+        :requestID => request_id,
+        :responseTo => 0
+      },
+      :database => database_name,
+      :collection => collection_name,
+      :query => query,
+      :returnFieldSelector => fields,
+      :numberToReturn => num_to_return,
+      :flags => flags,
+      :numberToSkip => number_to_skip
+    }
+  end
+
+  def self.write_query(doc)
+    raise 'missing full collection name' unless doc[:database] && doc[:collection]
+    flags = (doc[:flags] or 0)
+    numberToSkip = (doc[:numberToSkip] or 0)
+    numberToReturn = (doc[:numberToReturn] or 4294967295)
+    query = (doc[:query] or {})
+    returnFieldSelector = (doc[:returnFieldSelector] or nil)
+
+    msg = [flags, build_full_collection(doc[:database], doc[:collection]), numberToSkip, numberToReturn].pack('VZ*VV')
+    msg << write_bson([query])
+    msg << write_bson([returnFieldSelector]) if returnFieldSelector
+
+    return msg
+  end
+
+  # OP_QUERY: 2005
+  #
+  # header :header - Message header.
+  # int32 - Empty.
+  # string :database.:collection - Database + collection name.
+  # int32 :numberToReturn - Limit for next results reply.
+  # int64 :cursorID - ID of cursor to consume more from.
+  def self.receive_get_more(chunk)
+    x = {}
+    _, full, x[:numberToReturn], x[:cursorID] = chunk.unpack('VZ*VQ<')
+    x[:database], x[:collection] = parse_full_collection(full)
+
+    return x
+  end
+
+  def self.build_get_more(request_id, response_to, database_name, collection_name, cursor_id, number_to_return = 0)
+    {
+      :header => {
+        :opCode => OP_GET_MORE,
+        :requestID => request_id,
+        :responseTo => response_to
+      },
+      :database => database_name,
+      :collection => collection_name,
+      :cursorID => cursor_id,
+      :numberToReturn => number_to_return
+    }
+  end
+
+  def self.write_get_more(doc)
+    raise 'missing full collection name' unless doc[:database] && doc[:collection]
+    raise 'missing cursorID' unless doc[:cursorID]
+
+    numberToReturn = (doc[:numberToReturn] or 0)
+    msg = [0, build_full_collection(doc[:database], doc[:collection]), numberToReturn, doc[:cursorID]].pack('VZ*VQ<')
+
+    return msg
+  end
+
+  # OP_DELETE: 2006
+  #
+  # header :header - Message header.
+  # int32 - Empty.
+  # string :database.:collection - Database + collection name.
+  # int32 :flags - Bit vector of delete-related flags.
+  # document selector - Selector for deletion
+  def self.receive_delete(chunk)
+    x = {}
+    _, full, x[:flags] = chunk.unpack('VZ*V')
+    docs, start = receive_bson(chunk, full.bytesize + 9, 1)
+    x[:database], x[:collection] = parse_full_collection(full)
+    x[:selector], _ = docs[0]
+
+    return x
+  end
+
+  def self.build_delete(request_id, database_name, collection_name, selector, opt = [])
+    flags = 0
+    flags = FLAGS_DELETE_MULTI if opt.include?(:multi)
+
+    {
+      :header => {
+        :opCode => OP_DELETE,
+        :requestID => request_id,
+        :responseTo => 0
+      },
+      :database => database_name,
+      :collection => collection_name,
+      :selector => selector,
+      :flags => flags
+    }
+  end
+
+  def self.write_delete(doc)
+    raise 'missing full collection name' unless doc[:database] && doc[:collection]
+    raise 'missing selector' unless doc[:selector]
+    flags = (doc[:flags] or 0)
+
+    msg = [0, build_full_collection(doc[:database], doc[:collection]), flags].pack('VZ*V')
+    msg << write_bson(doc[:selector])
+
+    return msg
+  end
+
+  # OP_KILL_CURSORS: 2007
+  # Message to explicitly delete a cursor. Only sent from the client in very
+  # specific circumstances, cursors can also time out.
+  #
+  # header :header - Message header.
+  # int32 - Empty.
+  # int32 :numberOfCursorIDs - Number of cursors to kill.
+  # int64[] :cursorIDs - Array of cursor ids to kill.
+  def self.receive_kill_cursors(chunk)
+    x = {}
+    _, n = chunk.unpack('VV')
+    x[:cursorIDs] = chunk[8..-1].unpack("Q<#{n}")
+    return x
+  end
+
+  def self.build_kill_cursors(request_id, response_to, cursor_ids)
+    return {
+      :cursorIDs => cursor_ids,
+      :header => {
+        :opCode => OP_KILL_CURSORS,
+        :requestID => request_id,
+        :responseTo => response_to
+      }
+    }
+  end
+
+  def self.write_kill_cursors(doc)
+    raise 'missing cursorIDs' unless doc[:cursorIDs]
+    return ([0, doc[:cursorIDs].size] + doc[:cursorIDs]).pack("VVQ<*")
+  end
+
+  @@hash_getmore_history = {}
+
+  def self.hash doc
+    if doc[:header][:requestID]
+      temp_req_id = doc[:header][:requestID]
+      doc[:header][:requestID] = 1234
+    end
+    if doc[:header][:responseTo]
+      temp_response_to = doc[:header][:responseTo]
+      doc[:header][:responseTo] = 4321
+    end
+
+    if doc[:header][:opCode] == OP_GET_MORE
+      key = build_full_collection(doc[:database], doc[:collection]) + doc[:cursorID].to_s
+      @@hash_getmore_history[key] ||= 0
+      doc[:header][:requestID] = @@hash_getmore_history[key]
+      temp_cursor_id = doc[:cursorID]
+      doc[:cursorID] = 0
+      @@hash_getmore_history[key] += 1
+    end
+
+    x = Digest::SHA1.hexdigest(write doc)
+    
+    doc[:header][:requestID] = temp_req_id if temp_req_id
+    doc[:header][:responseTo] = temp_response_to if temp_response_to
+    doc[:cursorID] = temp_cursor_id if temp_cursor_id
+
+    return x
+  end
+end
+