diff --git a/src/comsock.c b/src/comsock.c index fea50574c..c55670120 100644 --- a/src/comsock.c +++ b/src/comsock.c @@ -78,6 +78,60 @@ natsSock_SetCommonTcpOptions(natsSock fd) return NATS_OK; } +void +natsSock_ShuffleIPs(natsSockCtx *ctx, struct addrinfo **tmp, int tmpSize, struct addrinfo **ipListHead, int count) +{ + struct addrinfo **ips = NULL; + struct addrinfo *p = NULL; + bool doFree = false; + int i, j; + + if (ctx->noRandomize || (ipListHead == NULL) || (*ipListHead == NULL) || count <= 1) + return; + + if (count > tmpSize) + { + ips = (struct addrinfo**) NATS_CALLOC(count, sizeof(struct addrinfo*)); + // Let's not fail the caller, simply don't shuffle. + if (ips == NULL) + return; + doFree = true; + } + else + { + ips = tmp; + } + p = *ipListHead; + // Put them in a array + for (i=0; iai_next; + } + // Shuffle the array + for (i=0; iai_next = ips[i+1]; + else + ips[i]->ai_next = NULL; + } + // Update the list head with the first in the array. + *ipListHead = ips[0]; + + // If we allocated the array, free it. + if (doFree) + NATS_FREE(ips); +} + #define MAX_HOST_NAME (256) natsStatus @@ -97,6 +151,7 @@ natsSock_ConnectTcp(natsSockCtx *ctx, const char *phost, int port) int64_t start = 0; int64_t totalTimeout = 0; int64_t timeoutPerIP = 0; + struct addrinfo *tmpStorage[64]; if (phost == NULL) return nats_setError(NATS_ADDRESS_MISSING, "%s", "No host specified"); @@ -124,6 +179,7 @@ natsSock_ConnectTcp(natsSockCtx *ctx, const char *phost, int port) { struct addrinfo hints; struct addrinfo *servinfo = NULL; + int count = 0; struct addrinfo *p; memset(&hints,0,sizeof(hints)); @@ -144,9 +200,14 @@ natsSock_ConnectTcp(natsSockCtx *ctx, const char *phost, int port) gai_strerror(res)); continue; } - servInfos[numServInfo++] = servinfo; + servInfos[numServInfo] = servinfo; for (p = servinfo; (p != NULL); p = p->ai_next) + { + count++; numIPs++; + } + natsSock_ShuffleIPs(ctx, tmpStorage, sizeof(tmpStorage), &(servInfos[numServInfo]), count); + numServInfo++; } // If we got a getaddrinfo() and there is no servInfos to try to connect to // bail out now. diff --git a/src/comsock.h b/src/comsock.h index 96417ca89..73c01fadf 100644 --- a/src/comsock.h +++ b/src/comsock.h @@ -23,6 +23,9 @@ natsSock_Init(natsSockCtx *ctx); natsStatus natsSock_WaitReady(int waitMode, natsSockCtx *ctx); +void +natsSock_ShuffleIPs(natsSockCtx *ctx, struct addrinfo **tmp, int tmpSize, struct addrinfo **ipListHead, int count); + natsStatus natsSock_ConnectTcp(natsSockCtx *ctx, const char *host, int port); diff --git a/src/conn.c b/src/conn.c index c4d10d9cc..9942e78f1 100644 --- a/src/conn.c +++ b/src/conn.c @@ -395,6 +395,9 @@ _createConn(natsConnection *nc) // Set the IP resolution order nc->sockCtx.orderIP = nc->opts->orderIP; + // Set ctx.noRandomize based on public NoRandomize option. + nc->sockCtx.noRandomize = nc->opts->noRandomize; + s = natsSock_ConnectTcp(&(nc->sockCtx), nc->cur->url->host, nc->cur->url->port); if (s == NATS_OK) nc->sockCtx.fdActive = true; diff --git a/src/natsp.h b/src/natsp.h index 30876efcf..03a4f9828 100644 --- a/src/natsp.h +++ b/src/natsp.h @@ -460,6 +460,10 @@ typedef struct __natsSockCtx int orderIP; // possible values: 0,4,6,46,64 + // By default, the list of IPs returned by the hostname resolution will + // be shuffled. This option, if `true`, will disable the shuffling. + bool noRandomize; + } natsSockCtx; typedef struct __respInfo diff --git a/test/list.txt b/test/list.txt index 17f3c1fed..b2cfac0e8 100644 --- a/test/list.txt +++ b/test/list.txt @@ -21,6 +21,7 @@ natsStrHash natsInbox natsOptions natsSock_ConnectTcp +natsSock_ShuffleIPs natsSock_IPOrder natsSock_ReadLine natsJSON diff --git a/test/test.c b/test/test.c index 25599fbaf..d18e48472 100644 --- a/test/test.c +++ b/test/test.c @@ -4859,6 +4859,12 @@ test_natsSock_ConnectTcp(void) _stopServer(serverPid); serverPid = NATS_INVALID_PID; + test("Check connect tcp hostname: "); + serverPid = _startServer("nats://localhost:4222", "-p 4222", true); + testCond(serverPid != NATS_INVALID_PID); + _stopServer(serverPid); + serverPid = NATS_INVALID_PID; + test("Check connect tcp (force server to listen to IPv4): "); serverPid = _startServer("nats://127.0.0.1:4222", "-a 127.0.0.1 -p 4222", true); testCond(serverPid != NATS_INVALID_PID); @@ -4866,6 +4872,95 @@ test_natsSock_ConnectTcp(void) serverPid = NATS_INVALID_PID; } +static bool +listOrder(struct addrinfo *head, bool ordered) +{ + struct addrinfo *p; + int i; + + p = head; + for (i=0; i<10; i++) + { + if (ordered && (p->ai_flags != (i+1))) + return false; + p = p->ai_next; + } + return true; +} + +static void +test_natsSock_ShuffleIPs(void) +{ + struct addrinfo *tmp[10]; + struct addrinfo *head = NULL; + struct addrinfo *tail = NULL; + struct addrinfo *list = NULL; + struct addrinfo *p; + natsSockCtx ctx; + int i=0; + + // Create a fake list that has `ai_flags` set to 1 to 10. + // We will use that to check that the list is shuffled or not. + for (i=0; i<10; i++) + { + p = calloc(1, sizeof(struct addrinfo)); + p->ai_flags = (i+1); + if (head == NULL) + head=p; + if (tail != NULL) + tail->ai_next = p; + tail = p; + } + + test("No randomize, so no shuffling: "); + natsSock_Init(&ctx); + ctx.noRandomize = true; + list = head; + natsSock_ShuffleIPs(&ctx, tmp, sizeof(tmp), &list, 10); + testCond((list == head) && listOrder(list, true)); + + test("Shuffling bad args 2: "); + natsSock_Init(&ctx); + list = head; + natsSock_ShuffleIPs(&ctx, tmp, sizeof(tmp), NULL, 10); + testCond((list == head) && listOrder(list, true)); + + test("Shuffling bad args 1: "); + natsSock_Init(&ctx); + list = head; + natsSock_ShuffleIPs(&ctx, tmp, sizeof(tmp), &list, 0); + testCond((list == head) && listOrder(list, true)); + + test("No shuffling count==1: "); + natsSock_Init(&ctx); + list = head; + natsSock_ShuffleIPs(&ctx, tmp, sizeof(tmp), &list, 1); + testCond((list == head) && listOrder(list, true)); + + test("Shuffling: "); + natsSock_Init(&ctx); + list = head; + natsSock_ShuffleIPs(&ctx, tmp, sizeof(tmp), &list, 10); + testCond(listOrder(list, false)); + + // Reorder the list, and we will try with a tmp buffer too small, + // so API is going to allocate memory. + p = list; + for (i=0; i<10; i++) + p->ai_flags = (i+1); + head = list; + test("Shuffling mem alloc: "); + natsSock_Init(&ctx); + natsSock_ShuffleIPs(&ctx, tmp, 5, &list, 10); + testCond(listOrder(list, false)); + + for (p = list; p != NULL; p = list) + { + list = list->ai_next; + free(p); + } +} + static natsOptions* _createReconnectOptions(void) { @@ -5241,6 +5336,7 @@ test_ServersRandomize(void) natsStatus s; natsOptions *opts = NULL; natsConnection *nc = NULL; + natsPid pid = NATS_INVALID_PID; int serversCount; serversCount = sizeof(testServers) / sizeof(char *); @@ -5330,7 +5426,43 @@ test_ServersRandomize(void) testCond(s == NATS_OK); natsConn_release(nc); + nc = NULL; + + pid = _startServer("nats://127.0.0.1:4222", NULL, true); + CHECK_SERVER_STARTED(pid); + + test("NoRandomize==true passed to context: "); + s = natsOptions_SetNoRandomize(opts, true); + IFOK(s, natsOptions_SetURL(opts, NATS_DEFAULT_URL)); + IFOK(s, natsConnection_Connect(&nc, opts)); + if (s == NATS_OK) + { + natsConn_Lock(nc); + if (!nc->sockCtx.noRandomize) + s = NATS_ERR; + natsConn_Unlock(nc); + } + testCond(s == NATS_OK); + + natsConnection_Destroy(nc); + nc = NULL; + + test("NoRandomize==false passed to context: "); + s = natsOptions_SetNoRandomize(opts, false); + IFOK(s, natsOptions_SetURL(opts, NATS_DEFAULT_URL)); + IFOK(s, natsConnection_Connect(&nc, opts)); + if (s == NATS_OK) + { + natsConn_Lock(nc); + if (nc->sockCtx.noRandomize) + s = NATS_ERR; + natsConn_Unlock(nc); + } + testCond(s == NATS_OK); + + natsConnection_Destroy(nc); natsOptions_Destroy(opts); + _stopServer(pid); } static void @@ -21897,6 +22029,7 @@ static testInfo allTests[] = {"natsInbox", test_natsInbox}, {"natsOptions", test_natsOptions}, {"natsSock_ConnectTcp", test_natsSock_ConnectTcp}, + {"natsSock_ShuffleIPs", test_natsSock_ShuffleIPs}, {"natsSock_IPOrder", test_natsSock_IPOrder}, {"natsSock_ReadLine", test_natsSock_ReadLine}, {"natsJSON", test_natsJSON},