Skip to content

Commit

Permalink
Merge pull request nats-io#442 from nats-io/ips_shuffling
Browse files Browse the repository at this point in the history 
[CHANGED] Shuffling of IPs after hostname resolution
  • Loading branch information
@kozlovic
kozlovic authored Jul 21, 2021
2 parents 267e1b6 + 83c421c commit 3556a90
Show file tree
Hide file tree
Showing 6 changed files with 206 additions and 1 deletion.
63 changes: 62 additions & 1 deletion src/comsock.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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; i<count; i++)
{
ips[i] = p;
p = p->ai_next;
}
// Shuffle the array
for (i=0; i<count; i++)
{
j = rand() % (i + 1);
p = ips[i];
ips[i] = ips[j];
ips[j] = p;
}
// Relink them
for (i=0; i<count; i++)
{
if (i < count-1)
ips[i]->ai_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
Expand All @@ -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");
Expand Down Expand Up @@ -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));
Expand All @@ -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.
Expand Down
3 changes: 3 additions & 0 deletions src/comsock.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
3 changes: 3 additions & 0 deletions src/conn.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
4 changes: 4 additions & 0 deletions src/natsp.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions test/list.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ natsStrHash
natsInbox
natsOptions
natsSock_ConnectTcp
natsSock_ShuffleIPs
natsSock_IPOrder
natsSock_ReadLine
natsJSON
Expand Down
133 changes: 133 additions & 0 deletions test/test.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4859,13 +4859,108 @@ 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);
_stopServer(serverPid);
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)
{
Expand Down Expand Up @@ -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 *);
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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},
Expand Down

0 comments on commit 3556a90

Please sign in to comment.