Merge pull request #897 from nats-io/fix_870

[ADDED] `natsOptions_SetProxyConnHandler` option

Author: kozlovic

src/comsock.c

@@ -1,4 +1,4 @@
-// Copyright 2015-2021 The NATS Authors
+// Copyright 2015-2025 The NATS Authors
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
@@ -134,6 +134,15 @@ natsSock_ShuffleIPs(natsSockCtx *ctx, struct addrinfo **tmp, int tmpSize, struct
 
 #define MAX_HOST_NAME   (256)
 
+static void
+_resetDeadline(natsSockCtx* ctx, int64_t start, int64_t totalTimeout)
+{
+    int64_t used = nats_Now() - start;
+    int64_t left = totalTimeout - used;
+
+    natsDeadline_Init(&(ctx->writeDeadline), (left > 0 ? left : 0));
+}
+
 natsStatus
 natsSock_ConnectTcp(natsSockCtx *ctx, const char *phost, int port)
 {
@@ -170,12 +179,25 @@ natsSock_ConnectTcp(natsSockCtx *ctx, const char *phost, int port)
 
     snprintf(sport, sizeof(sport), "%d", port);
 
-    if ((ctx->orderIP == 46) || (ctx->orderIP == 64))
-        max = 2;
-
     start = nats_Now();
 
-    for (i=0; i<max; i++)
+    // Call the proxy connect callback if provided
+    if (ctx->proxyConnectCb != NULL)
+    {
+        // Invoke the proxy connect callback.
+        s = ctx->proxyConnectCb(&ctx->fd, host, port, ctx->proxyConnectClosure);
+
+        // If there was a deadline, reset the deadline with whatever is left.
+        if (totalTimeout > 0)
+            _resetDeadline(ctx, start, totalTimeout);
+
+        return NATS_UPDATE_ERR_STACK(s);
+    }
+
+    if ((ctx->orderIP == 46) || (ctx->orderIP == 64))
+	    max = 2;
+
+    for (i = 0; i < max; i++)
     {
         struct addrinfo hints;
         struct addrinfo *servinfo = NULL;
@@ -306,12 +328,7 @@ natsSock_ConnectTcp(natsSockCtx *ctx, const char *phost, int port)
 
     // If there was a deadline, reset the deadline with whatever is left.
     if (totalTimeout > 0)
-    {
-        int64_t used = nats_Now() - start;
-        int64_t left = totalTimeout - used;
-
-        natsDeadline_Init(&(ctx->writeDeadline), (left > 0 ? left : 0));
-    }
+        _resetDeadline(ctx, start, totalTimeout);
 
     return NATS_UPDATE_ERR_STACK(s);
 }

src/conn.c

@@ -1,4 +1,4 @@
-// Copyright 2015-2024 The NATS Authors
+// Copyright 2015-2025 The NATS Authors
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
@@ -403,6 +403,10 @@ _createConn(natsConnection *nc)
     // Set ctx.noRandomize based on public NoRandomize option.
     nc->sockCtx.noRandomize = nc->opts->noRandomize;
 
+    // Set the proxy connect callback and closure.
+    nc->sockCtx.proxyConnectCb      = nc->opts->proxyConnectCb;
+    nc->sockCtx.proxyConnectClosure = nc->opts->proxyConnectClosure;
+
     s = natsSock_ConnectTcp(&(nc->sockCtx), nc->cur->url->host, nc->cur->url->port);
     if (s == NATS_OK)
         nc->sockCtx.fdActive = true;

src/nats.h

@@ -1672,6 +1672,20 @@ typedef void (*natsMsgHandler)(
 typedef void (*natsConnectionHandler)(
         natsConnection  *nc, void *closure);
 
+/** \brief Callback used to handle connections via proxy.
+ *
+ * This callback is used to handle connections via proxy.
+ * It creates a socket, uses it for proxy verification, and returns it in `fd`
+ * to be used for the NATS connection.
+ *
+ * The callback should return #NATS_OK if it has created a socket and returned it
+ * in `fd`, or any other error, such as #NATS_ERR if not.
+ *
+ * @see natsOptions_SetProxyConnHandler
+ */
+typedef natsStatus (*natsProxyConnHandler)(
+    natsSock *fd, char *host, int port, void *closure);
+
 /** \brief Callback used to notify the user of errors encountered while processing
  *         inbound messages.
  *
@@ -2951,6 +2965,21 @@ natsOptions_SetMaxPendingMsgs(natsOptions *opts, int maxPending);
 NATS_EXTERN natsStatus
 natsOptions_SetMaxPendingBytes(natsOptions* opts, int64_t maxPending);
 
+/** \brief Sets the proxy connection handler.
+ *
+ * Specifies the callback to invoke for proxy connection returning the socket to use.
+ *
+ * @see natsProxyConnHandler
+ *
+ * @param opts the pointer to the #natsOptions object.
+ * @param proxyConnHandler the proxy connection handler callback.
+ * @param closure a pointer to an user object that will be passed to
+ * the callback. `closure` can be `NULL`.
+ */
+NATS_EXTERN natsStatus
+natsOptions_SetProxyConnHandler(natsOptions *opts, natsProxyConnHandler proxyConnHandler,
+                                void *closure);
+
 /** \brief Sets the error handler for asynchronous events.
  *
  * Specifies the callback to invoke when an asynchronous error

src/natsp.h

@@ -277,6 +277,9 @@ struct __natsOptions
     natsConnectionHandler   microClosedCb;
     natsErrHandler          microAsyncErrCb;
 
+    natsProxyConnHandler    proxyConnectCb;
+    void                    *proxyConnectClosure;
+
     int64_t                 pingInterval;
     int                     maxPingsOut;
     int                     maxPendingMsgs;
@@ -653,22 +656,25 @@ typedef struct __natsPongList
 
 typedef struct __natsSockCtx
 {
-    natsSock        fd;
-    bool            fdActive;
+    natsSock                fd;
+    bool                    fdActive;
 
-    natsDeadline    readDeadline;
-    natsDeadline    writeDeadline;
+    natsDeadline            readDeadline;
+    natsDeadline            writeDeadline;
 
-    SSL             *ssl;
+    SSL                     *ssl;
 
     // This is true when we are using an external event loop (such as libuv).
-    bool            useEventLoop;
+    bool                    useEventLoop;
 
-    int             orderIP; // possible values: 0,4,6,46,64
+    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;
+    bool                    noRandomize;
+
+    natsProxyConnHandler    proxyConnectCb;
+    void                    *proxyConnectClosure;
 
 } natsSockCtx;
 

src/opts.c

@@ -1,4 +1,4 @@
-// Copyright 2015-2024 The NATS Authors
+// Copyright 2015-2025 The NATS Authors
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
@@ -989,6 +989,19 @@ natsOptions_SetMaxPendingBytes(natsOptions* opts, int64_t maxPending)
     return NATS_OK;
 }
 
+natsStatus
+natsOptions_SetProxyConnHandler(natsOptions* opts, natsProxyConnHandler proxyConnHandler, void *closure)
+{
+    LOCK_AND_CHECK_OPTIONS(opts, 0)
+
+    opts->proxyConnectCb        = proxyConnHandler;
+    opts->proxyConnectClosure   = closure;
+
+    UNLOCK_OPTS(opts);
+
+    return NATS_OK;
+}
+
 natsStatus
 natsOptions_SetErrorHandler(natsOptions *opts, natsErrHandler errHandler,
                             void *closure)

test/list_test.txt

@@ -220,6 +220,7 @@ _test(PingReconnect)
 _test(ProcessMsgArgs)
 _test(ProperFalloutAfterMaxAttempts)
 _test(ProperReconnectDelay)
+_test(ProxyConnectCb)
 _test(PublishMsg)
 _test(PubSubWithReply)
 _test(QueueSubscriber)

test/test.c

@@ -2705,6 +2705,21 @@ _dummyTokenHandler(void *closure)
     return "token";
 }
 
+static natsStatus
+_dummyProxyConnHandler(natsSock *fd, char *host, int port, void *closure)
+{
+    struct threadArg *args = (struct threadArg*) closure;
+
+    if (closure != NULL)
+    {
+        natsMutex_Lock(args->m);
+        args->sum++;
+        natsMutex_Unlock(args->m);
+    }
+
+    return NATS_ERR;
+}
+
 static void
 _dummyErrHandler(natsConnection *nc, natsSubscription *sub, natsStatus err,
                  void *closure)
@@ -3000,6 +3015,16 @@ void test_natsOptions(void)
     s = natsOptions_SetMaxPendingBytes(opts, 1000000);
     testCond((s == NATS_OK) && (opts->maxPendingBytes == 1000000))
 
+    test("Set Proxy Connection Handler: ")
+    s = natsOptions_SetProxyConnHandler(opts, _dummyProxyConnHandler, (void*) 1);
+    testCond((s == NATS_OK) && (opts->proxyConnectCb == _dummyProxyConnHandler)
+                && (opts->proxyConnectClosure == (void*) 1));
+
+    test("Remove Proxy Connection Handler: ")
+    s = natsOptions_SetProxyConnHandler(opts, NULL, NULL);
+    testCond((s == NATS_OK) && (opts->proxyConnectCb == NULL)
+                && (opts->proxyConnectClosure == NULL));
+
     test("Set Error Handler: ");
     s = natsOptions_SetErrorHandler(opts, _dummyErrHandler, NULL);
     testCond((s == NATS_OK) && (opts->asyncErrCb == _dummyErrHandler));
@@ -5952,6 +5977,88 @@ void test_ReconnectServerStats(void)
     _destroyDefaultThreadArgs(&args);
 }
 
+
+static natsStatus
+_proxyConnectCb(natsSock *fd, char *host, int port, void *closure)
+{
+    struct threadArg    *arg    = (struct threadArg*) closure;
+    natsStatus          s       = NATS_OK;
+    natsSockCtx         ctx;
+
+    natsMutex_Lock(arg->m);
+    arg->sum++;
+    natsMutex_Unlock(arg->m);
+
+    s = natsSock_Init(&ctx);
+    IFOK(s, natsSock_ConnectTcp(&ctx, host, port))
+
+    if (s == NATS_OK)
+        *fd = ctx.fd;
+
+    return s;
+}
+
+void test_ProxyConnectCb(void)
+{
+    natsStatus          s;
+    natsConnection      *nc         = NULL;
+    natsOptions         *opts       = NULL;
+    natsPid             serverPid   = NATS_INVALID_PID;
+    const char          *servers[]  = { "nats://127.0.0.1:4222", "nats://127.0.0.1:4223" };
+    struct threadArg    args;
+
+    s = _createDefaultThreadArgsForCbTests(&args);
+    if (s != NATS_OK)
+        FAIL("Unable to setup test")
+
+    serverPid = _startServer("nats://127.0.0.1:4222", NULL, true);
+    CHECK_SERVER_STARTED(serverPid)
+
+    s = natsOptions_Create(&opts);
+    IFOK(s, natsOptions_SetServers(opts, servers, 2));
+    IFOK(s, natsOptions_SetNoRandomize(opts, true));
+    if (s != NATS_OK)
+        FAIL("Unable to create options")
+
+    test("Set connectCb that returns failure: ");
+    s = natsOptions_SetProxyConnHandler(opts, _dummyProxyConnHandler, (void*) &args);
+    testCond(s == NATS_OK);
+
+    test("Connect with connectCb returning failure: ");
+    s = natsConnection_Connect(&nc, opts);
+    testCond(s == NATS_NO_SERVER);
+    nats_clearLastError();
+
+    test("Check invoked properly: ");
+    natsMutex_Lock(args.m);
+    s = (args.sum == 2 ? NATS_OK : NATS_ERR);
+    args.sum = 0;
+    natsMutex_Unlock(args.m);
+    testCond(s == NATS_OK);
+
+    test("Set connectCb that returns success: ");
+    s = natsOptions_SetProxyConnHandler(opts, _proxyConnectCb, (void*) &args);
+    testCond(s == NATS_OK);
+
+    test("Connect with connectCb returning success: ");
+    s = natsConnection_Connect(&nc, opts);
+    testCond(s == NATS_OK);
+
+    test("Check invoked properly: ");
+    natsMutex_Lock(args.m);
+    s = (args.sum == 1 ? NATS_OK : NATS_ERR);
+    args.sum = 0;
+    natsMutex_Unlock(args.m);
+    testCond(s == NATS_OK);
+
+    natsConnection_Destroy(nc);
+    natsOptions_Destroy(opts);
+
+    _stopServer(serverPid);
+
+    _destroyDefaultThreadArgs(&args);
+}
+
 static void
 _disconnectedCb(natsConnection *nc, void *closure)
 {