docs

credentials/oauth/oauth_argument.go

@@ -7,26 +7,33 @@ import (
 )
 
 // OAuthArgument is an interface that provides the necessary information to
-// authenticate with PersistentAuth.
+// authenticate with PersistentAuth. Implementations of this interface must
+// implement either the WorkspaceOAuthArgument or AccountOAuthArgument
+// interface.
 type OAuthArgument interface {
 	// GetCacheKey returns a unique key for the OAuthArgument. This key is used
 	// to store and retrieve the token from the token cache.
 	GetCacheKey(ctx context.Context) string
 }
 
+// WorkspaceOAuthArgument is an interface that provides the necessary information
+// to authenticate using OAuth to a specific workspace.
 type WorkspaceOAuthArgument interface {
 	OAuthArgument
 
 	// GetWorkspaceHost returns the host of the workspace to authenticate to.
 	GetWorkspaceHost(ctx context.Context) string
 }
 
+// BasicWorkspaceOAuthArgument is a basic implementation of the WorkspaceOAuthArgument
+// interface that links each host with exactly one OAuth token.
 type BasicWorkspaceOAuthArgument struct {
 	// host is the host of the workspace to authenticate to. This must start
 	// with "https://" and must not have a trailing slash.
 	host string
 }
 
+// NewBasicWorkspaceOAuthArgument creates a new BasicWorkspaceOAuthArgument.
 func NewBasicWorkspaceOAuthArgument(host string) (BasicWorkspaceOAuthArgument, error) {
 	if !strings.HasPrefix(host, "https://") {
 		return BasicWorkspaceOAuthArgument{}, fmt.Errorf("host must start with 'https://': %s", host)
@@ -37,13 +44,12 @@ func NewBasicWorkspaceOAuthArgument(host string) (BasicWorkspaceOAuthArgument, e
 	return BasicWorkspaceOAuthArgument{host: host}, nil
 }
 
+// GetWorkspaceHost returns the host of the workspace to authenticate to.
 func (a BasicWorkspaceOAuthArgument) GetWorkspaceHost(ctx context.Context) string {
 	return a.host
 }
 
-// key is currently used for two purposes: OIDC URL prefix and token cache key.
-// once we decide to start storing scopes in the token cache, we should change
-// this approach.
+// GetCacheKey returns a unique key for caching the OAuth token for the workspace.
 func (a BasicWorkspaceOAuthArgument) GetCacheKey(ctx context.Context) string {
 	a.host = strings.TrimSuffix(a.host, "/")
 	if !strings.HasPrefix(a.host, "http") {
@@ -54,6 +60,8 @@ func (a BasicWorkspaceOAuthArgument) GetCacheKey(ctx context.Context) string {
 
 var _ WorkspaceOAuthArgument = BasicWorkspaceOAuthArgument{}
 
+// AccountOAuthArgument is an interface that provides the necessary information
+// to authenticate using OAuth to a specific account.
 type AccountOAuthArgument interface {
 	OAuthArgument
 
@@ -64,13 +72,16 @@ type AccountOAuthArgument interface {
 	GetAccountId(ctx context.Context) string
 }
 
+// BasicAccountOAuthArgument is a basic implementation of the AccountOAuthArgument
+// interface that links each account with exactly one OAuth token.
 type BasicAccountOAuthArgument struct {
 	accountHost string
 	accountID   string
 }
 
 var _ AccountOAuthArgument = BasicAccountOAuthArgument{}
 
+// NewBasicAccountOAuthArgument creates a new BasicAccountOAuthArgument.
 func NewBasicAccountOAuthArgument(accountsHost, accountID string) (BasicAccountOAuthArgument, error) {
 	if !strings.HasPrefix(accountsHost, "https://") {
 		return BasicAccountOAuthArgument{}, fmt.Errorf("accountsHost must start with 'https://': %s", accountsHost)
@@ -81,17 +92,17 @@ func NewBasicAccountOAuthArgument(accountsHost, accountID string) (BasicAccountO
 	return BasicAccountOAuthArgument{accountHost: accountsHost, accountID: accountID}, nil
 }
 
+// GetAccountHost returns the host of the account to authenticate to.
 func (a BasicAccountOAuthArgument) GetAccountHost(ctx context.Context) string {
 	return a.accountHost
 }
 
+// GetAccountId returns the account ID of the account to authenticate to.
 func (a BasicAccountOAuthArgument) GetAccountId(ctx context.Context) string {
 	return a.accountID
 }
 
-// key is currently used for two purposes: OIDC URL prefix and token cache key.
-// once we decide to start storing scopes in the token cache, we should change
-// this approach.
+// GetCacheKey returns a unique key for caching the OAuth token for the account.
 func (a BasicAccountOAuthArgument) GetCacheKey(ctx context.Context) string {
 	return fmt.Sprintf("%s/oidc/accounts/%s", a.accountHost, a.accountID)
 }

credentials/oauth/persistent_auth.go

@@ -160,6 +160,8 @@ func (a *PersistentAuth) Load(ctx context.Context, arg OAuthArgument) (t *oauth2
 	a.locker.Lock()
 	defer a.locker.Unlock()
 
+	a.validateArg(arg)
+
 	// TODO: remove this listener after several releases.
 	err = a.startListener(ctx)
 	if err != nil {
@@ -223,6 +225,8 @@ func (a *PersistentAuth) refresh(ctx context.Context, arg OAuthArgument, oldToke
 func (a *PersistentAuth) Challenge(ctx context.Context, arg OAuthArgument) error {
 	a.locker.Lock()
 	defer a.locker.Unlock()
+
+	a.validateArg(arg)
 	err := a.startListener(ctx)
 	if err != nil {
 		return fmt.Errorf("starting listener: %w", err)
@@ -281,18 +285,30 @@ func (a *PersistentAuth) Close() error {
 	return a.ln.Close()
 }
 
+func (a *PersistentAuth) validateArg(arg OAuthArgument) error {
+	_, isWorkspaceArg := arg.(WorkspaceOAuthArgument)
+	_, isAccountArg := arg.(AccountOAuthArgument)
+	if !isWorkspaceArg && !isAccountArg {
+		return fmt.Errorf("unsupported OAuthArgument type: %T, must implement either WorkspaceOAuthArgument or AccountOAuthArgument interface", arg)
+	}
+	return nil
+}
+
 func (a *PersistentAuth) oauth2Config(ctx context.Context, arg OAuthArgument) (*oauth2.Config, error) {
 	scopes := []string{
 		"offline_access", // ensures OAuth token includes refresh token
 		"all-apis",       // ensures OAuth token has access to all control-plane APIs
 	}
 	var endpoints *OAuthAuthorizationServer
 	var err error
-	if workspaceArg, ok := arg.(WorkspaceOAuthArgument); ok {
-		endpoints, err = a.client.GetWorkspaceOAuthEndpoints(ctx, workspaceArg.GetWorkspaceHost(ctx))
-	} else if accountArg, ok := arg.(AccountOAuthArgument); ok {
+	switch argg := arg.(type) {
+	case WorkspaceOAuthArgument:
+		endpoints, err = a.client.GetWorkspaceOAuthEndpoints(ctx, argg.GetWorkspaceHost(ctx))
+	case AccountOAuthArgument:
 		endpoints, err = a.client.GetAccountOAuthEndpoints(
-			ctx, accountArg.GetAccountHost(ctx), accountArg.GetAccountId(ctx))
+			ctx, argg.GetAccountHost(ctx), argg.GetAccountId(ctx))
+	default:
+		return nil, fmt.Errorf("unsupported OAuthArgument type: %T, must implement either WorkspaceOAuthArgument or AccountOAuthArgument interface", arg)
 	}
 	if err != nil {
 		return nil, fmt.Errorf("fetching OAuth endpoints: %w", err)