SNOW-1859664 honour OCSP check settings in cloud storage clients too

snowflakedb · Jan 9, 2025 · c2219f0 · c2219f0
1 parent 8257f91
commit c2219f0
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 12 deletions.
diff --git a/azure_storage_client.go b/azure_storage_client.go
@@ -45,14 +45,15 @@ func (util *snowflakeAzureClient) createClient(info *execResponseStageInfo, _ bo
 	if err != nil {
 		return nil, err
 	}
+	transport := getTransport(util.cfg)
 	client, err := azblob.NewClientWithNoCredential(u.String(), &azblob.ClientOptions{
 		ClientOptions: azcore.ClientOptions{
 			Retry: policy.RetryOptions{
 				MaxRetries: 60,
 				RetryDelay: 2 * time.Second,
 			},
 			Transport: &http.Client{
-				Transport: SnowflakeTransport,
+				Transport: transport,
 			},
 		},
 	})
@@ -74,7 +75,7 @@ func (util *snowflakeAzureClient) getFileHeader(meta *fileMetadata, filename str
 		return nil, err
 	}
 	path := azureLoc.path + strings.TrimLeft(filename, "/")
-	containerClient, err := createContainerClient(client.URL())
+	containerClient, err := createContainerClient(client.URL(), util.cfg)
 	if err != nil {
 		return nil, &SnowflakeError{
 			Message: "failed to create container client",
@@ -188,7 +189,7 @@ func (util *snowflakeAzureClient) uploadFile(
 			Message: "failed to cast to azure client",
 		}
 	}
-	containerClient, err := createContainerClient(client.URL())
+	containerClient, err := createContainerClient(client.URL(), util.cfg)
 
 	if err != nil {
 		return &SnowflakeError{
@@ -273,7 +274,7 @@ func (util *snowflakeAzureClient) nativeDownloadFile(
 			Message: "failed to cast to azure client",
 		}
 	}
-	containerClient, err := createContainerClient(client.URL())
+	containerClient, err := createContainerClient(client.URL(), util.cfg)
 	if err != nil {
 		return &SnowflakeError{
 			Message: "failed to create container client",
@@ -348,10 +349,11 @@ func (util *snowflakeAzureClient) detectAzureTokenExpireError(resp *http.Respons
 		strings.Contains(errStr, "Server failed to authenticate the request")
 }
 
-func createContainerClient(clientURL string) (*container.Client, error) {
+func createContainerClient(clientURL string, cfg *Config) (*container.Client, error) {
+	transport := getTransport(cfg)
 	return container.NewClientWithNoCredential(clientURL, &container.ClientOptions{ClientOptions: azcore.ClientOptions{
 		Transport: &http.Client{
-			Transport: SnowflakeTransport,
+			Transport: transport,
 		},
 	}})
 }
diff --git a/gcs_storage_client.go b/gcs_storage_client.go
@@ -73,7 +73,7 @@ func (util *snowflakeGcsClient) getFileHeader(meta *fileMetadata, filename strin
 			for k, v := range gcsHeaders {
 				req.Header.Add(k, v)
 			}
-			client := newGcsClient()
+			client := newGcsClient(util.cfg)
 			// for testing only
 			if meta.mockGcsClient != nil {
 				client = meta.mockGcsClient
@@ -221,7 +221,7 @@ func (util *snowflakeGcsClient) uploadFile(
 		for k, v := range gcsHeaders {
 			req.Header.Add(k, v)
 		}
-		client := newGcsClient()
+		client := newGcsClient(util.cfg)
 		// for testing only
 		if meta.mockGcsClient != nil {
 			client = meta.mockGcsClient
@@ -302,7 +302,7 @@ func (util *snowflakeGcsClient) nativeDownloadFile(
 		for k, v := range gcsHeaders {
 			req.Header.Add(k, v)
 		}
-		client := newGcsClient()
+		client := newGcsClient(util.cfg)
 		// for testing only
 		if meta.mockGcsClient != nil {
 			client = meta.mockGcsClient
@@ -404,9 +404,10 @@ func (util *snowflakeGcsClient) isTokenExpired(resp *http.Response) bool {
 	return resp.StatusCode == 401
 }
 
-func newGcsClient() gcsAPI {
+func newGcsClient(cfg *Config) gcsAPI {
+	transport := getTransport(cfg)
 	return &http.Client{
-		Transport: SnowflakeTransport,
+		Transport: transport,
 	}
 }
 

diff --git a/s3_storage_client.go b/s3_storage_client.go
@@ -49,6 +49,7 @@ func (util *snowflakeS3Client) createClient(info *execResponseStageInfo, useAcce
 	stageCredentials := info.Creds
 	s3Logger := logging.LoggerFunc(s3LoggingFunc)
 	endPoint := getS3CustomEndpoint(info)
+	transport := getTransport(util.cfg)
 
 	return s3.New(s3.Options{
 		Region: info.Region,
@@ -59,7 +60,7 @@ func (util *snowflakeS3Client) createClient(info *execResponseStageInfo, useAcce
 		BaseEndpoint:  endPoint,
 		UseAccelerate: useAccelerateEndpoint,
 		HTTPClient: &http.Client{
-			Transport: SnowflakeTransport,
+			Transport: transport,
 		},
 		ClientLogMode: S3LoggingMode,
 		Logger:        s3Logger,

diff --git a/util.go b/util.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io"
 	"math/rand"
+	"net/http"
 	"os"
 	"strings"
 	"sync"
@@ -348,3 +349,12 @@ func findByPrefix(in []string, prefix string) int {
 	}
 	return -1
 }
+
+func getTransport(cfg *Config) *http.Transport {
+	if cfg.DisableOCSPChecks || cfg.InsecureMode {
+		logger.Debug("getTransport: won't perform OCSP validation")
+		return snowflakeInsecureTransport
+	}
+	logger.Debug("getTransport: will perform OCSP validation")
+	return SnowflakeTransport
+}