Add SetToken() to KeyFlow

core/auth/auth.go

@@ -150,6 +150,19 @@ func KeyAuth(cfg *config.Configuration) (http.RoundTripper, error) {
 		return nil, fmt.Errorf("configuring key authentication: private key could not be found: %w", err)
 	}
 
+	if cfg.TokenCustomUrl == "" {
+		tokenCustomUrl, tokenUrlSet := os.LookupEnv("STACKIT_TOKEN_BASEURL")
+		if tokenUrlSet {
+			cfg.TokenCustomUrl = tokenCustomUrl
+		}
+	}
+	if cfg.JWKSCustomUrl == "" {
+		jwksCustomUrl, jwksUrlSet := os.LookupEnv("STACKIT_JWKS_BASEURL")
+		if jwksUrlSet {
+			cfg.JWKSCustomUrl = jwksCustomUrl
+		}
+	}
+
 	keyCfg := clients.KeyFlowConfig{
 		ServiceAccountKey: cfg.ServiceAccountKey,
 		PrivateKey:        cfg.PrivateKey,

core/clients/key_flow.go

@@ -9,7 +9,6 @@ import (
 	"io"
 	"net/http"
 	"net/url"
-	"os"
 	"strings"
 	"time"
 
@@ -25,11 +24,12 @@ const (
 	PrivateKey            = "STACKIT_PRIVATE_KEY"
 	ServiceAccountKeyPath = "STACKIT_SERVICE_ACCOUNT_KEY_PATH"
 	PrivateKeyPath        = "STACKIT_PRIVATE_KEY_PATH"
+	tokenAPI              = "https://service-account.api.stackit.cloud/token" //nolint:gosec // linter false positive
+	jwksAPI               = "https://service-account.api.stackit.cloud/.well-known/jwks.json"
+	defaultTokenType      = "Bearer"
+	defaultScope          = ""
 )
 
-var tokenAPI = "https://service-account.api.stackit.cloud/token" //nolint:gosec // linter false positive
-var jwksAPI = "https://service-account.api.stackit.cloud/.well-known/jwks.json"
-
 // KeyFlow handles auth with SA key
 type KeyFlow struct {
 	client        *http.Client
@@ -108,21 +108,13 @@ func (c *KeyFlow) Init(cfg *KeyFlowConfig) error {
 	c.token = &TokenResponseBody{}
 	c.config = cfg
 	c.doer = Do
+
+	// set defaults if no custom token and jwks url are provided
 	if c.config.TokenUrl == "" {
-		tokenCustomUrl, tokenUrlSet := os.LookupEnv("STACKIT_TOKEN_BASEURL")
-		if !tokenUrlSet || tokenCustomUrl == "" {
-			c.config.TokenUrl = tokenAPI
-		} else {
-			c.config.TokenUrl = tokenCustomUrl
-		}
+		c.config.TokenUrl = tokenAPI
 	}
 	if c.config.JWKSUrl == "" {
-		jwksCustomUrl, jwksUrlSet := os.LookupEnv("STACKIT_JWKS_BASEURL")
-		if !jwksUrlSet || jwksCustomUrl == "" {
-			c.config.JWKSUrl = jwksAPI
-		} else {
-			c.config.TokenUrl = jwksCustomUrl
-		}
+		c.config.JWKSUrl = jwksAPI
 	}
 	c.configureHTTPClient()
 	if c.config.ClientRetry == nil {
@@ -131,6 +123,30 @@ func (c *KeyFlow) Init(cfg *KeyFlowConfig) error {
 	return c.validate()
 }
 
+// SetToken can be used to set an access and refresh token manually in the client
+// the other fields in the token field are determined by inspecting the token or setting default values
+func (c *KeyFlow) SetToken(accessToken, refreshToken string) error {
+	// We can safely use ParseUnverified because we are not authenticating the user,
+	// We are parsing the token just to get the expiration time claim
+	parsedAccessToken, _, err := jwt.NewParser().ParseUnverified(accessToken, &jwt.RegisteredClaims{})
+	if err != nil {
+		return fmt.Errorf("parse access token to read expiration time: %w", err)
+	}
+	exp, err := parsedAccessToken.Claims.GetExpirationTime()
+	if err != nil {
+		return fmt.Errorf("get expiration time from access token: %w", err)
+	}
+
+	c.token = &TokenResponseBody{
+		AccessToken:  accessToken,
+		ExpiresIn:    int(exp.Time.Unix()),
+		Scope:        defaultScope,
+		RefreshToken: refreshToken,
+		TokenType:    defaultTokenType,
+	}
+	return nil
+}
+
 // Clone creates a clone of the client
 func (c *KeyFlow) Clone() interface{} {
 	sc := *c

core/clients/key_flow_test.go

@@ -12,7 +12,9 @@ import (
 	"reflect"
 	"strings"
 	"testing"
+	"time"
 
+	"github.com/golang-jwt/jwt/v5"
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/uuid"
 )
@@ -34,7 +36,10 @@ const saKeyStrPattern = `{
 	"validUntil": "2024-03-22T18:05:41Z"
 }`
 
-var saKey = fmt.Sprintf(saKeyStrPattern, uuid.New().String(), uuid.New().String(), uuid.New().String())
+var (
+	saKey          = fmt.Sprintf(saKeyStrPattern, uuid.New().String(), uuid.New().String(), uuid.New().String())
+	testSigningKey = []byte("Test")
+)
 
 func generatePrivateKey() ([]byte, error) {
 	// Generate a new RSA key pair with a size of 2048 bits
@@ -113,6 +118,67 @@ func TestKeyFlowInit(t *testing.T) {
 	}
 }
 
+type MyCustomClaims struct {
+	Foo string `json:"foo"`
+}
+
+func TestSetToken(t *testing.T) {
+	tests := []struct {
+		name         string
+		tokenInvalid bool
+		refreshToken string
+		wantErr      bool
+	}{
+		{
+			name:         "ok",
+			tokenInvalid: false,
+			refreshToken: "refresh_token",
+			wantErr:      false,
+		},
+		{
+			name:         "invalid_token",
+			tokenInvalid: true,
+			refreshToken: "",
+			wantErr:      true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var accessToken string
+			var err error
+			if tt.tokenInvalid {
+				accessToken = "foo"
+			} else {
+				accessTokenJWT := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.RegisteredClaims{
+					ExpiresAt: jwt.NewNumericDate(time.Now().Add(24 * time.Hour))})
+				accessToken, err = accessTokenJWT.SignedString(testSigningKey)
+				if err != nil {
+					t.Fatalf("get test access token as string: %s", err)
+				}
+			}
+
+			c := &KeyFlow{}
+			err = c.SetToken(accessToken, tt.refreshToken)
+
+			if (err != nil) != tt.wantErr {
+				t.Errorf("KeyFlow.SetToken() error = %v, wantErr %v", err, tt.wantErr)
+			}
+			if err == nil {
+				expectedKeyFlowToken := &TokenResponseBody{
+					AccessToken:  accessToken,
+					ExpiresIn:    int(time.Now().Add(24 * time.Hour).Unix()),
+					RefreshToken: tt.refreshToken,
+					Scope:        defaultScope,
+					TokenType:    defaultTokenType,
+				}
+				if !cmp.Equal(expectedKeyFlowToken, c.token) {
+					t.Errorf("The returned result is wrong. Expected %+v, got %+v", expectedKeyFlowToken, c.token)
+				}
+			}
+		})
+	}
+}
+
 func TestKeyClone(t *testing.T) {
 	c := &KeyFlow{
 		client: &http.Client{},