From fc2a57a344576854ca66fb5160346683a46fa2fd Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Thu, 31 Aug 2023 15:38:58 +0200 Subject: [PATCH 1/3] [DECO-2484] Handle Azure authentication when WorkspaceResourceID is provided --- config/auth_azure_cli.go | 46 +++++++++++++++++++++++++++++------ config/auth_azure_cli_test.go | 17 +++++++++++++ 2 files changed, 56 insertions(+), 7 deletions(-) diff --git a/config/auth_azure_cli.go b/config/auth_azure_cli.go index e1b061403..f273e1117 100644 --- a/config/auth_azure_cli.go +++ b/config/auth_azure_cli.go @@ -95,17 +95,49 @@ type internalCliToken struct { ExpiresOn string `json:"expiresOn"` } +func (ts *azureCliTokenSource) getSubscription() string { + if ts.workspaceResourceId == "" { + return "" + } + components := strings.Split(ts.workspaceResourceId, "/") + if len(components) < 3 { + logger.Warnf(context.Background(), "Invalid azure workspace resource ID") + return "" + } + return components[2] +} + func (ts *azureCliTokenSource) Token() (*oauth2.Token, error) { - out, err := exec.Command("az", "account", "get-access-token", "--resource", - ts.resource, "--output", "json").Output() - if ee, ok := err.(*exec.ExitError); ok { - return nil, fmt.Errorf("cannot get access token: %s", string(ee.Stderr)) + subscription := ts.getSubscription() + var out []byte + args := []string{"account", "get-access-token", "--resource", + ts.resource, "--output", "json"} + if subscription != "" { + extendedArgs := make([]string, len(args)) + copy(extendedArgs, args) + extendedArgs = append(extendedArgs, "--subscription", subscription) + // This will fail if the user has access to the workspace, but not to the subscription + // itself. + // In such case, we fall back to not using the subscription. + result, err := exec.Command("az", extendedArgs...).Output() + if err != nil { + logger.Warnf(context.Background(), "Failed to get token for subscription. Using resource only token.") + } else { + out = result + } } - if err != nil { - return nil, fmt.Errorf("cannot get access token: %v", err) + if out == nil { + result, err := exec.Command("az", args...).Output() + if ee, ok := err.(*exec.ExitError); ok { + return nil, fmt.Errorf("cannot get access token: %s", string(ee.Stderr)) + } + if err != nil { + return nil, fmt.Errorf("cannot get access token: %v", err) + } + out = result } var it internalCliToken - err = json.Unmarshal(out, &it) + err := json.Unmarshal(out, &it) if err != nil { return nil, fmt.Errorf("cannot unmarshal CLI result: %w", err) } diff --git a/config/auth_azure_cli_test.go b/config/auth_azure_cli_test.go index 120e5315f..8a6917410 100644 --- a/config/auth_azure_cli_test.go +++ b/config/auth_azure_cli_test.go @@ -15,6 +15,7 @@ import ( var azDummy = &Config{Host: "https://adb-xyz.c.azuredatabricks.net/"} var azDummyWithResourceId = &Config{Host: "https://adb-xyz.c.azuredatabricks.net/", AzureResourceID: "/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123"} +var azDummyWitInvalidResourceId = &Config{Host: "https://adb-xyz.c.azuredatabricks.net/", AzureResourceID: "invalidResourceId"} // testdataPath returns the PATH to use for the duration of a test. // It must only return absolute directories because Go refuses to run @@ -104,6 +105,22 @@ func TestAzureCliCredentials_ValidWithAzureResourceId(t *testing.T) { assert.Equal(t, azDummyWithResourceId.AzureResourceID, r.Header.Get("X-Databricks-Azure-Workspace-Resource-Id")) } +func TestAzureCliCredentials_Fallback(t *testing.T) { + env.CleanupEnvironment(t) + os.Setenv("PATH", testdataPath()) + aa := AzureCliCredentials{} + visitor, err := aa.Configure(context.Background(), azDummyWitInvalidResourceId) + assert.NoError(t, err) + + r := &http.Request{Header: http.Header{}} + err = visitor(r) + assert.NoError(t, err) + + assert.Equal(t, "Bearer ...", r.Header.Get("Authorization")) + assert.Equal(t, azDummyWitInvalidResourceId.AzureResourceID, r.Header.Get("X-Databricks-Azure-Workspace-Resource-Id")) + assert.Equal(t, "...", r.Header.Get("X-Databricks-Azure-SP-Management-Token")) +} + func TestAzureCliCredentials_AlwaysExpired(t *testing.T) { env.CleanupEnvironment(t) os.Setenv("PATH", testdataPath()) From 287fc1c88756ab89ae1f829829726a0a4462cfb7 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 5 Sep 2023 11:22:58 +0200 Subject: [PATCH 2/3] Extract getTokenBytes function --- config/auth_azure_cli.go | 61 +++++++++++++++++++++------------------- 1 file changed, 32 insertions(+), 29 deletions(-) diff --git a/config/auth_azure_cli.go b/config/auth_azure_cli.go index f273e1117..44114f33e 100644 --- a/config/auth_azure_cli.go +++ b/config/auth_azure_cli.go @@ -108,36 +108,12 @@ func (ts *azureCliTokenSource) getSubscription() string { } func (ts *azureCliTokenSource) Token() (*oauth2.Token, error) { - subscription := ts.getSubscription() - var out []byte - args := []string{"account", "get-access-token", "--resource", - ts.resource, "--output", "json"} - if subscription != "" { - extendedArgs := make([]string, len(args)) - copy(extendedArgs, args) - extendedArgs = append(extendedArgs, "--subscription", subscription) - // This will fail if the user has access to the workspace, but not to the subscription - // itself. - // In such case, we fall back to not using the subscription. - result, err := exec.Command("az", extendedArgs...).Output() - if err != nil { - logger.Warnf(context.Background(), "Failed to get token for subscription. Using resource only token.") - } else { - out = result - } - } - if out == nil { - result, err := exec.Command("az", args...).Output() - if ee, ok := err.(*exec.ExitError); ok { - return nil, fmt.Errorf("cannot get access token: %s", string(ee.Stderr)) - } - if err != nil { - return nil, fmt.Errorf("cannot get access token: %v", err) - } - out = result + tokenBytes, err := ts.getTokenBytes() + if err != nil { + return nil, err } var it internalCliToken - err := json.Unmarshal(out, &it) + err = json.Unmarshal(tokenBytes, &it) if err != nil { return nil, fmt.Errorf("cannot unmarshal CLI result: %w", err) } @@ -149,7 +125,7 @@ func (ts *azureCliTokenSource) Token() (*oauth2.Token, error) { ts.resource, it.ExpiresOn) var extra map[string]interface{} - err = json.Unmarshal(out, &extra) + err = json.Unmarshal(tokenBytes, &extra) if err != nil { return nil, fmt.Errorf("cannot unmarshal extra: %w", err) } @@ -160,3 +136,30 @@ func (ts *azureCliTokenSource) Token() (*oauth2.Token, error) { Expiry: expiresOn, }).WithExtra(extra), nil } + +func (ts *azureCliTokenSource) getTokenBytes() ([]byte, error) { + subscription := ts.getSubscription() + args := []string{"account", "get-access-token", "--resource", + ts.resource, "--output", "json"} + if subscription != "" { + extendedArgs := make([]string, len(args)) + copy(extendedArgs, args) + extendedArgs = append(extendedArgs, "--subscription", subscription) + // This will fail if the user has access to the workspace, but not to the subscription + // itself. + // In such case, we fall back to not using the subscription. + result, err := exec.Command("az", extendedArgs...).Output() + if err == nil { + return result, nil + } + logger.Warnf(context.Background(), "Failed to get token for subscription. Using resource only token.") + } + result, err := exec.Command("az", args...).Output() + if ee, ok := err.(*exec.ExitError); ok { + return nil, fmt.Errorf("cannot get access token: %s", string(ee.Stderr)) + } + if err != nil { + return nil, fmt.Errorf("cannot get access token: %v", err) + } + return result, nil +} From 81656ea49381aae1acc21acbd03fddd72e5f390c Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 5 Sep 2023 11:53:04 +0200 Subject: [PATCH 3/3] Fix error wrapping --- config/auth_azure_cli.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/config/auth_azure_cli.go b/config/auth_azure_cli.go index 44114f33e..08b433672 100644 --- a/config/auth_azure_cli.go +++ b/config/auth_azure_cli.go @@ -159,7 +159,7 @@ func (ts *azureCliTokenSource) getTokenBytes() ([]byte, error) { return nil, fmt.Errorf("cannot get access token: %s", string(ee.Stderr)) } if err != nil { - return nil, fmt.Errorf("cannot get access token: %v", err) + return nil, fmt.Errorf("cannot get access token: %w", err) } return result, nil }