From baf8587637f7cf3e90fe70b57dd4615b286ee0d4 Mon Sep 17 00:00:00 2001
From: John Yun <140559986+sfc-gh-ext-simba-jy@users.noreply.github.com>
Date: Fri, 25 Oct 2024 14:21:36 -0700
Subject: [PATCH 1/6] implementd privatelink

---
 include/snowflake/client.h    |  2 ++
 lib/client.c                  | 18 +++++++++++
 lib/connection.c              | 13 ++++++++
 lib/connection.h              |  2 ++
 tests/CMakeLists.txt          |  1 +
 tests/test_unit_privatelink.c | 59 +++++++++++++++++++++++++++++++++++
 6 files changed, 95 insertions(+)
 create mode 100644 tests/test_unit_privatelink.c

diff --git a/include/snowflake/client.h b/include/snowflake/client.h
index 1dd5f53b30..dab7f9e44a 100644
--- a/include/snowflake/client.h
+++ b/include/snowflake/client.h
@@ -93,6 +93,8 @@ extern "C" {
  */
 #define SF_JWT_CNXN_WAIT_TIME 10
 
+#define PRIVATELINK_HOSTNAME_SUFFIX ".privatelink.snowflakecomputing."
+
 /**
  * Snowflake Data types
  *
diff --git a/lib/client.c b/lib/client.c
index 587cdb1692..0c633cd8c5 100644
--- a/lib/client.c
+++ b/lib/client.c
@@ -28,6 +28,7 @@
 
 #define curl_easier_escape(curl, string) curl_easy_escape(curl, string, 0)
 
+
 // Define internal constants
 sf_bool DISABLE_VERIFY_PEER;
 char *CA_BUNDLE_FILE;
@@ -475,15 +476,32 @@ _snowflake_check_connection_parameters(SF_CONNECT *sf) {
     }
 
     char* top_domain = strrchr(sf->host, '.');
+    char host_without_top_domain[4096];
     if (top_domain)
     {
         top_domain++;
+        size_t length = strlen(sf->host) - strlen(top_domain);
+        sf_strncpy(host_without_top_domain, sizeof(host_without_top_domain), sf->host,length);
     }
     else
     {
         // It's basically impossible not finding top domain in host.
         // Log the entire host just in case.
         top_domain = sf->host;
+        host_without_top_domain[0] = NULL;
+    }
+
+    //check privatelink
+    if (end_with(host_without_top_domain, PRIVATELINK_HOSTNAME_SUFFIX))
+    {
+        char urlbuf[4096];
+        sf_sprintf(urlbuf, sizeof(urlbuf), "http://ocsp.%s/%s",
+           sf->host, "ocsp_response_cache.json");
+        log_trace(
+            "sf", "Connection", "connect",
+            "Setting SF_OCSP_RESPONSE_CACHE_SERVER_URL to %s",
+            urlbuf);
+        sf_setenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL", urlbuf);
     }
 
     log_info("Connecting to %s Snowflake domain", (strcasecmp(top_domain, "cn") == 0) ? "CHINA" : "GLOBAL");
diff --git a/lib/connection.c b/lib/connection.c
index 94661ac6d9..51c065bd57 100644
--- a/lib/connection.c
+++ b/lib/connection.c
@@ -1344,4 +1344,17 @@ int64 get_retry_timeout(SF_CONNECT *sf)
 int8 get_login_retry_count(SF_CONNECT *sf)
 {
   return (int8)get_less_one(sf->retry_on_connect_count, sf->retry_count);
+}
+
+sf_bool end_with(char* a, char* b)
+{
+    if (strlen(b) > strlen(a))
+    {
+        return SF_BOOLEAN_FALSE;
+    }
+
+    char as[4096];
+    sf_strncpy(as, sizeof(as), a + (strlen(a) - strlen(b)), strlen(b));
+
+    return sf_strncasecmp(as, b, strlen(b)) == 0 ? SF_BOOLEAN_TRUE : SF_BOOLEAN_FALSE;
 }
\ No newline at end of file
diff --git a/lib/connection.h b/lib/connection.h
index fe605526d0..af54b87a65 100644
--- a/lib/connection.h
+++ b/lib/connection.h
@@ -624,6 +624,8 @@ int64 get_retry_timeout(SF_CONNECT *sf);
 */
 uint64 sf_get_current_time_millis();
 
+sf_bool end_with(char* a, char* b);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 615a859af1..77aa622bdf 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -44,6 +44,7 @@ SET(TESTS_C
         test_unit_oauth
         test_unit_mfa_auth
         test_ocsp_fail_open
+        test_unit_privatelink
 # FEATURE_INCREASED_MAX_LOB_SIZE_IN_MEMORY is internal switch
 # will enable lob test when the change on server side will be published
 #        test_lob
diff --git a/tests/test_unit_privatelink.c b/tests/test_unit_privatelink.c
new file mode 100644
index 0000000000..f95d798b00
--- /dev/null
+++ b/tests/test_unit_privatelink.c
@@ -0,0 +1,59 @@
+/*
+ * Copyright (c) 2018-2024 Snowflake Computing, Inc. All rights reserved.
+ */
+
+#include <string.h>
+#include "utils/test_setup.h"
+#include "connection.h"
+#include "memory.h"
+
+/**
+ * Test json body is properly updated.
+ */
+void test_private_link_core(void** unused)
+{
+    char* original_env = getenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL");
+    SF_CONNECT* sf = (SF_CONNECT*)SF_CALLOC(1, sizeof(SF_CONNECT));
+    sf->account = "testaccount";
+    sf->user = "testuser";
+    sf->password = "testpassword";
+    sf->authenticator = SF_AUTHENTICATOR_DEFAULT;
+
+    _snowflake_check_connection_parameters(sf);
+    assert_int_equal(NULL, getenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL"));
+    sf_unsetenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL");
+
+    sf->host = "account.privateLINK.snowflakecomputING.com";
+    _snowflake_check_connection_parameters(sf);
+    assert_string_equal("http://ocsp.account.privateLINK.snowflakecomputING.com/ocsp_response_cache.json", getenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL"));
+    sf_unsetenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL");
+
+    sf->host = "account.snowflakecomputing.com";
+    _snowflake_check_connection_parameters(sf);
+    assert_int_equal(NULL, getenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL"));
+    sf_unsetenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL");
+
+    sf->host = "account.privatelink.snowflakecomputing.cn";
+    _snowflake_check_connection_parameters(sf);
+    assert_string_equal("http://ocsp.account.privatelink.snowflakecomputing.cn/ocsp_response_cache.json", getenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL"));
+    sf_unsetenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL");
+
+
+    if (original_env) {
+        sf_setenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL", original_env);
+    }
+    else {
+        sf_unsetenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL");
+    }
+}
+
+int main(void)
+{
+    initialize_test(SF_BOOLEAN_FALSE);
+    const struct CMUnitTest tests[] = {
+        cmocka_unit_test(test_private_link_core),
+    };
+    int ret = cmocka_run_group_tests(tests, NULL, NULL);
+    snowflake_global_term();
+    return ret;
+}

From cc225bc8b000186d8f755a5c9c8153626a0a0e32 Mon Sep 17 00:00:00 2001
From: John Yun <140559986+sfc-gh-ext-simba-jy@users.noreply.github.com>
Date: Fri, 25 Oct 2024 15:39:34 -0700
Subject: [PATCH 2/6] fix

---
 include/snowflake/client.h | 3 +++
 lib/client.c               | 5 ++---
 lib/connection.c           | 2 +-
 lib/connection.h           | 3 +++
 4 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/include/snowflake/client.h b/include/snowflake/client.h
index dab7f9e44a..864d1900d8 100644
--- a/include/snowflake/client.h
+++ b/include/snowflake/client.h
@@ -93,6 +93,9 @@ extern "C" {
  */
 #define SF_JWT_CNXN_WAIT_TIME 10
 
+/**
+ * Privatelink host suffix.
+ */
 #define PRIVATELINK_HOSTNAME_SUFFIX ".privatelink.snowflakecomputing."
 
 /**
diff --git a/lib/client.c b/lib/client.c
index 0c633cd8c5..f19c7fd722 100644
--- a/lib/client.c
+++ b/lib/client.c
@@ -28,7 +28,6 @@
 
 #define curl_easier_escape(curl, string) curl_easy_escape(curl, string, 0)
 
-
 // Define internal constants
 sf_bool DISABLE_VERIFY_PEER;
 char *CA_BUNDLE_FILE;
@@ -476,7 +475,7 @@ _snowflake_check_connection_parameters(SF_CONNECT *sf) {
     }
 
     char* top_domain = strrchr(sf->host, '.');
-    char host_without_top_domain[4096];
+    char host_without_top_domain[1024];
     if (top_domain)
     {
         top_domain++;
@@ -488,7 +487,7 @@ _snowflake_check_connection_parameters(SF_CONNECT *sf) {
         // It's basically impossible not finding top domain in host.
         // Log the entire host just in case.
         top_domain = sf->host;
-        host_without_top_domain[0] = NULL;
+        host_without_top_domain[0] = '\0';
     }
 
     //check privatelink
diff --git a/lib/connection.c b/lib/connection.c
index 51c065bd57..bc19511b51 100644
--- a/lib/connection.c
+++ b/lib/connection.c
@@ -1353,7 +1353,7 @@ sf_bool end_with(char* a, char* b)
         return SF_BOOLEAN_FALSE;
     }
 
-    char as[4096];
+    char as[1024];
     sf_strncpy(as, sizeof(as), a + (strlen(a) - strlen(b)), strlen(b));
 
     return sf_strncasecmp(as, b, strlen(b)) == 0 ? SF_BOOLEAN_TRUE : SF_BOOLEAN_FALSE;
diff --git a/lib/connection.h b/lib/connection.h
index af54b87a65..90160fe8f9 100644
--- a/lib/connection.h
+++ b/lib/connection.h
@@ -624,6 +624,9 @@ int64 get_retry_timeout(SF_CONNECT *sf);
 */
 uint64 sf_get_current_time_millis();
 
+/**
+ * Validate that the value of a ends with the value of b 
+ */
 sf_bool end_with(char* a, char* b);
 
 #ifdef __cplusplus

From cdc644402d614b3411a8f4191ff90db1dfb2a9ac Mon Sep 17 00:00:00 2001
From: John Yun <140559986+sfc-gh-ext-simba-jy@users.noreply.github.com>
Date: Mon, 28 Oct 2024 09:58:04 -0700
Subject: [PATCH 3/6] change lint

---
 lib/client.c | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/lib/client.c b/lib/client.c
index f19c7fd722..90b5398411 100644
--- a/lib/client.c
+++ b/lib/client.c
@@ -493,14 +493,14 @@ _snowflake_check_connection_parameters(SF_CONNECT *sf) {
     //check privatelink
     if (end_with(host_without_top_domain, PRIVATELINK_HOSTNAME_SUFFIX))
     {
-        char urlbuf[4096];
-        sf_sprintf(urlbuf, sizeof(urlbuf), "http://ocsp.%s/%s",
+        char url_buf[4096];
+        sf_sprintf(url_buf, sizeof(url_buf), "http://ocsp.%s/%s",
            sf->host, "ocsp_response_cache.json");
         log_trace(
             "sf", "Connection", "connect",
             "Setting SF_OCSP_RESPONSE_CACHE_SERVER_URL to %s",
-            urlbuf);
-        sf_setenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL", urlbuf);
+            url_buf);
+        sf_setenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL", url_buf);
     }
 
     log_info("Connecting to %s Snowflake domain", (strcasecmp(top_domain, "cn") == 0) ? "CHINA" : "GLOBAL");

From 8803af5d392e80acfb338add3ceffd3c4a3f1c16 Mon Sep 17 00:00:00 2001
From: John Yun <140559986+sfc-gh-ext-simba-jy@users.noreply.github.com>
Date: Mon, 28 Oct 2024 11:43:52 -0700
Subject: [PATCH 4/6] updated PR based on comment

---
 lib/client.c     | 2 +-
 lib/connection.c | 3 +--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/lib/client.c b/lib/client.c
index 90b5398411..bba136e9cb 100644
--- a/lib/client.c
+++ b/lib/client.c
@@ -479,7 +479,7 @@ _snowflake_check_connection_parameters(SF_CONNECT *sf) {
     if (top_domain)
     {
         top_domain++;
-        size_t length = strlen(sf->host) - strlen(top_domain);
+        size_t length = top_domain - sf->host;
         sf_strncpy(host_without_top_domain, sizeof(host_without_top_domain), sf->host,length);
     }
     else
diff --git a/lib/connection.c b/lib/connection.c
index bc19511b51..b05ad731cb 100644
--- a/lib/connection.c
+++ b/lib/connection.c
@@ -1353,8 +1353,7 @@ sf_bool end_with(char* a, char* b)
         return SF_BOOLEAN_FALSE;
     }
 
-    char as[1024];
-    sf_strncpy(as, sizeof(as), a + (strlen(a) - strlen(b)), strlen(b));
+    char* as = a + (strlen(a) - strlen(b));
 
     return sf_strncasecmp(as, b, strlen(b)) == 0 ? SF_BOOLEAN_TRUE : SF_BOOLEAN_FALSE;
 }
\ No newline at end of file

From 81d7d8307485edc11562ba44618d4df53fd74897 Mon Sep 17 00:00:00 2001
From: John Yun <140559986+sfc-gh-ext-simba-jy@users.noreply.github.com>
Date: Wed, 30 Oct 2024 14:56:26 -0700
Subject: [PATCH 5/6] updated PR based on comments

---
 CMakeLists.txt   |  2 ++
 lib/client.c     | 12 ++++++++++--
 lib/connection.c | 12 ------------
 lib/connection.h |  6 ------
 lib/util.c       | 16 ++++++++++++++++
 lib/util.h       | 19 +++++++++++++++++++
 6 files changed, 47 insertions(+), 20 deletions(-)
 create mode 100644 lib/util.c
 create mode 100644 lib/util.h

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 529f322af2..e29c7d45c1 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -124,6 +124,8 @@ set(SOURCE_FILES
         lib/basic_types.c
         lib/error.h
         lib/error.c
+        lib/util.c
+        lib/util.h
         lib/client_int.h
         lib/chunk_downloader.h
         lib/chunk_downloader.c
diff --git a/lib/client.c b/lib/client.c
index bba136e9cb..4d00f11478 100644
--- a/lib/client.c
+++ b/lib/client.c
@@ -19,6 +19,7 @@
 #include "chunk_downloader.h"
 #include "authenticator.h"
 #include "query_context_cache.h"
+#include "util.h"
 
 #ifdef _WIN32
 #include <Shellapi.h>
@@ -491,11 +492,18 @@ _snowflake_check_connection_parameters(SF_CONNECT *sf) {
     }
 
     //check privatelink
-    if (end_with(host_without_top_domain, PRIVATELINK_HOSTNAME_SUFFIX))
+    if (ends_with(host_without_top_domain, PRIVATELINK_HOSTNAME_SUFFIX))
     {
         char url_buf[4096];
-        sf_sprintf(url_buf, sizeof(url_buf), "http://ocsp.%s/%s",
+        sf_snprintf(url_buf, sizeof(url_buf), sizeof(url_buf)-1, "http://ocsp.%s/%s",
            sf->host, "ocsp_response_cache.json");
+        if (getenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL")) 
+        {
+            log_warn(
+                "sf", "Connection", "connect",
+                "Replace SF_OCSP_RESPONSE_CACHE_SERVER_URL from %s to %s",
+                getenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL"),url_buf);
+        }
         log_trace(
             "sf", "Connection", "connect",
             "Setting SF_OCSP_RESPONSE_CACHE_SERVER_URL to %s",
diff --git a/lib/connection.c b/lib/connection.c
index b05ad731cb..e5a71e9d8a 100644
--- a/lib/connection.c
+++ b/lib/connection.c
@@ -1345,15 +1345,3 @@ int8 get_login_retry_count(SF_CONNECT *sf)
 {
   return (int8)get_less_one(sf->retry_on_connect_count, sf->retry_count);
 }
-
-sf_bool end_with(char* a, char* b)
-{
-    if (strlen(b) > strlen(a))
-    {
-        return SF_BOOLEAN_FALSE;
-    }
-
-    char* as = a + (strlen(a) - strlen(b));
-
-    return sf_strncasecmp(as, b, strlen(b)) == 0 ? SF_BOOLEAN_TRUE : SF_BOOLEAN_FALSE;
-}
\ No newline at end of file
diff --git a/lib/connection.h b/lib/connection.h
index 90160fe8f9..6425d3255e 100644
--- a/lib/connection.h
+++ b/lib/connection.h
@@ -623,12 +623,6 @@ int64 get_retry_timeout(SF_CONNECT *sf);
 * Get current time since epoch in milliseconds
 */
 uint64 sf_get_current_time_millis();
-
-/**
- * Validate that the value of a ends with the value of b 
- */
-sf_bool end_with(char* a, char* b);
-
 #ifdef __cplusplus
 }
 #endif
diff --git a/lib/util.c b/lib/util.c
new file mode 100644
index 0000000000..de8a218075
--- /dev/null
+++ b/lib/util.c
@@ -0,0 +1,16 @@
+#include "util.h"
+#include <string.h>
+
+sf_bool ends_with(char* str, char* suffix)
+{
+    size_t str_length = strlen(str);
+    size_t suffix_length = strlen(suffix);
+    if (suffix_length > str_length)
+    {
+        return SF_BOOLEAN_FALSE;
+    }
+
+    char* str_suffix = str + (str_length - suffix_length);
+
+    return sf_strncasecmp(str_suffix, suffix, suffix_length) == 0;
+}
\ No newline at end of file
diff --git a/lib/util.h b/lib/util.h
new file mode 100644
index 0000000000..12f11b233e
--- /dev/null
+++ b/lib/util.h
@@ -0,0 +1,19 @@
+#ifndef SNOWFLAKE_UTIL_H
+#define SNOWFLAKE_UTIL_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include "../include/snowflake/basic_types.h"
+
+/**
+ * Validate str ends with the suffix
+ */
+sf_bool ends_with(char* str, char* suffix);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif //SNOWFLAKE_CONNECTION_H

From 8b7a2972c9c46be5c08db2bc1850f740ca96dd25 Mon Sep 17 00:00:00 2001
From: John Yun <140559986+sfc-gh-ext-simba-jy@users.noreply.github.com>
Date: Mon, 4 Nov 2024 14:38:56 -0800
Subject: [PATCH 6/6] fix error

---
 lib/util.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/util.h b/lib/util.h
index 12f11b233e..aa24e2ec71 100644
--- a/lib/util.h
+++ b/lib/util.h
@@ -16,4 +16,4 @@ sf_bool ends_with(char* str, char* suffix);
 }
 #endif
 
-#endif //SNOWFLAKE_CONNECTION_H
+#endif //SNOWFLAKE_UTIL_H