CDRIVER-4056 Send LB in handshake and reject non-LB connections

src/libmongoc/doc/errors.rst

@@ -31,6 +31,8 @@ Many C Driver functions report errors by returning ``false`` or -1 and filling o
 +-----------------------------------------+---------------------------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 |                                         | ``MONGOC_ERROR_CLIENT_INVALID_ENCRYPTION_STATE``                                                                                | Failure related to Client-Side Field Level Encryption.                                                                                                                                                                                                                                                                                     |
 +-----------------------------------------+---------------------------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+|                                         | ``MONGOC_ERROR_CLIENT_INVALID_LOAD_BALANCER``                                                                                   | You attempted to connect to a MongoDB server behind a load balancer, but the server does not advertize load balanced support.                                                                                                                                                                                                              |
++-----------------------------------------+---------------------------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | ``MONGOC_ERROR_STREAM``                 | ``MONGOC_ERROR_STREAM_NAME_RESOLUTION``                                                                                         | DNS failure.                                                                                                                                                                                                                                                                                                                               |
 +-----------------------------------------+---------------------------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 |                                         | ``MONGOC_ERROR_STREAM_SOCKET``                                                                                                  | Timeout communicating with server, or connection closed.                                                                                                                                                                                                                                                                                   |

src/libmongoc/src/mongoc/mongoc-cluster.c

@@ -2278,8 +2278,30 @@ _mongoc_cluster_stream_for_server (mongoc_cluster_t *cluster,
       }
       bson_mutex_unlock (&topology->mutex);
       _mongoc_bson_init_with_transient_txn_error (cs, reply);
+      return NULL;
    }
 
+   /* If this is a load balanced topology and the server stream does not have a
+    * service id, disconnect and return an error. */
+   bson_mutex_lock (&topology->mutex);
+   if (topology->description.type == MONGOC_TOPOLOGY_LOAD_BALANCED) {
+      bson_oid_t service_id;
+
+      if (!mongoc_server_description_service_id (server_stream->sd,
+                                                 &service_id)) {
+         bson_set_error (error,
+                         MONGOC_ERROR_CLIENT,
+                         MONGOC_ERROR_CLIENT_INVALID_LOAD_BALANCER,
+                         "Driver attempted to initialize in load balancing "
+                         "mode, but the server does not support this mode.");
+         mongoc_server_stream_cleanup (server_stream);
+         mongoc_cluster_disconnect_node (cluster, server_id);
+         bson_mutex_unlock (&topology->mutex);
+         return NULL;
+      }
+   }
+   bson_mutex_unlock (&topology->mutex);
+
    RETURN (server_stream);
 }
 

src/libmongoc/src/mongoc/mongoc-error.h

@@ -131,7 +131,9 @@ typedef enum {
    MONGOC_ERROR_CLIENT_API_ALREADY_SET,
    MONGOC_ERROR_CLIENT_API_FROM_POOL,
    MONGOC_ERROR_POOL_API_ALREADY_SET,
-   MONGOC_ERROR_POOL_API_TOO_LATE
+   MONGOC_ERROR_POOL_API_TOO_LATE,
+
+   MONGOC_ERROR_CLIENT_INVALID_LOAD_BALANCER,
 } mongoc_error_code_t;
 
 MONGOC_EXPORT (bool)

src/libmongoc/src/mongoc/mongoc-server-description-private.h

@@ -112,6 +112,7 @@ struct _mongoc_server_description_t {
    pre-4.2 server.
    */
    uint32_t generation;
+   bson_oid_t service_id;
 };
 
 void
@@ -183,4 +184,11 @@ void
 mongoc_server_description_set_topology_version (mongoc_server_description_t *sd,
                                                 const bson_t *tv);
 
+/* If a service_id is set on the topology description, copies it to @oid and
+ * returns true. Otherwise returns false and zeros out oid.
+ */
+bool
+mongoc_server_description_service_id (
+   const mongoc_server_description_t *description, bson_oid_t *oid);
+
 #endif

src/libmongoc/src/mongoc/mongoc-server-description.c

@@ -92,6 +92,7 @@ mongoc_server_description_reset (mongoc_server_description_t *sd)
    sd->current_primary = NULL;
    sd->set_version = MONGOC_NO_SET_VERSION;
    bson_oid_copy_unsafe (&kObjectIdZero, &sd->election_id);
+   bson_oid_copy_unsafe (&kObjectIdZero, &sd->service_id);
 }
 
 /*
@@ -723,6 +724,10 @@ mongoc_server_description_handle_hello (mongoc_server_description_t *sd,
          mongoc_server_description_set_topology_version (
             sd, &incoming_topology_version);
          bson_destroy (&incoming_topology_version);
+      } else if (strcmp ("serviceId", bson_iter_key (&iter)) == 0) {
+          if (!BSON_ITER_HOLDS_OID (&iter))
+            goto failure;
+         bson_oid_copy_unsafe (bson_iter_oid (&iter), &sd->service_id);
       }
    }
 
@@ -798,6 +803,7 @@ mongoc_server_description_new_copy (
    bson_init (&copy->tags);
    bson_init (&copy->compressors);
    bson_copy_to (&description->topology_version, &copy->topology_version);
+   bson_oid_copy (&description->service_id, &copy->service_id);
 
    if (description->has_hello_response) {
       /* calls mongoc_server_description_reset */
@@ -1239,3 +1245,14 @@ mongoc_server_description_set_topology_version (mongoc_server_description_t *sd,
    bson_destroy (&sd->topology_version);
    bson_copy_to (tv, &sd->topology_version);
 }
+
+bool
+mongoc_server_description_service_id (const mongoc_server_description_t *description, bson_oid_t *oid) {
+   bson_oid_copy (&description->service_id, oid);
+
+   if (0 == bson_oid_compare (oid, &kObjectIdZero)) {
+      /* serviceID is unset. */
+      return false;
+   }
+   return true;
+}

src/libmongoc/src/mongoc/mongoc-topology-scanner-private.h

@@ -118,6 +118,7 @@ typedef struct mongoc_topology_scanner {
    bool speculative_authentication;
 
    mongoc_server_api_t *api;
+   bool loadbalanced;
 } mongoc_topology_scanner_t;
 
 mongoc_topology_scanner_t *
@@ -241,6 +242,9 @@ void
 _mongoc_topology_scanner_set_server_api (mongoc_topology_scanner_t *ts,
                                          const mongoc_server_api_t *api);
 
+void
+_mongoc_topology_scanner_set_loadbalanced (mongoc_topology_scanner_t *ts, bool val);
+
 /* for testing. */
 mongoc_stream_t *
 _mongoc_topology_scanner_tcp_initiate (mongoc_async_cmd_t *acmd);

src/libmongoc/src/mongoc/mongoc-topology-scanner.c

@@ -278,6 +278,10 @@ _build_handshake_cmd (mongoc_topology_scanner_t *ts)
    }
    bson_append_array_end (doc, &subdoc);
 
+   if (ts->loadbalanced) {
+      BSON_APPEND_BOOL (doc, "loadBalanced", true);
+   }
+
    /* Return whether the handshake doc fit the size limit */
    return res;
 }
@@ -1394,3 +1398,13 @@ _mongoc_topology_scanner_set_server_api (mongoc_topology_scanner_t *ts,
    ts->api = mongoc_server_api_copy (api);
    _reset_hello (ts);
 }
+
+/* This must be called before the handshake command is constructed. Caller does
+ * not need to lock the topology->mutex. */
+void
+_mongoc_topology_scanner_set_loadbalanced (mongoc_topology_scanner_t *ts,
+                                           bool val)
+{
+   BSON_ASSERT (bson_empty (&ts->handshake_cmd));
+   ts->loadbalanced = true;
+}

src/libmongoc/src/mongoc/mongoc-topology.c

@@ -439,6 +439,7 @@ mongoc_topology_new (const mongoc_uri_t *uri, bool single_threaded)
           * should occur. */
          _mongoc_topology_bypass_cooldown (topology);
       }
+      _mongoc_topology_scanner_set_loadbalanced (topology->scanner, true);
    } else if (service && !has_directconnection) {
       init_type = MONGOC_TOPOLOGY_UNKNOWN;
    } else if (has_directconnection) {

src/libmongoc/tests/test-mongoc-loadbalanced.c

@@ -20,6 +20,10 @@
 #include "test-libmongoc.h"
 #include "TestSuite.h"
 
+#include "mock_server/future-functions.h"
+#include "mock_server/mock-server.h"
+#include "mock_server/request.h"
+
 typedef struct {
    int server_changed_events;
    int server_opening_events;
@@ -437,6 +441,114 @@ test_loadbalanced_cooldown_is_bypassed_single (void *unused)
    free_and_assert_stats (stats);
 }
 
+/* Tests:
+ * - loadBalanced: true is added to the handshake
+ * - serviceId is set in the server description.
+ */
+#define LB_HELLO                                                               \
+   "{'ismaster': true, 'maxWireVersion': 13, 'msg': 'isdbgrid', 'serviceId': " \
+   "{'$oid': 'AAAAAAAAAAAAAAAAAAAAAAAA'}}"
+static void
+test_loadbalanced_handshake_sends_loadbalanced (void)
+{
+   mock_server_t *server;
+   mongoc_client_t *client;
+   mongoc_uri_t *uri;
+   request_t *request;
+   future_t *future;
+   bson_error_t error;
+   mongoc_server_description_t *monitor_sd;
+   mongoc_server_description_t *handshake_sd;
+   bson_oid_t expected;
+   bson_oid_t actual;
+
+   server = mock_server_new ();
+   mock_server_run (server);
+   mock_server_auto_endsessions (server);
+   uri = mongoc_uri_copy (mock_server_get_uri (server));
+   mongoc_uri_set_option_as_bool (uri, MONGOC_URI_LOADBALANCED, true);
+   client = mongoc_client_new_from_uri (uri);
+
+   future = future_client_command_simple (client,
+                                          "admin",
+                                          tmp_bson ("{'ping': 1}"),
+                                          NULL /* read prefs */,
+                                          NULL /* reply */,
+                                          &error);
+   request =
+      mock_server_receives_legacy_hello (server, "{'loadBalanced': true}");
+   mock_server_replies_simple (request, LB_HELLO);
+   request_destroy (request);
+
+   request = mock_server_receives_msg (server, 0, tmp_bson ("{'ping': 1}"));
+   mock_server_replies_ok_and_destroys (request);
+
+   ASSERT_OR_PRINT (future_get_bool (future), error);
+   future_destroy (future);
+
+   monitor_sd = mongoc_client_select_server (
+      client, true /* for writes */, NULL /* read prefs */, &error);
+   ASSERT_OR_PRINT (monitor_sd, error);
+   handshake_sd = mongoc_client_get_handshake_description (
+      client, 1, NULL /* opts */, &error);
+   ASSERT_OR_PRINT (handshake_sd, error);
+
+   bson_oid_init_from_string (&expected, "AAAAAAAAAAAAAAAAAAAAAAAA");
+   BSON_ASSERT (mongoc_server_description_service_id (handshake_sd, &actual));
+   ASSERT_CMPOID (&actual, &expected);
+
+   mongoc_server_description_destroy (handshake_sd);
+   mongoc_server_description_destroy (monitor_sd);
+   mongoc_uri_destroy (uri);
+   mongoc_client_destroy (client);
+   mock_server_destroy (server);
+}
+
+/* Tests that a connection is rejected if the handshake reply does not include a
+ * serviceID field. */
+#define NON_LB_HELLO \
+   "{'ismaster': true, 'maxWireVersion': 13, 'msg': 'isdbgrid'}"
+static void
+test_loadbalanced_handshake_rejects_non_loadbalanced (void)
+{
+   mock_server_t *server;
+   mongoc_client_t *client;
+   mongoc_uri_t *uri;
+   request_t *request;
+   future_t *future;
+   bson_error_t error;
+
+   server = mock_server_new ();
+   mock_server_run (server);
+   mock_server_auto_endsessions (server);
+   uri = mongoc_uri_copy (mock_server_get_uri (server));
+   mongoc_uri_set_option_as_bool (uri, MONGOC_URI_LOADBALANCED, true);
+   client = mongoc_client_new_from_uri (uri);
+
+   future = future_client_command_simple (client,
+                                          "admin",
+                                          tmp_bson ("{'ping': 1}"),
+                                          NULL /* read prefs */,
+                                          NULL /* reply */,
+                                          &error);
+   request =
+      mock_server_receives_legacy_hello (server, "{'loadBalanced': true}");
+   mock_server_replies_simple (request, NON_LB_HELLO);
+   request_destroy (request);
+   BSON_ASSERT (!future_get_bool (future));
+   future_destroy (future);
+
+   ASSERT_ERROR_CONTAINS (error,
+                          MONGOC_ERROR_CLIENT,
+                          MONGOC_ERROR_CLIENT_INVALID_LOAD_BALANCER,
+                          "Driver attempted to initialize in load balancing "
+                          "mode, but the server does not support this mode");
+
+   mongoc_uri_destroy (uri);
+   mongoc_client_destroy (client);
+   mock_server_destroy (server);
+}
+
 static int
 skip_if_not_loadbalanced (void)
 {
@@ -499,4 +611,13 @@ test_loadbalanced_install (TestSuite *suite)
                       NULL /* ctx */,
                       skip_if_not_loadbalanced,
                       test_framework_skip_if_no_failpoint);
+
+   TestSuite_AddMockServerTest (suite,
+                                "/loadbalanced/handshake_sends_loadbalanced",
+                                test_loadbalanced_handshake_sends_loadbalanced);
+
+   TestSuite_AddMockServerTest (
+      suite,
+      "/loadbalanced/handshake_rejects_non_loadbalanced",
+      test_loadbalanced_handshake_rejects_non_loadbalanced);
 }