Repository: snowflakedb/gosnowflake
Pull request #442 (2021-06-21 18:57:41)

SNOW-365916 Add Privatelink Revocation Check
### Description
Add privatelink revocation check when building connection.

### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing
- [x] Extended the README / documentation, if necessary


sfc-gh-hchaturvedi: Please make sure the tests pass before merging.

Patch:
diff --git a/connection.go b/connection.go
index a9a221b..faa07d8 100644
--- a/connection.go
+++ b/connection.go
@@ -12,6 +12,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"os"
 	"regexp"
 	"strconv"
 	"strings"
@@ -755,6 +756,18 @@ func (sc *snowflakeConn) buildRowsForRunningQuery(ctx context.Context, qid strin
 	return rows, err
 }
 
+func (sc *snowflakeConn) setupOCSPPrivatelink(app string, host string) error {
+	ocspCacheServer := fmt.Sprintf("http://ocsp.%v/ocsp_response_cache.json", host)
+	if err := os.Setenv(cacheServerURLEnv, ocspCacheServer); err != nil {
+		return err
+	}
+	ocspRetryHost := fmt.Sprintf("http://ocsp.%v/retry/", host) + "%v/%v"
+	if err := os.Setenv(ocspRetryURLEnv, ocspRetryHost); err != nil {
+		return err
+	}
+	return nil
+}
+
 func isAsyncMode(ctx context.Context) bool {
 	val := ctx.Value(asyncMode)
 	if val == nil {
@@ -987,6 +1000,15 @@ func buildSnowflakeConn(ctx context.Context, config Config) (*snowflakeConn, err
 		// use the custom transport
 		st = sc.cfg.Transporter
 	}
+	if strings.HasSuffix(sc.cfg.Host, "privatelink.snowflakecomputing.com") {
+		if err := sc.setupOCSPPrivatelink(sc.cfg.Application, sc.cfg.Host); err != nil {
+			return nil, err
+		}
+	} else {
+		if _, set := os.LookupEnv(cacheServerURLEnv); set {
+			os.Unsetenv(cacheServerURLEnv)
+		}
+	}
 	var tokenAccessor TokenAccessor
 	if sc.cfg.TokenAccessor != nil {
 		tokenAccessor = sc.cfg.TokenAccessor
diff --git a/ocsp.go b/ocsp.go
index e3816e4..5ac27b7 100644
--- a/ocsp.go
+++ b/ocsp.go
@@ -32,22 +32,21 @@ import (
 	"time"
 )
 
-// caRoot includes the CA certificates.
-var caRoot map[string]*x509.Certificate
-
-// certPOol includes the CA certificates.
-var certPool *x509.CertPool
-
-// cacheDir is the location of OCSP response cache file
-var cacheDir = ""
-
-// cacheFileName is the file name of OCSP response cache file
-var cacheFileName = ""
-
-// cacheUpdated is true if the memory cache is updated
-var cacheUpdated = true
+var (
+	// caRoot includes the CA certificates.
+	caRoot map[string]*x509.Certificate
+	// certPOol includes the CA certificates.
+	certPool *x509.CertPool
+	// cacheDir is the location of OCSP response cache file
+	cacheDir = ""
+	// cacheFileName is the file name of OCSP response cache file
+	cacheFileName = ""
+	// cacheUpdated is true if the memory cache is updated
+	cacheUpdated = true
+)
 
-// OCSPFailOpenMode is OCSP fail open mode. OCSPFailOpenTrue by default and may set to ocspModeFailClosed for fail closed mode
+// OCSPFailOpenMode is OCSP fail open mode. OCSPFailOpenTrue by default and may
+// set to ocspModeFailClosed for fail closed mode
 type OCSPFailOpenMode uint32
 
 const (
@@ -82,6 +81,7 @@ const (
 	cacheServerEnabledEnv = "SF_OCSP_RESPONSE_CACHE_SERVER_ENABLED"
 	cacheServerURLEnv     = "SF_OCSP_RESPONSE_CACHE_SERVER_URL"
 	cacheDirEnv           = "SF_OCSP_RESPONSE_CACHE_DIR"
+	ocspRetryURLEnv       = "SF_OCSP_RESPONSE_RETRY_URL"
 )
 
 const (
@@ -89,7 +89,7 @@ const (
 	ocspTestInjectUnknownStatusEnv        = "SF_OCSP_TEST_INJECT_UNKNOWN_STATUS"
 	ocspTestResponseCacheServerTimeoutEnv = "SF_OCSP_TEST_OCSP_RESPONSE_CACHE_SERVER_TIMEOUT"
 	ocspTestResponderTimeoutEnv           = "SF_OCSP_TEST_OCSP_RESPONDER_TIMEOUT"
-	ocspTestResponderURLEnv               = "SF_OCSP_RESPONDER_URL"
+	ocspTestResponderURLEnv               = "SF_OCSP_TEST_RESPONDER_URL"
 	ocspTestNoOCSPURLEnv                  = "SF_OCSP_TEST_NO_OCSP_RESPONDER_URL"
 )
 
@@ -489,7 +489,12 @@ func getRevocationStatus(ctx context.Context, subject, issuer *x509.Certificate)
 		}
 	}
 	hostnameStr := os.Getenv(ocspTestResponderURLEnv)
-	hostname := u.Hostname()
+	var hostname string
+	if retryURL := os.Getenv(ocspRetryURLEnv); retryURL != "" {
+		hostname = fmt.Sprintf(retryURL, u.Hostname(), base64.StdEncoding.EncodeToString(ocspReq))
+	} else {
+		hostname = u.Hostname()
+	}
 	if hostnameStr != "" {
 		u0, err := url.Parse(hostnameStr)
 		if err == nil {


Test patch:
diff --git a/connection_test.go b/connection_test.go
index e609545..bded99a 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -11,6 +11,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/url"
+	"os"
 	"strings"
 	"sync"
 	"testing"
@@ -24,10 +25,12 @@ const (
 	serviceNameAppend = "a"
 )
 
-// postQueryMock would generate a response based on the X-Snowflake-Service header, to generate a response
-// with the SERVICE_NAME field appending a character at the end of the header
-// This way it could test both the send and receive logic
-func postQueryMock(_ context.Context, _ *snowflakeRestful, _ *url.Values, headers map[string]string, _ []byte, _ time.Duration, _ uuid.UUID, _ *Config) (*execResponse, error) {
+// postQueryMock generates a response based on the X-Snowflake-Service header,
+// to generate a response with the SERVICE_NAME field appending a character at
+// the end of the header. This way it could test both the send and receive logic
+func postQueryMock(_ context.Context, _ *snowflakeRestful, _ *url.Values,
+	headers map[string]string, _ []byte, _ time.Duration, _ uuid.UUID,
+	_ *Config) (*execResponse, error) {
 	var serviceName string
 	if serviceHeader, ok := headers[httpHeaderServiceName]; ok {
 		serviceName = serviceHeader + serviceNameAppend
@@ -48,7 +51,9 @@ func postQueryMock(_ context.Context, _ *snowflakeRestful, _ *url.Values, header
 
 func TestExecWithEmptyRequestID(t *testing.T) {
 	ctx := WithRequestID(context.Background(), uuid.Nil)
-	postQueryMock := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid.UUID, _ *Config) (*execResponse, error) {
+	postQueryMock := func(_ context.Context, _ *snowflakeRestful,
+		_ *url.Values, _ map[string]string, _ []byte, _ time.Duration,
+		requestID uuid.UUID, _ *Config) (*execResponse, error) {
 		// ensure the same requestID from context is used
 		if len(requestID) == 0 {
 			t.Fatal("requestID is empty")
@@ -70,7 +75,8 @@ func TestExecWithEmptyRequestID(t *testing.T) {
 		cfg:  &Config{Params: map[string]*string{}},
 		rest: sr,
 	}
-	if _, err := sc.exec(ctx, "", false /* noResult */, false /* isInternal */, false /* describeOnly */, nil); err != nil {
+	if _, err := sc.exec(ctx, "", false /* noResult */, false /* isInternal */,
+		false /* describeOnly */, nil); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 }
@@ -79,7 +85,8 @@ func TestGetQueryResultUsesTokenFromTokenAccessor(t *testing.T) {
 	ta := getSimpleTokenAccessor()
 	token := "snowflake-test-token"
 	ta.SetTokens(token, "", 1)
-	funcGetMock := func(_ context.Context, _ *snowflakeRestful, _ *url.URL, headers map[string]string, _ time.Duration) (*http.Response, error) {
+	funcGetMock := func(_ context.Context, _ *snowflakeRestful, _ *url.URL,
+		headers map[string]string, _ time.Duration) (*http.Response, error) {
 		if headers[headerAuthorizationKey] != fmt.Sprintf(headerSnowflakeToken, token) {
 			t.Fatalf("header authorization key is not correct: %v", headers[headerAuthorizationKey])
 		}
@@ -115,7 +122,9 @@ func TestGetQueryResultUsesTokenFromTokenAccessor(t *testing.T) {
 func TestExecWithSpecificRequestID(t *testing.T) {
 	origRequestID := uuid.New()
 	ctx := WithRequestID(context.Background(), origRequestID)
-	postQueryMock := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid.UUID, _ *Config) (*execResponse, error) {
+	postQueryMock := func(_ context.Context, _ *snowflakeRestful,
+		_ *url.Values, _ map[string]string, _ []byte, _ time.Duration,
+		requestID uuid.UUID, _ *Config) (*execResponse, error) {
 		// ensure the same requestID from context is used
 		if requestID != origRequestID {
 			t.Fatal("requestID doesn't match")
@@ -137,7 +146,8 @@ func TestExecWithSpecificRequestID(t *testing.T) {
 		cfg:  &Config{Params: map[string]*string{}},
 		rest: sr,
 	}
-	if _, err := sc.exec(ctx, "", false /* noResult */, false /* isInternal */, false /* describeOnly */, nil); err != nil {
+	if _, err := sc.exec(ctx, "", false /* noResult */, false /* isInternal */,
+		false /* describeOnly */, nil); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 }
@@ -158,10 +168,12 @@ func TestServiceName(t *testing.T) {
 
 	expectServiceName := serviceNameStub
 	for i := 0; i < 5; i++ {
-		sc.exec(context.TODO(), "", false /* noResult */, false /* isInternal */, false /* describeOnly */, nil)
+		sc.exec(context.TODO(), "", false /* noResult */,
+			false /* isInternal */, false /* describeOnly */, nil)
 		if actualServiceName, ok := sc.cfg.Params[serviceName]; ok {
 			if *actualServiceName != expectServiceName {
-				t.Errorf("service name mis-match. expected %v, actual %v", expectServiceName, actualServiceName)
+				t.Errorf("service name mis-match. expected %v, actual %v",
+					expectServiceName, actualServiceName)
 			}
 		} else {
 			t.Error("No service name in the response")
@@ -227,13 +239,13 @@ func TestFetchRunningQueryByID(t *testing.T) {
 }
 
 func TestFetchErrorQueryByID(t *testing.T) {
-	err := &SnowflakeError{
-		Number: ErrQueryReportedError}
-	fetchResultByQueryID(t, returnQueryIsErrStatus, err)
+	fetchResultByQueryID(t, returnQueryIsErrStatus, &SnowflakeError{
+		Number: ErrQueryReportedError})
 }
 
 func customGetQuery(ctx context.Context, rest *snowflakeRestful, url *url.URL,
-	vals map[string]string, _ time.Duration, jsonStr string) (*http.Response, error) {
+	vals map[string]string, _ time.Duration, jsonStr string) (
+	*http.Response, error) {
 	if strings.Contains(url.Path, "/monitoring/queries/") {
 		return &http.Response{
 			StatusCode: http.StatusOK,
@@ -245,21 +257,27 @@ func customGetQuery(ctx context.Context, rest *snowflakeRestful, url *url.URL,
 
 func returnQueryIsRunningStatus(ctx context.Context, rest *snowflakeRestful, fullURL *url.URL,
 	vals map[string]string, duration time.Duration) (*http.Response, error) {
-	var jsonStr = `{"data" : { "queries" : [{"status" : "RUNNING", "state" : "FILE_SET_INITIALIZATION",
-        "errorCode" : 0, "errorMessage" : null}] }, "code" : null, "message" : null, "success" : true }`
+	var jsonStr = `{"data" : { "queries" : [{"status" : "RUNNING", "state" :
+		"FILE_SET_INITIALIZATION", "errorCode" : 0, "errorMessage" : null}] },
+		"code" : null, "message" : null, "success" : true }`
 	return customGetQuery(ctx, rest, fullURL, vals, duration, jsonStr)
 }
 
 func returnQueryIsErrStatus(ctx context.Context, rest *snowflakeRestful, fullURL *url.URL,
 	vals map[string]string, duration time.Duration) (*http.Response, error) {
-	var jsonStr = `{"data" : { "queries" : [{"status" : "FAILED_WITH_ERROR", "errorCode" : 0, "errorMessage" : ""}] },
-       "code" : null, "message" : null, "success" : true }`
+	var jsonStr = `{"data" : { "queries" : [{"status" : "FAILED_WITH_ERROR",
+		"errorCode" : 0, "errorMessage" : ""}] }, "code" : null,
+		"message" : null, "success" : true }`
 	return customGetQuery(ctx, rest, fullURL, vals, duration, jsonStr)
 }
 
 // this function is going to: 1, create a table, 2, query on this table,
-//      3, fetch result of query in step 2, mock running status and error status of that query.
-func fetchResultByQueryID(t *testing.T, customget FuncGetType, expectedFetchErr *SnowflakeError) error {
+// 3, fetch result of query in step 2, mock running status and error status
+// of that query.
+func fetchResultByQueryID(
+	t *testing.T,
+	customget FuncGetType,
+	expectedFetchErr *SnowflakeError) error {
 	config, _ := ParseDSN(dsn)
 	ctx := context.Background()
 	sc, err := buildSnowflakeConn(ctx, *config)
@@ -314,3 +332,24 @@ func fetchResultByQueryID(t *testing.T, customget FuncGetType, expectedFetchErr
 	}
 	return nil
 }
+
+func TestPrivateLink(t *testing.T) {
+	if _, err := buildSnowflakeConn(context.Background(), Config{
+		Account:  "testaccount",
+		User:     "testuser",
+		Password: "testpassword",
+		Host:     "testaccount.us-east-1.privatelink.snowflakecomputing.com",
+	}); err != nil {
+		t.Error(err)
+	}
+	ocspURL := os.Getenv(cacheServerURLEnv)
+	expectedURL := "http://ocsp.testaccount.us-east-1.privatelink.snowflakecomputing.com/ocsp_response_cache.json"
+	if ocspURL != expectedURL {
+		t.Errorf("expected: %v, got: %v", expectedURL, ocspURL)
+	}
+	retryURL := os.Getenv(ocspRetryURLEnv)
+	expectedURL = "http://ocsp.testaccount.us-east-1.privatelink.snowflakecomputing.com/retry/%v/%v"
+	if retryURL != expectedURL {
+		t.Errorf("expected: %v, got: %v", expectedURL, retryURL)
+	}
+}