Golang bson.Marshal函数代码示例

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
			}
		}
	}
}