Golang bson.Unmarshal函数代码示例

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