magiclink.go

package gomagiclink

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base32"
	"errors"
	"fmt"
	"log/slog"
	"slices"
	"strconv"
	"strings"
	"time"

	"github.com/google/uuid"
)

const sesionIdSplitChar = "_"

type RecordWithID interface {
	GetID() uuid.UUID
}

type RecordWithKeyName interface {
	GetKeyName() string
}

// When a new storage provider is created, it implements this interface.
// See the provided storage provided in the `storage` package.
type UserAuthDatabase interface {
	UserExistsByEmail(email string) bool
	StoreUser(user *AuthUserRecord) error
	GetUserById(id uuid.UUID) (*AuthUserRecord, error)
	GetUserByEmail(email string) (*AuthUserRecord, error)
	GetUserCount() (int, error) // Slow
	UsersExist() (bool, error)  // Fast
}

const challengeSignature = "9"
const sessionIdSignature = "S"
const saltLength = 8

var ErrUserAlreadyExists = errors.New("user already exists")
var ErrUserNotFound = errors.New("user not found")
var ErrUserDisabled = errors.New("user disabled")
var ErrSecretKeyTooShort = errors.New("secret Key too short (min 16 bytes)")
var ErrInvalidChallenge = errors.New("invalid challenge")
var ErrBrokenChallenge = errors.New("broken challenge")
var ErrExpiredChallenge = errors.New("expired challenge")
var ErrInvalidSessionId = errors.New("invalid session id")
var ErrBrokenSessionId = errors.New("broken session id")
var ErrExpiredSessionId = errors.New("expired session id")

// All functionalities needed to implement the Magic Link login system is available
// through the AuthMagicLinkController.
type AuthMagicLinkController struct {
	secretKeyHash        []byte
	challengeExpDuration time.Duration
	sessionExpDuration   time.Duration
	db                   UserAuthDatabase
}

// NewAuthMagicLinkController configures and creates a new instance of the AuthMagicLinkController.
// The secretKey needs to be kept safe. To provide your own storage mechanism for the magic
// link data, implement the UserAuthDatabase interface. There are file system and SQL database
// implementations provided.
func NewAuthMagicLinkController(secretKey []byte, challengeExpDuration time.Duration, sessionExpDuration time.Duration, db UserAuthDatabase) (mlc *AuthMagicLinkController, err error) {
	if len(secretKey) < 16 {
		return nil, ErrSecretKeyTooShort
	}
	keyHash := sha256.Sum256(secretKey)
	return &AuthMagicLinkController{
		secretKeyHash:        keyHash[:],
		challengeExpDuration: challengeExpDuration,
		sessionExpDuration:   sessionExpDuration,
		db:                   db,
	}, nil
}

func (mlc *AuthMagicLinkController) makeHMAC(payload []byte) []byte {
	mac := hmac.New(sha256.New, mlc.secretKeyHash)
	mac.Write(payload)
	return mac.Sum(nil)
}

func (mlc *AuthMagicLinkController) GetUserByEmail(email string) (*AuthUserRecord, error) {
	return mlc.db.GetUserByEmail(email)
}

func (mlc *AuthMagicLinkController) StoreUser(user *AuthUserRecord) error {
	return mlc.db.StoreUser(user)
}

func (mlc *AuthMagicLinkController) UserExistsByEmail(email string) bool {
	return mlc.db.UserExistsByEmail(email)
}

func (mlc *AuthMagicLinkController) GetUserCount() (int, error) {
	return mlc.db.GetUserCount()
}

func (mlc *AuthMagicLinkController) UsersExist() (bool, error) {
	return mlc.db.UsersExist()
}

// GenerateChallenge creates a challenge string to be used for constructing the magic link.
// This challenge string needs to be verified by VerifyChallenge()
func (mlc *AuthMagicLinkController) GenerateChallenge(email string) (challenge string, err error) {
	// Challenge is in the format:
	// SALT-EMAIL-EXPTIME-HMAC(SALT || EMAIL || EXPTIME, secredKeyHash)
	email = NormalizeEmail(email)
	salt := make([]byte, saltLength)
	_, err = rand.Read(salt)
	if err != nil {
		return
	}
	expTime := time.Now().Add(mlc.challengeExpDuration).Unix()
	hmac := mlc.makeHMAC(slices.Concat(salt, []byte{0}, []byte(email), []byte{0}, []byte(strconv.Itoa(int(expTime)))))
	challenge = fmt.Sprintf("%s%s-%s-%d-%s", challengeSignature, encodeToString(salt), encodeToString([]byte(email)), expTime, encodeToString(hmac))
	return challenge, nil
}

// VerifyChallenge verifies the challenge string generated by GenerateChallenge(),
// and returns the AuthUserRecord corresponding to the user for which the challenge
// was created (identifying them by their email address).
func (mlc *AuthMagicLinkController) VerifyChallenge(challenge string) (user *AuthUserRecord, err error) {
	if !strings.HasPrefix(challenge, challengeSignature) {
		return nil, ErrInvalidChallenge
	}
	challenge = challenge[len(challengeSignature):]
	parts := strings.Split(challenge, "-")
	if len(parts) != 4 {
		return nil, ErrInvalidChallenge
	}

	salt, err := decodeFromString(parts[0])
	if err != nil {
		return nil, ErrInvalidChallenge
	}
	email, err := decodeFromString(parts[1])
	if err != nil {
		return nil, ErrInvalidChallenge
	}
	expTime, err := strconv.Atoi(parts[2])
	if err != nil {
		return nil, ErrInvalidChallenge
	}
	if expTime < int(time.Now().Unix()) {
		return nil, ErrExpiredChallenge
	}
	hmac1, err := decodeFromString(parts[3])
	if err != nil {
		return nil, ErrInvalidChallenge
	}
	hmac2 := mlc.makeHMAC(slices.Concat(salt, []byte{0}, []byte(email), []byte{0}, []byte(strconv.Itoa(int(expTime)))))
	if !hmac.Equal(hmac1, hmac2) {
		return nil, ErrBrokenChallenge
	}
	// We've verified the challenge, so assume the user is real.
	// Now either create a new AuthUserRecord or load an existing one.
	user, err = mlc.db.GetUserByEmail(string(email))
	if err != nil {
		if err == ErrUserNotFound {
			user, err = NewAuthUserRecord(string(email))
		}
	}

	if user != nil {
		if !user.Enabled {
			return nil, ErrUserDisabled
		}
		user.RecentLoginTime = time.Now()
	}
	return
}

// GenerateSessionId generates a session id suitable for using as a cookie
// in a web app.
func (mlc *AuthMagicLinkController) GenerateSessionId(user *AuthUserRecord) (sessionId string, err error) {
	// Session ID is in the format:
	// SALT-USER_ID-EXPTIME-HMAC(SALT || USER_ID || EXPTIME, secretKeyHash)
	salt := make([]byte, saltLength)
	_, err = rand.Read(salt)
	if err != nil {
		return
	}
	userId := user.ID.String()
	expTime := 0
	if mlc.sessionExpDuration > 0 {
		expTime = int(time.Now().Add(mlc.sessionExpDuration).Unix())
	}
	expTimeStr := strconv.Itoa(expTime)

	userIDBytes, err := user.ID.MarshalBinary()
	if err != nil {
		return
	}

	hmac := mlc.makeHMAC(slices.Concat(salt, []byte{0}, userIDBytes, []byte{0}, []byte(expTimeStr)))

	return strings.Join([]string{
		sessionIdSignature + encodeToString(salt),
		userId,
		expTimeStr,
		encodeToString(hmac),
	}, sesionIdSplitChar), nil
}

// VerifySessionId verifies the session ID generated by GenerateSessionId() and if it's valid,
// returns the AuthUserRecord of the associated user.
func (mlc *AuthMagicLinkController) VerifySessionId(sessionId string) (user *AuthUserRecord, err error) {
	if !strings.HasPrefix(sessionId, sessionIdSignature) {
		slog.Error("Error finding sessionId prefix")
		return nil, ErrInvalidSessionId
	}
	sessionId = sessionId[len(sessionIdSignature):]
	parts := strings.Split(sessionId, sesionIdSplitChar)
	if len(parts) != 4 {
		slog.Error("Error in splitting sessionId", "parts", parts, "sessionId", sessionId)
		return nil, ErrInvalidSessionId
	}

	salt, err := decodeFromString(parts[0])
	if err != nil {
		slog.Error("Error decoding part 0", "error", err)
		return nil, ErrInvalidSessionId
	}
	userId, err := uuid.Parse(parts[1])
	if err != nil {
		slog.Error("Error parsing UUID", "error", err)
		return nil, ErrInvalidSessionId
	}
	expTime, err := strconv.Atoi(parts[2])
	if err != nil {
		slog.Error("Error decoding expTime", "error", err)
		return nil, ErrInvalidSessionId
	}
	if expTime < int(time.Now().Unix()) {
		slog.Error("Session ID expired")
		return nil, ErrExpiredSessionId
	}
	hmac1, err := decodeFromString(parts[3])
	if err != nil {
		slog.Error("Error decoding part 3", "error", err)
		return nil, ErrInvalidSessionId
	}
	userIdBinary, err := userId.MarshalBinary()
	if err != nil {
		slog.Error("Error marshaling userID to binary", "error", err)
		return nil, ErrInvalidSessionId
	}
	hmac2 := mlc.makeHMAC(slices.Concat(salt, []byte{0}, userIdBinary, []byte{0}, []byte(parts[2])))
	if !hmac.Equal(hmac1, hmac2) {
		return nil, ErrBrokenSessionId
	}
	// Now we're sure the session Id is validated, so the userId should be valid
	user, err = mlc.db.GetUserById(userId)
	if !user.Enabled {
		return nil, ErrUserDisabled
	}
	user.RecentLoginTime = time.Now()
	return
}

// AuthUser represents user data
type AuthUserRecord struct {
	ID              uuid.UUID         `json:"id"` // Unique identifier
	Enabled         bool              `json:"enabled"`
	Email           string            `json:"email"` // Also must be unique
	AccessLevel     int               `json:"access_level"`
	FirstLoginTime  time.Time         `json:"first_login_time"`
	RecentLoginTime time.Time         `json:"recent_login_time"`
	CustomData      map[string]string `json:"custom_data"` // Apps can attach custom data to the user record
}

// NewAuthUserRecords constructs a new AuthUserRecord. This function isn't normally
// directly called by the users of this package.
func NewAuthUserRecord(email string) (aur *AuthUserRecord, err error) {
	newID, err := uuid.NewV7()
	if err != nil {
		return
	}
	now := time.Now()
	aur = &AuthUserRecord{
		ID:              newID,
		Email:           NormalizeEmail(email),
		Enabled:         true,
		FirstLoginTime:  now,
		RecentLoginTime: now,
		CustomData:      nil,
	}
	return aur, nil
}

// Returns the user ID.
func (aur *AuthUserRecord) GetID() uuid.UUID {
	if aur.ID == uuid.Nil {
		var err error
		aur.ID, err = uuid.NewV7()
		if err != nil {
			panic(err)
		}
	}
	return aur.ID
}

// Returns the Key name suitable for key-value databases.
func (aur *AuthUserRecord) GetKeyName() string {
	if aur.ID == uuid.Nil {
		var err error
		aur.ID, err = uuid.NewV7()
		if err != nil {
			panic(err)
		}
	}
	return fmt.Sprintf("_%s_%s", aur.ID.String(), aur.Email)
}

// Binary-string encoding
func encodeToString(b []byte) string {
	return strings.TrimRight(base32.StdEncoding.EncodeToString(b), "=")
}

func decodeFromString(s string) ([]byte, error) {
	s = s + strings.Repeat("=", 8-(len(s)%8))
	return base32.StdEncoding.DecodeString(s)
}