Skip to content

Commit 4919d4f

Browse files
kevinAlbskevinAlbs
authored
CDRIVER-4056 Send LB in handshake and reject non-LB connections (#828)
1 parent f9db623 commit 4919d4f

9 files changed

+192
-1
lines changed

‎src/libmongoc/doc/errors.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@ Many C Driver functions report errors by returning ``false`` or -1 and filling o
3131
+-----------------------------------------+---------------------------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
3232
| | ``MONGOC_ERROR_CLIENT_INVALID_ENCRYPTION_STATE`` | Failure related to Client-Side Field Level Encryption. |
3333
+-----------------------------------------+---------------------------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
34+
| | ``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. |
35+
+-----------------------------------------+---------------------------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
3436
| ``MONGOC_ERROR_STREAM`` | ``MONGOC_ERROR_STREAM_NAME_RESOLUTION`` | DNS failure. |
3537
+-----------------------------------------+---------------------------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
3638
| | ``MONGOC_ERROR_STREAM_SOCKET`` | Timeout communicating with server, or connection closed. |

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2278,8 +2278,30 @@ _mongoc_cluster_stream_for_server (mongoc_cluster_t *cluster,
22782278
}
22792279
bson_mutex_unlock (&topology->mutex);
22802280
_mongoc_bson_init_with_transient_txn_error (cs, reply);
2281+
return NULL;
22812282
}
22822283

2284+
/* If this is a load balanced topology and the server stream does not have a
2285+
* service id, disconnect and return an error. */
2286+
bson_mutex_lock (&topology->mutex);
2287+
if (topology->description.type == MONGOC_TOPOLOGY_LOAD_BALANCED) {
2288+
bson_oid_t service_id;
2289+
2290+
if (!mongoc_server_description_service_id (server_stream->sd,
2291+
&service_id)) {
2292+
bson_set_error (error,
2293+
MONGOC_ERROR_CLIENT,
2294+
MONGOC_ERROR_CLIENT_INVALID_LOAD_BALANCER,
2295+
"Driver attempted to initialize in load balancing "
2296+
"mode, but the server does not support this mode.");
2297+
mongoc_server_stream_cleanup (server_stream);
2298+
mongoc_cluster_disconnect_node (cluster, server_id);
2299+
bson_mutex_unlock (&topology->mutex);
2300+
return NULL;
2301+
}
2302+
}
2303+
bson_mutex_unlock (&topology->mutex);
2304+
22832305
RETURN (server_stream);
22842306
}
22852307

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,9 @@ typedef enum {
131131
MONGOC_ERROR_CLIENT_API_ALREADY_SET,
132132
MONGOC_ERROR_CLIENT_API_FROM_POOL,
133133
MONGOC_ERROR_POOL_API_ALREADY_SET,
134-
MONGOC_ERROR_POOL_API_TOO_LATE
134+
MONGOC_ERROR_POOL_API_TOO_LATE,
135+
136+
MONGOC_ERROR_CLIENT_INVALID_LOAD_BALANCER,
135137
} mongoc_error_code_t;
136138

137139
MONGOC_EXPORT (bool)

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@ struct _mongoc_server_description_t {
112112
pre-4.2 server.
113113
*/
114114
uint32_t generation;
115+
bson_oid_t service_id;
115116
};
116117

117118
void
@@ -183,4 +184,11 @@ void
183184
mongoc_server_description_set_topology_version (mongoc_server_description_t *sd,
184185
const bson_t *tv);
185186

187+
/* If a service_id is set on the topology description, copies it to @oid and
188+
* returns true. Otherwise returns false and zeros out oid.
189+
*/
190+
bool
191+
mongoc_server_description_service_id (
192+
const mongoc_server_description_t *description, bson_oid_t *oid);
193+
186194
#endif

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@ mongoc_server_description_reset (mongoc_server_description_t *sd)
9292
sd->current_primary = NULL;
9393
sd->set_version = MONGOC_NO_SET_VERSION;
9494
bson_oid_copy_unsafe (&kObjectIdZero, &sd->election_id);
95+
bson_oid_copy_unsafe (&kObjectIdZero, &sd->service_id);
9596
}
9697

9798
/*
@@ -723,6 +724,10 @@ mongoc_server_description_handle_hello (mongoc_server_description_t *sd,
723724
mongoc_server_description_set_topology_version (
724725
sd, &incoming_topology_version);
725726
bson_destroy (&incoming_topology_version);
727+
} else if (strcmp ("serviceId", bson_iter_key (&iter)) == 0) {
728+
if (!BSON_ITER_HOLDS_OID (&iter))
729+
goto failure;
730+
bson_oid_copy_unsafe (bson_iter_oid (&iter), &sd->service_id);
726731
}
727732
}
728733

@@ -798,6 +803,7 @@ mongoc_server_description_new_copy (
798803
bson_init (&copy->tags);
799804
bson_init (&copy->compressors);
800805
bson_copy_to (&description->topology_version, &copy->topology_version);
806+
bson_oid_copy (&description->service_id, &copy->service_id);
801807

802808
if (description->has_hello_response) {
803809
/* calls mongoc_server_description_reset */
@@ -1239,3 +1245,14 @@ mongoc_server_description_set_topology_version (mongoc_server_description_t *sd,
12391245
bson_destroy (&sd->topology_version);
12401246
bson_copy_to (tv, &sd->topology_version);
12411247
}
1248+
1249+
bool
1250+
mongoc_server_description_service_id (const mongoc_server_description_t *description, bson_oid_t *oid) {
1251+
bson_oid_copy (&description->service_id, oid);
1252+
1253+
if (0 == bson_oid_compare (oid, &kObjectIdZero)) {
1254+
/* serviceID is unset. */
1255+
return false;
1256+
}
1257+
return true;
1258+
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@ typedef struct mongoc_topology_scanner {
118118
bool speculative_authentication;
119119

120120
mongoc_server_api_t *api;
121+
bool loadbalanced;
121122
} mongoc_topology_scanner_t;
122123

123124
mongoc_topology_scanner_t *
@@ -241,6 +242,9 @@ void
241242
_mongoc_topology_scanner_set_server_api (mongoc_topology_scanner_t *ts,
242243
const mongoc_server_api_t *api);
243244

245+
void
246+
_mongoc_topology_scanner_set_loadbalanced (mongoc_topology_scanner_t *ts, bool val);
247+
244248
/* for testing. */
245249
mongoc_stream_t *
246250
_mongoc_topology_scanner_tcp_initiate (mongoc_async_cmd_t *acmd);

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -278,6 +278,10 @@ _build_handshake_cmd (mongoc_topology_scanner_t *ts)
278278
}
279279
bson_append_array_end (doc, &subdoc);
280280

281+
if (ts->loadbalanced) {
282+
BSON_APPEND_BOOL (doc, "loadBalanced", true);
283+
}
284+
281285
/* Return whether the handshake doc fit the size limit */
282286
return res;
283287
}
@@ -1394,3 +1398,13 @@ _mongoc_topology_scanner_set_server_api (mongoc_topology_scanner_t *ts,
13941398
ts->api = mongoc_server_api_copy (api);
13951399
_reset_hello (ts);
13961400
}
1401+
1402+
/* This must be called before the handshake command is constructed. Caller does
1403+
* not need to lock the topology->mutex. */
1404+
void
1405+
_mongoc_topology_scanner_set_loadbalanced (mongoc_topology_scanner_t *ts,
1406+
bool val)
1407+
{
1408+
BSON_ASSERT (bson_empty (&ts->handshake_cmd));
1409+
ts->loadbalanced = true;
1410+
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -439,6 +439,7 @@ mongoc_topology_new (const mongoc_uri_t *uri, bool single_threaded)
439439
* should occur. */
440440
_mongoc_topology_bypass_cooldown (topology);
441441
}
442+
_mongoc_topology_scanner_set_loadbalanced (topology->scanner, true);
442443
} else if (service && !has_directconnection) {
443444
init_type = MONGOC_TOPOLOGY_UNKNOWN;
444445
} else if (has_directconnection) {

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

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,10 @@
2020
#include "test-libmongoc.h"
2121
#include "TestSuite.h"
2222

23+
#include "mock_server/future-functions.h"
24+
#include "mock_server/mock-server.h"
25+
#include "mock_server/request.h"
26+
2327
typedef struct {
2428
int server_changed_events;
2529
int server_opening_events;
@@ -437,6 +441,114 @@ test_loadbalanced_cooldown_is_bypassed_single (void *unused)
437441
free_and_assert_stats (stats);
438442
}
439443

444+
/* Tests:
445+
* - loadBalanced: true is added to the handshake
446+
* - serviceId is set in the server description.
447+
*/
448+
#define LB_HELLO \
449+
"{'ismaster': true, 'maxWireVersion': 13, 'msg': 'isdbgrid', 'serviceId': " \
450+
"{'$oid': 'AAAAAAAAAAAAAAAAAAAAAAAA'}}"
451+
static void
452+
test_loadbalanced_handshake_sends_loadbalanced (void)
453+
{
454+
mock_server_t *server;
455+
mongoc_client_t *client;
456+
mongoc_uri_t *uri;
457+
request_t *request;
458+
future_t *future;
459+
bson_error_t error;
460+
mongoc_server_description_t *monitor_sd;
461+
mongoc_server_description_t *handshake_sd;
462+
bson_oid_t expected;
463+
bson_oid_t actual;
464+
465+
server = mock_server_new ();
466+
mock_server_run (server);
467+
mock_server_auto_endsessions (server);
468+
uri = mongoc_uri_copy (mock_server_get_uri (server));
469+
mongoc_uri_set_option_as_bool (uri, MONGOC_URI_LOADBALANCED, true);
470+
client = mongoc_client_new_from_uri (uri);
471+
472+
future = future_client_command_simple (client,
473+
"admin",
474+
tmp_bson ("{'ping': 1}"),
475+
NULL /* read prefs */,
476+
NULL /* reply */,
477+
&error);
478+
request =
479+
mock_server_receives_legacy_hello (server, "{'loadBalanced': true}");
480+
mock_server_replies_simple (request, LB_HELLO);
481+
request_destroy (request);
482+
483+
request = mock_server_receives_msg (server, 0, tmp_bson ("{'ping': 1}"));
484+
mock_server_replies_ok_and_destroys (request);
485+
486+
ASSERT_OR_PRINT (future_get_bool (future), error);
487+
future_destroy (future);
488+
489+
monitor_sd = mongoc_client_select_server (
490+
client, true /* for writes */, NULL /* read prefs */, &error);
491+
ASSERT_OR_PRINT (monitor_sd, error);
492+
handshake_sd = mongoc_client_get_handshake_description (
493+
client, 1, NULL /* opts */, &error);
494+
ASSERT_OR_PRINT (handshake_sd, error);
495+
496+
bson_oid_init_from_string (&expected, "AAAAAAAAAAAAAAAAAAAAAAAA");
497+
BSON_ASSERT (mongoc_server_description_service_id (handshake_sd, &actual));
498+
ASSERT_CMPOID (&actual, &expected);
499+
500+
mongoc_server_description_destroy (handshake_sd);
501+
mongoc_server_description_destroy (monitor_sd);
502+
mongoc_uri_destroy (uri);
503+
mongoc_client_destroy (client);
504+
mock_server_destroy (server);
505+
}
506+
507+
/* Tests that a connection is rejected if the handshake reply does not include a
508+
* serviceID field. */
509+
#define NON_LB_HELLO \
510+
"{'ismaster': true, 'maxWireVersion': 13, 'msg': 'isdbgrid'}"
511+
static void
512+
test_loadbalanced_handshake_rejects_non_loadbalanced (void)
513+
{
514+
mock_server_t *server;
515+
mongoc_client_t *client;
516+
mongoc_uri_t *uri;
517+
request_t *request;
518+
future_t *future;
519+
bson_error_t error;
520+
521+
server = mock_server_new ();
522+
mock_server_run (server);
523+
mock_server_auto_endsessions (server);
524+
uri = mongoc_uri_copy (mock_server_get_uri (server));
525+
mongoc_uri_set_option_as_bool (uri, MONGOC_URI_LOADBALANCED, true);
526+
client = mongoc_client_new_from_uri (uri);
527+
528+
future = future_client_command_simple (client,
529+
"admin",
530+
tmp_bson ("{'ping': 1}"),
531+
NULL /* read prefs */,
532+
NULL /* reply */,
533+
&error);
534+
request =
535+
mock_server_receives_legacy_hello (server, "{'loadBalanced': true}");
536+
mock_server_replies_simple (request, NON_LB_HELLO);
537+
request_destroy (request);
538+
BSON_ASSERT (!future_get_bool (future));
539+
future_destroy (future);
540+
541+
ASSERT_ERROR_CONTAINS (error,
542+
MONGOC_ERROR_CLIENT,
543+
MONGOC_ERROR_CLIENT_INVALID_LOAD_BALANCER,
544+
"Driver attempted to initialize in load balancing "
545+
"mode, but the server does not support this mode");
546+
547+
mongoc_uri_destroy (uri);
548+
mongoc_client_destroy (client);
549+
mock_server_destroy (server);
550+
}
551+
440552
static int
441553
skip_if_not_loadbalanced (void)
442554
{
@@ -499,4 +611,13 @@ test_loadbalanced_install (TestSuite *suite)
499611
NULL /* ctx */,
500612
skip_if_not_loadbalanced,
501613
test_framework_skip_if_no_failpoint);
614+
615+
TestSuite_AddMockServerTest (suite,
616+
"/loadbalanced/handshake_sends_loadbalanced",
617+
test_loadbalanced_handshake_sends_loadbalanced);
618+
619+
TestSuite_AddMockServerTest (
620+
suite,
621+
"/loadbalanced/handshake_rejects_non_loadbalanced",
622+
test_loadbalanced_handshake_rejects_non_loadbalanced);
502623
}

0 commit comments

Comments
 (0)