Copy Key Vault internal module to sdk/security/keyvault

sdk/security/keyvault/internal/CHANGELOG.md

@@ -0,0 +1,12 @@
+# Release History
+
+## 0.6.0 (Unreleased)
+* Initial release from this path. The prior history of this module is in the
+  changelog for `github.com/Azure/azure-sdk-for-go/sdk/keyvault/internal`.
+
+### Breaking Changes
+* `ParseID()` no longer appends a trailing slash to vault URLs
+
+### Bugs Fixed
+
+### Other Changes

sdk/security/keyvault/internal/LICENSE.txt

@@ -0,0 +1,21 @@
+    MIT License
+
+    Copyright (c) Microsoft Corporation. All rights reserved.
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE

sdk/security/keyvault/internal/README.md

@@ -0,0 +1,23 @@
+# Key Vault Internal Module for Go
+
+[![PkgGoDev](https://pkg.go.dev/badge/github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/internal)](https://pkg.go.dev/github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/internal)
+
+This module contains shared code for all the Key Vault SDKs, mainly the challenge authentication policy.
+
+## Contributing
+This project welcomes contributions and suggestions. Most contributions require
+you to agree to a Contributor License Agreement (CLA) declaring that you have
+the right to, and actually do, grant us the rights to use your contribution.
+For details, visit [https://cla.microsoft.com](https://cla.microsoft.com).
+
+When you submit a pull request, a CLA-bot will automatically determine whether
+you need to provide a CLA and decorate the PR appropriately (e.g., label,
+comment). Simply follow the instructions provided by the bot. You will only
+need to do this once across all repos using our CLA.
+
+This project has adopted the
+[Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
+For more information, see the
+[Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/)
+or contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any
+additional questions or comments.

sdk/security/keyvault/internal/challenge_policy.go

@@ -0,0 +1,228 @@
+//go:build go1.18
+// +build go1.18
+
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+package internal
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/streaming"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	"github.com/Azure/azure-sdk-for-go/sdk/internal/errorinfo"
+	"github.com/Azure/azure-sdk-for-go/sdk/internal/temporal"
+)
+
+const headerAuthorization = "Authorization"
+const bearerHeader = "Bearer "
+
+type KeyVaultChallengePolicy struct {
+	// mainResource is the resource to be retrieved using the tenant specified in the credential
+	mainResource *temporal.Resource[azcore.AccessToken, acquiringResourceState]
+	cred         azcore.TokenCredential
+	scope        *string
+	tenantID     *string
+}
+
+func NewKeyVaultChallengePolicy(cred azcore.TokenCredential) *KeyVaultChallengePolicy {
+	return &KeyVaultChallengePolicy{
+		cred:         cred,
+		mainResource: temporal.NewResource(acquire),
+	}
+}
+
+func (k *KeyVaultChallengePolicy) Do(req *policy.Request) (*http.Response, error) {
+	as := acquiringResourceState{
+		p:   k,
+		req: req,
+	}
+
+	if k.scope == nil || k.tenantID == nil {
+		// First request, get both to get the token
+		challengeReq, err := k.getChallengeRequest(*req)
+		if err != nil {
+			return nil, err
+		}
+
+		resp, err := challengeReq.Next()
+		if err != nil {
+			return nil, err
+		}
+
+		if resp.StatusCode > 399 && resp.StatusCode != http.StatusUnauthorized {
+			// the request failed for some other reason, don't try any further
+			return resp, nil
+		}
+		err = k.findScopeAndTenant(resp)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	tk, err := k.mainResource.Get(as)
+	if err != nil {
+		return nil, err
+	}
+
+	req.Raw().Header.Set(
+		headerAuthorization,
+		fmt.Sprintf("%s%s", bearerHeader, tk.Token),
+	)
+
+	// send a copy of the request
+	cloneReq := req.Clone(req.Raw().Context())
+	resp, cloneReqErr := cloneReq.Next()
+	if cloneReqErr != nil {
+		return nil, cloneReqErr
+	}
+
+	// If it fails and has a 401, try it with a new token
+	if resp.StatusCode == 401 {
+		// Force a new token
+		k.mainResource.Expire()
+
+		// Find the scope and tenant again in case they have changed
+		err := k.findScopeAndTenant(resp)
+		if err != nil {
+			// Error parsing challenge, doomed to fail. Return
+			return resp, cloneReqErr
+		}
+
+		tk, err := k.mainResource.Get(as)
+		if err != nil {
+			return resp, err
+		}
+
+		req.Raw().Header.Set(
+			headerAuthorization,
+			bearerHeader+tk.Token,
+		)
+
+		// send the original request now
+		return req.Next()
+	}
+
+	return resp, err
+}
+
+// parses Tenant ID from auth challenge
+// https://login.microsoftonline.com/00000000-0000-0000-0000-000000000000
+func parseTenant(url string) *string {
+	if url == "" {
+		return to.Ptr("")
+	}
+	parts := strings.Split(url, "/")
+	tenant := parts[3]
+	tenant = strings.ReplaceAll(tenant, ",", "")
+	return &tenant
+}
+
+type challengePolicyError struct {
+	err error
+}
+
+func (c *challengePolicyError) Error() string {
+	return c.err.Error()
+}
+
+func (*challengePolicyError) NonRetriable() {
+	// marker method
+}
+
+func (c *challengePolicyError) Unwrap() error {
+	return c.err
+}
+
+var _ errorinfo.NonRetriable = (*challengePolicyError)(nil)
+
+// sets the k.scope and k.tenantID from the WWW-Authenticate header
+func (k *KeyVaultChallengePolicy) findScopeAndTenant(resp *http.Response) error {
+	authHeader := resp.Header.Get("WWW-Authenticate")
+	if authHeader == "" {
+		return &challengePolicyError{err: errors.New("response has no WWW-Authenticate header for challenge authentication")}
+	}
+
+	// Strip down to auth and resource
+	// Format is "Bearer authorization=\"<site>\" resource=\"<site>\"" OR
+	// "Bearer authorization=\"<site>\" scope=\"<site>\" resource=\"<resource>\""
+	authHeader = strings.ReplaceAll(authHeader, "Bearer ", "")
+
+	parts := strings.Split(authHeader, " ")
+
+	vals := map[string]string{}
+	for _, part := range parts {
+		subParts := strings.Split(part, "=")
+		if len(subParts) == 2 {
+			stripped := strings.ReplaceAll(subParts[1], "\"", "")
+			stripped = strings.TrimSuffix(stripped, ",")
+			vals[subParts[0]] = stripped
+		}
+	}
+
+	k.tenantID = parseTenant(vals["authorization"])
+	if scope, ok := vals["scope"]; ok {
+		k.scope = &scope
+	} else if resource, ok := vals["resource"]; ok {
+		if !strings.HasSuffix(resource, "/.default") {
+			resource += "/.default"
+		}
+		k.scope = &resource
+	} else {
+		return &challengePolicyError{err: errors.New("could not find a valid resource in the WWW-Authenticate header")}
+	}
+
+	return nil
+}
+
+func (k KeyVaultChallengePolicy) getChallengeRequest(orig policy.Request) (*policy.Request, error) {
+	req, err := runtime.NewRequest(orig.Raw().Context(), orig.Raw().Method, orig.Raw().URL.String())
+	if err != nil {
+		return nil, &challengePolicyError{err: err}
+	}
+
+	req.Raw().Header = orig.Raw().Header
+	req.Raw().Header.Set("Content-Length", "0")
+	req.Raw().ContentLength = 0
+
+	copied := orig.Clone(orig.Raw().Context())
+	copied.Raw().Body = req.Body()
+	copied.Raw().ContentLength = 0
+	copied.Raw().Header.Set("Content-Length", "0")
+	err = copied.SetBody(streaming.NopCloser(bytes.NewReader([]byte{})), "application/json")
+	if err != nil {
+		return nil, &challengePolicyError{err: err}
+	}
+	copied.Raw().Header.Del("Content-Type")
+
+	return copied, err
+}
+
+type acquiringResourceState struct {
+	req *policy.Request
+	p   *KeyVaultChallengePolicy
+}
+
+// acquire acquires or updates the resource; only one
+// thread/goroutine at a time ever calls this function
+func acquire(state acquiringResourceState) (newResource azcore.AccessToken, newExpiration time.Time, err error) {
+	tk, err := state.p.cred.GetToken(
+		state.req.Raw().Context(),
+		policy.TokenRequestOptions{
+			Scopes: []string{*state.p.scope},
+		},
+	)
+	if err != nil {
+		return azcore.AccessToken{}, time.Time{}, err
+	}
+	return tk, tk.ExpiresOn, nil
+}