API Endpoint from Session Endpoint (#239)

client/session/api_token.go

@@ -18,8 +18,8 @@ import (
 const timePadding = 30 * time.Second
 
 // FromAPIToken creates a session from a ready API token.
-func FromAPIToken(_ context.Context, client *http.Client) func(string) (Session, error) {
-	return func(token string) (Session, error) {
+func FromAPIToken(_ context.Context, client *http.Client) func(string, string) (Session, error) {
+	return func(endpoint, token string) (Session, error) {
 		var claims jwt.RegisteredClaims
 
 		_, _, err := (&jwt.Parser{}).ParseUnverified(token, &claims)
@@ -31,9 +31,14 @@ func FromAPIToken(_ context.Context, client *http.Client) func(string) (Session,
 			return nil, fmt.Errorf("unexpected audience: %v", claims.Audience)
 		}
 
+		apiEndpoint := claims.Audience[0]
+		if endpoint != "" {
+			apiEndpoint = endpoint
+		}
+
 		return &apiToken{
 			client:          client,
-			endpoint:        claims.Audience[0],
+			endpoint:        apiEndpoint,
 			jwt:             token,
 			tokenValidUntil: claims.ExpiresAt.Time,
 			timer:           time.Now,

client/session/from_environment.go

@@ -50,7 +50,7 @@ func FromEnvironment(ctx context.Context, client *http.Client) func(func(string)
 		}
 
 		if token, ok := lookup(EnvSpaceliftAPIToken); ok && token != "" {
-			return FromAPIToken(ctx, client)(token)
+			return FromAPIToken(ctx, client)("", token)
 		}
 
 		endpoint, ok := lookup(EnvSpaceliftAPIKeyEndpoint)

client/session/stored_credentials.go

@@ -49,7 +49,7 @@ func (s *StoredCredentials) Session(ctx context.Context, client *http.Client) (S
 	case CredentialsTypeGitHubToken:
 		return FromGitHubToken(ctx, client)(s.Endpoint, s.AccessToken)
 	case CredentialsTypeAPIToken:
-		return FromAPIToken(ctx, client)(s.AccessToken)
+		return FromAPIToken(ctx, client)(s.Endpoint, s.AccessToken)
 	default:
 		return nil, fmt.Errorf("unexpected credentials type: %d", s.Type)
 	}