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)
		}
	}
}
Exemple #2
0
// 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

}
Exemple #3
0
// 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
}
Exemple #4
0
// 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
}
Exemple #5
0
// 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
}
Exemple #6
0
// 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)
	}
}
Exemple #8
0
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
}
Exemple #9
0
// 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
}
Exemple #10
0
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
}
Exemple #11
0
// 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
}
Exemple #12
0
// 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
}
Exemple #13
0
// 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
}
Exemple #14
0
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)
	}
}
Exemple #16
0
// 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)
		}
	}
}
Exemple #19
0
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
			}
		}
	}
}
Exemple #22
0
// 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
}