Repository: snowflakedb/gosnowflake
Pull request #1337 (2025-03-20 08:51:30)

SNOW-1958566 Do fallback for OCSP GET request only when malformed response for POST is returned
### Description

SNOW-1958566 PR provides change to not fallback to GET OCSP request unless POST returns malformed response. PR provides also simple OCSP test server.

### Checklist
- [x] Created tests which fail without the change (if possible)
- [x] Extended the README / documentation, if necessary



Patch:
diff --git a/ocsp.go b/ocsp.go
index 3929c2c..14f68ac 100644
--- a/ocsp.go
+++ b/ocsp.go
@@ -478,9 +478,18 @@ func retryOCSP(
 	}
 	ocspRes, err = ocsp.ParseResponse(ocspResBytes, issuer)
 	if err != nil {
-		logger.WithContext(ctx).Warnf("error when parsing ocsp response: %v", err)
-		logger.WithContext(ctx).Warnf("performing GET fallback request to OCSP")
-		return fallbackRetryOCSPToGETRequest(ctx, client, req, ocspHost, headers, issuer, totalTimeout)
+		_, ok1 := err.(asn1.StructuralError)
+		_, ok2 := err.(asn1.SyntaxError)
+		if ok1 || ok2 {
+			logger.WithContext(ctx).Warnf("error when parsing ocsp response: %v", err)
+			logger.WithContext(ctx).Warnf("performing GET fallback request to OCSP")
+			return fallbackRetryOCSPToGETRequest(ctx, client, req, ocspHost, headers, issuer, totalTimeout)
+		}
+		logger.Warnf("Unknown response status from OCSP responder: %v", err)
+		return nil, nil, &ocspStatus{
+			code: ocspStatusUnknown,
+			err:  err,
+		}
 	}
 
 	logger.WithContext(ctx).Debugf("OCSP Status from server: %v", printStatus(ocspRes))
@@ -580,7 +589,8 @@ func getRevocationStatus(ctx context.Context, subject, issuer *x509.Certificate)
 	}
 	logger.WithContext(ctx).Infof("cache missed")
 	logger.WithContext(ctx).Infof("OCSP Server: %v", subject.OCSPServer)
-	if len(subject.OCSPServer) == 0 || isTestNoOCSPURL() {
+	testResponderURL := os.Getenv(ocspTestResponderURLEnv)
+	if (len(subject.OCSPServer) == 0 || isTestNoOCSPURL()) && testResponderURL == "" {
 		return &ocspStatus{
 			code: ocspNoServer,
 			err: &SnowflakeError{
@@ -590,7 +600,10 @@ func getRevocationStatus(ctx context.Context, subject, issuer *x509.Certificate)
 			},
 		}
 	}
-	ocspHost := subject.OCSPServer[0]
+	ocspHost := testResponderURL
+	if ocspHost == "" && len(subject.OCSPServer) > 0 {
+		ocspHost = subject.OCSPServer[0]
+	}
 	u, err := url.Parse(ocspHost)
 	if err != nil {
 		return &ocspStatus{
@@ -598,7 +611,6 @@ func getRevocationStatus(ctx context.Context, subject, issuer *x509.Certificate)
 			err:  fmt.Errorf("failed to parse OCSP server host. %v", ocspHost),
 		}
 	}
-	hostnameStr := os.Getenv(ocspTestResponderURLEnv)
 	var hostname string
 	if retryURL := os.Getenv(ocspRetryURLEnv); retryURL != "" {
 		hostname = fmt.Sprintf(retryURL, fullOCSPURL(u), base64.StdEncoding.EncodeToString(ocspReq))
@@ -610,13 +622,6 @@ func getRevocationStatus(ctx context.Context, subject, issuer *x509.Certificate)
 	} else {
 		hostname = fullOCSPURL(u)
 	}
-	if hostnameStr != "" {
-		u0, err := url.Parse(hostnameStr)
-		if err == nil {
-			hostname = u0.Hostname()
-			u = u0
-		}
-	}
 
 	logger.WithContext(ctx).Debugf("Fetching OCSP response from server: %v", u)
 	logger.WithContext(ctx).Debugf("Host in headers: %v", hostname)
@@ -1078,7 +1083,7 @@ func init() {
 }
 
 // snowflakeNoOcspTransport is the transport object that doesn't do certificate revocation check with OCSP.
-var snowflakeNoOcspTransport = &http.Transport{
+var snowflakeNoOcspTransport http.RoundTripper = &http.Transport{
 	MaxIdleConns:    10,
 	IdleConnTimeout: 30 * time.Minute,
 	Proxy:           http.ProxyFromEnvironment,


Test patch:
diff --git a/driver_ocsp_test.go b/driver_ocsp_test.go
index 96b8f6f..df0ba38 100644
--- a/driver_ocsp_test.go
+++ b/driver_ocsp_test.go
@@ -3,8 +3,10 @@
 package gosnowflake
 
 import (
+	"crypto/tls"
 	"crypto/x509"
 	"database/sql"
+	"net/http"
 	"net/url"
 	"os"
 	"strings"
@@ -51,6 +53,7 @@ func cleanup() {
 	unsetenv(ocspTestNoOCSPURLEnv)
 	unsetenv(ocspRetryURLEnv)
 	unsetenv(cacheDirEnv)
+	ocspFailOpen = OCSPFailOpenTrue
 }
 
 // TestOCSPFailOpen just confirms OCSPFailOpenTrue works.
@@ -860,3 +863,114 @@ func TestOCSPFailClosedNoOCSPURL(t *testing.T) {
 		t.Fatalf("should failed to connect %v", err)
 	}
 }
+
+func TestOCSPUnexpectedResponses(t *testing.T) {
+	cleanup()
+	defer cleanup()
+
+	cfg := wiremockHTTPS.connectionConfig()
+	testCertPool := x509.NewCertPool()
+	caBytes, err := os.ReadFile("ci/scripts/ca.der")
+	assertNilF(t, err)
+	certificate, err := x509.ParseCertificate(caBytes)
+	assertNilF(t, err)
+	testCertPool.AddCert(certificate)
+	customCertPoolTransporter := &http.Transport{
+		TLSClientConfig: &tls.Config{
+			RootCAs:               testCertPool,
+			VerifyPeerCertificate: verifyPeerCertificateSerial,
+		},
+		DisableKeepAlives: true,
+	}
+
+	countingRoundTripper := newCountingRoundTripper(customCertPoolTransporter)
+	originalNoOcspTransport := snowflakeNoOcspTransport
+	defer func() {
+		snowflakeNoOcspTransport = originalNoOcspTransport
+	}()
+	snowflakeNoOcspTransport = countingRoundTripper
+	cfg.Transporter = countingRoundTripper
+
+	runSampleQuery := func(cfg *Config) {
+		connector := NewConnector(SnowflakeDriver{}, *cfg)
+		db := sql.OpenDB(connector)
+		rows, err := db.Query("SELECT 1")
+		if err != nil {
+			println(err.Error())
+		}
+		assertNilF(t, err)
+		defer rows.Close()
+		var v int
+		next := rows.Next()
+		assertTrueF(t, next)
+		err = rows.Scan(&v)
+		if err != nil {
+			println(err)
+		}
+		assertNilF(t, err)
+		assertEqualE(t, v, 1)
+	}
+
+	t.Run("should retry when OCSP is not reachable", func(t *testing.T) {
+		countingRoundTripper.reset()
+		assertNilF(t, os.Setenv(ocspTestResponderURLEnv, "http://localhost:56734")) // not existing port
+		wiremock.registerMappings(t, wiremockMapping{filePath: "select1.json"},
+			wiremockMapping{filePath: "auth/password/successful_flow.json"},
+		)
+		runSampleQuery(cfg)
+		assertTrueE(t, countingRoundTripper.postReqCount["http://localhost:56734"] > 1)
+		assertEqualE(t, countingRoundTripper.getReqCount["http://localhost:56734"], 0)
+	})
+
+	t.Run("should fallback to GET when POST returns malformed response", func(t *testing.T) {
+		countingRoundTripper.reset()
+		assertNilF(t, os.Setenv(ocspTestResponderURLEnv, wiremock.baseURL()))
+		wiremock.registerMappings(t, wiremockMapping{filePath: "ocsp/malformed.json"},
+			wiremockMapping{filePath: "select1.json"},
+			wiremockMapping{filePath: "auth/password/successful_flow.json"},
+		)
+		runSampleQuery(cfg)
+		assertEqualE(t, countingRoundTripper.postReqCount[wiremock.baseURL()], 3)
+		assertEqualE(t, countingRoundTripper.getReqCount[wiremock.baseURL()], 3)
+	})
+
+	t.Run("should not fallback to GET when for POST unauthorized is returned", func(t *testing.T) {
+		countingRoundTripper.reset()
+		assertNilF(t, os.Setenv(ocspTestResponderURLEnv, wiremock.baseURL()))
+		wiremock.registerMappings(t, wiremockMapping{filePath: "ocsp/unauthorized.json"},
+			wiremockMapping{filePath: "select1.json"},
+			wiremockMapping{filePath: "auth/password/successful_flow.json"},
+		)
+		runSampleQuery(cfg)
+		assertEqualE(t, countingRoundTripper.postReqCount[wiremock.baseURL()], 3)
+		assertEqualE(t, countingRoundTripper.getReqCount[wiremock.baseURL()], 0)
+	})
+}
+
+type countingRoundTripper struct {
+	delegate     http.RoundTripper
+	postReqCount map[string]int
+	getReqCount  map[string]int
+}
+
+func newCountingRoundTripper(delegate http.RoundTripper) *countingRoundTripper {
+	return &countingRoundTripper{
+		delegate:     delegate,
+		postReqCount: make(map[string]int),
+		getReqCount:  make(map[string]int),
+	}
+}
+
+func (crt *countingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	if req.Method == "GET" {
+		crt.getReqCount[req.URL.String()]++
+	} else if req.Method == "POST" {
+		crt.postReqCount[req.URL.String()]++
+	}
+	return crt.delegate.RoundTrip(req)
+}
+
+func (crt *countingRoundTripper) reset() {
+	crt.getReqCount = make(map[string]int)
+	crt.postReqCount = make(map[string]int)
+}
diff --git a/ocsp_test.go b/ocsp_test.go
index 9a1eb89..cb67ebc 100644
--- a/ocsp_test.go
+++ b/ocsp_test.go
@@ -33,7 +33,7 @@ func TestOCSP(t *testing.T) {
 		"https://sfcdev2.blob.core.windows.net/",
 	}
 
-	transports := []*http.Transport{
+	transports := []http.RoundTripper{
 		snowflakeNoOcspTransport,
 		SnowflakeTransport,
 	}
diff --git a/test_data/wiremock/mappings/auth/password/successful_flow.json b/test_data/wiremock/mappings/auth/password/successful_flow.json
index a69d9e0..5030f08 100644
--- a/test_data/wiremock/mappings/auth/password/successful_flow.json
+++ b/test_data/wiremock/mappings/auth/password/successful_flow.json
@@ -1,9 +1,6 @@
 {
   "mappings": [
     {
-      "scenarioName": "Successful PAT authentication flow",
-      "requiredScenarioState": "Started",
-      "newScenarioState": "Authenticated",
       "request": {
         "urlPathPattern": "/session/v1/login-request.*",
         "method": "POST",
diff --git a/test_data/wiremock/mappings/ocsp/malformed.json b/test_data/wiremock/mappings/ocsp/malformed.json
new file mode 100644
index 0000000..545a826
--- /dev/null
+++ b/test_data/wiremock/mappings/ocsp/malformed.json
@@ -0,0 +1,13 @@
+{
+  "mappings": [
+    {
+      "request": {
+        "urlPathPattern": "/"
+      },
+      "response": {
+        "status": 200,
+        "base64Body": "AQID"
+      }
+    }
+  ]
+}
\ No newline at end of file
diff --git a/test_data/wiremock/mappings/ocsp/unauthorized.json b/test_data/wiremock/mappings/ocsp/unauthorized.json
new file mode 100644
index 0000000..3fbe2f6
--- /dev/null
+++ b/test_data/wiremock/mappings/ocsp/unauthorized.json
@@ -0,0 +1,14 @@
+{
+  "mappings": [
+    {
+      "request": {
+        "urlPathPattern": "/",
+        "method": "POST"
+      },
+      "response": {
+        "status": 200,
+        "base64Body": "MAMKAQY="
+      }
+    }
+  ]
+}
\ No newline at end of file
diff --git a/wiremock_test.go b/wiremock_test.go
index e12d4d6..9cf128c 100644
--- a/wiremock_test.go
+++ b/wiremock_test.go
@@ -156,6 +156,10 @@ func (wm *wiremockClient) mappingsURL() string {
 	return fmt.Sprintf("http://%v:%v/__admin/mappings", wm.host, wm.adminPort)
 }
 
+func (wm *wiremockClient) baseURL() string {
+	return fmt.Sprintf("http://%v:%v", wm.host, wm.port)
+}
+
 func TestQueryViaHttps(t *testing.T) {
 	wiremockHTTPS.registerMappings(t,
 		wiremockMapping{filePath: "auth/password/successful_flow.json"},