CLOUDP-282641: Use oauth2 token source (#471)

Co-authored-by: Gustavo Bazan <[email protected]>
Co-authored-by: Jeroen Vervaeke <[email protected]>

Author: 3 people

admin/atlas_client.go

@@ -1,12 +1,14 @@
 package admin // import "go.mongodb.org/atlas-sdk/v20241023002/admin"
 
 import (
+	"context"
 	"errors"
 	"net/http"
 	"strings"
 
 	"github.com/mongodb-forks/digest"
-	"go.mongodb.org/atlas-sdk/v20241023002/auth/credentials"
+	"go.mongodb.org/atlas-sdk/v20241023002/auth"
+	"go.mongodb.org/atlas-sdk/v20241023002/auth/clientcredentials"
 	"go.mongodb.org/atlas-sdk/v20241023002/internal/core"
 )
 
@@ -35,7 +37,7 @@ func NewClient(modifiers ...ClientModifier) (*APIClient, error) {
 	return NewAPIClient(defaultConfig), nil
 }
 
-// ClientModifiers lets you create function that controls configuration before creating client.
+// ClientModifier lets you create function that controls configuration before creating client.
 type ClientModifier func(*Configuration) error
 
 // UseDigestAuth provides Digest authentication for Go SDK.
@@ -53,30 +55,26 @@ func UseDigestAuth(apiKey, apiSecret string) ClientModifier {
 	}
 }
 
-// UseOAuthAuth provides OAuthAuth authentication for Go SDK.
+// UseOAuthAuth provides OAuth authentication for Go SDK.
 // Method is provided as helper to create a default HTTP client that supports OAuth (Service Accounts) authentication.
-// credentials.LocalTokenCache can be supplied to reuse OAuth Token across application restarts.
-// Warning: for advanced use cases please use credentials.NewTokenSource directly in your code pass it to UseHTTPClient method.
+//
 // Warning: any previously set httpClient will be overwritten. To fully customize HttpClient use UseHTTPClient method.
-func UseOAuthAuth(clientID, clientSecret string, tokenCache credentials.LocalTokenCache) ClientModifier {
+func UseOAuthAuth(ctx context.Context, clientID, clientSecret string) ClientModifier {
+	ctx2 := ctx
+	if hc := ctx.Value(auth.HTTPClient); hc == nil {
+		client := http.DefaultClient
+		client.Transport = &clientcredentials.Transport{
+			Base: http.DefaultTransport, UserAgent: core.DefaultUserAgent,
+		}
+		ctx2 = context.WithValue(ctx, auth.HTTPClient, client)
+	}
+	oauth := clientcredentials.NewConfig(clientID, clientSecret)
 	return func(c *Configuration) error {
-		var tokenSource credentials.TokenSource
-		if tokenCache != nil {
-			tokenSource = credentials.NewTokenSourceWithOptions(credentials.AtlasTokenSourceOptions{
-				ClientID:     clientID,
-				ClientSecret: clientSecret,
-				TokenCache:   tokenCache,
-				BaseURL:      &c.Servers[0].URL,
-			})
-		} else {
-			tokenSource = credentials.NewTokenSourceWithOptions(credentials.AtlasTokenSourceOptions{
-				ClientID:     clientID,
-				ClientSecret: clientSecret,
-				BaseURL:      &c.Servers[0].URL,
-			})
+		if len(c.Servers) > 0 {
+			oauth.TokenURL = c.Servers[0].URL + clientcredentials.TokenAPIPath
+			oauth.RevokeURL = c.Servers[0].URL + clientcredentials.RevokeAPIPath
 		}
-		httpClient := credentials.NewHTTPClientWithOAuthToken(tokenSource)
-		c.HTTPClient = httpClient
+		c.HTTPClient = oauth.Client(ctx2)
 		return nil
 	}
 }
@@ -86,7 +84,7 @@ func UseOAuthAuth(clientID, clientSecret string, tokenCache credentials.LocalTok
 // UseHTTPClient set custom http client implementation.
 //
 // Warning: UseHTTPClient overrides any previously set httpClient including the one set by UseDigestAuth.
-// To set a custom http client with HTTP diggest support use:
+// To set a custom http client with HTTP digest support use:
 //
 //	transport := digest.NewTransportWithHTTPRoundTripper(apiKey, apiSecret, yourHttpTransport)
 //	client := UseHTTPClient(transport.Client())

auth/auth.go

@@ -0,0 +1,34 @@
+package auth
+
+import "golang.org/x/oauth2"
+
+// oauth2 alias
+var (
+	// HTTPClient is the context key to use with golang.org/x/net/context's
+	// WithValue function to associate an *http.Client value with a context.
+	HTTPClient = oauth2.HTTPClient
+	// ReuseTokenSource returns a TokenSource which repeatedly returns the
+	// same token as long as it's valid, starting with t.
+	// When its cached token is invalid, a new token is obtained from src.
+	//
+	// ReuseTokenSource is typically used to reuse tokens from a cache
+	// (such as a file on disk) between runs of a program, rather than
+	// obtaining new tokens unnecessarily.
+	//
+	// The initial token t may be nil, in which case the TokenSource is
+	// wrapped in a caching version if it isn't one already. This also
+	// means it's always safe to wrap ReuseTokenSource around any other
+	// TokenSource without adverse effects.
+	ReuseTokenSource = oauth2.ReuseTokenSource
+	// NewClient creates an *http.Client from a Context and TokenSource.
+	// The returned client is not valid beyond the lifetime of the context.
+	//
+	// Note that if a custom *http.Client is provided via the Context it
+	// is used only for token acquisition and is not used to configure the
+	// *http.Client returned from NewClient.
+	//
+	// As a special case, if src is nil, a non-OAuth2 client is returned
+	// using the provided context. This exists to support related OAuth2
+	// packages.
+	NewClient = oauth2.NewClient
+)

auth/clientcredentials/clientcredentials.go

@@ -0,0 +1,106 @@
+package clientcredentials
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"golang.org/x/oauth2"
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+
+	"go.mongodb.org/atlas-sdk/v20241023002/auth"
+	"go.mongodb.org/atlas-sdk/v20241023002/internal/core"
+	"golang.org/x/oauth2/clientcredentials"
+)
+
+const (
+	// TokenAPIPath for getting OAuth Access Token from server
+	TokenAPIPath = "/api/oauth/token" //nolint:gosec //url only
+	// serverTokenURL for Token Atlas API
+	serverTokenURL = core.DefaultCloudURL + TokenAPIPath
+	// RevokeAPIPath for revoking OAuth Access Token from server
+	RevokeAPIPath = "/api/oauth/revoke"
+
+	// serverRevokeURL for Revoke Atlas API
+	serverRevokeURL = core.DefaultCloudURL + RevokeAPIPath
+	userAgent       = "User-Agent"
+)
+
+func NewConfig(clientID, clientSecret string) *Config {
+	c := &Config{}
+	c.ClientID = clientID
+	c.ClientSecret = clientSecret
+	c.RevokeURL = serverRevokeURL
+	c.TokenURL = serverTokenURL
+	c.AuthStyle = oauth2.AuthStyleInHeader
+	c.userAgent = core.DefaultUserAgent
+
+	return c
+}
+
+// Config describes a 2-legged OAuth2 flow, with both the
+// client application information and the server's endpoint URLs.
+//
+// NOTE: Config values are used only internally
+// and should not be overridden by clients
+type Config struct {
+	clientcredentials.Config
+	RevokeURL string
+	userAgent string
+}
+
+func (c *Config) Client(ctx context.Context) *http.Client {
+	client := c.Config.Client(ctx)
+	client.Transport = &Transport{
+		Base:      client.Transport,
+		UserAgent: core.DefaultUserAgent,
+	}
+	return client
+}
+
+// RevokeToken revokes OAuth Token
+func (c *Config) RevokeToken(ctx context.Context, t *auth.Token) error {
+	if c.RevokeURL == "" {
+		return errors.New("endpoint missing RevokeURL")
+	}
+	if !t.Valid() {
+		return nil // nothing to do
+	}
+	v := url.Values{
+		"token":           {t.AccessToken},
+		"token_type_hint": {"access_token"},
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.RevokeURL, strings.NewReader(v.Encode()))
+	if err != nil {
+		return err
+	}
+	req.SetBasicAuth(url.QueryEscape(c.ClientID), url.QueryEscape(c.ClientSecret))
+	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+	req.Header.Set(userAgent, c.userAgent)
+
+	client := http.DefaultClient
+	resp, err := client.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		if resp.StatusCode == http.StatusTooManyRequests {
+			msg, _ := io.ReadAll(resp.Body)
+			formattedMessage := fmt.Sprintf("%s %s: HTTP %v Detail: %v Reason: %v",
+				http.MethodPost, c.RevokeURL, resp.StatusCode,
+				"Token Revocation request was rate limited", string(msg))
+			return errors.New(formattedMessage)
+		}
+		formattedMessage := fmt.Sprintf("%s %s: HTTP %v Detail: %v Reason: %v",
+			http.MethodPost, c.RevokeURL, resp.StatusCode,
+			"Failed to revoke Access Token when fetching new OAuth Token from remote server",
+			resp.Header.Get("www-authenticate"))
+		return errors.New(formattedMessage)
+	}
+	return nil
+}

auth/clientcredentials/clientcredentials_test.go

@@ -0,0 +1,57 @@
+package clientcredentials
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"go.mongodb.org/atlas-sdk/v20241023002/auth"
+)
+
+// mockOAuthRevokeEndpoint creates a mock OAuth revoke endpoint,
+// that simulates token revocation responses.
+func mockOAuthRevokeEndpoint(statusCode int) *httptest.Server {
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost || r.FormValue("token") == "" {
+			http.Error(w, "invalid request", http.StatusBadRequest)
+			return
+		}
+		w.WriteHeader(statusCode)
+	})
+	return httptest.NewServer(handler)
+}
+
+// Test OAuthTokenSource_RevokeToken_Success tests successful token revocation.
+func TestOAuthTokenSource_RevokeToken_Success(t *testing.T) {
+	mockServer := mockOAuthRevokeEndpoint(http.StatusOK)
+	defer mockServer.Close()
+
+	config := NewConfig("clientID", "clientSecret")
+	config.RevokeURL = mockServer.URL
+	expiry := time.Now().Add(1 * time.Hour)
+	err := config.RevokeToken(context.Background(), &auth.Token{
+		AccessToken: "test",
+		Expiry:      expiry,
+		ExpiresIn:   expiry.Unix(),
+	})
+	assert.NoError(t, err)
+}
+
+// TestOAuthTokenSource_RevokeToken_Failure tests token revocation failure due to unauthorized access.
+func TestOAuthTokenSource_RevokeToken_Failure(t *testing.T) {
+	mockServer := mockOAuthRevokeEndpoint(http.StatusUnauthorized)
+	defer mockServer.Close()
+
+	config := NewConfig("clientID", "clientSecret")
+	expiry := time.Now().Add(1 * time.Hour)
+	err := config.RevokeToken(context.Background(), &auth.Token{
+		AccessToken: "test",
+		Expiry:      expiry,
+		ExpiresIn:   expiry.Unix(),
+	})
+	assert.Error(t, err)
+	assert.ErrorContains(t, err, "Failed to revoke Access Token when fetching new OAuth Token from remote server")
+}

auth/credentials/doc.go renamed to auth/clientcredentials/doc.go

@@ -13,6 +13,6 @@
 // limitations under the License.
 
 /*
-Package credentials provides an SDK internal client_credentials grant implementation https://datatracker.ietf.org/doc/html/rfc6749#section-1.3.4
+Package credentials provide an SDK internal client_credentials grant implementation https://datatracker.ietf.org/doc/html/rfc6749#section-1.3.4
 */
-package credentials
+package clientcredentials

auth/clientcredentials/transport.go

@@ -0,0 +1,25 @@
+package clientcredentials
+
+import (
+	"net/http"
+)
+
+// Transport supplies custom user agent to token requests
+type Transport struct {
+	Base      http.RoundTripper
+	UserAgent string
+}
+
+func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
+	if t.UserAgent != "" {
+		req.Header.Set("User-Agent", t.UserAgent)
+	}
+	return t.base().RoundTrip(req)
+}
+
+func (t *Transport) base() http.RoundTripper {
+	if t.Base != nil {
+		return t.Base
+	}
+	return http.DefaultTransport
+}