llink/go/internal/auth/service.go

package auth

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"time"

	firebaseauth "firebase.google.com/go/v4/auth"
	pbaero "github.com/flowy-live/llink/genproto/aero"
	"github.com/flowy-live/llink/internal/utils"
	"github.com/redis/go-redis/v9"
	"go.jetify.com/typeid"
)

const (
	sessionExpiry          = time.Hour * 24 * 15
	extendSessionThreshold = time.Hour * 24 * 3
	codeExpiry             = time.Minute * 15
)

var (
	ErrInvalidCode     = errors.New("invalid code")
	ErrSessionNotFound = errors.New("session not found")
)

type sessionTokenPrefix struct{}

func (sessionTokenPrefix) Prefix() string { return "session" }

type sessionToken struct {
	typeid.TypeID[sessionTokenPrefix]
}

func newSessionToken() (sessionToken, error) {
	return typeid.New[sessionToken]()
}

type Session struct {
	Email   string `json:"email"`
	HumanId string `json:"human_id"`
}

type AuthService interface {
	SessionReader

	// RequestSignInCode generates a code and emails it to the provided email.
	// To retrieve a session, client must verify with VerifySignInCode.
	RequestSignInCode(ctx context.Context, email string) error
	// VerifySignInCode returns ErrInvalidCode if incorrect code, otherwise creates a session.
	// humanId is stored in the session alongside the email.
	VerifySignInCode(ctx context.Context, email, code, humanId string) (sessionToken string, err error)
	// ExtendSession returns ErrSessionNotFound if no valid session
	ExtendSession(ctx context.Context, sessionToken string) error
	SignOut(ctx context.Context, sessionToken string) error

	// MintFirebaseCustomToken returns a Firebase custom token with uid=humanId and no custom claims.
	MintFirebaseCustomToken(ctx context.Context, humanId string) (string, error)

	IsSystemAdmin(ctx context.Context, email string) bool
}

type authServiceImpl struct {
	*sessionReaderImpl
	aeroSvc pbaero.PrimaryClient
	fbAuth  *firebaseauth.Client
}

func NewAuthService(redisClient *redis.Client, aeroSvc pbaero.PrimaryClient, fbAuth *firebaseauth.Client) AuthService {
	return &authServiceImpl{
		sessionReaderImpl: newSessionReader(redisClient),
		aeroSvc:           aeroSvc,
		fbAuth:            fbAuth,
	}
}

func (a *authServiceImpl) MintFirebaseCustomToken(ctx context.Context, humanId string) (string, error) {
	if humanId == "" {
		return "", errors.New("humanId is required")
	}
	return a.fbAuth.CustomToken(ctx, humanId)
}

func (a *authServiceImpl) IsSystemAdmin(ctx context.Context, email string) bool {
	formattedEmail, err := utils.NormalizeEmail(email)
	if err != nil {
		slog.Error("problem validating email", "error", err)
		return false
	}

	if strings.Contains(formattedEmail, "@flowylabs.ai") {
		return true
	}
	return false
}

func (a *authServiceImpl) RequestSignInCode(ctx context.Context, email string) error {
	if email == "" {
		return errors.New("email is required")
	}

	code := utils.RandomStringNumbers(4)

	formattedEmail, err := utils.NormalizeEmail(email)
	if err != nil {
		return errors.New("email is not valid")
	}

	err = a.redisClient.Set(ctx, formattedEmail, code, codeExpiry).Err()
	if err != nil {
		slog.Error("error setting code in redis", "error", err)
		return fmt.Errorf("error storing sign-in code: %w", err)
	}

	message := fmt.Sprintf("Here is your one-time code for signing into Flowy: %s\n\nPlease do not share this with anyone.\n\nBest, \nFlowy Team", code)
	subject := fmt.Sprint("Sign In - Your One-Time Code for Flowy.llink")
	_, err = a.aeroSvc.ShootEmail(ctx, &pbaero.ShootEmailRequest{
		ToEmails: []string{formattedEmail},
		Subject:  subject,
		TemplateData: &pbaero.ShootEmailRequest_SimpleTextData{
			SimpleTextData: &pbaero.SimpleTextData{
				Message: message,
			},
		},
	})
	if err != nil {
		return fmt.Errorf("an error occurred while sending the email: %w", err)
	}

	slog.Info("sent sign in code", "email", formattedEmail)
	return nil
}

func (a *authServiceImpl) VerifySignInCode(ctx context.Context, email, code, humanId string) (string, error) {
	formattedEmail, err := utils.NormalizeEmail(email)
	if err != nil {
		return "", fmt.Errorf("invalid email: %w", err)
	}

	storedCode, err := a.redisClient.Get(ctx, formattedEmail).Result()
	if err != nil {
		if errors.Is(err, redis.Nil) {
			return "", ErrInvalidCode
		}
		slog.Error("error getting code from redis", "error", err)
		return "", fmt.Errorf("error verifying code: %w", err)
	}

	if storedCode != code {
		return "", ErrInvalidCode
	}

	if err := a.redisClient.Del(ctx, formattedEmail).Err(); err != nil {
		slog.Error("error deleting code from redis", "error", err)
	}

	token, err := a.createSession(ctx, formattedEmail, humanId)
	if err != nil {
		return "", err
	}

	return token, nil
}

func (a *authServiceImpl) ExtendSession(ctx context.Context, token string) error {
	ttl, err := a.redisClient.TTL(ctx, token).Result()
	if err != nil {
		return fmt.Errorf("error checking session TTL: %w", err)
	}

	if ttl < 0 {
		return ErrSessionNotFound
	}

	if ttl < extendSessionThreshold {
		if err := a.redisClient.Expire(ctx, token, sessionExpiry).Err(); err != nil {
			return fmt.Errorf("error extending session: %w", err)
		}
	}

	return nil
}

func (a *authServiceImpl) SignOut(ctx context.Context, token string) error {
	if err := a.redisClient.Del(ctx, token).Err(); err != nil {
		return fmt.Errorf("error deleting session: %w", err)
	}
	return nil
}

func (a *authServiceImpl) createSession(ctx context.Context, email, humanId string) (string, error) {
	formattedEmail, err := utils.NormalizeEmail(email)
	if err != nil {
		return "", fmt.Errorf("invalid email: %w", err)
	}

	token, err := newSessionToken()
	if err != nil {
		return "", fmt.Errorf("error generating session token: %w", err)
	}

	session := Session{Email: formattedEmail, HumanId: humanId}
	data, err := json.Marshal(session)
	if err != nil {
		return "", fmt.Errorf("error marshaling session: %w", err)
	}

	if err := a.redisClient.Set(ctx, token.String(), data, sessionExpiry).Err(); err != nil {
		return "", fmt.Errorf("error storing session: %w", err)
	}

	return token.String(), nil
}