Validate SPOG org-id query values

Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>

Author: msrathore-db

connector.go

@@ -139,7 +139,10 @@ func NewConnector(options ...ConnOption) (driver.Connector, error) {
 
 // clusterPathOrgIDPattern matches the workspace ID inside an all-purpose-compute
 // Thrift path of the form [/]sql/protocolv1/o/<workspace-id>/<cluster-id>[/...].
-var clusterPathOrgIDPattern = regexp.MustCompile(`(?:^|/)sql/protocolv1/o/(\d+)/[^/?]+`)
+var (
+	orgIDPattern            = regexp.MustCompile(`^[0-9]+$`)
+	clusterPathOrgIDPattern = regexp.MustCompile(`^/?sql/protocolv1/o/([0-9]+)/[^/?]+`)
+)
 
 // extractSpogHeaders inspects httpPath for the workspace ID and returns it as an
 // x-databricks-org-id header dict for SPOG routing.
@@ -170,10 +173,15 @@ func extractSpogHeaders(httpPath string) map[string]string {
 				"SPOG header extraction: malformed query string in httpPath, falling back to path inspection: %s",
 				err)
 		} else if orgID := params.Get("o"); orgID != "" {
-			logger.Debug().Msgf(
-				"SPOG header extraction: injecting x-databricks-org-id=%s (extracted from ?o= in httpPath)",
-				orgID)
-			return map[string]string{"x-databricks-org-id": orgID}
+			if !orgIDPattern.MatchString(orgID) {
+				logger.Debug().Msg(
+					"SPOG header extraction: ignoring non-numeric ?o= value in httpPath, falling back to path inspection")
+			} else {
+				logger.Debug().Msgf(
+					"SPOG header extraction: injecting x-databricks-org-id=%s (extracted from ?o= in httpPath)",
+					orgID)
+				return map[string]string{"x-databricks-org-id": orgID}
+			}
 		}
 	}
 

connector_spog_test.go

@@ -48,9 +48,24 @@ func TestExtractSpogHeaders(t *testing.T) {
 			want:     map[string]string{"x-databricks-org-id": "12345"},
 		},
 		{
-			name:     "first o= wins when duplicated",
-			httpPath: "/sql/1.0/warehouses/abc?o=first&o=second",
-			want:     map[string]string{"x-databricks-org-id": "first"},
+			name:     "first numeric o= wins when duplicated",
+			httpPath: "/sql/1.0/warehouses/abc?o=111&o=222",
+			want:     map[string]string{"x-databricks-org-id": "111"},
+		},
+		{
+			name:     "non-numeric o= value returns nil",
+			httpPath: "/sql/1.0/warehouses/abc?o=abc123",
+			want:     nil,
+		},
+		{
+			name:     "control-character o= value returns nil",
+			httpPath: "/sql/1.0/warehouses/abc?o=123%0D%0AX-Injected:%20yes",
+			want:     nil,
+		},
+		{
+			name:     "invalid o= falls back to valid cluster path segment",
+			httpPath: "sql/protocolv1/o/6051921418418893/0528-220959-uzmcn1qt?o=abc123",
+			want:     map[string]string{"x-databricks-org-id": "6051921418418893"},
 		},
 		{
 			name:     "just ? with nothing after returns nil",
@@ -75,6 +90,21 @@ func TestExtractSpogHeaders(t *testing.T) {
 			httpPath: "sql/protocolv1/o/111/0528-220959-uzmcn1qt?o=222",
 			want:     map[string]string{"x-databricks-org-id": "222"},
 		},
+		{
+			name:     "nested cluster path prefix returns nil",
+			httpPath: "evil/sql/protocolv1/o/999/0528-220959-uzmcn1qt",
+			want:     nil,
+		},
+		{
+			name:     "incomplete cluster path returns nil",
+			httpPath: "sql/protocolv1/o/999/",
+			want:     nil,
+		},
+		{
+			name:     "warehouse path containing cluster-looking suffix returns nil",
+			httpPath: "/sql/1.0/warehouses/sql/protocolv1/o/999/cluster-id",
+			want:     nil,
+		},
 		{
 			// Regression guard: the new cluster-path regex must not match
 			// warehouse paths (which never embed the workspace ID).