func TestCommandOp(t *testing.T) {
generator := newRecordedOpGenerator()
op := CommandOp{}
op.Database = "foo"
op.CommandName = "query"
metadata := bson.D{{"metadata", 1}}
op.Metadata = metadata
change := bson.D{{"updated", true}}
commandArgs := bson.D{{"$set", change}}
op.CommandArgs = commandArgs
inputDocs := []interface{}{}
for i := 0; i < 5; i++ {
inputDocs = append(inputDocs, &bson.D{{"inputDoc", 1}})
}
op.InputDocs = inputDocs
t.Logf("Generated CommandOp: %#v\n", op.CommandOp)
result, err := generator.fetchRecordedOpsFromConn(&op.CommandOp)
if err != nil {
t.Error(err)
}
receivedOp, err := result.RawOp.Parse()
if err != nil {
t.Error(err)
}
commandOp := receivedOp.(*CommandOp)
metadataAsBytes, _ := bson.Marshal(metadata)
metadataRaw := &bson.Raw{}
bson.Unmarshal(metadataAsBytes, metadataRaw)
commandArgsAsBytes, _ := bson.Marshal(commandArgs)
commandArgsRaw := &bson.Raw{}
bson.Unmarshal(commandArgsAsBytes, commandArgsRaw)
t.Log("Comparing parsed Command to original Command")
switch {
case commandOp.Database != op.Database:
t.Errorf("Databases not equal. Saw %v -- Expected %v\n", commandOp.Database, op.Database)
case commandOp.CommandName != op.CommandName:
t.Errorf("CommandNames not equal. Saw %v -- Expected %v\n", commandOp.CommandName, op.CommandName)
case !reflect.DeepEqual(commandOp.Metadata, metadataRaw):
t.Errorf("Metadata not equal. Saw %v -- Expected %v\n", commandOp.Metadata, metadataRaw)
case !reflect.DeepEqual(commandOp.CommandArgs, commandArgsRaw):
t.Errorf("CommandArgs not equal. Saw %v -- Expected %v\n", commandOp.CommandArgs, commandArgsRaw)
}
for i, doc := range commandOp.InputDocs {
marshaledAsBytes, _ := bson.Marshal(inputDocs[i])
unmarshaled := &bson.Raw{}
bson.Unmarshal(marshaledAsBytes, unmarshaled)
if !reflect.DeepEqual(unmarshaled, doc) {
t.Errorf("Document from InputDocs not matched. Saw %v -- Expected %v\n", unmarshaled, doc)
}
}
}
// Execute performs the CommandOp on a given session, yielding the reply when
// successful (and an error otherwise).
func (op *CommandOp) Execute(session *mgo.Session) (Replyable, error) {
session.SetSocketTimeout(0)
before := time.Now()
metadata, commandReply, replyData, resultReply, err := mgo.ExecOpWithReply(session, &op.CommandOp)
after := time.Now()
if err != nil {
return nil, err
}
mgoCommandReplyOp, ok := resultReply.(*mgo.CommandReplyOp)
if !ok {
panic("reply from execution was not the correct type")
}
commandReplyOp := &CommandReplyOp{
CommandReplyOp: *mgoCommandReplyOp,
}
commandReplyOp.Metadata = &bson.Raw{}
err = bson.Unmarshal(metadata, commandReplyOp.Metadata)
if err != nil {
return nil, err
}
commandReplyAsRaw := &bson.Raw{}
err = bson.Unmarshal(commandReply, commandReplyAsRaw)
if err != nil {
return nil, err
}
commandReplyOp.CommandReply = commandReplyAsRaw
doc := &struct {
Cursor struct {
FirstBatch []bson.Raw `bson:"firstBatch"`
NextBatch []bson.Raw `bson:"nextBatch"`
} `bson:"cursor"`
}{}
err = commandReplyAsRaw.Unmarshal(&doc)
if err != nil {
return nil, err
}
if doc.Cursor.FirstBatch != nil {
commandReplyOp.Docs = doc.Cursor.FirstBatch
} else if doc.Cursor.NextBatch != nil {
commandReplyOp.Docs = doc.Cursor.NextBatch
}
for _, d := range replyData {
dataDoc := &bson.Raw{}
err = bson.Unmarshal(d, &dataDoc)
if err != nil {
return nil, err
}
commandReplyOp.OutputDocs = append(commandReplyOp.OutputDocs, dataDoc)
}
commandReplyOp.Latency = after.Sub(before)
return commandReplyOp, nil
}
// FromReader extracts data from a serialized OpCommand into its concrete
// structure.
func (op *CommandOp) FromReader(r io.Reader) error {
database, err := readCStringFromReader(r)
if err != nil {
return err
}
op.Database = string(database)
commandName, err := readCStringFromReader(r)
if err != nil {
return err
}
op.CommandName = string(commandName)
commandArgsAsSlice, err := ReadDocument(r)
if err != nil {
return err
}
op.CommandArgs = &bson.Raw{}
err = bson.Unmarshal(commandArgsAsSlice, op.CommandArgs)
if err != nil {
return err
}
metadataAsSlice, err := ReadDocument(r)
if err != nil {
return err
}
op.Metadata = &bson.Raw{}
err = bson.Unmarshal(metadataAsSlice, op.Metadata)
if err != nil {
return err
}
lengthRead := len(database) + 1 + len(commandName) + 1 + len(commandArgsAsSlice) + len(metadataAsSlice)
op.InputDocs = make([]interface{}, 0)
docLen := 0
for lengthRead+docLen < int(op.Header.MessageLength)-MsgHeaderLen {
docAsSlice, err := ReadDocument(r)
doc := &bson.Raw{}
err = bson.Unmarshal(docAsSlice, doc)
if err != nil {
return err
}
docLen += len(docAsSlice)
op.InputDocs = append(op.InputDocs, doc)
}
return nil
}
// FromReader extracts data from a serialized ReplyOp into its concrete structure.
func (op *ReplyOp) FromReader(r io.Reader) error {
var b [20]byte
if _, err := io.ReadFull(r, b[:]); err != nil {
return err
}
op.Flags = uint32(getInt32(b[:], 0))
op.CursorId = getInt64(b[:], 4)
op.FirstDoc = getInt32(b[:], 12)
op.ReplyDocs = getInt32(b[:], 16)
op.Docs = []bson.Raw{}
// read as many docs as we can from the reader
for {
docBytes, err := ReadDocument(r)
if err != nil {
if err != io.EOF {
// Broken BSON in reply data. TODO log something here?
return err
}
break
}
if len(docBytes) == 0 {
break
}
nextDoc := bson.Raw{}
err = bson.Unmarshal(docBytes, &nextDoc)
if err != nil {
// Unmarshaling []byte to bson.Raw should never ever fail.
panic("failed to unmarshal []byte to Raw")
}
op.Docs = append(op.Docs, nextDoc)
}
return nil
}
// FromReader extracts data from a serialized DeleteOp into its concrete
// structure.
func (op *DeleteOp) FromReader(r io.Reader) error {
var b [4]byte
_, err := io.ReadFull(r, b[:]) //skip ZERO
if err != nil {
return err
}
name, err := readCStringFromReader(r)
if err != nil {
return err
}
op.Collection = string(name)
_, err = io.ReadFull(r, b[:]) //Grab the flags
if err != nil {
return err
}
op.Flags = uint32(getInt32(b[:], 0))
selectorAsSlice, err := ReadDocument(r)
if err != nil {
return err
}
op.Selector = &bson.D{}
err = bson.Unmarshal(selectorAsSlice, op.Selector)
if err != nil {
return err
}
return nil
}
// Execute performs the GetMoreOp on a given session, yielding the reply when
// successful (and an error otherwise).
func (op *GetMoreOp) Execute(session *mgo.Session) (Replyable, error) {
session.SetSocketTimeout(0)
before := time.Now()
_, _, data, resultReply, err := mgo.ExecOpWithReply(session, &op.GetMoreOp)
after := time.Now()
mgoReply, ok := resultReply.(*mgo.ReplyOp)
if !ok {
panic("reply from execution was not the correct type")
}
reply := &ReplyOp{
ReplyOp: *mgoReply,
Docs: make([]bson.Raw, 0, len(data)),
}
for _, d := range data {
dataDoc := bson.Raw{}
err = bson.Unmarshal(d, &dataDoc)
if err != nil {
return nil, err
}
reply.Docs = append(reply.Docs, dataDoc)
}
reply.Latency = after.Sub(before)
return reply, nil
}
func TestOpCommandReplyGetCursorID(t *testing.T) {
testCursorID := int64(123)
doc := &struct {
Cursor struct {
ID int64 `bson:"id"`
} `bson:"cursor"`
}{}
doc.Cursor.ID = testCursorID
asByte, err := bson.Marshal(doc)
if err != nil {
t.Errorf("could not marshal bson: %v", err)
}
asRaw := &bson.Raw{}
bson.Unmarshal(asByte, asRaw)
commandReplyOp := &CommandReplyOp{}
commandReplyOp.CommandReply = asRaw
cursorID, err := commandReplyOp.getCursorID()
if err != nil {
t.Errorf("error fetching cursor %v", err)
}
if cursorID != testCursorID {
t.Errorf("cursorID did not match expected. Found: %v --- Expected: %v", cursorID, testCursorID)
}
t.Log("Ensuring cursorID consistent between multiple calls")
cursorID, err = commandReplyOp.getCursorID()
if err != nil {
t.Errorf("error fetching cursor %v", err)
}
if cursorID != testCursorID {
t.Errorf("cursorID did not match expected. Found: %v --- Expected: %v", cursorID, testCursorID)
}
}
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
}
// FromReader extracts data from a serialized InsertOp into its concrete
// structure.
func (op *InsertOp) FromReader(r io.Reader) error {
var b [4]byte
_, err := io.ReadFull(r, b[:])
if err != nil {
return err
}
op.Flags = uint32(getInt32(b[:], 0))
name, err := readCStringFromReader(r)
if err != nil {
return err
}
op.Collection = string(name)
op.Documents = make([]interface{}, 0)
docLen := 0
for len(name)+1+4+docLen < int(op.Header.MessageLength)-MsgHeaderLen {
docAsSlice, err := ReadDocument(r)
doc := &bson.D{}
err = bson.Unmarshal(docAsSlice, doc)
if err != nil {
return err
}
docLen += len(docAsSlice)
op.Documents = append(op.Documents, doc)
}
return nil
}
func (socket *MongoSocket) loginRun(db string, query, result interface{}, f func() error) error {
var mutex sync.Mutex
var replyErr error
mutex.Lock()
op := QueryOp{}
op.Query = query
op.Collection = db + ".$cmd"
op.Limit = -1
op.replyFunc = func(err error, rfl *replyFuncLegacyArgs, rfc *replyFuncCommandArgs) {
defer mutex.Unlock()
if err != nil {
replyErr = err
return
}
err = bson.Unmarshal(rfl.docData, result)
if err != nil {
replyErr = err
} else {
// Must handle this within the read loop for the socket, so
// that concurrent login requests are properly ordered.
replyErr = f()
}
}
err := socket.Query(&op)
if err != nil {
return err
}
mutex.Lock() // Wait.
return replyErr
}
// GetMeta unmarshals the optional "metadata" field associated with the
// file into the result parameter. The meaning of keys under that field
// is user-defined. For example:
//
// result := struct{ INode int }{}
// err = file.GetMeta(&result)
// if err != nil {
// panic(err.String())
// }
// fmt.Printf("inode: %d\n", result.INode)
//
func (file *GridFile) GetMeta(result interface{}) (err error) {
file.m.Lock()
if file.doc.Metadata != nil {
err = bson.Unmarshal(file.doc.Metadata.Data, result)
}
file.m.Unlock()
return
}
// FromReader extracts data from a serialized QueryOp into its concrete
// structure.
func (op *QueryOp) FromReader(r io.Reader) error {
var b [8]byte
if _, err := io.ReadFull(r, b[:4]); err != nil {
return err
}
op.Flags = mgo.QueryOpFlags(getInt32(b[:], 0))
name, err := readCStringFromReader(r)
if err != nil {
return err
}
op.Collection = string(name)
if _, err := io.ReadFull(r, b[:]); err != nil {
return err
}
op.Skip = getInt32(b[:], 0)
op.Limit = getInt32(b[:], 4)
queryAsSlice, err := ReadDocument(r)
if err != nil {
return err
}
op.Query = &bson.Raw{}
err = bson.Unmarshal(queryAsSlice, op.Query)
if err != nil {
return err
}
currentRead := len(queryAsSlice) + len(op.Collection) + 1 + 12 + MsgHeaderLen
if int(op.Header.MessageLength) > currentRead {
selectorAsSlice, err := ReadDocument(r)
if err != nil {
return err
}
op.Selector = &bson.D{}
err = bson.Unmarshal(selectorAsSlice, op.Selector)
if err != nil {
return err
}
}
return nil
}
// FromReader extracts data from a serialized CommandReplyOp into its
// concrete structure.
func (op *CommandReplyOp) FromReader(r io.Reader) error {
commandReplyAsSlice, err := ReadDocument(r)
if err != nil {
return err
}
op.CommandReply = &bson.Raw{}
err = bson.Unmarshal(commandReplyAsSlice, op.CommandReply)
if err != nil {
return err
}
metadataAsSlice, err := ReadDocument(r)
if err != nil {
return err
}
op.Metadata = &bson.Raw{}
err = bson.Unmarshal(metadataAsSlice, op.Metadata)
if err != nil {
return err
}
op.OutputDocs = make([]interface{}, 0)
for {
docAsSlice, err := ReadDocument(r)
if err != nil {
if err != io.EOF {
// Broken BSON in reply data. TODO log something here?
return err
}
break
}
if len(docAsSlice) == 0 {
break
}
doc := &bson.Raw{}
err = bson.Unmarshal(docAsSlice, doc)
if err != nil {
return err
}
op.OutputDocs = append(op.OutputDocs, doc)
}
return nil
}
func objToDoc(obj interface{}) (d bson.D, err error) {
data, err := bson.Marshal(obj)
if err != nil {
return nil, err
}
err = bson.Unmarshal(data, &d)
if err != nil {
return nil, err
}
return d, err
}
func TestLegacyOpReplyGetCursorID(t *testing.T) {
testCursorID := int64(123)
doc := &struct {
Cursor struct {
ID int64 `bson:"id"`
} `bson:"cursor"`
}{}
doc.Cursor.ID = testCursorID
asByte, err := bson.Marshal(doc)
if err != nil {
t.Errorf("could not marshal bson: %v", err)
}
asRaw := bson.Raw{}
bson.Unmarshal(asByte, &asRaw)
reply := &ReplyOp{}
reply.Docs = []bson.Raw{asRaw}
t.Log("Retrieving cursorID from reply docs")
cursorID, err := reply.getCursorID()
if err != nil {
t.Errorf("error fetching cursor %v", err)
}
if cursorID != testCursorID {
t.Errorf("cursorID did not match expected. Found: %v --- Expected: %v", cursorID, testCursorID)
}
t.Log("Ensuring cursorID consistent between multiple calls")
cursorID, err = reply.getCursorID()
if err != nil {
t.Errorf("error fetching cursor %v", err)
}
if cursorID != testCursorID {
t.Errorf("cursorID did not match expected. Found: %v --- Expected: %v", cursorID, testCursorID)
}
reply2 := &ReplyOp{}
reply2.CursorId = testCursorID
t.Log("Retrieving cursorID from reply field")
cursorID, err = reply.getCursorID()
if err != nil {
t.Errorf("error fetching cursor %v", err)
}
if cursorID != testCursorID {
t.Errorf("cursorID did not match expected. Found: %v --- Expected: %v", cursorID, testCursorID)
}
}
// NextRecordedOp iterates through the PlaybackFileReader to yield the next
// RecordedOp. It returns io.EOF when successfully complete.
func (file *PlaybackFileReader) NextRecordedOp() (*RecordedOp, error) {
buf, err := ReadDocument(file)
if err != nil {
if err != io.EOF {
err = fmt.Errorf("ReadDocument Error: %v", err)
}
return nil, err
}
doc := new(RecordedOp)
err = bson.Unmarshal(buf, doc)
if err != nil {
return nil, fmt.Errorf("Unmarshal RecordedOp Error: %v\n", err)
}
return doc, nil
}
func TestPreciseTimeMarshal(t *testing.T) {
t1 := time.Date(2015, 4, 8, 15, 16, 23, 651387237, time.UTC)
preciseTime := &PreciseTime{t1}
asBson, err := bson.Marshal(preciseTime)
if err != nil {
t.Error(err)
}
result := &PreciseTime{}
err = bson.Unmarshal(asBson, result)
if err != nil {
t.Error(err)
}
if t1 != result.Time {
t.Errorf("Times not equal. Input: %v -- Result: %v", t1, result.Time)
}
}
func TestInsertOp(t *testing.T) {
generator := newRecordedOpGenerator()
op := InsertOp{}
op.Collection = "mongoreplay_test.test"
op.Flags = 7
documents := []interface{}(nil)
for i := 0; i < 10; i++ {
insertDoc := &testDoc{
DocumentNumber: i,
Success: true,
}
documents = append(documents, insertDoc)
}
op.Documents = documents
t.Logf("Generated Insert: %#v\n", op.InsertOp)
result, err := generator.fetchRecordedOpsFromConn(&op.InsertOp)
if err != nil {
t.Error(err)
}
receivedOp, err := result.RawOp.Parse()
if err != nil {
t.Error(err)
}
insertOp := receivedOp.(*InsertOp)
t.Log("Comparing parsed Insert to original Insert")
switch {
case insertOp.Collection != "mongoreplay_test.test":
t.Errorf("Collection not matched. Saw %v -- Expected %v\n", insertOp.Collection, "mongoreplay_test.test")
case insertOp.Flags != 7:
t.Errorf("Flags not matched. Saw %v -- Expected %v\n", insertOp.Flags, 7)
}
for i, doc := range insertOp.Documents {
marshaled, _ := bson.Marshal(documents[i])
unmarshaled := &bson.D{}
bson.Unmarshal(marshaled, unmarshaled)
if !reflect.DeepEqual(unmarshaled, doc) {
t.Errorf("Document not matched. Saw %v -- Expected %v\n", unmarshaled, doc)
}
}
}
func (socket *MongoSocket) resetNonce() {
debugf("Socket %p to %s: requesting a new nonce", socket, socket.addr)
op := &QueryOp{}
op.Query = &getNonceCmd{GetNonce: 1}
op.Collection = "admin.$cmd"
op.Limit = -1
op.replyFunc = func(err error, rfl *replyFuncLegacyArgs, rfc *replyFuncCommandArgs) {
if err != nil {
socket.kill(errors.New("getNonce: "+err.Error()), true)
return
}
result := &getNonceResult{}
err = bson.Unmarshal(rfl.docData, &result)
if err != nil {
socket.kill(errors.New("Failed to unmarshal nonce: "+err.Error()), true)
return
}
debugf("Socket %p to %s: nonce unmarshalled: %#v", socket, socket.addr, result)
if result.Code == 13390 {
// mongos doesn't yet support auth (see http://j.mp/mongos-auth)
result.Nonce = "mongos"
} else if result.Nonce == "" {
var msg string
if result.Err != "" {
msg = fmt.Sprintf("Got an empty nonce: %s (%d)", result.Err, result.Code)
} else {
msg = "Got an empty nonce"
}
socket.kill(errors.New(msg), true)
return
}
socket.Lock()
if socket.cachedNonce != "" {
socket.Unlock()
panic("resetNonce: nonce already cached")
}
socket.cachedNonce = result.Nonce
socket.gotNonce.Signal()
socket.Unlock()
}
err := socket.Query(op)
if err != nil {
socket.kill(errors.New("resetNonce: "+err.Error()), true)
}
}
func TestShortenCommandReply(t *testing.T) {
generator := newRecordedOpGenerator()
op := CommandReplyOp{}
op.Metadata = &testDoc{
Name: "Metadata",
DocumentNumber: 100000,
Success: true,
}
doc1 := testDoc{
Name: "Op Raw Short Reply Test 1",
DocumentNumber: 1,
Success: true,
}
doc2 := testDoc{
Name: "Op Raw Short Reply Test 2",
DocumentNumber: 2,
Success: true,
}
batch := []interface{}{doc1, doc2}
cursorDocIn := cursorDoc{
batch, 12345678, "test"}
op.CommandReply = findReply{cursorDocIn, 1}
op.OutputDocs = []interface{}{}
result, err := generator.fetchRecordedOpsFromConn(&op.CommandReplyOp)
// reply should be functional and parseable
parsed, err := result.RawOp.Parse()
if err != nil {
t.Errorf("error parsing op: %#v", err)
}
t.Logf("parsed Op: %v", parsed)
fullReply, ok := parsed.(*CommandReplyOp)
if !ok {
t.Errorf("parsed op was wrong type")
}
commandReplyCheckRaw, ok := fullReply.CommandReply.(*bson.Raw)
if !ok {
t.Errorf("comamndReply not bson.Raw")
}
commandReplyCheck := &findReply{
Cursor: cursorDoc{},
}
err = bson.Unmarshal(commandReplyCheckRaw.Data, commandReplyCheck)
if err != nil {
t.Errorf("error unmarshaling commandReply %v", err)
}
// ensure that the reply now has 2 document
if !(len(commandReplyCheck.Cursor.Batch) == 2) {
t.Errorf("parsed reply has wrong number of docs: %d", len(commandReplyCheck.Cursor.Batch))
}
// shorten the reply
result.ShortenReply()
parsed, err = result.RawOp.Parse()
if err != nil {
t.Errorf("error parsing op: %v", err)
}
fullReply, ok = parsed.(*CommandReplyOp)
if !ok {
t.Errorf("parsed op was wrong type")
}
commandReplyRaw, ok := fullReply.CommandReply.(*bson.Raw)
if !ok {
t.Errorf("comamndReply not bson.Raw")
}
commandReplyOut := &findReply{
Cursor: cursorDoc{},
}
err = bson.Unmarshal(commandReplyRaw.Data, commandReplyOut)
if err != nil {
t.Errorf("error unmarshaling commandReply %v", err)
}
// ensure that the reply now has 0 documents
if !(len(commandReplyOut.Cursor.Batch) == 0) {
t.Errorf("parsed reply has wrong number of docs: %d", len(commandReplyOut.Cursor.Batch))
}
}
func TestCommandOpGetMoreCursorsRewriteable(t *testing.T) {
oldCursorID := int64(1234)
newCursorID := int64(5678)
commandGM := &CommandGetMore{
CommandOp: CommandOp{},
}
doc := &bson.D{{"getMore", oldCursorID}}
asByte, err := bson.Marshal(doc)
if err != nil {
t.Errorf("could not marshal bson: %v", err)
}
asRaw := &bson.Raw{}
bson.Unmarshal(asByte, asRaw)
commandGM.CommandOp.CommandArgs = asRaw
t.Log("fetching getmore cursorID")
cursorIDs, err := commandGM.getCursorIDs()
if err != nil {
t.Errorf("error fetching cursorIDs: %v", err)
}
if len(cursorIDs) != 1 {
t.Errorf("differing number of cursorIDs found in commandlgetmore. Expected: %v --- Found: %v", 1, len(cursorIDs))
} else {
if oldCursorID != cursorIDs[0] {
t.Errorf("cursorIDs not matched when retrieved. Expected: %v --- Found: %v", oldCursorID, cursorIDs[0])
}
}
t.Log("setting getmore cursorID")
err = commandGM.setCursorIDs([]int64{newCursorID})
if err != nil {
t.Errorf("error setting cursorIDs: %v", err)
}
t.Log("fetching new getmore cursorID")
cursorIDs, err = commandGM.getCursorIDs()
if err != nil {
t.Errorf("error fetching cursorIDs: %v", err)
}
if len(cursorIDs) != 1 {
t.Errorf("differing number of cursorIDs found in killcursors. Expected: %v --- Found: %v", 1, len(cursorIDs))
} else {
if newCursorID != cursorIDs[0] {
t.Errorf("cursorIDs not matched when retrieved. Expected: %v --- Found: %v", newCursorID, cursorIDs[0])
}
}
commandArgs, ok := commandGM.CommandOp.CommandArgs.(*bson.D)
if !ok {
t.Errorf("commandArgs not a *bson.D")
} else {
for _, bsonDoc := range *commandArgs {
if bsonDoc.Name == "getMore" {
getmoreID, ok := bsonDoc.Value.(int64)
if !ok {
t.Errorf("cursorID in command is not int64")
}
if newCursorID != getmoreID {
t.Errorf("cursorIDs not matched when retrieved. Expected: %v --- Found: %v", newCursorID, getmoreID)
}
break
}
}
}
}
// ShortReplyFromReader reads an op from the given reader. It only holds on
// to header-related information and the first document.
func (op *RawOp) ShortenReply() error {
if op.Header.MessageLength < MsgHeaderLen {
return fmt.Errorf("expected message header to have length: %d bytes but was %d bytes", MsgHeaderLen, op.Header.MessageLength)
}
if op.Header.MessageLength > MaxMessageSize {
return fmt.Errorf("wire message size, %v, was greater then the maximum, %v bytes", op.Header.MessageLength, MaxMessageSize)
}
switch op.Header.OpCode {
case OpCodeReply:
if op.Header.MessageLength <= 20+MsgHeaderLen {
//there are no reply docs
return nil
}
firstDocSize := getInt32(op.Body, 20+MsgHeaderLen)
if 20+MsgHeaderLen+int(firstDocSize) > len(op.Body) || firstDocSize > maxBSONSize {
return fmt.Errorf("the size of the first document is greater then the size of the message")
}
op.Body = op.Body[0:(20 + MsgHeaderLen + firstDocSize)]
case OpCodeCommandReply:
// unmarshal the needed fields for replacing into the buffer
commandReply := &CommandReplyStruct{}
err := bson.Unmarshal(op.Body[MsgHeaderLen:], commandReply)
if err != nil {
return fmt.Errorf("unmarshaling op to shorten: %v", err)
}
switch {
case commandReply.Cursor.FirstBatch.Data != nil:
commandReply.Cursor.FirstBatch.Data, _ = bson.Marshal([0]byte{})
case commandReply.Cursor.NextBatch.Data != nil:
commandReply.Cursor.NextBatch.Data, _ = bson.Marshal([0]byte{})
default:
// it's not a findReply so we don't care about it
return nil
}
out, err := bson.Marshal(commandReply)
if err != nil {
return err
}
// calculate the new sizes for offsets into the new buffer
commandReplySize := getInt32(op.Body, MsgHeaderLen)
newCommandReplySize := getInt32(out, 0)
sizeDiff := commandReplySize - newCommandReplySize
newSize := op.Header.MessageLength - sizeDiff
newBody := make([]byte, newSize)
// copy the new data into a buffer that will replace the old buffer
copy(newBody, op.Body[:MsgHeaderLen])
copy(newBody[MsgHeaderLen:], out)
copy(newBody[MsgHeaderLen+newCommandReplySize:], op.Body[MsgHeaderLen+commandReplySize:])
// update the size of this message in the headers
SetInt32(newBody, 0, newSize)
op.Header.MessageLength = newSize
op.Body = newBody
default:
return fmt.Errorf("unexpected op type : %v", op.Header.OpCode)
}
return nil
}
