vendor/github.com/mongodb/mongo-go-driver/x/network/wiremessage/reply.go

// Copyright (C) MongoDB, Inc. 2017-present.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0

package wiremessage

import (
	"errors"
	"fmt"
	"strings"

	"github.com/mongodb/mongo-go-driver/bson"
	"github.com/mongodb/mongo-go-driver/x/bsonx"
)

// Reply represents the OP_REPLY message of the MongoDB wire protocol.
type Reply struct {
	MsgHeader      Header
	ResponseFlags  ReplyFlag
	CursorID       int64
	StartingFrom   int32
	NumberReturned int32
	Documents      []bson.Raw
}

// MarshalWireMessage implements the Marshaler and WireMessage interfaces.
//
// See AppendWireMessage for a description of the rules this method follows.
func (r Reply) MarshalWireMessage() ([]byte, error) {
	b := make([]byte, 0, r.Len())
	return r.AppendWireMessage(b)
}

// ValidateWireMessage implements the Validator and WireMessage interfaces.
func (r Reply) ValidateWireMessage() error {
	if int(r.MsgHeader.MessageLength) != r.Len() {
		return errors.New("incorrect header: message length is not correct")
	}
	if r.MsgHeader.OpCode != OpReply {
		return errors.New("incorrect header: op code is not OpReply")
	}

	return nil
}

// AppendWireMessage implements the Appender and WireMessage interfaces.
//
// AppendWireMessage will set the MessageLength property of the MsgHeader
// if it is zero. It will also set the OpCode to OpQuery if the OpCode is
// zero. If either of these properties are non-zero and not correct, this
// method will return both the []byte with the wire message appended to it
// and an invalid header error.
func (r Reply) AppendWireMessage(b []byte) ([]byte, error) {
	var err error
	err = r.MsgHeader.SetDefaults(r.Len(), OpReply)

	b = r.MsgHeader.AppendHeader(b)
	b = appendInt32(b, int32(r.ResponseFlags))
	b = appendInt64(b, r.CursorID)
	b = appendInt32(b, r.StartingFrom)
	b = appendInt32(b, r.NumberReturned)
	for _, d := range r.Documents {
		b = append(b, d...)
	}
	return b, err
}

// String implements the fmt.Stringer interface.
func (r Reply) String() string {
	return fmt.Sprintf(
		`OP_REPLY{MsgHeader: %s, ResponseFlags: %s, CursorID: %d, StartingFrom: %d, NumberReturned: %d, Documents: %v}`,
		r.MsgHeader, r.ResponseFlags, r.CursorID, r.StartingFrom, r.NumberReturned, r.Documents,
	)
}

// Len implements the WireMessage interface.
func (r Reply) Len() int {
	// Header + Flags + CursorID + StartingFrom + NumberReturned + Length of Length of Documents
	docsLen := 0
	for _, d := range r.Documents {
		docsLen += len(d)
	}
	return 16 + 4 + 8 + 4 + 4 + docsLen
}

// UnmarshalWireMessage implements the Unmarshaler interface.
func (r *Reply) UnmarshalWireMessage(b []byte) error {
	var err error
	r.MsgHeader, err = ReadHeader(b, 0)
	if err != nil {
		return err
	}
	if r.MsgHeader.MessageLength < 36 {
		return errors.New("invalid OP_REPLY: header length too small")
	}
	if len(b) < int(r.MsgHeader.MessageLength) {
		return errors.New("invalid OP_REPLY: []byte too small")
	}

	r.ResponseFlags = ReplyFlag(readInt32(b, 16))
	r.CursorID = readInt64(b, 20)
	r.StartingFrom = readInt32(b, 28)
	r.NumberReturned = readInt32(b, 32)
	pos := 36
	for pos < len(b) {
		rdr, size, err := readDocument(b, int32(pos))
		if err.Message != "" {
			err.Type = ErrOpReply
			return err
		}
		r.Documents = append(r.Documents, rdr)
		pos += size
	}

	return nil
}

// GetMainLegacyDocument constructs and returns a BSON document for this reply.
func (r *Reply) GetMainLegacyDocument(fullCollectionName string) (bsonx.Doc, error) {
	if r.ResponseFlags&CursorNotFound > 0 {
		fmt.Println("cursor not found err")
		return bsonx.Doc{
			{"ok", bsonx.Int32(0)},
		}, nil
	}
	if r.ResponseFlags&QueryFailure > 0 {
		firstDoc := r.Documents[0]
		return bsonx.Doc{
			{"ok", bsonx.Int32(0)},
			{"errmsg", bsonx.String(firstDoc.Lookup("$err").StringValue())},
			{"code", bsonx.Int32(firstDoc.Lookup("code").Int32())},
		}, nil
	}

	doc := bsonx.Doc{
		{"ok", bsonx.Int32(1)},
	}

	batchStr := "firstBatch"
	if r.StartingFrom != 0 {
		batchStr = "nextBatch"
	}

	batchArr := make([]bsonx.Val, len(r.Documents))
	for i, docRaw := range r.Documents {
		doc, err := bsonx.ReadDoc(docRaw)
		if err != nil {
			return nil, err
		}

		batchArr[i] = bsonx.Document(doc)
	}

	cursorDoc := bsonx.Doc{
		{"id", bsonx.Int64(r.CursorID)},
		{"ns", bsonx.String(fullCollectionName)},
		{batchStr, bsonx.Array(batchArr)},
	}

	doc = doc.Append("cursor", bsonx.Document(cursorDoc))
	return doc, nil
}

// GetMainDocument returns the main BSON document for this reply.
func (r *Reply) GetMainDocument() (bsonx.Doc, error) {
	return bsonx.ReadDoc([]byte(r.Documents[0]))
}

// ReplyFlag represents the flags of an OP_REPLY message.
type ReplyFlag int32

// These constants represent the individual flags of an OP_REPLY message.
const (
	CursorNotFound ReplyFlag = 1 << iota
	QueryFailure
	ShardConfigStale
	AwaitCapable
)

// String implements the fmt.Stringer interface.
func (rf ReplyFlag) String() string {
	strs := make([]string, 0)
	if rf&CursorNotFound == CursorNotFound {
		strs = append(strs, "CursorNotFound")
	}
	if rf&QueryFailure == QueryFailure {
		strs = append(strs, "QueryFailure")
	}
	if rf&ShardConfigStale == ShardConfigStale {
		strs = append(strs, "ShardConfigStale")
	}
	if rf&AwaitCapable == AwaitCapable {
		strs = append(strs, "AwaitCapable")
	}
	str := "["
	str += strings.Join(strs, ", ")
	str += "]"
	return str
}