OAuth: Use the attached external session data in OAuthToken and OAuthTokenSync (#96655)

* wip

* wip + tests

* wip

* wip opt2

* Use authn.Identity struct's SessionToken

* Merge fixes

* Handle disabling the feature flag correctly

* Fix test

* Cleanup

* Remove HasOAuthEntry from the OAuthTokenService interface

* Remove unused function

pkg/services/auth/authimpl/auth_token.go

@@ -248,14 +248,18 @@ func (s *UserAuthTokenService) GetTokenByExternalSessionID(ctx context.Context,
 	return &userToken, err
 }
 
-func (s *UserAuthTokenService) GetExternalSession(ctx context.Context, extSessionID int64) (*auth.ExternalSession, error) {
-	return s.externalSessionStore.Get(ctx, extSessionID)
+func (s *UserAuthTokenService) GetExternalSession(ctx context.Context, externalSessionID int64) (*auth.ExternalSession, error) {
+	return s.externalSessionStore.Get(ctx, externalSessionID)
 }
 
 func (s *UserAuthTokenService) FindExternalSessions(ctx context.Context, query *auth.ListExternalSessionQuery) ([]*auth.ExternalSession, error) {
 	return s.externalSessionStore.List(ctx, query)
 }
 
+func (s *UserAuthTokenService) UpdateExternalSession(ctx context.Context, externalSessionID int64, cmd *auth.UpdateExternalSessionCommand) error {
+	return s.externalSessionStore.Update(ctx, externalSessionID, cmd)
+}
+
 func (s *UserAuthTokenService) RotateToken(ctx context.Context, cmd auth.RotateCommand) (*auth.UserToken, error) {
 	if cmd.UnHashedToken == "" {
 		return nil, auth.ErrInvalidSessionToken

pkg/services/auth/authimpl/external_session_store.go

@@ -27,11 +27,11 @@ func provideExternalSessionStore(sqlStore db.DB, secretService secrets.Service,
 	}
 }
 
-func (s *store) Get(ctx context.Context, extSessionID int64) (*auth.ExternalSession, error) {
+func (s *store) Get(ctx context.Context, ID int64) (*auth.ExternalSession, error) {
 	ctx, span := s.tracer.Start(ctx, "externalsession.Get")
 	defer span.End()
 
-	externalSession := &auth.ExternalSession{ID: extSessionID}
+	externalSession := &auth.ExternalSession{ID: ID}
 
 	err := s.sqlStore.WithDbSession(ctx, func(sess *db.Session) error {
 		found, err := sess.Get(externalSession)
@@ -150,6 +150,45 @@ func (s *store) Create(ctx context.Context, extSession *auth.ExternalSession) er
 	return nil
 }
 
+func (s *store) Update(ctx context.Context, ID int64, cmd *auth.UpdateExternalSessionCommand) error {
+	ctx, span := s.tracer.Start(ctx, "externalsession.Update")
+	defer span.End()
+
+	var err error
+	externalSession := &auth.ExternalSession{}
+
+	externalSession.AccessToken, err = s.encryptAndEncode(cmd.Token.AccessToken)
+	if err != nil {
+		return err
+	}
+
+	externalSession.RefreshToken, err = s.encryptAndEncode(cmd.Token.RefreshToken)
+	if err != nil {
+		return err
+	}
+
+	var secretIdToken string
+	if idToken, ok := cmd.Token.Extra("id_token").(string); ok && idToken != "" {
+		secretIdToken, err = s.encryptAndEncode(idToken)
+		if err != nil {
+			return err
+		}
+		externalSession.IDToken = secretIdToken
+	}
+
+	externalSession.ExpiresAt = cmd.Token.Expiry
+
+	err = s.sqlStore.WithDbSession(ctx, func(sess *db.Session) error {
+		_, err := sess.ID(ID).Cols("access_token", "refresh_token", "id_token", "expires_at").Update(externalSession)
+		return err
+	})
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func (s *store) Delete(ctx context.Context, ID int64) error {
 	ctx, span := s.tracer.Start(ctx, "externalsession.Delete")
 	defer span.End()