SNOW-1538117: Support privatelink environments (#769)

CMakeLists.txt

@@ -57,6 +57,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

include/snowflake/client.h

@@ -88,6 +88,11 @@ extern "C" {
  */
 #define SF_JWT_CNXN_WAIT_TIME 10
 
+/**
+ * Privatelink host suffix.
+ */
+#define PRIVATELINK_HOSTNAME_SUFFIX ".privatelink.snowflakecomputing."
+
 /**
  * Snowflake Data types
  *

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>
@@ -482,15 +483,39 @@ _snowflake_check_connection_parameters(SF_CONNECT *sf) {
     }
 
     char* top_domain = strrchr(sf->host, '.');
+    char host_without_top_domain[1024] = { 0 };
     if (top_domain)
     {
         top_domain++;
+        size_t length = top_domain - sf->host;
+        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] = '\0';
+    }
+
+    //check privatelink
+    if (ends_with(host_without_top_domain, PRIVATELINK_HOSTNAME_SUFFIX))
+    {
+        char url_buf[4096];
+        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",
+            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");

lib/connection.c

@@ -1364,4 +1364,4 @@ 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);
-}
+}

lib/connection.h

@@ -644,7 +644,6 @@ int64 get_retry_timeout(SF_CONNECT *sf);
 * Get current time since epoch in milliseconds
 */
 uint64 sf_get_current_time_millis();
-
 #ifdef __cplusplus
 }
 #endif

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;
+}

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_UTIL_H

tests/CMakeLists.txt

@@ -44,6 +44,7 @@ SET(TESTS_C
         test_unit_oauth
         test_unit_mfa_auth
         test_ocsp_fail_open
+        test_unit_privatelink
         test_multiple_statements
 # FEATURE_INCREASED_MAX_LOB_SIZE_IN_MEMORY is internal switch
 # will enable lob test when the change on server side will be published

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;
+}