GODRIVER-2935 Use OP_QUERY in connection handshakes (#1377)

mongo/integration/client_test.go

@@ -731,8 +731,8 @@ func TestClient(t *testing.T) {
 		pair := msgPairs[0]
 		assert.Equal(mt, handshake.LegacyHello, pair.CommandName, "expected command name %s at index 0, got %s",
 			handshake.LegacyHello, pair.CommandName)
-		assert.Equal(mt, wiremessage.OpMsg, pair.Sent.OpCode,
-			"expected 'OP_MSG' OpCode in wire message, got %q", pair.Sent.OpCode.String())
+		assert.Equal(mt, wiremessage.OpQuery, pair.Sent.OpCode,
+			"expected 'OP_QUERY' OpCode in wire message, got %q", pair.Sent.OpCode.String())
 
 		// Look for a saslContinue in the remaining proxied messages and assert that it uses the OP_MSG OpCode, as wire
 		// version is now known to be >= 6.

mongo/integration/mtest/sent_message.go

@@ -37,6 +37,8 @@ type sentMsgParseFn func([]byte) (*SentMessage, error)
 
 func getSentMessageParser(opcode wiremessage.OpCode) (sentMsgParseFn, bool) {
 	switch opcode {
+	case wiremessage.OpQuery:
+		return parseOpQuery, true
 	case wiremessage.OpMsg:
 		return parseSentOpMsg, true
 	case wiremessage.OpCompressed:
@@ -46,6 +48,69 @@ func getSentMessageParser(opcode wiremessage.OpCode) (sentMsgParseFn, bool) {
 	}
 }
 
+func parseOpQuery(wm []byte) (*SentMessage, error) {
+	var ok bool
+
+	if _, wm, ok = wiremessage.ReadQueryFlags(wm); !ok {
+		return nil, errors.New("failed to read query flags")
+	}
+	if _, wm, ok = wiremessage.ReadQueryFullCollectionName(wm); !ok {
+		return nil, errors.New("failed to read full collection name")
+	}
+	if _, wm, ok = wiremessage.ReadQueryNumberToSkip(wm); !ok {
+		return nil, errors.New("failed to read number to skip")
+	}
+	if _, wm, ok = wiremessage.ReadQueryNumberToReturn(wm); !ok {
+		return nil, errors.New("failed to read number to return")
+	}
+
+	query, wm, ok := wiremessage.ReadQueryQuery(wm)
+	if !ok {
+		return nil, errors.New("failed to read query")
+	}
+
+	// If there is no read preference document, the command document is query.
+	// Otherwise, query is in the format {$query: <command document>, $readPreference: <read preference document>}.
+	commandDoc := query
+	var rpDoc bsoncore.Document
+
+	dollarQueryVal, err := query.LookupErr("$query")
+	if err == nil {
+		commandDoc = dollarQueryVal.Document()
+
+		rpVal, err := query.LookupErr("$readPreference")
+		if err != nil {
+			return nil, fmt.Errorf("query %s contains $query but not $readPreference fields", query)
+		}
+		rpDoc = rpVal.Document()
+	}
+
+	// For OP_QUERY, inserts, updates, and deletes are sent as a BSON array of documents inside the main command
+	// document. Pull these sequences out into an ArrayStyle DocumentSequence.
+	var docSequence *bsoncore.DocumentSequence
+	cmdElems, _ := commandDoc.Elements()
+	for _, elem := range cmdElems {
+		switch elem.Key() {
+		case "documents", "updates", "deletes":
+			docSequence = &bsoncore.DocumentSequence{
+				Style: bsoncore.ArrayStyle,
+				Data:  elem.Value().Array(),
+			}
+		}
+		if docSequence != nil {
+			// There can only be one of these arrays in a well-formed command, so we exit the loop once one is found.
+			break
+		}
+	}
+
+	sm := &SentMessage{
+		Command:          commandDoc,
+		ReadPreference:   rpDoc,
+		DocumentSequence: docSequence,
+	}
+	return sm, nil
+}
+
 func parseSentMessage(wm []byte) (*SentMessage, error) {
 	// Re-assign the wire message to "remaining" so "wm" continues to point to the entire message after parsing.
 	_, requestID, _, opcode, remaining, ok := wiremessage.ReadHeader(wm)

x/mongo/driver/auth/speculative_scram_test.go

@@ -93,7 +93,7 @@ func TestSpeculativeSCRAM(t *testing.T) {
 				// Assert that the driver sent hello with the speculative authentication message.
 				assert.Equal(t, len(tc.payloads), len(conn.Written), "expected %d wire messages to be sent, got %d",
 					len(tc.payloads), (conn.Written))
-				helloCmd, err := drivertest.GetCommandFromMsgWireMessage(<-conn.Written)
+				helloCmd, err := drivertest.GetCommandFromQueryWireMessage(<-conn.Written)
 				assert.Nil(t, err, "error parsing hello command: %v", err)
 				assertCommandName(t, helloCmd, handshake.LegacyHello)
 
@@ -177,7 +177,7 @@ func TestSpeculativeSCRAM(t *testing.T) {
 
 				assert.Equal(t, numResponses, len(conn.Written), "expected %d wire messages to be sent, got %d",
 					numResponses, len(conn.Written))
-				hello, err := drivertest.GetCommandFromMsgWireMessage(<-conn.Written)
+				hello, err := drivertest.GetCommandFromQueryWireMessage(<-conn.Written)
 				assert.Nil(t, err, "error parsing hello command: %v", err)
 				assertCommandName(t, hello, handshake.LegacyHello)
 				_, err = hello.LookupErr("speculativeAuthenticate")

x/mongo/driver/auth/speculative_x509_test.go

@@ -58,7 +58,7 @@ func TestSpeculativeX509(t *testing.T) {
 
 		assert.Equal(t, numResponses, len(conn.Written), "expected %d wire messages to be sent, got %d",
 			numResponses, len(conn.Written))
-		hello, err := drivertest.GetCommandFromMsgWireMessage(<-conn.Written)
+		hello, err := drivertest.GetCommandFromQueryWireMessage(<-conn.Written)
 		assert.Nil(t, err, "error parsing hello command: %v", err)
 		assertCommandName(t, hello, handshake.LegacyHello)
 
@@ -103,7 +103,7 @@ func TestSpeculativeX509(t *testing.T) {
 
 		assert.Equal(t, numResponses, len(conn.Written), "expected %d wire messages to be sent, got %d",
 			numResponses, len(conn.Written))
-		hello, err := drivertest.GetCommandFromMsgWireMessage(<-conn.Written)
+		hello, err := drivertest.GetCommandFromQueryWireMessage(<-conn.Written)
 		assert.Nil(t, err, "error parsing hello command: %v", err)
 		assertCommandName(t, hello, handshake.LegacyHello)
 		_, err = hello.LookupErr("speculativeAuthenticate")

x/mongo/driver/drivertest/channel_conn.go

@@ -99,6 +99,38 @@ func MakeReply(doc bsoncore.Document) []byte {
 	return bsoncore.UpdateLength(dst, idx, int32(len(dst[idx:])))
 }
 
+// GetCommandFromQueryWireMessage returns the command sent in an OP_QUERY wire message.
+func GetCommandFromQueryWireMessage(wm []byte) (bsoncore.Document, error) {
+	var ok bool
+	_, _, _, _, wm, ok = wiremessage.ReadHeader(wm)
+	if !ok {
+		return nil, errors.New("could not read header")
+	}
+	_, wm, ok = wiremessage.ReadQueryFlags(wm)
+	if !ok {
+		return nil, errors.New("could not read flags")
+	}
+	_, wm, ok = wiremessage.ReadQueryFullCollectionName(wm)
+	if !ok {
+		return nil, errors.New("could not read fullCollectionName")
+	}
+	_, wm, ok = wiremessage.ReadQueryNumberToSkip(wm)
+	if !ok {
+		return nil, errors.New("could not read numberToSkip")
+	}
+	_, wm, ok = wiremessage.ReadQueryNumberToReturn(wm)
+	if !ok {
+		return nil, errors.New("could not read numberToReturn")
+	}
+
+	var query bsoncore.Document
+	query, wm, ok = wiremessage.ReadQueryQuery(wm)
+	if !ok {
+		return nil, errors.New("could not read query")
+	}
+	return query, nil
+}
+
 // GetCommandFromMsgWireMessage returns the command document sent in an OP_MSG wire message.
 func GetCommandFromMsgWireMessage(wm []byte) (bsoncore.Document, error) {
 	var ok bool

x/mongo/driver/legacy.go

@@ -19,4 +19,5 @@ const (
 	LegacyKillCursors
 	LegacyListCollections
 	LegacyListIndexes
+	LegacyHandshake
 )

x/mongo/driver/operation.go

@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/bsontype"
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/event"
 	"go.mongodb.org/mongo-driver/internal/csot"
@@ -629,7 +630,7 @@ func (op Operation) Execute(ctx context.Context) error {
 		}
 
 		var startedInfo startedInformation
-		*wm, startedInfo, err = op.createMsgWireMessage(ctx, maxTimeMS, (*wm)[:0], desc, conn, requestID)
+		*wm, startedInfo, err = op.createWireMessage(ctx, maxTimeMS, (*wm)[:0], desc, conn, requestID)
 
 		if err != nil {
 			return err
@@ -1103,6 +1104,85 @@ func (op Operation) addBatchArray(dst []byte) []byte {
 	return dst
 }
 
+func (op Operation) createLegacyHandshakeWireMessage(
+	maxTimeMS uint64,
+	dst []byte,
+	desc description.SelectedServer,
+) ([]byte, startedInformation, error) {
+	var info startedInformation
+	flags := op.secondaryOK(desc)
+	var wmindex int32
+	info.requestID = wiremessage.NextRequestID()
+	wmindex, dst = wiremessage.AppendHeaderStart(dst, info.requestID, 0, wiremessage.OpQuery)
+	dst = wiremessage.AppendQueryFlags(dst, flags)
+
+	dollarCmd := [...]byte{'.', '$', 'c', 'm', 'd'}
+
+	// FullCollectionName
+	dst = append(dst, op.Database...)
+	dst = append(dst, dollarCmd[:]...)
+	dst = append(dst, 0x00)
+	dst = wiremessage.AppendQueryNumberToSkip(dst, 0)
+	dst = wiremessage.AppendQueryNumberToReturn(dst, -1)
+
+	wrapper := int32(-1)
+	rp, err := op.createReadPref(desc, true)
+	if err != nil {
+		return dst, info, err
+	}
+	if len(rp) > 0 {
+		wrapper, dst = bsoncore.AppendDocumentStart(dst)
+		dst = bsoncore.AppendHeader(dst, bsontype.EmbeddedDocument, "$query")
+	}
+	idx, dst := bsoncore.AppendDocumentStart(dst)
+	dst, err = op.CommandFn(dst, desc)
+	if err != nil {
+		return dst, info, err
+	}
+
+	if op.Batches != nil && len(op.Batches.Current) > 0 {
+		dst = op.addBatchArray(dst)
+	}
+
+	dst, err = op.addReadConcern(dst, desc)
+	if err != nil {
+		return dst, info, err
+	}
+
+	dst, err = op.addWriteConcern(dst, desc)
+	if err != nil {
+		return dst, info, err
+	}
+
+	dst, err = op.addSession(dst, desc)
+	if err != nil {
+		return dst, info, err
+	}
+
+	dst = op.addClusterTime(dst, desc)
+	dst = op.addServerAPI(dst)
+	// If maxTimeMS is greater than 0 append it to wire message. A maxTimeMS value of 0 only explicitly
+	// specifies the default behavior of no timeout server-side.
+	if maxTimeMS > 0 {
+		dst = bsoncore.AppendInt64Element(dst, "maxTimeMS", int64(maxTimeMS))
+	}
+
+	dst, _ = bsoncore.AppendDocumentEnd(dst, idx)
+	// Command monitoring only reports the document inside $query
+	info.cmd = dst[idx:]
+
+	if len(rp) > 0 {
+		var err error
+		dst = bsoncore.AppendDocumentElement(dst, "$readPreference", rp)
+		dst, err = bsoncore.AppendDocumentEnd(dst, wrapper)
+		if err != nil {
+			return dst, info, err
+		}
+	}
+
+	return bsoncore.UpdateLength(dst, wmindex, int32(len(dst[wmindex:]))), info, nil
+}
+
 func (op Operation) createMsgWireMessage(
 	ctx context.Context,
 	maxTimeMS uint64,
@@ -1191,6 +1271,29 @@ func (op Operation) createMsgWireMessage(
 	return bsoncore.UpdateLength(dst, wmindex, int32(len(dst[wmindex:]))), info, nil
 }
 
+// isLegacyHandshake returns True if the operation is the first message of
+// the initial handshake and should use a legacy hello.
+func isLegacyHandshake(op Operation, desc description.SelectedServer) bool {
+	isInitialHandshake := desc.WireVersion == nil || desc.WireVersion.Max == 0
+
+	return op.Legacy == LegacyHandshake && isInitialHandshake
+}
+
+func (op Operation) createWireMessage(
+	ctx context.Context,
+	maxTimeMS uint64,
+	dst []byte,
+	desc description.SelectedServer,
+	conn Connection,
+	requestID int32,
+) ([]byte, startedInformation, error) {
+	if isLegacyHandshake(op, desc) {
+		return op.createLegacyHandshakeWireMessage(maxTimeMS, dst, desc)
+	}
+
+	return op.createMsgWireMessage(ctx, maxTimeMS, dst, desc, conn, requestID)
+}
+
 // addCommandFields adds the fields for a command to the wire message in dst. This assumes that the start of the document
 // has already been added and does not add the final 0 byte.
 func (op Operation) addCommandFields(ctx context.Context, dst []byte, desc description.SelectedServer) ([]byte, error) {

x/mongo/driver/operation/hello.go

@@ -537,8 +537,16 @@ func (h *Hello) StreamResponse(ctx context.Context, conn driver.StreamerConnecti
 	return h.createOperation().ExecuteExhaust(ctx, conn)
 }
 
+// isLegacyHandshake returns True if server API version is not requested and
+// loadBalanced is False. If this is the case, then the drivers MUST use legacy
+// hello for the first message of the initial handshake with the OP_QUERY
+// protocol
+func isLegacyHandshake(srvAPI *driver.ServerAPIOptions, deployment driver.Deployment) bool {
+	return srvAPI == nil && deployment.Kind() != description.LoadBalanced
+}
+
 func (h *Hello) createOperation() driver.Operation {
-	return driver.Operation{
+	op := driver.Operation{
 		Clock:      h.clock,
 		CommandFn:  h.command,
 		Database:   "admin",
@@ -549,23 +557,36 @@ func (h *Hello) createOperation() driver.Operation {
 		},
 		ServerAPI: h.serverAPI,
 	}
+
+	if isLegacyHandshake(h.serverAPI, h.d) {
+		op.Legacy = driver.LegacyHandshake
+	}
+
+	return op
 }
 
 // GetHandshakeInformation performs the MongoDB handshake for the provided connection and returns the relevant
 // information about the server. This function implements the driver.Handshaker interface.
 func (h *Hello) GetHandshakeInformation(ctx context.Context, _ address.Address, c driver.Connection) (driver.HandshakeInformation, error) {
-	err := driver.Operation{
+	deployment := driver.SingleConnectionDeployment{C: c}
+
+	op := driver.Operation{
 		Clock:      h.clock,
 		CommandFn:  h.handshakeCommand,
-		Deployment: driver.SingleConnectionDeployment{C: c},
+		Deployment: deployment,
 		Database:   "admin",
 		ProcessResponseFn: func(info driver.ResponseInfo) error {
 			h.res = info.ServerResponse
 			return nil
 		},
 		ServerAPI: h.serverAPI,
-	}.Execute(ctx)
-	if err != nil {
+	}
+
+	if isLegacyHandshake(h.serverAPI, deployment) {
+		op.Legacy = driver.LegacyHandshake
+	}
+
+	if err := op.Execute(ctx); err != nil {
 		return driver.HandshakeInformation{}, err
 	}
 
@@ -578,6 +599,9 @@ func (h *Hello) GetHandshakeInformation(ctx context.Context, _ address.Address,
 	if serverConnectionID, ok := h.res.Lookup("connectionId").AsInt64OK(); ok {
 		info.ServerConnectionID = &serverConnectionID
 	}
+
+	var err error
+
 	// Cast to bson.Raw to lookup saslSupportedMechs to avoid converting from bsoncore.Value to bson.RawValue for the
 	// StringSliceFromRawValue call.
 	if saslSupportedMechs, lookupErr := bson.Raw(h.res).LookupErr("saslSupportedMechs"); lookupErr == nil {