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 }