ultisuite-backend/internal/migration/oauth.go

package migration

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"time"

	"github.com/redis/go-redis/v9"
	"golang.org/x/oauth2"
)

const pendingKeyPrefix = "migration_oauth_pending:"
const pendingTTL = 15 * time.Minute

var ErrUnknownState = errors.New("migration oauth state expired or unknown")
var ErrProviderDisabled = errors.New("migration oauth provider not configured")

type Provider string

const (
	ProviderGoogle    Provider = "google"
	ProviderMicrosoft Provider = "microsoft"
)

type PendingOAuth struct {
	UserID    string `json:"user_id"`
	ProjectID string `json:"project_id"`
	Provider  string `json:"provider"`
	InviteToken string `json:"invite_token,omitempty"`
	PKCEVerifier string `json:"pkce_verifier"`
}

type OAuthConfig struct {
	GoogleClientID     string
	GoogleClientSecret string
	MicrosoftClientID  string
	MicrosoftSecret    string
	MicrosoftTenant    string
	RedirectURL        string
}

type OAuthService struct {
	cfg OAuthConfig
	rdb *redis.Client
}

func NewOAuthService(cfg OAuthConfig, rdb *redis.Client) *OAuthService {
	return &OAuthService{cfg: cfg, rdb: rdb}
}

func (s *OAuthService) EnabledProviders() []string {
	var out []string
	if s.providerConfig(ProviderGoogle) != nil {
		out = append(out, string(ProviderGoogle))
	}
	if s.providerConfig(ProviderMicrosoft) != nil {
		out = append(out, string(ProviderMicrosoft))
	}
	return out
}

func (s *OAuthService) Start(ctx context.Context, pending PendingOAuth, provider Provider) (authURL, state string, err error) {
	oauthCfg := s.providerConfig(provider)
	if oauthCfg == nil {
		return "", "", ErrProviderDisabled
	}
	verifier, challenge, err := newPKCE()
	if err != nil {
		return "", "", err
	}
	state, err = randomState()
	if err != nil {
		return "", "", err
	}
	pending.Provider = string(provider)
	pending.PKCEVerifier = verifier
	if err := s.savePending(ctx, state, pending); err != nil {
		return "", "", err
	}
	authURL = oauthCfg.AuthCodeURL(state,
		oauth2.AccessTypeOffline,
		oauth2.SetAuthURLParam("code_challenge", challenge),
		oauth2.SetAuthURLParam("code_challenge_method", "S256"),
		oauth2.SetAuthURLParam("prompt", "consent"),
	)
	return authURL, state, nil
}

func (s *OAuthService) Exchange(ctx context.Context, state, code string) (PendingOAuth, *oauth2.Token, []string, error) {
	pending, err := s.loadPending(ctx, state)
	if err != nil {
		return PendingOAuth{}, nil, nil, err
	}
	oauthCfg := s.providerConfig(Provider(pending.Provider))
	if oauthCfg == nil {
		return PendingOAuth{}, nil, nil, ErrProviderDisabled
	}
	token, err := oauthCfg.Exchange(ctx, code, oauth2.SetAuthURLParam("code_verifier", pending.PKCEVerifier))
	if err != nil {
		return PendingOAuth{}, nil, nil, fmt.Errorf("token exchange: %w", err)
	}
	_ = s.rdb.Del(ctx, pendingKeyPrefix+state).Err()
	return pending, token, oauthCfg.Scopes, nil
}

func (s *OAuthService) Refresh(ctx context.Context, provider Provider, refreshToken string) (*oauth2.Token, error) {
	oauthCfg := s.providerConfig(provider)
	if oauthCfg == nil {
		return nil, ErrProviderDisabled
	}
	if strings.TrimSpace(refreshToken) == "" {
		return nil, fmt.Errorf("refresh token required")
	}
	token, err := oauthCfg.TokenSource(ctx, &oauth2.Token{RefreshToken: refreshToken}).Token()
	if err != nil {
		return nil, fmt.Errorf("token refresh: %w", err)
	}
	return token, nil
}

// AdminConsentURL returns the Microsoft tenant admin consent URL for migration scopes.
func (s *OAuthService) AdminConsentURL(tenant, state string) (string, error) {
	if s.cfg.MicrosoftClientID == "" || s.cfg.RedirectURL == "" {
		return "", ErrProviderDisabled
	}
	tenant = strings.TrimSpace(tenant)
	if tenant == "" {
		tenant = s.cfg.MicrosoftTenant
	}
	if tenant == "" {
		tenant = "common"
	}
	values := url.Values{}
	values.Set("client_id", s.cfg.MicrosoftClientID)
	values.Set("redirect_uri", s.cfg.RedirectURL)
	if state = strings.TrimSpace(state); state != "" {
		values.Set("state", state)
	}
	return fmt.Sprintf("https://login.microsoftonline.com/%s/adminconsent?%s", url.PathEscape(tenant), values.Encode()), nil
}

func (s *OAuthService) MicrosoftClientID() string {
	return s.cfg.MicrosoftClientID
}

func (s *OAuthService) providerConfig(provider Provider) *oauth2.Config {
	switch provider {
	case ProviderGoogle:
		if s.cfg.GoogleClientID == "" || s.cfg.GoogleClientSecret == "" || s.cfg.RedirectURL == "" {
			return nil
		}
		return &oauth2.Config{
			ClientID:     s.cfg.GoogleClientID,
			ClientSecret: s.cfg.GoogleClientSecret,
			RedirectURL:  s.cfg.RedirectURL,
			Scopes: []string{
				"https://www.googleapis.com/auth/gmail.readonly",
				"https://www.googleapis.com/auth/drive.readonly",
				"https://www.googleapis.com/auth/calendar.readonly",
				"https://www.googleapis.com/auth/contacts.readonly",
			},
			Endpoint: oauth2.Endpoint{
				AuthURL:  "https://accounts.google.com/o/oauth2/v2/auth",
				TokenURL: "https://oauth2.googleapis.com/token",
			},
		}
	case ProviderMicrosoft:
		if s.cfg.MicrosoftClientID == "" || s.cfg.MicrosoftSecret == "" || s.cfg.RedirectURL == "" {
			return nil
		}
		tenant := s.cfg.MicrosoftTenant
		if tenant == "" {
			tenant = "common"
		}
		return &oauth2.Config{
			ClientID:     s.cfg.MicrosoftClientID,
			ClientSecret: s.cfg.MicrosoftSecret,
			RedirectURL:  s.cfg.RedirectURL,
			Scopes: []string{
				"offline_access",
				"User.Read",
				"Mail.Read",
				"Files.Read.All",
				"Calendars.Read",
				"Contacts.Read",
			},
			Endpoint: oauth2.Endpoint{
				AuthURL:  fmt.Sprintf("https://login.microsoftonline.com/%s/oauth2/v2.0/authorize", tenant),
				TokenURL: fmt.Sprintf("https://login.microsoftonline.com/%s/oauth2/v2.0/token", tenant),
			},
		}
	default:
		return nil
	}
}

func (s *OAuthService) savePending(ctx context.Context, state string, pending PendingOAuth) error {
	if s.rdb == nil {
		return errors.New("oauth state store unavailable")
	}
	raw, err := json.Marshal(pending)
	if err != nil {
		return err
	}
	return s.rdb.Set(ctx, pendingKeyPrefix+state, raw, pendingTTL).Err()
}

func (s *OAuthService) loadPending(ctx context.Context, state string) (PendingOAuth, error) {
	if s.rdb == nil {
		return PendingOAuth{}, errors.New("oauth state store unavailable")
	}
	raw, err := s.rdb.Get(ctx, pendingKeyPrefix+state).Bytes()
	if err != nil {
		if errors.Is(err, redis.Nil) {
			return PendingOAuth{}, ErrUnknownState
		}
		return PendingOAuth{}, err
	}
	var pending PendingOAuth
	if err := json.Unmarshal(raw, &pending); err != nil {
		return PendingOAuth{}, err
	}
	return pending, nil
}