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)
}
}
}
func TestRepeatGeneration(t *testing.T) {
recOp := &RecordedOp{
Seen: &PreciseTime{time.Now()},
}
bsonBytes, err := bson.Marshal(recOp)
if err != nil {
t.Errorf("couldn't marshal %v", err)
}
playbackReader := &PlaybackFileReader{bytes.NewReader(bsonBytes)}
repeat := 2
opChan, errChan := NewOpChanFromFile(playbackReader, repeat)
op1, ok := <-opChan
if !ok {
t.Errorf("read of 0-generation op failed")
}
if op1.Generation != 0 {
t.Errorf("generation of 0 generation op is %v", op1.Generation)
}
op2, ok := <-opChan
if !ok {
t.Errorf("read of 1-generation op failed")
}
if op2.Generation != 1 {
t.Errorf("generation of 1 generation op is %v", op2.Generation)
}
_, ok = <-opChan
if ok {
t.Errorf("Successfully read past end of op chan")
}
err = <-errChan
if err != io.EOF {
t.Errorf("should have eof at end, but got %v", err)
}
}
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 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
}
// Record writes pcap data into a playback file
func Record(ctx *packetHandlerContext,
playbackWriter *PlaybackWriter,
noShortenReply bool) error {
ch := make(chan error)
go func() {
defer close(ch)
var fail error
for op := range ctx.mongoOpStream.Ops {
// since we don't currently have a way to shutdown packetHandler.Handle()
// continue to read from ctx.mongoOpStream.Ops even after a faltal error
if fail != nil {
toolDebugLogger.Logvf(DebugHigh, "not recording op because of record error %v", fail)
continue
}
if (op.Header.OpCode == OpCodeReply || op.Header.OpCode == OpCodeCommandReply) &&
!noShortenReply {
err := op.ShortenReply()
if err != nil {
userInfoLogger.Logvf(DebugLow, "stream %v problem shortening reply: %v", op.SeenConnectionNum, err)
continue
}
}
bsonBytes, err := bson.Marshal(op)
if err != nil {
userInfoLogger.Logvf(DebugLow, "stream %v error marshaling message: %v", op.SeenConnectionNum, err)
continue
}
_, err = playbackWriter.Write(bsonBytes)
if err != nil {
fail = fmt.Errorf("error writing message: %v", err)
userInfoLogger.Logvf(Always, "%v", err)
continue
}
}
ch <- fail
}()
if err := ctx.packetHandler.Handle(ctx.mongoOpStream, -1); err != nil {
return fmt.Errorf("record: error handling packet stream: %s", err)
}
stats, err := ctx.pcapHandle.Stats()
if err != nil {
toolDebugLogger.Logvf(Always, "Warning: got err %v getting pcap handle stats", err)
} else {
toolDebugLogger.Logvf(Info, "PCAP stats: %#v", stats)
}
err = <-ch
if err == nil && stats != nil && stats.PacketsDropped != 0 {
err = ErrPacketsDropped{stats.PacketsDropped}
}
return err
}
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
}
// SetMeta changes the optional "metadata" field associated with the
// file. The meaning of keys under that field is user-defined.
// For example:
//
// file.SetMeta(bson.M{"inode": inode})
//
// It is a runtime error to call this function when the file is not open
// for writing.
func (file *GridFile) SetMeta(metadata interface{}) {
file.assertMode(gfsWriting)
data, err := bson.Marshal(metadata)
file.m.Lock()
if err != nil && file.err == nil {
file.err = err
} else {
file.doc.Metadata = &bson.Raw{Data: data}
}
file.m.Unlock()
}
func TestPlayOpEOF(t *testing.T) {
ops := []RecordedOp{{
Seen: &PreciseTime{time.Now()},
}, {
Seen: &PreciseTime{time.Now()},
EOF: true,
}}
var buf bytes.Buffer
for _, op := range ops {
bsonBytes, err := bson.Marshal(op)
if err != nil {
t.Errorf("couldn't marshal op %v", err)
}
buf.Write(bsonBytes)
}
playbackReader := &PlaybackFileReader{bytes.NewReader(buf.Bytes())}
repeat := 2
opChan, errChan := NewOpChanFromFile(playbackReader, repeat)
op1, ok := <-opChan
if !ok {
t.Errorf("read of op1 failed")
}
if op1.EOF {
t.Errorf("op1 should not be an EOF op")
}
op2, ok := <-opChan
if !ok {
t.Errorf("read op2 failed")
}
if op2.EOF {
t.Errorf("op2 should not be an EOF op")
}
op3, ok := <-opChan
if !ok {
t.Errorf("read of op3 failed")
}
if !op3.EOF {
t.Errorf("op3 is not an EOF op")
}
_, ok = <-opChan
if ok {
t.Errorf("Successfully read past end of op chan")
}
err := <-errChan
if err != io.EOF {
t.Errorf("should have eof at end, but got %v", 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)
}
}
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)
}
}
}
// Record writes pcap data into a playback file
func Record(ctx *packetHandlerContext,
playbackWriter *PlaybackWriter,
noShortenReply bool) error {
ch := make(chan error)
go func() {
defer close(ch)
for op := range ctx.mongoOpStream.Ops {
if (op.Header.OpCode == OpCodeReply || op.Header.OpCode == OpCodeCommandReply) &&
!noShortenReply {
op.ShortenReply()
}
bsonBytes, err := bson.Marshal(op)
if err != nil {
ch <- fmt.Errorf("error marshaling message: %v", err)
return
}
_, err = playbackWriter.Write(bsonBytes)
if err != nil {
ch <- fmt.Errorf("error writing message: %v", err)
return
}
}
ch <- nil
}()
if err := ctx.packetHandler.Handle(ctx.mongoOpStream, -1); err != nil {
return fmt.Errorf("record: error handling packet stream: %s", err)
}
stats, err := ctx.pcapHandle.Stats()
if err != nil {
toolDebugLogger.Logvf(Always, "Warning: got err %v getting pcap handle stats", err)
} else {
toolDebugLogger.Logvf(Info, "PCAP stats: %#v", stats)
}
err = <-ch
if err == nil && stats != nil && stats.PacketsDropped != 0 {
err = ErrPacketsDropped{stats.PacketsDropped}
}
return err
}
func (file *GridFile) insertChunk(data []byte) {
n := file.chunk
file.chunk++
debugf("GridFile %p: adding to checksum: %q", file, string(data))
file.wsum.Write(data)
for file.doc.ChunkSize*file.wpending >= 1024*1024 {
// Hold on.. we got a MB pending.
file.c.Wait()
if file.err != nil {
return
}
}
file.wpending++
debugf("GridFile %p: inserting chunk %d with %d bytes", file, n, len(data))
// We may not own the memory of data, so rather than
// simply copying it, we'll marshal the document ahead of time.
data, err := bson.Marshal(gfsChunk{bson.NewObjectId(), file.doc.Id, n, data})
if err != nil {
file.err = err
return
}
go func() {
err := file.gfs.Chunks.Insert(bson.Raw{Data: data})
file.m.Lock()
file.wpending--
if err != nil && file.err == nil {
file.err = err
}
file.c.Broadcast()
file.m.Unlock()
}()
}
func TestShortenLegacyReply(t *testing.T) {
generator := newRecordedOpGenerator()
op := ReplyOp{}
op.ReplyDocs = 2
result, err := generator.fetchRecordedOpsFromConn(&op.ReplyOp)
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,
}
asByte1, err := bson.Marshal(doc1)
if err != nil {
t.Errorf("could not marshal bson: %v", err)
}
asByte2, err := bson.Marshal(doc2)
if err != nil {
t.Errorf("could not marshal bson: %v", err)
}
// add the two docs as the docs from the reply
result.RawOp.Body = append(result.RawOp.Body, asByte1...)
result.RawOp.Body = append(result.RawOp.Body, asByte2...)
result.Header.MessageLength = int32(len(result.RawOp.Body))
// reply should be functional and parseable
parsed, err := result.RawOp.Parse()
if err != nil {
t.Errorf("error parsing op: %v", err)
}
fullReply, ok := parsed.(*ReplyOp)
if !ok {
t.Errorf("parsed op was wrong type")
}
if !(len(fullReply.Docs) == 2) {
t.Errorf("parsed reply has wrong number of docs: %d", len(fullReply.Docs))
}
// shorten the reply
result.ShortenReply()
parsed, err = result.RawOp.Parse()
if err != nil {
t.Errorf("error parsing op: %v", err)
}
fullReply, ok = parsed.(*ReplyOp)
if !ok {
t.Errorf("parsed op was wrong type")
}
// ensure that the reply now has only 1 document
if !(len(fullReply.Docs) == 1) {
t.Errorf("parsed reply has wrong number of docs: %d", len(fullReply.Docs))
}
}
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
}