pocket-id/backend/internal/service/one_time_access_service.go

package service

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/url"
	"strings"
	"time"

	"github.com/pocket-id/pocket-id/backend/internal/appconfig"
	"github.com/pocket-id/pocket-id/backend/internal/common"
	"github.com/pocket-id/pocket-id/backend/internal/model"
	datatype "github.com/pocket-id/pocket-id/backend/internal/model/types"
	"github.com/pocket-id/pocket-id/backend/internal/utils"
	"github.com/pocket-id/pocket-id/backend/internal/utils/email"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

type OneTimeAccessService struct {
	db              *gorm.DB
	userService     *UserService
	jwtService      *JwtService
	auditLogService *AuditLogService
	emailService    *EmailService
}

func NewOneTimeAccessService(db *gorm.DB, userService *UserService, jwtService *JwtService, auditLogService *AuditLogService, emailService *EmailService) *OneTimeAccessService {
	return &OneTimeAccessService{
		db:              db,
		userService:     userService,
		jwtService:      jwtService,
		auditLogService: auditLogService,
		emailService:    emailService,
	}
}

func (s *OneTimeAccessService) RequestOneTimeAccessEmailAsAdmin(ctx context.Context, dbConfig *appconfig.AppConfigModel, userID string, ttl time.Duration) error {
	if !dbConfig.EmailOneTimeAccessAsAdminEnabled.IsTrue() {
		return &common.OneTimeAccessDisabledError{}
	}

	_, err := s.requestOneTimeAccessEmailInternal(ctx, userID, "", ttl, false, dbConfig)
	return err
}

func (s *OneTimeAccessService) RequestOneTimeAccessEmailAsUnauthenticatedUser(ctx context.Context, dbConfig *appconfig.AppConfigModel, userID, redirectPath string) (string, error) {
	if !dbConfig.EmailOneTimeAccessAsUnauthenticatedEnabled.IsTrue() {
		return "", &common.OneTimeAccessDisabledError{}
	}

	var userId string
	err := s.db.Model(&model.User{}).Select("id").Where("email = ?", userID).First(&userId).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		// Do not return error if user not found to prevent email enumeration
		return "", nil
	} else if err != nil {
		return "", err
	}

	deviceToken, err := s.requestOneTimeAccessEmailInternal(ctx, userId, redirectPath, 15*time.Minute, true, dbConfig)
	if err != nil {
		return "", err
	} else if deviceToken == nil {
		return "", errors.New("device token expected but not returned")
	}

	return *deviceToken, nil
}

func (s *OneTimeAccessService) requestOneTimeAccessEmailInternal(ctx context.Context, userID, redirectPath string, ttl time.Duration, withDeviceToken bool, dbConfig *appconfig.AppConfigModel) (*string, error) {
	tx := s.db.Begin()
	defer func() {
		tx.Rollback()
	}()

	user, err := s.userService.getUserInternal(ctx, userID, tx)
	if err != nil {
		return nil, err
	}

	if user.Email == nil {
		return nil, &common.UserEmailNotSetError{}
	}

	oneTimeAccessToken, deviceToken, err := s.createOneTimeAccessTokenInternal(ctx, user.ID, ttl, withDeviceToken, tx)
	if err != nil {
		return nil, err
	}
	err = tx.Commit().Error
	if err != nil {
		return nil, err
	}

	go func() {
		// This runs in background, so use a context without cancellation (or it would be stopped when the request ends)
		// We still want to have a context derived from the request's to carry over tracing info
		innerCtx := context.WithoutCancel(ctx)

		link := common.EnvConfig.AppURL + "/lc"
		linkWithCode := link + "/" + oneTimeAccessToken

		// Add redirect path to the link
		if strings.HasPrefix(redirectPath, "/") {
			encodedRedirectPath := url.QueryEscape(redirectPath)
			linkWithCode = linkWithCode + "?redirect=" + encodedRedirectPath
		}

		errInternal := SendEmail(innerCtx, s.emailService, dbConfig, email.Address{
			Name:  user.FullName(),
			Email: *user.Email,
		}, OneTimeAccessTemplate, &OneTimeAccessTemplateData{
			Code:              oneTimeAccessToken,
			LoginLink:         link,
			LoginLinkWithCode: linkWithCode,
			ExpirationString:  utils.DurationToString(ttl),
		})
		if errInternal != nil {
			slog.ErrorContext(innerCtx, "Failed to send one-time access token email", slog.Any("error", errInternal), slog.String("address", *user.Email))
			return
		}
	}()

	return deviceToken, nil
}

func (s *OneTimeAccessService) CreateOneTimeAccessToken(ctx context.Context, userID string, ttl time.Duration) (token string, err error) {
	tx := s.db.Begin()
	defer func() {
		tx.Rollback()
	}()

	// Load the user to ensure it exists
	_, err = s.userService.getUserInternal(ctx, userID, tx)
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return "", &common.UserNotFoundError{}
	} else if err != nil {
		return "", err
	}

	// Create the one-time access token
	token, _, err = s.createOneTimeAccessTokenInternal(ctx, userID, ttl, false, tx)
	if err != nil {
		return "", err
	}

	// Commit
	err = tx.Commit().Error
	if err != nil {
		return "", fmt.Errorf("error committing transaction: %w", err)
	}

	return token, nil
}

func (s *OneTimeAccessService) createOneTimeAccessTokenInternal(ctx context.Context, userID string, ttl time.Duration, withDeviceToken bool, tx *gorm.DB) (token string, deviceToken *string, err error) {
	oneTimeAccessToken, err := NewOneTimeAccessToken(userID, ttl, withDeviceToken)
	if err != nil {
		return "", nil, err
	}

	err = tx.WithContext(ctx).Create(oneTimeAccessToken).Error
	if err != nil {
		return "", nil, err
	}

	return oneTimeAccessToken.Token, oneTimeAccessToken.DeviceToken, nil
}

func (s *OneTimeAccessService) ExchangeOneTimeAccessToken(ctx context.Context, dbConfig *appconfig.AppConfigModel, token, deviceToken, ipAddress, userAgent string) (model.User, string, error) {
	tx := s.db.Begin()
	defer func() {
		tx.Rollback()
	}()

	var oneTimeAccessToken model.OneTimeAccessToken
	err := tx.
		WithContext(ctx).
		Where("token = ? AND expires_at > ?", token, datatype.DateTime(time.Now())).
		Preload("User").
		Clauses(clause.Locking{Strength: "UPDATE"}).
		First(&oneTimeAccessToken).
		Error
	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return model.User{}, "", &common.TokenInvalidOrExpiredError{}
		}
		return model.User{}, "", err
	}
	if oneTimeAccessToken.DeviceToken != nil && deviceToken != *oneTimeAccessToken.DeviceToken {
		return model.User{}, "", &common.DeviceCodeInvalid{}
	}
	if oneTimeAccessToken.User.Disabled {
		return model.User{}, "", &common.UserDisabledError{}
	}

	accessToken, err := s.jwtService.GenerateAccessToken(
		oneTimeAccessToken.User,
		AuthenticationMethodOneTimePassword,
		dbConfig.SessionDuration.AsDurationMinutes(),
	)
	if err != nil {
		return model.User{}, "", err
	}

	err = tx.
		WithContext(ctx).
		Delete(&oneTimeAccessToken).
		Error
	if err != nil {
		return model.User{}, "", err
	}

	s.auditLogService.Create(
		ctx, model.AuditLogEventOneTimeAccessTokenSignIn,
		ipAddress, userAgent,
		oneTimeAccessToken.User.ID, model.AuditLogData{},
		tx,
	)

	err = tx.Commit().Error
	if err != nil {
		return model.User{}, "", fmt.Errorf("error committing transaction: %w", err)
	}

	return oneTimeAccessToken.User, accessToken, nil
}

func NewOneTimeAccessToken(userID string, ttl time.Duration, withDeviceToken bool) (*model.OneTimeAccessToken, error) {
	// If expires at is less than 15 minutes, use a 6-character token instead of 16
	tokenLength := 16
	if ttl <= 15*time.Minute {
		tokenLength = 6
	}

	token, err := utils.GenerateRandomUnambiguousString(tokenLength)
	if err != nil {
		return nil, err
	}

	var deviceToken *string
	if withDeviceToken {
		dt, err := utils.GenerateRandomAlphanumericString(16)
		if err != nil {
			return nil, err
		}
		deviceToken = &dt
	}

	now := time.Now().Round(time.Second)
	o := &model.OneTimeAccessToken{
		UserID:      userID,
		ExpiresAt:   datatype.DateTime(now.Add(ttl)),
		Token:       token,
		DeviceToken: deviceToken,
	}

	return o, nil
}