// Copyright (C) MongoDB, Inc. 2015-present.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
//
// Based on gopkg.io/mgo.v2 by Gustavo Niemeyer.
// See THIRD-PARTY-NOTICES for original license terms.

package mgo

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"github.com/10gen/llmgo/bson"
)

const (
	opInvalid      = 0
	opReply        = 1
	dbMsg          = 1000
	dbUpdate       = 2001
	dbInsert       = 2002
	dbQuery        = 2004
	dbGetMore      = 2005
	dbDelete       = 2006
	dbKillCursors  = 2007
	dbCommand      = 2010
	dbCommandReply = 2011
	dbCompressed   = 2012
	dbMessage      = 2013
)

const (
	MsgFlagChecksumPresent = 1
	MsgFlagMoreToCome      = (1 << 1)
	MsgFlagExhaustAllowed  = (1 << 16)
)

const (
	MsgPayload0 = uint8(0)
	MsgPayload1 = uint8(1)
)

type replyFunc func(err error, rfl *replyFuncLegacyArgs, rfc *replyFuncCommandArgs, rfm *replyFuncMsgArgs)

type MongoSocket struct {
	sync.Mutex
	server        *MongoServer // nil when cached
	Conn          net.Conn
	timeout       time.Duration
	addr          string // For debugging only.
	nextRequestId uint32
	replyFuncs    map[uint32]replyFunc
	references    int
	creds         []Credential
	logout        []Credential
	cachedNonce   string
	gotNonce      sync.Cond
	dead          error
	serverInfo    *mongoServerInfo
}

type QueryOpFlags uint32

const (
	_ QueryOpFlags = 1 << iota
	flagTailable
	flagSlaveOk
	flagLogReplay
	flagNoCursorTimeout
	flagAwaitData
)

type QueryOp struct {
	Collection string
	Query      interface{}
	Skip       int32
	Limit      int32
	Selector   interface{}
	Flags      QueryOpFlags
	replyFunc  replyFunc

	mode       Mode
	Options    QueryWrapper
	HasOptions bool
	ServerTags []bson.D
}

func (op *QueryOp) SetReplyFunc(reply replyFunc) {
	op.replyFunc = reply
}

type QueryWrapper struct {
	Query          interface{} "$query"
	OrderBy        interface{} "$orderby,omitempty"
	Hint           interface{} "$hint,omitempty"
	Explain        bool        "$explain,omitempty"
	Snapshot       bool        "$snapshot,omitempty"
	ReadPreference bson.D      "$readPreference,omitempty"
	MaxScan        int         "$maxScan,omitempty"
	MaxTimeMS      int         "$maxTimeMS,omitempty"
	Comment        string      "$comment,omitempty"
}

func (op *QueryOp) finalQuery(socket *MongoSocket) interface{} {
	if op.Flags&flagSlaveOk != 0 && socket.ServerInfo().Mongos {
		var modeName string
		switch op.mode {
		case Strong:
			modeName = "primary"
		case Monotonic, Eventual:
			modeName = "secondaryPreferred"
		case PrimaryPreferred:
			modeName = "primaryPreferred"
		case Secondary:
			modeName = "secondary"
		case SecondaryPreferred:
			modeName = "secondaryPreferred"
		case Nearest:
			modeName = "nearest"
		default:
			panic(fmt.Sprintf("unsupported read mode: %d", op.mode))
		}
		op.HasOptions = true
		op.Options.ReadPreference = make(bson.D, 0, 2)
		op.Options.ReadPreference = append(op.Options.ReadPreference, bson.DocElem{"mode", modeName})
		if len(op.ServerTags) > 0 {
			op.Options.ReadPreference = append(op.Options.ReadPreference, bson.DocElem{"tags", op.ServerTags})
		}
	}
	if op.HasOptions {
		if op.Query == nil {
			var empty bson.D
			op.Options.Query = empty
		} else {
			op.Options.Query = op.Query
		}
		debugf("final query is %#v\n", &op.Options)
		return &op.Options
	}
	return op.Query
}

type GetMoreOp struct {
	Collection string
	Limit      int32
	CursorId   int64
	replyFunc  replyFunc
}

type ReplyOp struct {
	Flags     uint32
	CursorId  int64
	FirstDoc  int32
	ReplyDocs int32
}

type InsertOp struct {
	Collection string        // "database.collection"
	Documents  []interface{} // One or more documents to insert
	Flags      uint32
}

type UpdateOp struct {
	Collection string      `bson:"-"` // "database.collection"
	Selector   interface{} `bson:"q"`
	Update     interface{} `bson:"u"`
	Flags      uint32      `bson:"-"`
	Multi      bool        `bson:"multi,omitempty"`
	Upsert     bool        `bson:"upsert,omitempty"`
}

type DeleteOp struct {
	Collection string      `bson:"-"` // "database.collection"
	Selector   interface{} `bson:"q"`
	Flags      uint32      `bson:"-"`
	Limit      int         `bson:"limit"`
}

type KillCursorsOp struct {
	CursorIds []int64
}

type CommandOp struct {
	Database    string
	CommandName string
	Metadata    interface{}
	CommandArgs interface{}
	InputDocs   []interface{}
	replyFunc   replyFunc
}

type CommandReplyOp struct {
	Metadata     interface{}
	CommandReply interface{}
	OutputDocs   []interface{}
}

type MsgOp struct {
	Flags     uint32
	Sections  []MsgSection
	Checksum  uint32
	replyFunc replyFunc
}

type MsgSection struct {
	PayloadType uint8
	Data        interface{}
}

// PayloadType1 is a container for the OP_MSG payload data of type 1.
// There is no definition of the type 0 payload because that is simply a
// bson document.
type PayloadType1 struct {
	Size       int32
	Identifier string
	Docs       []interface{}
}

func (p *PayloadType1) CalculateSize() (int32, error) {
	docsLen := 0
	for _, d := range p.Docs {
		docAsSlice, err := bson.Marshal(d)
		if err != nil {
			return 0, err
		}
		docsLen += len(docAsSlice)
	}

	// This math is as follows:
	// (4 bytes for int32 length)  +
	// (length of the identifier string) +
	// (1 for the null byte added to c strings) +
	// (length of the marshalled documents)
	return int32(4 + len(p.Identifier) + 1 + docsLen), nil
}

// replyFuncCommandArgs contains the arguments needed by the replyFunc to complete a CommandReplyOp.
type replyFuncCommandArgs struct {
	// op is the newly generated CommandReplyOp
	op *CommandReplyOp

	// bytesLeft is the number of bytes that remain to be read by the readLoop.
	// This indicates if there is more data or not so that the reply can decide
	// whether or not to release its lock.
	bytesLeft int

	// metadata is a slice containing the unread bson of the CommandReplyOp's metadata field.
	metadata []byte

	// commandReply is a slice containing the unread bson of the CommandReplyOp's commandReply field.
	commandReply []byte

	// outputDoc is a slice of bytes containing the unread bson of a reply document being handed to the
	// replyFunc
	outputDoc []byte
}

// replyFuncLegacyArgs contains the arguments needed by the replyFunc to complete a ReplyOp.
type replyFuncLegacyArgs struct {
	// op is the newly generated ReplyOp
	op *ReplyOp

	//docNum is the number of the current document being handed to the reply func.
	// This indicates how many docs have been read so the replyFunc can determine if
	// it can release its lock.
	docNum int

	// docData is a slice of bytes containing the unread bson of reply document being handed to the
	// replyFunc
	docData []byte
}

// replyFuncMsgArgs contains the arguments needed by the replyFunc to complete a MsgOp
// as a reply.
type replyFuncMsgArgs struct {
	// op is the newly generated MsgOp. It will contain the Flags and Checksum but
	// not the data of the parsed ops.
	op *MsgOp

	// docData is a slice of bytes containing the unread bson of reply document being handed to the
	// replyFunc
	sectionsData []byte
}

func (op *GetMoreOp) SetReplyFunc(reply replyFunc) {
	op.replyFunc = reply
}
func (op *CommandOp) SetReplyFunc(reply replyFunc) {
	op.replyFunc = reply
}
func (op *MsgOp) SetReplyFunc(reply replyFunc) {
	op.replyFunc = reply
}

type OpWithReply interface {
	SetReplyFunc(reply replyFunc)
}

type requestInfo struct {
	bufferPos int
	replyFunc replyFunc
}

func NewSocket(server *MongoServer, conn net.Conn, timeout time.Duration) *MongoSocket {
	socket := &MongoSocket{
		Conn:       conn,
		addr:       server.Addr,
		server:     server,
		replyFuncs: make(map[uint32]replyFunc),
	}
	socket.gotNonce.L = &socket.Mutex
	if err := socket.InitialAcquire(server.Info(), timeout); err != nil {
		panic("newSocket: InitialAcquire returned error: " + err.Error())
	}
	stats.socketsAlive(+1)
	debugf("Socket %p to %s: initialized", socket, socket.addr)
	socket.resetNonce()
	go socket.readLoop()
	return socket
}

func NewDumbSocket(conn net.Conn) *MongoSocket {
	server := &MongoServer{}
	return &MongoSocket{
		server:     server,
		addr:       server.Addr,
		Conn:       conn,
		replyFuncs: make(map[uint32]replyFunc),
	}
}

// Server returns the server that the socket is associated with.
// It returns nil while the socket is cached in its respective server.
func (socket *MongoSocket) Server() *MongoServer {
	socket.Lock()
	server := socket.server
	socket.Unlock()
	return server
}

// ServerInfo returns details for the server at the time the socket
// was initially acquired.
func (socket *MongoSocket) ServerInfo() *mongoServerInfo {
	socket.Lock()
	serverInfo := socket.serverInfo
	socket.Unlock()
	return serverInfo
}

// InitialAcquire obtains the first reference to the socket, either
// right after the connection is made or once a recycled socket is
// being put back in use.
func (socket *MongoSocket) InitialAcquire(serverInfo *mongoServerInfo, timeout time.Duration) error {
	socket.Lock()
	if socket.references > 0 {
		panic("Socket acquired out of cache with references")
	}
	if socket.dead != nil {
		dead := socket.dead
		socket.Unlock()
		return dead
	}
	socket.references++
	socket.serverInfo = serverInfo
	socket.timeout = timeout
	stats.socketsInUse(+1)
	stats.socketRefs(+1)
	socket.Unlock()
	return nil
}

// Acquire obtains an additional reference to the socket.
// The socket will only be recycled when it's released as many
// times as it's been acquired.
func (socket *MongoSocket) Acquire() (info *mongoServerInfo) {
	socket.Lock()
	if socket.references == 0 {
		panic("Socket got non-initial acquire with references == 0")
	}
	// We'll track references to dead sockets as well.
	// Caller is still supposed to release the socket.
	socket.references++
	stats.socketRefs(+1)
	serverInfo := socket.serverInfo
	socket.Unlock()
	return serverInfo
}

// Release decrements a socket reference. The socket will be
// recycled once its released as many times as it's been acquired.
func (socket *MongoSocket) Release() {
	socket.Lock()
	if socket.references == 0 {
		panic("socket.Release() with references == 0")
	}
	socket.references--
	stats.socketRefs(-1)
	if socket.references == 0 {
		stats.socketsInUse(-1)
		server := socket.server
		socket.Unlock()
		socket.LogoutAll()
		// If the socket is dead server is nil.
		if server != nil {
			server.RecycleSocket(socket)
		}
	} else {
		socket.Unlock()
	}
}

// SetTimeout changes the timeout used on socket operations.
func (socket *MongoSocket) SetTimeout(d time.Duration) {
	socket.Lock()
	socket.timeout = d
	socket.Unlock()
}

type deadlineType int

const (
	readDeadline  deadlineType = 1
	writeDeadline deadlineType = 2
)

func (socket *MongoSocket) updateDeadline(which deadlineType) {
	var when time.Time
	if socket.timeout > 0 {
		when = time.Now().Add(socket.timeout)
	}
	whichstr := ""
	switch which {
	case readDeadline | writeDeadline:
		whichstr = "read/write"
		socket.Conn.SetDeadline(when)
	case readDeadline:
		whichstr = "read"
		socket.Conn.SetReadDeadline(when)
	case writeDeadline:
		whichstr = "write"
		socket.Conn.SetWriteDeadline(when)
	default:
		panic("invalid parameter to updateDeadline")
	}
	debugf("Socket %p to %s: updated %s deadline to %s ahead (%s)", socket, socket.addr, whichstr, socket.timeout, when)
}

// Close terminates the socket use.
func (socket *MongoSocket) Close() {
	socket.kill(errors.New("Closed explicitly"), false)
}

func (socket *MongoSocket) kill(err error, abend bool) {
	socket.Lock()
	if socket.dead != nil {
		debugf("Socket %p to %s: killed again: %s (previously: %s)", socket, socket.addr, err.Error(), socket.dead.Error())
		socket.Unlock()
		return
	}
	logf("Socket %p to %s: closing: %s (abend=%v)", socket, socket.addr, err.Error(), abend)
	socket.dead = err
	socket.Conn.Close()
	stats.socketsAlive(-1)
	replyFuncs := socket.replyFuncs
	socket.replyFuncs = make(map[uint32]replyFunc)
	server := socket.server
	socket.server = nil
	socket.gotNonce.Broadcast()
	socket.Unlock()
	for _, replyFunc := range replyFuncs {
		logf("Socket %p to %s: notifying replyFunc of closed socket: %s", socket, socket.addr, err.Error())
		replyFunc(err, nil, nil, nil)
	}
	if abend {
		server.AbendSocket(socket)
	}
}

func (socket *MongoSocket) SimpleQuery(op *QueryOp) (data []byte, replyOp *ReplyOp, err error) {
	var wait, change sync.Mutex
	var replyDone bool
	var replyData []byte
	var replyErr error
	wait.Lock()
	op.replyFunc = func(err error, rfl *replyFuncLegacyArgs, rfc *replyFuncCommandArgs, rfm *replyFuncMsgArgs) {
		change.Lock()
		if !replyDone {
			replyDone = true
			replyErr = err
			if rfl != nil {
				replyOp = rfl.op
				if err == nil {
					replyData = rfl.docData
				}
			}
		}
		change.Unlock()
		wait.Unlock()
	}
	err = socket.Query(op)
	if err != nil {
		return nil, nil, err
	}
	wait.Lock()
	change.Lock()
	data = replyData
	err = replyErr
	change.Unlock()
	return data, replyOp, err
}

func (socket *MongoSocket) Query(ops ...interface{}) (err error) {
	if lops := socket.flushLogout(); len(lops) > 0 {
		ops = append(lops, ops...)
	}

	buf := make([]byte, 0, 256)

	// Serialize operations synchronously to avoid interrupting
	// other goroutines while we can't really be sending data.
	// Also, record id positions so that we can compute request
	// ids at once later with the lock already held.
	requests := make([]requestInfo, len(ops))
	requestCount := 0

	for _, op := range ops {
		debugf("Socket %p to %s: serializing op: %#v", socket, socket.addr, op)
		start := len(buf)
		var replyFunc replyFunc
		switch op := op.(type) {

		case *UpdateOp:
			buf = addHeader(buf, dbUpdate)
			buf = addInt32(buf, 0) // Reserved
			buf = addCString(buf, op.Collection)
			buf = addInt32(buf, int32(op.Flags))
			debugf("Socket %p to %s: serializing selector document: %#v", socket, socket.addr, op.Selector)
			buf, err = addBSON(buf, op.Selector)
			if err != nil {
				return err
			}
			debugf("Socket %p to %s: serializing update document: %#v", socket, socket.addr, op.Update)
			buf, err = addBSON(buf, op.Update)
			if err != nil {
				return err
			}

		case *InsertOp:
			buf = addHeader(buf, dbInsert)
			buf = addInt32(buf, int32(op.Flags))
			buf = addCString(buf, op.Collection)
			for _, doc := range op.Documents {
				debugf("Socket %p to %s: serializing document for insertion: %#v", socket, socket.addr, doc)
				buf, err = addBSON(buf, doc)
				if err != nil {
					return err
				}
			}

		case *QueryOp:
			buf = addHeader(buf, dbQuery)
			buf = addInt32(buf, int32(op.Flags))
			buf = addCString(buf, op.Collection)
			buf = addInt32(buf, op.Skip)
			buf = addInt32(buf, op.Limit)
			buf, err = addBSON(buf, op.finalQuery(socket))
			if err != nil {
				return err
			}
			if op.Selector != nil {
				buf, err = addBSON(buf, op.Selector)
				if err != nil {
					return err
				}
			}
			replyFunc = op.replyFunc

		case *GetMoreOp:
			buf = addHeader(buf, dbGetMore)
			buf = addInt32(buf, 0) // Reserved
			buf = addCString(buf, op.Collection)
			buf = addInt32(buf, op.Limit)
			buf = addInt64(buf, op.CursorId)
			replyFunc = op.replyFunc

		case *ReplyOp:
			buf = addHeader(buf, opReply)
			buf = addInt32(buf, int32(op.Flags))
			buf = addInt64(buf, op.CursorId)
			buf = addInt32(buf, op.FirstDoc)
			buf = addInt32(buf, op.ReplyDocs)

		case *DeleteOp:
			buf = addHeader(buf, dbDelete)
			buf = addInt32(buf, 0) // Reserved
			buf = addCString(buf, op.Collection)
			buf = addInt32(buf, int32(op.Flags))
			debugf("Socket %p to %s: serializing selector document: %#v", socket, socket.addr, op.Selector)
			buf, err = addBSON(buf, op.Selector)
			if err != nil {
				return err
			}

		case *KillCursorsOp:
			buf = addHeader(buf, dbKillCursors)
			buf = addInt32(buf, 0) // Reserved
			buf = addInt32(buf, int32(len(op.CursorIds)))
			for _, cursorId := range op.CursorIds {
				buf = addInt64(buf, cursorId)
			}
		case *CommandOp:
			buf = addHeader(buf, dbCommand)
			buf = addCString(buf, op.Database)
			buf = addCString(buf, op.CommandName)
			buf, err = addBSON(buf, op.CommandArgs)
			if err != nil {
				return err
			}
			buf, err = addBSON(buf, op.Metadata)
			if err != nil {
				return err
			}
			for _, doc := range op.InputDocs {
				debugf("Socket %p to %s: serializing document for opcommand: %#v", socket, socket.addr, doc)
				buf, err = addBSON(buf, doc)
				if err != nil {
					return err
				}
			}
			replyFunc = op.replyFunc
		case *CommandReplyOp:
			buf = addHeader(buf, dbCommandReply)
			buf, err = addBSON(buf, op.CommandReply)
			if err != nil {
				return err
			}
			buf, err = addBSON(buf, op.Metadata)
			if err != nil {
				return err
			}
			for _, doc := range op.OutputDocs {
				debugf("Socket %p to %s: serializing document for opcommand: %#v", socket, socket.addr, doc)
				buf, err = addBSON(buf, doc)
				if err != nil {
					return err
				}
			}
		case *MsgOp:
			buf = addHeader(buf, dbMessage)
			buf = addInt32(buf, int32(op.Flags))
			for _, section := range op.Sections {
				// Read the payload type
				switch section.PayloadType {
				// If it's the simpler type (type 0) write out its type and BSON
				case MsgPayload0:
					buf = append(buf, byte(0))
					buf, err = addBSON(buf, section.Data)
					if err != nil {
						return err
					}
				case MsgPayload1:
					// Write the payload type
					buf = append(buf, byte(1))
					payload, ok := section.Data.(PayloadType1)
					if !ok {
						panic("incorrect type given for payload")
					}

					//Write out the size
					buf = addInt32(buf, payload.Size)
					//Write out the identifier
					buf = addCString(buf, payload.Identifier)
					//Write out the docs
					for _, d := range payload.Docs {
						buf, err = addBSON(buf, d)
						if err != nil {
							return err
						}
					}
				default:
					return fmt.Errorf("unknown payload type: %d", section.PayloadType)
				}
			}
			// check if the checksum should be present
			if op.Flags&MsgFlagChecksumPresent != 0 {
				buf = addInt32(buf, int32(op.Checksum))
			}
			replyFunc = op.replyFunc

		default:
			panic("internal error: unknown operation type")
		}

		setInt32(buf, start, int32(len(buf)-start))

		if replyFunc != nil {
			request := &requests[requestCount]
			request.replyFunc = replyFunc
			request.bufferPos = start
			requestCount++
		}
	}

	// Buffer is ready for the pipe.  Lock, allocate ids, and enqueue.

	socket.Lock()
	if socket.dead != nil {
		dead := socket.dead
		socket.Unlock()
		debugf("Socket %p to %s: failing query, already closed: %s", socket, socket.addr, socket.dead.Error())
		// XXX This seems necessary in case the session is closed concurrently
		// with a query being performed, but it's not yet tested:
		for i := 0; i != requestCount; i++ {
			request := &requests[i]
			if request.replyFunc != nil {
				request.replyFunc(dead, nil, nil, nil)
			}
		}
		return dead
	}

	wasWaiting := len(socket.replyFuncs) > 0

	// Reserve id 0 for requests which should have no responses.
	requestId := socket.nextRequestId + 1
	if requestId == 0 {
		requestId++
	}
	socket.nextRequestId = requestId + uint32(requestCount)
	for i := 0; i != requestCount; i++ {
		request := &requests[i]
		setInt32(buf, request.bufferPos+4, int32(requestId))
		socket.replyFuncs[requestId] = request.replyFunc
		requestId++
	}

	debugf("Socket %p to %s: sending %d op(s) (%d bytes)", socket, socket.addr, len(ops), len(buf))
	stats.sentOps(len(ops))

	socket.updateDeadline(writeDeadline)
	_, err = socket.Conn.Write(buf)

	if !wasWaiting && requestCount > 0 {
		socket.updateDeadline(readDeadline)
	}
	socket.Unlock()
	return err
}

// Estimated minimum cost per socket: 1 goroutine + memory for the largest
// document ever seen.
func (socket *MongoSocket) readLoop() {
	headerBuf := make([]byte, 16) // 16 from header
	bodyBuf := make([]byte, 20)   // 20 from OP_REPLY fixed fields or other general uses
	for {
		var r io.Reader = socket.Conn

		// XXX Handle timeouts, , etc
		_, err := io.ReadFull(r, headerBuf)
		if err != nil {
			socket.kill(err, true)
			return
		}

		totalLen := getInt32(headerBuf, 0)
		responseTo := getInt32(headerBuf, 8)
		opCode := getInt32(headerBuf, 12)

		if opCode == dbCompressed {
			buf := bytes.NewBuffer(headerBuf)
			io.CopyN(buf, r, int64(totalLen-16))
			msg, err := DecompressMessage(buf.Bytes())
			if err != nil {
				socket.kill(err, true)
				return
			}
			r = bytes.NewBuffer(msg)

			_, err = io.ReadFull(r, headerBuf)
			if err != nil {
				socket.kill(err, true)
				return
			}
			opCode = getInt32(headerBuf, 12)
			if opCode == dbCompressed {
				err = fmt.Errorf("cannot recursively decompress messages")
				socket.kill(err, true)
				return
			}
		}

		// Don't use socket.server.Addr here.  socket is not
		// locked and socket.server may go away.
		debugf("Socket %p to %s: got reply (%d bytes)", socket, socket.addr, totalLen)

		socket.Lock()
		replyFunc, ok := socket.replyFuncs[uint32(responseTo)]
		if ok {
			delete(socket.replyFuncs, uint32(responseTo))
		}
		socket.Unlock()

		switch opCode {
		case opReply:
			_, err := io.ReadFull(r, bodyBuf)
			if err != nil {
				socket.kill(err, true)
				return
			}
			reply := ReplyOp{
				Flags:     uint32(getInt32(bodyBuf, 0)),
				CursorId:  getInt64(bodyBuf, 4),
				FirstDoc:  getInt32(bodyBuf, 12),
				ReplyDocs: getInt32(bodyBuf, 16),
			}
			stats.receivedOps(+1)
			stats.receivedDocs(int(reply.ReplyDocs))
			if replyFunc != nil && reply.ReplyDocs == 0 {
				rfl := replyFuncLegacyArgs{
					op:     &reply,
					docNum: -1,
				}
				replyFunc(nil, &rfl, nil, nil)
			} else {
				for i := 0; i != int(reply.ReplyDocs); i++ {
					b, err := readDocument(r)
					if err != nil {
						if replyFunc != nil {
							replyFunc(err, &replyFuncLegacyArgs{docNum: -1}, nil, nil)
						}
						socket.kill(err, true)
						return
					}

					if replyFunc != nil {
						rfl := replyFuncLegacyArgs{
							op:      &reply,
							docNum:  i,
							docData: b,
						}
						replyFunc(nil, &rfl, nil, nil)
					}
					// XXX Do bound checking against totalLen.
				}
			}
		case dbCommandReply:
			commandReplyAsSlice, err := readDocument(r)
			if err != nil {
				socket.kill(err, true)
				return
			}
			metadataAsSlice, err := readDocument(r)
			if err != nil {
				socket.kill(err, true)
				return
			}
			rfc := replyFuncCommandArgs{
				op:           &CommandReplyOp{},
				metadata:     metadataAsSlice,
				commandReply: commandReplyAsSlice,
			}
			lengthRead := len(commandReplyAsSlice) + len(metadataAsSlice)
			if replyFunc != nil && lengthRead+16 >= int(totalLen) {
				replyFunc(nil, nil, &rfc, nil)
			}

			docLen := 0
			for lengthRead+docLen < int(totalLen)-16 {
				documentBuf, err := readDocument(r)
				if err != nil {
					rfc.bytesLeft = 0
					if replyFunc != nil {
						replyFunc(err, nil, &rfc, nil)
					}
					socket.kill(err, true)
					return
				}
				rfc.outputDoc = documentBuf
				if replyFunc != nil {
					replyFunc(nil, nil, &rfc, nil)
				}
				docLen += len(documentBuf)
			}
		case dbMessage:
			rfm := replyFuncMsgArgs{}
			// READ THE FLAGS
			_, err := io.ReadFull(r, bodyBuf[:4])
			if err != nil {
				socket.kill(err, true)
				return
			}

			reply := &MsgOp{
				Flags: uint32(getInt32(bodyBuf, 0)),
			}
			rfm.op = reply
			lengthRead := 4
			checksumLength := 0
			checksumPresent := false
			if (reply.Flags & MsgFlagChecksumPresent) != 0 {
				checksumPresent = true
				checksumLength = 4
			}

			// read the sections into a big buffer.
			// 16 is used here for the length of the message header.
			sectionsLength := int(totalLen) - (checksumLength + lengthRead + 16)
			rfm.sectionsData = make([]byte, sectionsLength)
			_, err = io.ReadFull(r, rfm.sectionsData)
			if err != nil {
				socket.kill(err, true)
				return
			}
			if checksumPresent {
				//read the checksum
				_, err := io.ReadFull(r, bodyBuf[:4])
				if err != nil {
					socket.kill(err, true)
					return
				}
				rfm.op.Checksum = uint32(getInt32(bodyBuf, 0))
			}

			if replyFunc != nil {
				replyFunc(nil, nil, nil, &rfm)
			}

		default:
			socket.kill(errors.New("opcode != 1, 2011, or 2013, corrupted data?"), true)
			return
		}

		socket.Lock()
		if len(socket.replyFuncs) == 0 {
			// Nothing else to read for now. Disable deadline.
			socket.Conn.SetReadDeadline(time.Time{})
		} else {
			socket.updateDeadline(readDeadline)
		}
		socket.Unlock()

		// XXX Do bound checking against totalLen.
	}
}

func readDocument(r io.Reader) (docBuf []byte, err error) {
	sizeBuf := make([]byte, 4)
	_, err = io.ReadFull(r, sizeBuf)
	if err != nil {
		return
	}
	size := getInt32(sizeBuf, 0)
	docBuf = make([]byte, int(size))

	copy(docBuf, sizeBuf)

	_, err = io.ReadFull(r, docBuf[4:])
	if err != nil {
		return
	}
	if globalDebug && globalLogger != nil {
		m := bson.M{}
		if err := bson.Unmarshal(docBuf, m); err == nil {
			if conn, ok := r.(net.Conn); ok {
				debugf("Socket with addr '%s' received document: %#v", conn.RemoteAddr(), m)
			}
		}
	}
	return
}

var emptyHeader = []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}

func addHeader(b []byte, opcode int) []byte {
	i := len(b)
	b = append(b, emptyHeader...)
	// Enough for current opcodes.
	b[i+12] = byte(opcode)
	b[i+13] = byte(opcode >> 8)
	return b
}

func addInt32(b []byte, i int32) []byte {
	return append(b, byte(i), byte(i>>8), byte(i>>16), byte(i>>24))
}

func addInt64(b []byte, i int64) []byte {
	return append(b, byte(i), byte(i>>8), byte(i>>16), byte(i>>24),
		byte(i>>32), byte(i>>40), byte(i>>48), byte(i>>56))
}

func addCString(b []byte, s string) []byte {
	b = append(b, []byte(s)...)
	b = append(b, 0)
	return b
}

func addBSON(b []byte, doc interface{}) ([]byte, error) {
	if doc == nil {
		return append(b, 5, 0, 0, 0, 0), nil
	}
	data, err := bson.Marshal(doc)
	if err != nil {
		return b, err
	}
	return append(b, data...), nil
}

func setInt32(b []byte, pos int, i int32) {
	b[pos] = byte(i)
	b[pos+1] = byte(i >> 8)
	b[pos+2] = byte(i >> 16)
	b[pos+3] = byte(i >> 24)
}

func getInt32(b []byte, pos int) int32 {
	return (int32(b[pos+0])) |
		(int32(b[pos+1]) << 8) |
		(int32(b[pos+2]) << 16) |
		(int32(b[pos+3]) << 24)
}

func getInt64(b []byte, pos int) int64 {
	return (int64(b[pos+0])) |
		(int64(b[pos+1]) << 8) |
		(int64(b[pos+2]) << 16) |
		(int64(b[pos+3]) << 24) |
		(int64(b[pos+4]) << 32) |
		(int64(b[pos+5]) << 40) |
		(int64(b[pos+6]) << 48) |
		(int64(b[pos+7]) << 56)
}
