SNOW-1859664 use correct transport for calling cloud providers (#1288)

azure_storage_client.go

@@ -52,7 +52,7 @@ func (util *snowflakeAzureClient) createClient(info *execResponseStageInfo, _ bo
 				RetryDelay: 2 * time.Second,
 			},
 			Transport: &http.Client{
-				Transport: SnowflakeTransport,
+				Transport: getTransport(util.cfg),
 			},
 		},
 	})
@@ -74,7 +74,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 +188,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 +273,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 +348,10 @@ 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) {
 	return container.NewClientWithNoCredential(clientURL, &container.ClientOptions{ClientOptions: azcore.ClientOptions{
 		Transport: &http.Client{
-			Transport: SnowflakeTransport,
+			Transport: getTransport(cfg),
 		},
 	}})
 }

client_test.go

@@ -23,7 +23,7 @@ func (t *DummyTransport) RoundTrip(r *http.Request) (*http.Response, error) {
 		}
 		return &http.Response{StatusCode: 200}, nil
 	}
-	return snowflakeInsecureTransport.RoundTrip(r)
+	return snowflakeNoOcspTransport.RoundTrip(r)
 }
 
 func TestInternalClient(t *testing.T) {

connection.go

@@ -790,7 +790,7 @@ func buildSnowflakeConn(ctx context.Context, config Config) (*snowflakeConn, err
 	if sc.cfg.Transporter == nil {
 		if sc.cfg.DisableOCSPChecks || sc.cfg.InsecureMode {
 			// no revocation check with OCSP. Think twice when you want to enable this option.
-			st = snowflakeInsecureTransport
+			st = snowflakeNoOcspTransport
 		} else {
 			// set OCSP fail open mode
 			ocspResponseCacheLock.Lock()
@@ -856,3 +856,21 @@ func buildSnowflakeConn(ctx context.Context, config Config) (*snowflakeConn, err
 
 	return sc, nil
 }
+
+func getTransport(cfg *Config) http.RoundTripper {
+	if cfg == nil {
+		logger.Debug("getTransport: got nil Config, will perform OCSP validation for cloud storage")
+		return SnowflakeTransport
+	}
+	// if user configured a custom Transporter, prioritize that
+	if cfg.Transporter != nil {
+		logger.Debug("getTransport: using Transporter configured by the user")
+		return cfg.Transporter
+	}
+	if cfg.DisableOCSPChecks || cfg.InsecureMode {
+		logger.Debug("getTransport: skipping OCSP validation for cloud storage")
+		return snowflakeNoOcspTransport
+	}
+	logger.Debug("getTransport: will perform OCSP validation for cloud storage")
+	return SnowflakeTransport
+}

connection_test.go

@@ -826,3 +826,56 @@ func TestBeginCreatesTransaction(t *testing.T) {
 		}
 	})
 }
+
+type EmptyTransporter struct{}
+
+func (t EmptyTransporter) RoundTrip(*http.Request) (*http.Response, error) {
+	return nil, nil
+}
+
+func TestGetTransport(t *testing.T) {
+	testcases := []struct {
+		name      string
+		cfg       *Config
+		transport http.RoundTripper
+	}{
+		{
+			name:      "DisableOCSPChecks and InsecureMode false",
+			cfg:       &Config{Account: "one", DisableOCSPChecks: false, InsecureMode: false},
+			transport: SnowflakeTransport,
+		},
+		{
+			name:      "DisableOCSPChecks true and InsecureMode false",
+			cfg:       &Config{Account: "two", DisableOCSPChecks: true, InsecureMode: false},
+			transport: snowflakeNoOcspTransport,
+		},
+		{
+			name:      "DisableOCSPChecks false and InsecureMode true",
+			cfg:       &Config{Account: "three", DisableOCSPChecks: false, InsecureMode: true},
+			transport: snowflakeNoOcspTransport,
+		},
+		{
+			name:      "DisableOCSPChecks and InsecureMode missing from Config",
+			cfg:       &Config{Account: "four"},
+			transport: SnowflakeTransport,
+		},
+		{
+			name:      "whole Config is missing",
+			cfg:       nil,
+			transport: SnowflakeTransport,
+		},
+		{
+			name:      "Using custom Transporter",
+			cfg:       &Config{Account: "five", DisableOCSPChecks: true, InsecureMode: false, Transporter: EmptyTransporter{}},
+			transport: EmptyTransporter{},
+		},
+	}
+	for _, test := range testcases {
+		t.Run(test.name, func(t *testing.T) {
+			result := getTransport(test.cfg)
+			if test.transport != result {
+				t.Errorf("Failed to return the correct transport, input :%#v, expected: %v, got: %v", test.cfg, test.transport, result)
+			}
+		})
+	}
+}

driver_test.go

@@ -1993,7 +1993,7 @@ type CountingTransport struct {
 
 func (t *CountingTransport) RoundTrip(r *http.Request) (*http.Response, error) {
 	t.requests++
-	return snowflakeInsecureTransport.RoundTrip(r)
+	return snowflakeNoOcspTransport.RoundTrip(r)
 }
 
 func TestOpenWithTransport(t *testing.T) {

file_transfer_agent.go

@@ -597,15 +597,15 @@ type s3BucketAccelerateConfigGetter interface {
 
 type s3ClientCreator interface {
 	extractBucketNameAndPath(location string) (*s3Location, error)
-	createClient(info *execResponseStageInfo, useAccelerateEndpoint bool) (cloudClient, error)
+	createClientWithConfig(info *execResponseStageInfo, useAccelerateEndpoint bool, cfg *Config) (cloudClient, error)
 }
 
 func (sfa *snowflakeFileTransferAgent) transferAccelerateConfigWithUtil(s3Util s3ClientCreator) error {
 	s3Loc, err := s3Util.extractBucketNameAndPath(sfa.stageInfo.Location)
 	if err != nil {
 		return err
 	}
-	s3Cli, err := s3Util.createClient(sfa.stageInfo, false)
+	s3Cli, err := s3Util.createClientWithConfig(sfa.stageInfo, false, sfa.sc.cfg)
 	if err != nil {
 		return err
 	}

file_transfer_agent_test.go

@@ -53,15 +53,15 @@ func TestGetBucketAccelerateConfiguration(t *testing.T) {
 
 type s3ClientCreatorMock struct {
 	extract func(string) (*s3Location, error)
-	create  func(info *execResponseStageInfo, useAccelerateEndpoint bool) (cloudClient, error)
+	create  func(info *execResponseStageInfo, useAccelerateEndpoint bool, cfg *Config) (cloudClient, error)
 }
 
 func (mock *s3ClientCreatorMock) extractBucketNameAndPath(location string) (*s3Location, error) {
 	return mock.extract(location)
 }
 
-func (mock *s3ClientCreatorMock) createClient(info *execResponseStageInfo, useAccelerateEndpoint bool) (cloudClient, error) {
-	return mock.create(info, useAccelerateEndpoint)
+func (mock *s3ClientCreatorMock) createClientWithConfig(info *execResponseStageInfo, useAccelerateEndpoint bool, cfg *Config) (cloudClient, error) {
+	return mock.create(info, useAccelerateEndpoint, cfg)
 }
 
 type s3BucketAccelerateConfigGetterMock struct {
@@ -96,7 +96,7 @@ func TestGetBucketAccelerateConfigurationTooManyRetries(t *testing.T) {
 			extract: func(s string) (*s3Location, error) {
 				return &s3Location{bucketName: "test", s3Path: "test"}, nil
 			},
-			create: func(info *execResponseStageInfo, useAccelerateEndpoint bool) (cloudClient, error) {
+			create: func(info *execResponseStageInfo, useAccelerateEndpoint bool, cfg *Config) (cloudClient, error) {
 				return &s3BucketAccelerateConfigGetterMock{err: errors.New("testing")}, nil
 			},
 		})
@@ -146,7 +146,7 @@ func TestGetBucketAccelerateConfigurationFailedCreateClient(t *testing.T) {
 			extract: func(s string) (*s3Location, error) {
 				return &s3Location{bucketName: "test", s3Path: "test"}, nil
 			},
-			create: func(info *execResponseStageInfo, useAccelerateEndpoint bool) (cloudClient, error) {
+			create: func(info *execResponseStageInfo, useAccelerateEndpoint bool, cfg *Config) (cloudClient, error) {
 				return nil, errors.New("failed creation")
 			},
 		})
@@ -172,7 +172,7 @@ func TestGetBucketAccelerateConfigurationInvalidClient(t *testing.T) {
 			extract: func(s string) (*s3Location, error) {
 				return &s3Location{bucketName: "test", s3Path: "test"}, nil
 			},
-			create: func(info *execResponseStageInfo, useAccelerateEndpoint bool) (cloudClient, error) {
+			create: func(info *execResponseStageInfo, useAccelerateEndpoint bool, cfg *Config) (cloudClient, error) {
 				return 1, nil
 			},
 		})

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
@@ -226,7 +226,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
@@ -307,7 +307,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
@@ -409,9 +409,9 @@ func (util *snowflakeGcsClient) isTokenExpired(resp *http.Response) bool {
 	return resp.StatusCode == 401
 }
 
-func newGcsClient() gcsAPI {
+func newGcsClient(cfg *Config) gcsAPI {
 	return &http.Client{
-		Transport: SnowflakeTransport,
+		Transport: getTransport(cfg),
 	}
 }
 

ocsp.go

@@ -637,7 +637,7 @@ func getRevocationStatus(ctx context.Context, subject, issuer *x509.Certificate)
 	}
 	ocspClient := &http.Client{
 		Timeout:   timeout,
-		Transport: snowflakeInsecureTransport,
+		Transport: snowflakeNoOcspTransport,
 	}
 	ocspRes, ocspResBytes, ocspS := retryOCSP(
 		ctx, ocspClient, http.NewRequest, u, headers, ocspReq, issuer, timeout)
@@ -786,7 +786,7 @@ func downloadOCSPCacheServer() {
 	}
 	ocspClient := &http.Client{
 		Timeout:   timeout,
-		Transport: snowflakeInsecureTransport,
+		Transport: snowflakeNoOcspTransport,
 	}
 	ret, ocspStatus := checkOCSPCacheServer(context.Background(), ocspClient, http.NewRequest, u, timeout)
 	if ocspStatus.code != ocspSuccess {
@@ -1075,8 +1075,8 @@ func init() {
 	initOCSPCache()
 }
 
-// snowflakeInsecureTransport is the transport object that doesn't do certificate revocation check.
-var snowflakeInsecureTransport = &http.Transport{
+// snowflakeNoOcspTransport is the transport object that doesn't do certificate revocation check with OCSP.
+var snowflakeNoOcspTransport = &http.Transport{
 	MaxIdleConns:    10,
 	IdleConnTimeout: 30 * time.Minute,
 	Proxy:           http.ProxyFromEnvironment,

ocsp_test.go

@@ -34,7 +34,7 @@ func TestOCSP(t *testing.T) {
 	}
 
 	transports := []*http.Transport{
-		snowflakeInsecureTransport,
+		snowflakeNoOcspTransport,
 		SnowflakeTransport,
 	}
 

s3_storage_client.go

@@ -59,13 +59,20 @@ func (util *snowflakeS3Client) createClient(info *execResponseStageInfo, useAcce
 		BaseEndpoint:  endPoint,
 		UseAccelerate: useAccelerateEndpoint,
 		HTTPClient: &http.Client{
-			Transport: SnowflakeTransport,
+			Transport: getTransport(util.cfg),
 		},
 		ClientLogMode: S3LoggingMode,
 		Logger:        s3Logger,
 	}), nil
 }
 
+// to be used with S3 transferAccelerateConfigWithUtil
+func (util *snowflakeS3Client) createClientWithConfig(info *execResponseStageInfo, useAccelerateEndpoint bool, cfg *Config) (cloudClient, error) {
+	// copy snowflakeFileTransferAgent's config onto the cloud client so we could decide which Transport to use
+	util.cfg = cfg
+	return util.createClient(info, useAccelerateEndpoint)
+}
+
 func getS3CustomEndpoint(info *execResponseStageInfo) *string {
 	var endPoint *string
 	isRegionalURLEnabled := info.UseRegionalURL || info.UseS3RegionalURL