GODRIVER-1405 Send app name on heartbeat and application connections (#366)

Author: Divjot Arora

mongo/client.go

@@ -333,6 +333,13 @@ func (c *Client) configure(opts *options.ClientOptions) error {
 	var appName string
 	if opts.AppName != nil {
 		appName = *opts.AppName
+
+		serverOpts = append(serverOpts, topology.WithServerAppName(func(string) string {
+			return appName
+		}))
+		connOpts = append(connOpts, topology.WithConnectionAppName(func(string) string {
+			return appName
+		}))
 	}
 	// Compressors & ZlibLevel
 	var comps []string

mongo/integration/client_test.go

@@ -7,10 +7,15 @@
 package integration
 
 import (
+	"bytes"
+	"context"
 	"fmt"
+	"io"
+	"net"
 	"path"
 	"reflect"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -23,6 +28,8 @@ import (
 	"go.mongodb.org/mongo-driver/mongo/options"
 	"go.mongodb.org/mongo-driver/mongo/readpref"
 	"go.mongodb.org/mongo-driver/x/mongo/driver"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/drivertest"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/wiremessage"
 )
 
 const certificatesDir = "../../data/certificates"
@@ -333,4 +340,125 @@ func TestClient(t *testing.T) {
 			})
 		}
 	})
+
+	testAppName := "foo"
+	hosts := options.Client().ApplyURI(mt.ConnString()).Hosts
+	appNameProxyDialer := newProxyDialer()
+	appNameDialerOpts := options.Client().
+		SetDialer(appNameProxyDialer).
+		SetHosts(hosts[:1]).
+		SetDirect(true).
+		SetAppName(testAppName)
+	appNameMtOpts := mtest.NewOptions().
+		ClientOptions(appNameDialerOpts).
+		Topologies(mtest.Single)
+	mt.RunOpts("app name is always sent", appNameMtOpts, func(mt *mtest.T) {
+		err := mt.Client.Ping(mtest.Background, mtest.PrimaryRp)
+		assert.Nil(mt, err, "Ping error: %v", err)
+
+		sent := appNameProxyDialer.sent
+		assert.True(mt, len(sent) >= 2, "expected at least 2 events sent, got %v", len(sent))
+
+		// First two messages should be connection handshakes: one for the heartbeat connection and the other for the
+		// application connection.
+		for idx, wm := range sent[:2] {
+			cmd, err := drivertest.GetCommandFromQueryWireMessage(wm)
+			assert.Nil(mt, err, "GetCommandFromQueryWireMessage error at index %d: %v", idx, err)
+			heartbeatCmdName := cmd.Index(0).Key()
+			assert.Equal(mt, "isMaster", heartbeatCmdName,
+				"expected command name isMaster at index %d, got %v", idx, heartbeatCmdName)
+
+			appNameVal, err := cmd.LookupErr("client", "application", "name")
+			assert.Nil(mt, err, "expected command %s at index %d to contain app name", cmd, idx)
+			appName := appNameVal.StringValue()
+			assert.Equal(mt, testAppName, appName, "expected app name %v at index %d, got %v", testAppName, idx, appName)
+		}
+	})
+}
+
+// proxyDialer is a ContextDialer implementation that wraps a net.Dialer and records the messages sent and received
+// using connections created through it.
+type proxyDialer struct {
+	*net.Dialer
+	sync.Mutex
+	sent     []wiremessage.WireMessage
+	received []wiremessage.WireMessage
+}
+
+var _ options.ContextDialer = (*proxyDialer)(nil)
+
+func newProxyDialer() *proxyDialer {
+	return &proxyDialer{
+		Dialer: &net.Dialer{Timeout: 30 * time.Second},
+	}
+}
+
+// DialContext creates a new proxyConnection.
+func (p *proxyDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	netConn, err := p.Dialer.DialContext(ctx, network, address)
+	if err != nil {
+		return netConn, err
+	}
+
+	proxy := &proxyConn{
+		Conn:           netConn,
+		dialer:         p,
+		currentReading: bytes.NewBuffer(nil),
+	}
+	return proxy, nil
+}
+
+// storeSentMessage stores a copy of the wire message being sent to the server.
+func (p *proxyDialer) storeSentMessage(msg []byte) {
+	p.Lock()
+	defer p.Unlock()
+
+	msgCopy := make(wiremessage.WireMessage, len(msg))
+	copy(msgCopy, msg)
+	p.sent = append(p.sent, msgCopy)
+}
+
+// storeReceivedMessage stores a copy of the wire message being received from the server.
+func (p *proxyDialer) storeReceivedMessage(msg []byte) {
+	p.Lock()
+	defer p.Unlock()
+
+	msgCopy := make(wiremessage.WireMessage, len(msg))
+	copy(msgCopy, msg)
+	p.received = append(p.received, msgCopy)
+}
+
+// proxyConn is a net.Conn that wraps a network connection. All messages sent/received through a proxyConn are stored
+// in the associated proxyDialer and are forwarded over the wrapped connection.
+type proxyConn struct {
+	net.Conn
+	dialer         *proxyDialer
+	currentReading *bytes.Buffer // The current message being read.
+}
+
+// Write stores the given message in the proxyDialer associated with this connection and forwards the message to the
+// server.
+func (pc *proxyConn) Write(msg []byte) (n int, err error) {
+	pc.dialer.storeSentMessage(msg)
+	return pc.Conn.Write(msg)
+}
+
+// Read reads the message from the server into the given buffer and stores the read message in the proxyDialer
+// associated with this connection.
+func (pc *proxyConn) Read(buffer []byte) (int, error) {
+	n, err := pc.Conn.Read(buffer)
+	if err != nil {
+		return n, err
+	}
+
+	_, err = io.Copy(pc.currentReading, bytes.NewReader(buffer))
+	if err != nil {
+		return 0, fmt.Errorf("error copying to mock: %v", err)
+	}
+	if len(buffer) != 4 {
+		pc.dialer.storeReceivedMessage(pc.currentReading.Bytes())
+		pc.currentReading.Reset()
+	}
+
+	return n, err
 }