SERVER-61352: Test connections made to mongos over proxy protocol

5 files changed, 163 insertions, 0 deletions

diff --git a/etc/pip/components/resmoke.req b/etc/pip/components/resmoke.req
index 0e4a70cbb65..ed4d2a47dce 100644
--- a/etc/pip/components/resmoke.req
+++ b/etc/pip/components/resmoke.req
@@ -16,3 +16,4 @@ PyGithub == 1.53
 urllib3 >= 1.26.0
 distro == 1.5.0
 dnspython == 2.1.0
+proxy-protocol == 0.7.5
\ No newline at end of file
diff --git a/jstests/sharding/libs/proxy_protocol.js b/jstests/sharding/libs/proxy_protocol.js
new file mode 100644
index 00000000000..54847dbc595
--- /dev/null
+++ b/jstests/sharding/libs/proxy_protocol.js
@@ -0,0 +1,80 @@
+/**
+ * Control the proxy protocol server.
+ */
+class ProxyProtocolServer {
+    /**
+     * Create a new proxy protocol server.
+     */
+    constructor(ingress_port, egress_port, version) {
+        this.python = "python3";
+
+        if (_isWindows()) {
+            this.python = "python.exe";
+        }
+
+        print("Using python interpreter: " + this.python);
+        this.web_server_py = "jstests/sharding/libs/proxy_protocol_server.py";
+
+        this.pid = undefined;
+        this.ingress_port = ingress_port;
+        this.egress_port = egress_port;
+
+        assert(version === 1 || version === 2);
+        this.version = version;
+    }
+
+    /**
+     * Get the ingress port - the port over which a client wishing to appear proxied should connect.
+     *
+     * @return {number} ingress port number
+     */
+    getIngressPort() {
+        return ingress_port;
+    }
+
+    /**
+     * Get the egress port - the port that mongos should be listening on for proxied connections.
+     *
+     * @return {number} egress port number
+     */
+    getEgressPort() {
+        return egress_port;
+    }
+
+    /**
+     * Start the server.
+     */
+    start() {
+        print("Proxy protocol server is listening on port: " + this.ingress_port);
+        print("Proxy protocol server is proxying to port: " + this.egress_port);
+
+        let args = [
+            this.python,
+            "-u",
+            this.web_server_py,
+            "--service",
+            "localhost:" + this.ingress_port,
+            "localhost:" + this.egress_port + "?pp=v" + this.version
+        ];
+
+        clearRawMongoProgramOutput();
+
+        this.pid = _startMongoProgram({args: args});
+
+        assert(checkProgram(this.pid));
+
+        // Wait for the web server to start
+        assert.soon(function() {
+            return rawMongoProgramOutput().search("Starting proxy protocol server...") !== -1;
+        });
+
+        print("Proxy Protocol Server sucessfully started.");
+    }
+
+    /**
+     * Stop the server.
+     */
+    stop() {
+        stopMongoProgramByPid(this.pid);
+    }
+}
diff --git a/jstests/sharding/libs/proxy_protocol_server.py b/jstests/sharding/libs/proxy_protocol_server.py
new file mode 100644
index 00000000000..d6bf876bad1
--- /dev/null
+++ b/jstests/sharding/libs/proxy_protocol_server.py
@@ -0,0 +1,11 @@
+#! /usr/bin/env python3
+"""
+Python script to interact with proxy protocol server.
+"""
+
+from proxyprotocol.server.main import *
+import sys
+
+if __name__ == '__main__':
+    print("Starting proxy protocol server...")
+    sys.exit(main())
diff --git a/jstests/sharding/proxy_protocol_connect.js b/jstests/sharding/proxy_protocol_connect.js
new file mode 100644
index 00000000000..bb3b78c6cdb
--- /dev/null
+++ b/jstests/sharding/proxy_protocol_connect.js
@@ -0,0 +1,66 @@
+/**
+ * Validate we can connect over the proxy protocol port with the protocol appended.
+ * @tags: [requires_fcv_52]
+ */
+
+(function() {
+if (_isWindows()) {
+    // The proxy protocol python package currently doesn't support Windows.
+    return;
+}
+load("jstests/libs/logv2_helpers.js");
+load("jstests/sharding/libs/proxy_protocol.js");
+
+// Test that you can connect to the load balancer port over a proxy.
+function testProxyProtocolConnect(ingressPort, egressPort, version) {
+    'use strict';
+
+    let proxy_server = new ProxyProtocolServer(ingressPort, egressPort, version);
+    proxy_server.start();
+
+    let st = new ShardingTest({
+        shards: 1,
+        mongos: 1,
+        mongosOptions:
+            {setParameter: {"featureFlagLoadBalancer": true, "loadBalancerPort": egressPort}}
+    });
+
+    const uri = `mongodb://localhost:${ingressPort}/?loadBalanced=true`;
+    const conn = new Mongo(uri);
+    assert.neq(null, conn, 'Client was unable to connect to the load balancer port');
+    assert.commandWorked(conn.getDB('admin').runCommand({hello: 1}));
+    proxy_server.stop();
+    st.stop();
+}
+
+// Test that you can't connect to the load balancer port without being proxied.
+function testProxyProtocolConnectFailure(lbPort, sendLoadBalanced) {
+    'use strict';
+
+    let st = new ShardingTest({
+        shards: 1,
+        mongos: 1,
+        mongosOptions: {setParameter: {"featureFlagLoadBalancer": true, "loadBalancerPort": lbPort}}
+    });
+
+    const hostName = st.s.host.substring(0, st.s.host.indexOf(":"));
+    const uri = `mongodb://${hostName}:${lbPort}/?loadBalanced=${sendLoadBalanced}`;
+    try {
+        var conn = new Mongo(uri);
+        assert(false, 'Client was unable to connect to the load balancer port');
+    } catch (err) {
+        assert(checkLog.checkContainsOnceJsonStringMatch(
+                   st.s, 6067900, "msg", "Error while parsing proxy protocol header"),
+               "Connection failed for some reason other than lacking a proxy protocol header");
+    }
+    st.stop();
+}
+
+const ingressPort = 21234;
+const egressPort = 21235;
+
+testProxyProtocolConnect(ingressPort, egressPort, 1);
+testProxyProtocolConnect(ingressPort, egressPort, 2);
+testProxyProtocolConnectFailure(egressPort, "true");
+testProxyProtocolConnectFailure(egressPort, "false");
+})();
diff --git a/src/mongo/client/dbclient_connection.cpp b/src/mongo/client/dbclient_connection.cpp
index 3eab2b4fab6..55202e0c907 100644
--- a/src/mongo/client/dbclient_connection.cpp
+++ b/src/mongo/client/dbclient_connection.cpp
@@ -170,6 +170,11 @@ executor::RemoteCommandResponse initWireVersion(
         bob.append("helloOk", true);
     }
 
+    auto loadBalancedOpt = uri.getOption("loadBalanced");
+    if (loadBalancedOpt && (loadBalancedOpt.value() == "true")) {
+        bob.append("loadBalanced", true);
+    }
+
     *speculativeAuthType = auth::speculateAuth(&bob, uri, saslClientSession);
     if (!uri.getUser().empty()) {
         UserName user(uri.getUser(), uri.getAuthenticationDatabase());