GODRIVER-3544, GODRIVER-3653 Allow Client to Send Client Metadata On-Demand

internal/assert/assertbsoncore/assertions_bsoncore.go

@@ -0,0 +1,47 @@
+// Copyright (C) MongoDB, Inc. 2025-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 assertbsoncore
+
+import (
+	"errors"
+	"testing"
+
+	"go.mongodb.org/mongo-driver/v2/internal/assert"
+	"go.mongodb.org/mongo-driver/v2/internal/handshake"
+	"go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore"
+)
+
+// HandshakeClientMetadata compares the client metadata in two wire messages. It
+// extracts the client metadata document from each wire message and compares
+// them. If the document is not found, it assumes the wire message is just the
+// value of the client metadata document itself.
+func HandshakeClientMetadata(t testing.TB, expectedWM, actualWM []byte) bool {
+	gotCommand, err := handshake.ParseClientMetadata(actualWM)
+	if err != nil {
+		if errors.Is(err, bsoncore.ErrElementNotFound) {
+			// If the element is not found, the actual wire message may just be the
+			// client metadata document itself.
+			gotCommand = bsoncore.Document(actualWM)
+		} else {
+			return assert.Fail(t, "error parsing actual wire message: %v", err)
+		}
+	}
+
+	wantCommand, err := handshake.ParseClientMetadata(expectedWM)
+	if err != nil {
+		// If the element is not found, the expected wire message may just be the
+		// client metadata document itself.
+		if errors.Is(err, bsoncore.ErrElementNotFound) {
+			wantCommand = bsoncore.Document(expectedWM)
+		} else {
+			return assert.Fail(t, "error parsing expected wire message: %v", err)
+		}
+	}
+
+	return assert.Equal(t, wantCommand, gotCommand,
+		"expected: %v, got: %v", bsoncore.Document(wantCommand), bsoncore.Document(gotCommand))
+}

internal/handshake/handshake.go

@@ -6,8 +6,29 @@
 
 package handshake
 
+import (
+	"go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore"
+)
+
 // LegacyHello is the legacy version of the hello command.
 var LegacyHello = "isMaster"
 
 // LegacyHelloLowercase is the lowercase, legacy version of the hello command.
 var LegacyHelloLowercase = "ismaster"
+
+func ParseClientMetadata(msg []byte) ([]byte, error) {
+	command := bsoncore.Document(msg)
+
+	// Lookup the "client" field in the command document.
+	clientMetadataRaw, err := command.LookupErr("client")
+	if err != nil {
+		return nil, err
+	}
+
+	clientMetadata, ok := clientMetadataRaw.DocumentOK()
+	if !ok {
+		return nil, err
+	}
+
+	return clientMetadata, nil
+}

internal/integration/client_test.go

@@ -19,12 +19,13 @@ import (
 	"go.mongodb.org/mongo-driver/v2/bson"
 	"go.mongodb.org/mongo-driver/v2/event"
 	"go.mongodb.org/mongo-driver/v2/internal/assert"
+	"go.mongodb.org/mongo-driver/v2/internal/assert/assertbsoncore"
 	"go.mongodb.org/mongo-driver/v2/internal/eventtest"
 	"go.mongodb.org/mongo-driver/v2/internal/failpoint"
-	"go.mongodb.org/mongo-driver/v2/internal/handshake"
 	"go.mongodb.org/mongo-driver/v2/internal/integration/mtest"
 	"go.mongodb.org/mongo-driver/v2/internal/integtest"
 	"go.mongodb.org/mongo-driver/v2/internal/require"
+	"go.mongodb.org/mongo-driver/v2/internal/test"
 	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 	"go.mongodb.org/mongo-driver/v2/mongo/readpref"
@@ -456,26 +457,12 @@ func TestClient(t *testing.T) {
 		err := mt.Client.Ping(context.Background(), mtest.PrimaryRp)
 		assert.Nil(mt, err, "Ping error: %v", err)
 
-		msgPairs := mt.GetProxiedMessages()
-		assert.True(mt, len(msgPairs) >= 2, "expected at least 2 events sent, got %v", len(msgPairs))
+		want := test.EncodeClientMetadata(mt, test.WithClientMetadataAppName("foo"))
+		for i := 0; i < 2; i++ {
+			message := mt.GetProxyCapture().TryNext()
+			require.NotNil(mt, message, "expected handshake message, got nil")
 
-		// First two messages should be connection handshakes: one for the heartbeat connection and the other for the
-		// application connection.
-		for idx, pair := range msgPairs[:2] {
-			helloCommand := handshake.LegacyHello
-			//  Expect "hello" command name with API version.
-			if os.Getenv("REQUIRE_API_VERSION") == "true" {
-				helloCommand = "hello"
-			}
-			assert.Equal(mt, pair.CommandName, helloCommand, "expected command name %s at index %d, got %s", helloCommand, idx,
-				pair.CommandName)
-
-			sent := pair.Sent
-			appNameVal, err := sent.Command.LookupErr("client", "application", "name")
-			assert.Nil(mt, err, "expected command %s at index %d to contain app name", sent.Command, idx)
-			appName := appNameVal.StringValue()
-			assert.Equal(mt, testAppName, appName, "expected app name %v at index %d, got %v", testAppName, idx,
-				appName)
+			assertbsoncore.HandshakeClientMetadata(mt, want, message.Sent.Command)
 		}
 	})
 
@@ -604,24 +591,32 @@ func TestClient(t *testing.T) {
 		err := mt.Client.Ping(context.Background(), mtest.PrimaryRp)
 		assert.Nil(mt, err, "Ping error: %v", err)
 
-		msgPairs := mt.GetProxiedMessages()
-		assert.True(mt, len(msgPairs) >= 3, "expected at least 3 events, got %v", len(msgPairs))
+		proxyCapture := mt.GetProxyCapture()
 
 		// The first message should be a connection handshake.
-		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.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.
+		firstMessage := proxyCapture.TryNext()
+		require.NotNil(mt, firstMessage, "expected handshake message, got nil")
+
+		assert.True(t, firstMessage.IsHandshake())
+
+		opCode := firstMessage.Sent.OpCode
+		assert.Equal(mt, wiremessage.OpQuery, opCode,
+			"expected 'OP_MSG' OpCode in wire message, got %q", 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.
 		var saslContinueFound bool
-		for _, pair := range msgPairs[1:] {
-			if pair.CommandName == "saslContinue" {
+		for {
+			message := proxyCapture.TryNext()
+			if message == nil {
+				break
+			}
+
+			if message.CommandName == "saslContinue" {
 				saslContinueFound = true
-				assert.Equal(mt, wiremessage.OpMsg, pair.Sent.OpCode,
-					"expected 'OP_MSG' OpCode in wire message, got %s", pair.Sent.OpCode.String())
+				opCode := message.Sent.OpCode
+				assert.Equal(mt, wiremessage.OpMsg, opCode,
+					"expected 'OP_MSG' OpCode in wire message, got %q", opCode.String())
 				break
 			}
 		}
@@ -634,18 +629,18 @@ func TestClient(t *testing.T) {
 		err := mt.Client.Ping(context.Background(), mtest.PrimaryRp)
 		assert.Nil(mt, err, "Ping error: %v", err)
 
-		msgPairs := mt.GetProxiedMessages()
-		assert.True(mt, len(msgPairs) >= 3, "expected at least 3 events, got %v", len(msgPairs))
-
 		// First three messages should be connection handshakes: one for the heartbeat connection, another for the
 		// application connection, and a final one for the RTT monitor connection.
-		for idx, pair := range msgPairs[:3] {
-			assert.Equal(mt, "hello", pair.CommandName, "expected command name 'hello' at index %d, got %s", idx,
-				pair.CommandName)
+		for idx := 0; idx < 3; idx++ {
+			message := mt.GetProxyCapture().TryNext()
+			require.NotNil(mt, message, "expected handshake message, got nil")
+
+			assert.True(t, message.IsHandshake())
 
 			// Assert that appended OpCode is OP_MSG when API version is set.
-			assert.Equal(mt, wiremessage.OpMsg, pair.Sent.OpCode,
-				"expected 'OP_MSG' OpCode in wire message, got %q", pair.Sent.OpCode.String())
+			opCode := message.Sent.OpCode
+			assert.Equal(mt, wiremessage.OpMsg, opCode,
+				"expected 'OP_MSG' OpCode in wire message, got %q", opCode.String())
 		}
 	})