feat: add PAT auto-region configuration (#395)

j-luong · web-flow · commit a26581814257 · 2025-08-08T15:12:36.000+01:00
diff --git a/pkg/app/app.go b/pkg/app/app.go
@@ -81,6 +81,7 @@ func defaultFuncOrganization(engine workflow.Engine, config configuration.Config
 func defaultFuncApiUrl(_ configuration.Configuration, logger *zerolog.Logger) configuration.DefaultValueFunction {
 	callback := func(config configuration.Configuration, existingValue interface{}) (interface{}, error) {
 		urlString := constants.SNYK_DEFAULT_API_URL
+		authToken := config.GetString(configuration.AUTHENTICATION_TOKEN)
 
 		urlFromOauthToken, err := auth.GetAudienceClaimFromOauthToken(config.GetString(auth.CONFIG_KEY_OAUTH_TOKEN))
 		if err != nil {
@@ -89,6 +90,14 @@ func defaultFuncApiUrl(_ configuration.Configuration, logger *zerolog.Logger) co
 
 		if len(urlFromOauthToken) > 0 && len(urlFromOauthToken[0]) > 0 {
 			urlString = urlFromOauthToken[0]
+		} else if auth.IsAuthTypePAT(authToken) {
+			apiUrl, claimsErr := auth.GetApiUrlFromPAT(authToken)
+			if claimsErr != nil {
+				logger.Warn().Err(claimsErr).Msg("failed to get api url from pat")
+			}
+			if len(apiUrl) > 0 {
+				urlString = apiUrl
+			}
 		} else if existingValue != nil { // try the configured value as last resort
 			if temp, ok := existingValue.(string); ok {
 				urlString = temp
diff --git a/pkg/app/app_test.go b/pkg/app/app_test.go
@@ -1,6 +1,7 @@
 package app
 
 import (
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"log"
@@ -89,6 +90,40 @@ func Test_CreateAppEngine_config_replaceV1inApi(t *testing.T) {
 	assert.Equal(t, expectApiUrl, actualApiUrl)
 }
 
+func Test_CreateAppEngine_config_PAT_autoRegionDetection(t *testing.T) {
+	t.Run("default", func(t *testing.T) {
+		apiUrl := "api.snyk.io"
+		euPAT := createMockPAT(t, fmt.Sprintf(`{"h":"%s"}`, apiUrl))
+		engine := CreateAppEngine()
+		config := engine.GetConfiguration()
+		config.Set(configuration.AUTHENTICATION_TOKEN, euPAT)
+
+		actualApiUrl := config.GetString(configuration.API_URL)
+		assert.Equal(t, fmt.Sprintf("https://%s", apiUrl), actualApiUrl)
+	})
+
+	t.Run("eu", func(t *testing.T) {
+		apiUrl := "api.eu.snyk.io"
+		euPAT := createMockPAT(t, fmt.Sprintf(`{"h":"%s"}`, apiUrl))
+		engine := CreateAppEngine()
+		config := engine.GetConfiguration()
+		config.Set(configuration.AUTHENTICATION_TOKEN, euPAT)
+
+		actualApiUrl := config.GetString(configuration.API_URL)
+		assert.Equal(t, fmt.Sprintf("https://%s", apiUrl), actualApiUrl)
+	})
+
+	t.Run("invalid PAT reverts to default API URL", func(t *testing.T) {
+		patWithExtraSegments := "snyk_uat.12345678.payload.signature.extra"
+		engine := CreateAppEngine()
+		config := engine.GetConfiguration()
+		config.Set(configuration.AUTHENTICATION_TOKEN, patWithExtraSegments)
+
+		actualApiUrl := config.GetString(configuration.API_URL)
+		assert.Equal(t, constants.SNYK_DEFAULT_API_URL, actualApiUrl)
+	})
+}
+
 func Test_CreateAppEngine_config_OauthAudHasPrecedence(t *testing.T) {
 	config := configuration.New()
 	config.Set(auth.CONFIG_KEY_OAUTH_TOKEN,
@@ -368,3 +403,11 @@ func Test_initConfiguration_DEFAULT_TEMP_DIRECTORY(t *testing.T) {
 		assert.Equal(t, expected, actual)
 	})
 }
+
+func createMockPAT(t *testing.T, payload string) string {
+	t.Helper()
+
+	encodedPayload := base64.RawURLEncoding.EncodeToString([]byte(payload))
+	signature := "signature"
+	return fmt.Sprintf("snyk_uat.12345678.%s.%s", encodedPayload, signature)
+}
diff --git a/pkg/auth/tokenauthenticator.go b/pkg/auth/tokenauthenticator.go
@@ -1,9 +1,12 @@
 package auth
 
 import (
+	"encoding/base64"
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"regexp"
+	"strings"
 
 	"github.com/google/uuid"
 )
@@ -15,8 +18,15 @@ const (
 	CACHED_PAT_IS_VALID_KEY_PREFIX = "cached_pat_is_valid"
 	CONFIG_KEY_TOKEN               = "api"      // the snyk config key for api token
 	CONFIG_KEY_ENDPOINT            = "endpoint" // the snyk config key for api endpoint
+	PAT_REGEX                      = `(snyk_(?:uat|sat))\.([a-z0-9]{8}\.[a-zA-Z0-9-_]+\.[a-zA-Z0-9-_]+)`
 )
 
+// Claims represents the structure of the PATs claims, it does not represent all the claims; only the ones we need
+type Claims struct {
+	// Hostname PAT is valid for
+	Hostname string `json:"h,omitempty"`
+}
+
 var _ Authenticator = (*tokenAuthenticator)(nil)
 
 type tokenAuthenticator struct {
@@ -60,9 +70,40 @@ func IsAuthTypeToken(token string) bool {
 
 func IsAuthTypePAT(token string) bool {
 	// e.g. snyk_uat.1a2b3c4d.mySuperSecret_Token-Value.aChecksum_123-Value
-	patRegex := `^snyk_(?:uat|sat)\.[a-z0-9]{8}\.[a-zA-Z0-9-_]+\.[a-zA-Z0-9-_]+$`
-	if matched, err := regexp.MatchString(patRegex, token); err == nil && matched {
-		return matched
+	return regexp.MustCompile(fmt.Sprintf("^%s$", PAT_REGEX)).MatchString(token)
+}
+
+// extractClaimsFromPAT accepts a raw PAT string and returns the PAT claims
+// differs from the implementation in oauth.go as Snyk PATs do not strictly follow the JWT spec
+func extractClaimsFromPAT(raw string) (*Claims, error) {
+	parts := strings.Split(raw, ".")
+	if len(parts) != 4 {
+		return nil, fmt.Errorf("invalid number of segments: %d", len(parts))
 	}
-	return false
+
+	payload, err := base64.RawURLEncoding.DecodeString(parts[2])
+	if err != nil {
+		return nil, fmt.Errorf("failed to decode payload: %w", err)
+	}
+
+	var c Claims
+	if err = json.Unmarshal(payload, &c); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal payload: %w", err)
+	}
+
+	return &c, nil
+}
+
+func GetApiUrlFromPAT(pat string) (string, error) {
+	claims, err := extractClaimsFromPAT(pat)
+	if err != nil {
+		return "", err
+	}
+
+	hostname := claims.Hostname
+	if !strings.HasPrefix(hostname, "http") {
+		hostname = fmt.Sprintf("https://%s", hostname)
+	}
+
+	return hostname, nil
 }
diff --git a/pkg/auth/tokenauthenticator_test.go b/pkg/auth/tokenauthenticator_test.go
@@ -1,6 +1,8 @@
 package auth
 
 import (
+	"encoding/base64"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -17,3 +19,95 @@ func TestIsAuthTypePAT(t *testing.T) {
 	// legacy token format
 	assert.False(t, IsAuthTypePAT("f47ac10b-58cc-4372-a567-0e02b2c3d479"))
 }
+
+func TestExtractClaimsFromPAT(t *testing.T) {
+	t.Run("Valid PAT", func(t *testing.T) {
+		pat := createMockPAT(t, `{"h":"api.snyk.io"}`)
+
+		claims, err := extractClaimsFromPAT(pat)
+		assert.NoError(t, err)
+		assert.NotNil(t, claims)
+		assert.Equal(t, "api.snyk.io", claims.Hostname)
+	})
+
+	t.Run("Valid EU PAT", func(t *testing.T) {
+		pat := createMockPAT(t, `{"h":"api.eu.snyk.io"}`)
+
+		claims, err := extractClaimsFromPAT(pat)
+		assert.NoError(t, err)
+		assert.NotNil(t, claims)
+		assert.Equal(t, "api.eu.snyk.io", claims.Hostname)
+	})
+
+	t.Run("PAT with fewer than 4 segments", func(t *testing.T) {
+		pat := "snyk_test.12345678.payload"
+		claims, err := extractClaimsFromPAT(pat)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "invalid number of segments: 3")
+		assert.Nil(t, claims)
+	})
+
+	t.Run("PAT with more than 4 segments", func(t *testing.T) {
+		pat := "snyk_test.12345678.payload.signature.extra"
+		claims, err := extractClaimsFromPAT(pat)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "invalid number of segments: 5")
+		assert.Nil(t, claims)
+	})
+
+	t.Run("PAT with invalid base64 payload", func(t *testing.T) {
+		pat := "snyk_test.12345678.invalid-base64!@#$.signature"
+		claims, err := extractClaimsFromPAT(pat)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "failed to decode payload")
+		assert.Nil(t, claims)
+	})
+
+	t.Run("PAT with invalid JSON payload", func(t *testing.T) {
+		pat := createMockPAT(t, `{"j":"pat-id-123", "h":"api.snyk.io", "e":1678886400, "s":"sub-id-456`)
+
+		claims, err := extractClaimsFromPAT(pat)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "failed to unmarshal payload")
+		assert.Nil(t, claims)
+	})
+}
+
+func TestGetApiUrlFromPAT(t *testing.T) {
+	t.Run("Valid PAT", func(t *testing.T) {
+		pat := createMockPAT(t, `{"h":"api.snyk.io"}`)
+		apiUrl, err := GetApiUrlFromPAT(pat)
+		assert.NoError(t, err)
+		assert.Equal(t, "https://api.snyk.io", apiUrl)
+	})
+
+	t.Run("Valid EU PAT", func(t *testing.T) {
+		pat := createMockPAT(t, `{"h":"api.eu.snyk.io"}`)
+		apiUrl, err := GetApiUrlFromPAT(pat)
+		assert.NoError(t, err)
+		assert.Equal(t, "https://api.eu.snyk.io", apiUrl)
+	})
+
+	t.Run("PAT with scheme", func(t *testing.T) {
+		pat := createMockPAT(t, `{"h":"http://api.snyk.io"}`)
+		fmt.Println("pat", pat)
+		apiUrl, err := GetApiUrlFromPAT(pat)
+		assert.NoError(t, err)
+		assert.Equal(t, "http://api.snyk.io", apiUrl)
+	})
+
+	t.Run("Invalid PAT", func(t *testing.T) {
+		patTooManySegments := "snyk_test.12345678.payload.signature.extra"
+		apiUrl, err := GetApiUrlFromPAT(patTooManySegments)
+		assert.Error(t, err)
+		assert.Equal(t, "", apiUrl)
+	})
+}
+
+func createMockPAT(t *testing.T, payload string) string {
+	t.Helper()
+
+	encodedPayload := base64.RawURLEncoding.EncodeToString([]byte(payload))
+	signature := "signature"
+	return fmt.Sprintf("snyk_uat.12345678.%s.%s", encodedPayload, signature)
+}
diff --git a/pkg/local_workflows/auth_workflow.go b/pkg/local_workflows/auth_workflow.go
@@ -135,6 +135,8 @@ func entryPointDI(invocationCtx workflow.InvocationContext, logger *zerolog.Logg
 
 		logger.Print("Validating pat")
 		whoamiConfig := config.Clone()
+		// we don't want to use the cache here, so this is a workaround
+		whoamiConfig.ClearCache()
 		whoamiConfig.Set(configuration.FLAG_EXPERIMENTAL, true)
 		whoamiConfig.Set(configuration.AUTHENTICATION_TOKEN, pat)
 		_, whoamiErr := engine.InvokeWithConfig(workflow.NewWorkflowIdentifier("whoami"), whoamiConfig)
@@ -147,6 +149,8 @@ func entryPointDI(invocationCtx workflow.InvocationContext, logger *zerolog.Logg
 		}
 
 		logger.Print("Validation successful; set pat credentials in config")
+		// we don't want to use the cache here, so this is a workaround
+		engine.GetConfiguration().ClearCache()
 		engine.GetConfiguration().Set(auth.CONFIG_KEY_TOKEN, pat)
 
 		err = ui.DefaultUi().Output(auth.AUTHENTICATED_MESSAGE)
diff --git a/pkg/local_workflows/auth_workflow_test.go b/pkg/local_workflows/auth_workflow_test.go
@@ -103,12 +103,6 @@ func Test_auth_token(t *testing.T) {
 }
 
 func Test_pat(t *testing.T) {
-	const (
-		testPAT               = "snyk_pat.12345678.abcdefghijklmnopqrstuvwxyz123456"
-		mockedPatEndpoint     = "https://api.snyk.io"
-		expectedAPIKeyStorage = auth.CONFIG_KEY_TOKEN
-	)
-
 	mockCtl := gomock.NewController(t)
 	defer mockCtl.Finish()
 
@@ -118,23 +112,22 @@ func Test_pat(t *testing.T) {
 
 	engine := mocks.NewMockEngine(mockCtl)
 	authenticator := mocks.NewMockAuthenticator(mockCtl)
+	pat := "myPAT"
 
 	t.Run("happy", func(t *testing.T) {
-		config := configuration.New()
+		config := configuration.NewWithOpts()
 		config.Set(authTypeParameter, auth.AUTH_TYPE_PAT)
-		config.Set(ConfigurationNewAuthenticationToken, testPAT)
+		config.Set(ConfigurationNewAuthenticationToken, pat)
+
 		config.Set(auth.CONFIG_KEY_OAUTH_TOKEN, "some-oauth-token")
 		config.Set(configuration.AUTHENTICATION_TOKEN, "some-legacy-api-token")
 
-		config.Set(configuration.API_URL, []string{"https://api.snyk.io"})
-
 		mockInvocationContext := mocks.NewMockInvocationContext(mockCtl)
 		mockInvocationContext.EXPECT().GetConfiguration().Return(config).AnyTimes()
 		mockInvocationContext.EXPECT().GetEnhancedLogger().Return(&logger).AnyTimes()
 		mockInvocationContext.EXPECT().GetAnalytics().Return(analytics).Times(1)
 
-		engineConfig := configuration.New()
-		engine.EXPECT().GetConfiguration().Return(engineConfig).AnyTimes()
+		engine.EXPECT().GetConfiguration().Return(config).AnyTimes()
 		engine.EXPECT().InvokeWithConfig(gomock.Any(), gomock.Any())
 
 		err := entryPointDI(mockInvocationContext, &logger, engine, authenticator)
@@ -145,21 +138,20 @@ func Test_pat(t *testing.T) {
 	})
 
 	t.Run("invalid pat should fail", func(t *testing.T) {
-		config := configuration.New()
+		config := configuration.NewWithOpts()
 		config.Set(authTypeParameter, auth.AUTH_TYPE_PAT)
-		config.Set(ConfigurationNewAuthenticationToken, testPAT)
+		config.Set(ConfigurationNewAuthenticationToken, pat)
+
 		config.Set(auth.CONFIG_KEY_OAUTH_TOKEN, "some-oauth-token")
 		config.Set(configuration.AUTHENTICATION_TOKEN, "some-legacy-api-token")
 
-		config.Set(configuration.API_URL, []string{"https://api.snyk.io"})
-
 		mockInvocationContext := mocks.NewMockInvocationContext(mockCtl)
 		mockInvocationContext.EXPECT().GetConfiguration().Return(config).AnyTimes()
 		mockInvocationContext.EXPECT().GetEnhancedLogger().Return(&logger).AnyTimes()
 		mockInvocationContext.EXPECT().GetAnalytics().Return(analytics).Times(1)
 
-		engineConfig := configuration.New()
-		engine.EXPECT().GetConfiguration().Return(engineConfig).AnyTimes()
+		engine.EXPECT().GetConfiguration().Return(config).AnyTimes()
+
 		mockWhoAmIError := fmt.Errorf("mock whoami failure")
 		engine.EXPECT().InvokeWithConfig(gomock.Any(), gomock.Any()).Return(nil, mockWhoAmIError)
 
diff --git a/pkg/logging/scrubbingLogWriter.go b/pkg/logging/scrubbingLogWriter.go
@@ -194,6 +194,13 @@ func addMandatoryMasking(dict ScrubbingDict) ScrubbingDict {
 		regex:         regexp.MustCompile(s),
 	}
 
+	// Snyk PATs
+	s = auth.PAT_REGEX
+	dict[s] = scrubStruct{
+		groupToRedact: 2,
+		regex:         regexp.MustCompile(s),
+	}
+
 	// github
 	s = fmt.Sprintf(`(access_token[\\="\s:]+)(%s)&?`, charGroup)
 	dict[s] = scrubStruct{
diff --git a/pkg/logging/scrubbingLogWriter_test.go b/pkg/logging/scrubbingLogWriter_test.go
