Sanitize host field prior to auth flow (#1385)

Author: pietern

common/client.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"log"
 	"net/http"
+	"net/url"
 	"os"
 	"reflect"
 	"strings"
@@ -237,6 +238,11 @@ func (c *DatabricksClient) Authenticate(ctx context.Context) error {
 	if c.authVisitor != nil {
 		return nil
 	}
+	// Fix host prior to auth, because it may be used in the OIDC flow as "audience" field.
+	// If necessary, this function adds a scheme and strips a trailing slash.
+	if err := c.fixHost(); err != nil {
+		return err
+	}
 	type auth struct {
 		configure func(context.Context) (func(*http.Request) error, error)
 		name      string
@@ -326,12 +332,34 @@ func (c *DatabricksClient) niceAuthError(message string) error {
 	return fmt.Errorf("%s%s. Please check %s for details", message, info, docUrl)
 }
 
-func (c *DatabricksClient) fixHost() {
-	if c.Host != "" && !(strings.HasPrefix(c.Host, "https://") || strings.HasPrefix(c.Host, "http://")) {
-		// azurerm_databricks_workspace.*.workspace_url is giving URL without scheme
-		// so that is why this line is here
-		c.Host = "https://" + c.Host
+func (c *DatabricksClient) fixHost() error {
+	// Nothing to fix if the host isn't set.
+	if c.Host == "" {
+		return nil
+	}
+
+	u, err := url.Parse(c.Host)
+	if err != nil {
+		return err
+	}
+
+	// If the host is empty, assume the scheme wasn't included.
+	if u.Host == "" {
+		u, err = url.Parse("https://" + c.Host)
+		if err != nil {
+			return err
+		}
 	}
+
+	// Create new instance to ensure other fields are initialized as empty.
+	u = &url.URL{
+		Scheme: u.Scheme,
+		Host:   u.Host,
+	}
+
+	// Store sanitized version of c.Host.
+	c.Host = u.String()
+	return nil
 }
 
 func (c *DatabricksClient) configureWithPat(ctx context.Context) (func(*http.Request) error, error) {

common/client_test.go

@@ -297,3 +297,49 @@ func TestConfigAttributeSetNonsense(t *testing.T) {
 	}).Set(&DatabricksClient{}, 1)
 	assert.EqualError(t, err, "cannot set  of unknown type Chan")
 }
+
+func TestDatabricksClientFixHost(t *testing.T) {
+	hostForInput := func(in string) (string, error) {
+		client := &DatabricksClient{
+			Host: in,
+		}
+		if err := client.fixHost(); err != nil {
+			return "", err
+		}
+		return client.Host, nil
+	}
+
+	{
+		// Strip trailing slash.
+		out, err := hostForInput("https://accounts.gcp.databricks.com/")
+		assert.Nil(t, err)
+		assert.Equal(t, out, "https://accounts.gcp.databricks.com")
+	}
+
+	{
+		// Keep port.
+		out, err := hostForInput("https://accounts.gcp.databricks.com:443")
+		assert.Nil(t, err)
+		assert.Equal(t, out, "https://accounts.gcp.databricks.com:443")
+	}
+
+	{
+		// Default scheme.
+		out, err := hostForInput("accounts.gcp.databricks.com")
+		assert.Nil(t, err)
+		assert.Equal(t, out, "https://accounts.gcp.databricks.com")
+	}
+
+	{
+		// Default scheme with port.
+		out, err := hostForInput("accounts.gcp.databricks.com:443")
+		assert.Nil(t, err)
+		assert.Equal(t, out, "https://accounts.gcp.databricks.com:443")
+	}
+
+	{
+		// Return error.
+		_, err := hostForInput("://@@@accounts.gcp.databricks.com/")
+		assert.NotNil(t, err)
+	}
+}

provider/provider_test.go

@@ -198,7 +198,7 @@ func TestConfig_AzurePAT(t *testing.T) {
 		host:        "https://adb-xxx.y.azuredatabricks.net/",
 		token:       "y",
 		assertAzure: true,
-		assertHost:  "https://adb-xxx.y.azuredatabricks.net/",
+		assertHost:  "https://adb-xxx.y.azuredatabricks.net",
 		assertAuth:  "pat",
 	}.apply(t)
 }
@@ -244,7 +244,7 @@ func TestConfig_PatFromDatabricksCfg(t *testing.T) {
 		env: map[string]string{
 			"HOME": "../common/testdata",
 		},
-		assertHost: "https://dbc-XXXXXXXX-YYYY.cloud.databricks.com/",
+		assertHost: "https://dbc-XXXXXXXX-YYYY.cloud.databricks.com",
 		assertAuth: "databricks-cli",
 	}.apply(t)
 }