Authn: No longer hash service account token twice during authentication (#92598)

* APIKey: Only decode and hash token once during authentication

* Only update last used every 5 minutes

pkg/services/authn/clients/api_key.go

@@ -31,6 +31,11 @@ var _ authn.HookClient = new(APIKey)
 var _ authn.ContextAwareClient = new(APIKey)
 var _ authn.IdentityResolverClient = new(APIKey)
 
+const (
+	metaKeyID           = "keyID"
+	metaKeySkipLastUsed = "keySkipLastUsed"
+)
+
 func ProvideAPIKey(apiKeyService apikey.Service) *APIKey {
 	return &APIKey{
 		log:           log.New(authn.ClientAPIKey),
@@ -64,6 +69,14 @@ func (s *APIKey) Authenticate(ctx context.Context, r *authn.Request) (*authn.Ide
 		return nil, err
 	}
 
+	// Set keyID so we can use it in last used hook
+	r.SetMeta(metaKeyID, strconv.FormatInt(key.ID, 10))
+	if !shouldUpdateLastUsedAt(key) {
+		// Hack to just have some value, we will check this key in the hook
+		// and if its not an empty string we will not update last used.
+		r.SetMeta(metaKeySkipLastUsed, "true")
+	}
+
 	// if the api key don't belong to a service account construct the identity and return it
 	if key.ServiceAccountId == nil || *key.ServiceAccountId < 1 {
 		return newAPIKeyIdentity(key), nil
@@ -109,7 +122,6 @@ func (s *APIKey) getFromTokenLegacy(ctx context.Context, token string) (*apikey.
 	if err != nil {
 		return nil, err
 	}
-
 	// fetch key
 	keyQuery := apikey.GetByNameQuery{KeyName: decoded.Name, OrgID: decoded.OrgId}
 	key, err := s.apiKeyService.GetApiKeyByName(ctx, &keyQuery)
@@ -170,52 +182,32 @@ func (s *APIKey) ResolveIdentity(ctx context.Context, orgID int64, typ claims.Id
 }
 
 func (s *APIKey) Hook(ctx context.Context, identity *authn.Identity, r *authn.Request) error {
-	id, exists := s.getAPIKeyID(ctx, identity, r)
-
-	if !exists {
+	if r.GetMeta(metaKeySkipLastUsed) != "" {
 		return nil
 	}
 
-	go func(apikeyID int64) {
+	go func(keyID string) {
 		defer func() {
 			if err := recover(); err != nil {
 				s.log.Error("Panic during user last seen sync", "err", err)
 			}
 		}()
-		if err := s.apiKeyService.UpdateAPIKeyLastUsedDate(context.Background(), apikeyID); err != nil {
-			s.log.Warn("Failed to update last use date for api key", "id", apikeyID)
+
+		id, err := strconv.ParseInt(keyID, 10, 64)
+		if err != nil {
+			s.log.Warn("Invalid api key id", "id", keyID, "err", err)
+			return
 		}
-	}(id)
+
+		if err := s.apiKeyService.UpdateAPIKeyLastUsedDate(context.Background(), id); err != nil {
+			s.log.Warn("Failed to update last used date for api key", "id", keyID, "err", err)
+			return
+		}
+	}(r.GetMeta(metaKeyID))
 
 	return nil
 }
 
-func (s *APIKey) getAPIKeyID(ctx context.Context, id *authn.Identity, r *authn.Request) (apiKeyID int64, exists bool) {
-	internalId, err := id.GetInternalID()
-	if err != nil {
-		s.log.Warn("Failed to parse ID from identifier", "err", err)
-		return -1, false
-	}
-
-	if id.IsIdentityType(claims.TypeAPIKey) {
-		return internalId, true
-	}
-
-	if id.IsIdentityType(claims.TypeServiceAccount) {
-		// When the identity is service account, the ID in from the namespace is the service account ID.
-		// We need to fetch the API key in this scenario, as we could use it to uniquely identify a service account token.
-		apiKey, err := s.getAPIKey(ctx, getTokenFromRequest(r))
-		if err != nil {
-			s.log.Warn("Failed to fetch the API Key from request")
-			return -1, false
-		}
-
-		return apiKey.ID, true
-	}
-
-	return -1, false
-}
-
 func looksLikeApiKey(token string) bool {
 	return token != ""
 }
@@ -276,3 +268,7 @@ func newServiceAccountIdentity(key *apikey.APIKey) *authn.Identity {
 		ClientParams:    authn.ClientParams{FetchSyncedUser: true, SyncPermissions: true},
 	}
 }
+
+func shouldUpdateLastUsedAt(key *apikey.APIKey) bool {
+	return key.LastUsedAt == nil || time.Since(*key.LastUsedAt) > 5*time.Minute
+}

pkg/services/authn/clients/api_key_test.go

@@ -189,108 +189,6 @@ func TestAPIKey_Test(t *testing.T) {
 	}
 }
 
-func TestAPIKey_GetAPIKeyIDFromIdentity(t *testing.T) {
-	type TestCase struct {
-		desc             string
-		expectedKey      *apikey.APIKey
-		expectedIdentity *authn.Identity
-		expectedError    error
-		expectedKeyID    int64
-		expectedExists   bool
-	}
-
-	tests := []TestCase{
-		{
-			desc: "should return API Key ID for valid token that is connected to service account",
-			expectedKey: &apikey.APIKey{
-				ID:               1,
-				OrgID:            1,
-				Key:              hash,
-				ServiceAccountId: intPtr(1),
-			},
-			expectedIdentity: &authn.Identity{
-				ID:              "1",
-				Type:            claims.TypeServiceAccount,
-				OrgID:           1,
-				Name:            "test",
-				AuthenticatedBy: login.APIKeyAuthModule,
-			},
-			expectedKeyID:  1,
-			expectedExists: true,
-		},
-		{
-			desc: "should return API Key ID for valid token for API key",
-			expectedKey: &apikey.APIKey{
-				ID:    2,
-				OrgID: 1,
-				Key:   hash,
-			},
-			expectedIdentity: &authn.Identity{
-				ID:              "2",
-				Type:            claims.TypeAPIKey,
-				OrgID:           1,
-				Name:            "test",
-				AuthenticatedBy: login.APIKeyAuthModule,
-			},
-			expectedKeyID:  2,
-			expectedExists: true,
-		},
-		{
-			desc: "should not return any ID when the request is not made by API key or service account",
-			expectedKey: &apikey.APIKey{
-				ID:    2,
-				OrgID: 1,
-				Key:   hash,
-			},
-			expectedIdentity: &authn.Identity{
-				ID:              "2",
-				Type:            claims.TypeUser,
-				OrgID:           1,
-				Name:            "test",
-				AuthenticatedBy: login.APIKeyAuthModule,
-			},
-			expectedKeyID:  -1,
-			expectedExists: false,
-		},
-		{
-			desc: "should not return any ID when the can't fetch API Key",
-			expectedKey: &apikey.APIKey{
-				ID:    1,
-				OrgID: 1,
-				Key:   hash,
-			},
-			expectedIdentity: &authn.Identity{
-				ID:              "2",
-				Type:            claims.TypeServiceAccount,
-				OrgID:           1,
-				Name:            "test",
-				AuthenticatedBy: login.APIKeyAuthModule,
-			},
-			expectedError:  fmt.Errorf("invalid token"),
-			expectedKeyID:  -1,
-			expectedExists: false,
-		},
-	}
-
-	req := &authn.Request{HTTPRequest: &http.Request{
-		Header: map[string][]string{
-			"Authorization": {"Bearer " + secret},
-		},
-	}}
-
-	for _, tt := range tests {
-		t.Run(tt.desc, func(t *testing.T) {
-			c := ProvideAPIKey(&apikeytest.Service{
-				ExpectedError:  tt.expectedError,
-				ExpectedAPIKey: tt.expectedKey,
-			})
-			id, exists := c.getAPIKeyID(context.Background(), tt.expectedIdentity, req)
-			assert.Equal(t, tt.expectedExists, exists)
-			assert.Equal(t, tt.expectedKeyID, id)
-		})
-	}
-}
-
 func TestAPIKey_ResolveIdentity(t *testing.T) {
 	type testCase struct {
 		desc string