inventory-plus-plus/internal/domains/authentication/auth.go

package authentication

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

	"github.com/coreos/go-oidc/v3/oidc"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/jackc/pgx/v5/pgxpool"
	"golang.org/x/oauth2"

	"ruben/inventory2/internal/consts"
	"ruben/inventory2/internal/logging"
)

// TODO: move these to a config?
const (
	// The URL of our Auth0 Tenant Domain.
	// If you're using a Custom Domain, be sure to set this to that value instead.
	AUTH0_DOMAIN = "dev-uq3gqy5bdnwxmr6d.us.auth0.com"

	// Our Auth0 application"s Client ID.
	AUTH0_CLIENT_ID = "JEjrXTQ9fxlTLgp9RgTIACpUk8a2lqNT"

	// Our Auth0 application"s Client Secret.
	AUTH0_CLIENT_SECRET = "83U-iWdVaNnwk9XDzteo_2VMyOq_l1siKYqg1_2E7jCzgL8MnkaxlysPMcPMGlxA"

	// The Callback URL of our application.
	AUTH0_CALLBACK_URL = "https://inventory-plus-plus.com/api/auth/login/callback"
)

type (
	// Authenticator is used to authenticate our users.
	Authenticator struct {
		log *logging.Logger
		*oidc.Provider
		oauth2.Config
		db *pgxpool.Pool
	}

	// AccessTokenClaims is the claims Auth0 provides in access tokens
	AccessTokenClaims struct {
		Audience   string    `json:"aud"`
		Expires    int64     `json:"exp"`
		FamilyName string    `json:"family_name"`
		GivenName  string    `json:"given_name"`
		IssuedAt   int64     `json:"iat"`
		Issuer     string    `json:"iss"`
		Name       string    `json:"name"`
		Nickname   string    `json:"nickname"`
		Picture    string    `json:"picture"`
		SessionID  string    `json:"sid"`
		Subject    string    `json:"sub"`
		UpdatedAt  time.Time `json:"updated_at"`
	}
)

// New instantiates the *Authenticator.
func New(
	ctx context.Context,
	db *pgxpool.Pool,
	logger *logging.Logger,
) (*Authenticator, error) {
	provider, err := oidc.NewProvider(
		ctx,
		"https://"+AUTH0_DOMAIN+"/",
	)
	if err != nil {
		return nil, err
	}

	return &Authenticator{
		log:      logger,
		Provider: provider,
		Config: oauth2.Config{
			ClientID:     AUTH0_CLIENT_ID,
			ClientSecret: AUTH0_CLIENT_SECRET,
			RedirectURL:  AUTH0_CALLBACK_URL,
			Endpoint:     provider.Endpoint(),
			Scopes:       []string{oidc.ScopeOpenID, "profile"},
		},
		db: db,
	}, nil
}

func (a *Authenticator) RunBackgroundCleanup(ctx context.Context) error {
	for {
		if _, err := a.db.Exec(ctx, "DELETE FROM oauth_tokens WHERE expiry < NOW()"); err != nil {
			return fmt.Errorf("failed to delete all oauth_tokens rows that are expired: %w", err)
		}

		if _, err := a.db.Exec(ctx, "DELETE FROM oauth_login_states WHERE expiration < NOW()"); err != nil {
			return fmt.Errorf("failed to delete all oauth_login_states rows that are expired: %w", err)
		}

		select {
		case <-ctx.Done():
			return nil
		case <-time.After(time.Minute):
		}
	}
}

// Exchange exchanges an auth code for an access token.
func (a *Authenticator) Exchange(ctx context.Context, state, code string) (accessToken, targetURI string, expiration time.Time, err error) {
	// validate state

	exp, targetURI, err := a.GetStateExpirationAndURL(ctx, state)
	if errors.Is(err, consts.ErrNotFound) {
		return "", "", time.Time{}, fmt.Errorf("invalid state: %w", consts.ErrNotFound)
	} else if err != nil {
		return "", "", time.Time{}, fmt.Errorf("failed to load state expiration: %w", err)
	} else if exp.Before(time.Now()) {
		return "", "", time.Time{}, fmt.Errorf("invalid state: state expired")
	}

	// obtain token and profile

	tkn, err := a.Config.Exchange(ctx, code)
	if err != nil {
		return "", "", time.Time{}, fmt.Errorf("failed to exchange an authorization code for a token: %w", err)
	}

	idToken, claims, err := a.verifyIDTokenAndClaimsFromToken(ctx, tkn)
	if err != nil {
		return "", "", time.Time{}, err
	}

	claimsJSON, _ := json.Marshal(claims)

	// store the token, and the potentially new user

	if _, err := a.db.Exec(
		ctx,
		`
			WITH new_user AS (
				INSERT INTO oauth_users (
					user_id
				)
				VALUES (
					@id_token_subject
				)
				ON CONFLICT DO NOTHING
				RETURNING
					user_id
			), the_user AS (
				SELECT
					COALESCE(user_id, user_id_2) as user_id
				FROM
					new_user
				RIGHT JOIN
					(SELECT @id_token_subject as user_id_2)
					ON TRUE
			)
			INSERT INTO oauth_tokens (
				access_token,
				token_type,
				refresh_token,
				expiry,

				id_token_issuer,
				id_token_audience,
				id_token_subject,
				id_token_expiry,
				id_token_issued_at,
				id_token_nonce,
				id_token_access_token_hash,

				claims -- might want to open this up
			)
			SELECT
				@access_token,
				@token_type,
				@refresh_token,
				@expiry,

				@id_token_issuer,
				@id_token_audience,
				user_id,
				@id_token_expiry,
				@id_token_issued_at,
				@id_token_nonce,
				@id_token_access_token_hash,

				@claims
			FROM
				the_user
		`,
		pgx.NamedArgs{
			"access_token":  tkn.AccessToken,
			"token_type":    tkn.TokenType,
			"refresh_token": tkn.RefreshToken,
			"expiry":        tkn.Expiry,

			"id_token_issuer":            idToken.Issuer,
			"id_token_audience":          pgtype.FlatArray[string](idToken.Audience),
			"id_token_subject":           idToken.Subject,
			"id_token_expiry":            idToken.Expiry,
			"id_token_issued_at":         idToken.IssuedAt,
			"id_token_nonce":             idToken.Nonce,
			"id_token_access_token_hash": idToken.AccessTokenHash,

			"claims": json.RawMessage(claimsJSON),
		},
	); err != nil {
		return "", "", time.Time{}, fmt.Errorf("failed to perform query to save tokens: %w", err)
	}

	return tkn.AccessToken, targetURI, tkn.Expiry.UTC(), nil
}

// VerifyIDToken verifies that an *oauth2.Token is a valid *oidc.IDToken.
func (a *Authenticator) VerifyIDToken(ctx context.Context, token *oauth2.Token) (*oidc.IDToken, error) {
	rawIDToken, ok := token.Extra("id_token").(string)
	if !ok {
		return nil, errors.New("no id_token field in oauth2 token")
	}

	oidcConfig := &oidc.Config{
		ClientID: a.ClientID,
	}

	return a.Verifier(oidcConfig).Verify(ctx, rawIDToken)
}

func (a *Authenticator) GetLogoutURL(requestHost string) *url.URL {
	return &url.URL{
		Scheme: "https",
		Host:   AUTH0_DOMAIN,
		Path:   "/v2/logout",
		RawQuery: url.Values{
			"returnTo": {
				(&url.URL{
					Scheme: "https",
					Host:   requestHost,
				}).String(),
			},
			"client_id": {AUTH0_CLIENT_ID},
		}.Encode(),
	}
}

func (a *Authenticator) RefreshAccessToken(
	ctx context.Context,
	oldAccessToken string,
) (
	accessToken string,
	expiration time.Time,
	err error,
) {

	refreshToken, tokenType, err := a.getRefreshTokenForAccessToken(ctx, accessToken)
	if err != nil {
		return "", time.Time{}, fmt.Errorf("failed to load refresh token: %w", err)
	}

	tkn, err := a.TokenSource(ctx, &oauth2.Token{
		// AccessToken is the token that authorizes and authenticates
		// the requests.
		AccessToken: oldAccessToken,

		// TokenType is the type of token.
		// The Type method returns either this or "Bearer", the default.
		TokenType: tokenType,

		// RefreshToken is a token that's used by the application
		// (as opposed to the user) to refresh the access token
		// if it expires.
		RefreshToken: refreshToken,

		/*
			// Expiry is the optional expiration time of the access token.
			//
			// If zero, TokenSource implementations will reuse the same
			// token forever and RefreshToken or equivalent
			// mechanisms for that TokenSource will not be used.
			Expiry time.Time `json:"expiry,omitempty"`
		*/
	}).Token()
	if err != nil {
		return "", time.Time{}, fmt.Errorf("failed to fetch refresh token: %w", err)
	}

	idToken, claims, err := a.verifyIDTokenAndClaimsFromToken(ctx, tkn)
	if err != nil {
		return "", time.Time{}, err
	}

	claimsJSON, _ := json.Marshal(claims)

	if _, err = a.db.Exec(
		ctx,
		`
			WITH deleted_tokens AS (
				DELETE FROM
					oauth_tokens
				WHERE
					access_token = @old_access_token
				RETURNING
					access_token AS old_access_token
			), new_tokens AS (
				INSERT INTO oauth_tokens (
					access_token,
					token_type,
					refresh_token,
					expiry,

					id_token_issuer,
					id_token_audience,
					id_token_subject,
					id_token_expiry,
					id_token_issued_at,
					id_token_nonce,
					id_token_access_token_hash,

					claims -- might want to open this up
				)
				VALUES (
					@new_access_token,
					@token_type,
					@refresh_token,
					@expiry,

					@id_token_issuer,
					@id_token_audience,
					@id_token_subject,
					@id_token_expiry,
					@id_token_issued_at,
					@id_token_nonce,
					@id_token_access_token_hash,

					@claims
				)
				RETURNING
					access_token AS new_access_token
			)
			SELECT
				old_access_token,
				new_access_token
			FROM
				deleted_tokens
			FULL JOIN
				new_tokens
		`,
		pgx.NamedArgs{
			"old_access_token": oldAccessToken,

			"new_access_token": tkn.AccessToken,
			"token_type":       tkn.TokenType,
			"refresh_token":    tkn.RefreshToken,
			"expiry":           tkn.Expiry,

			"id_token_issuer":            idToken.Issuer,
			"id_token_audience":          pgtype.FlatArray[string](idToken.Audience),
			"id_token_subject":           idToken.Subject,
			"id_token_expiry":            idToken.Expiry,
			"id_token_issued_at":         idToken.IssuedAt,
			"id_token_nonce":             idToken.Nonce,
			"id_token_access_token_hash": idToken.AccessTokenHash,

			"claims": json.RawMessage(claimsJSON),
		},
	); err != nil {
		return "", time.Time{}, fmt.Errorf("failed to save new access token and delete old access token: %w", err)
	}

	return tkn.AccessToken, tkn.Expiry.UTC(), nil
}

func (a *Authenticator) verifyIDTokenAndClaimsFromToken(ctx context.Context, tkn *oauth2.Token) (idToken *oidc.IDToken, claims map[string]any, err error) {
	if idToken, err = a.VerifyIDToken(ctx, tkn); err != nil {
		return nil, nil, fmt.Errorf("failed to verify id Token: %w", err)
	}

	if err := idToken.Claims(&claims); err != nil {
		return nil, nil, fmt.Errorf("failed to obtain id token claims: %w", err)
	}

	return idToken, claims, nil
}