Use / as the default auth cookie path.

1 files changed, 63 insertions, 87 deletions

diff --git a/magic_link_login.go b/magic_link_login.go
index 9228a97..aa3ec27 100644
--- a/magic_link_login.go
+++ b/magic_link_login.go
@@ -1,21 +1,18 @@
 package kate
 
 import (
-	"errors"
 	"net/http"
-	"strings"
-	"time"
 )
 
 // MagicLinkConfig configures the magic link authentication handler behavior.
 type MagicLinkConfig[T any] struct {
-	// Mailer provides user data lookup and email sending
-	Mailer MagicLinkMailer[T]
+	// UserData provides user data lookup and email sending
+	UserData MagicLinkUserData[T]
 
 	// Redirects configures post-authentication redirect behavior
 	Redirects Redirects
 
-	// UsernameField is the form field name for username
+	// UsernameField is the POST form field name in the login handler for username
 	UsernameField string
 
 	// TokenField is the URL parameter name for the magic link token
@@ -24,20 +21,32 @@ type MagicLinkConfig[T any] struct {
 	// TokenLocation specifies where to retrieve the token from
 	TokenLocation TokenLocation
 
-	// TokenExpiry is how long the magic link token is valid
-	TokenExpiry time.Duration
+	// AfterSendHandler is run on the login request after the email has been sent
+	//
+	// optional - the default will write a simple 200 response.
+	AfterSendHandler http.Handler
+
+	// RejectHandler is run on the login verification request when authentication is rejected.
+	//
+	// optional - the default will write a simple 401 response with textual message.
+	RejectHandler http.Handler
 
 	// LogError is an optional function to log errors
 	LogError func(error)
 }
 
-// MagicLinkMailer provides user data lookup and email sending for magic link authentication.
-type MagicLinkMailer[T any] interface {
-	// Fetch retrieves user data by username.
+// MagicLinkUserData provides user data lookup and email sending for magic link authentication.
+type MagicLinkUserData[T any] interface {
+	// GenerateToken creates and stores a new login token.
+	GenerateToken(username, redirectPath string) (*MagicLinkToken[T], error)
+
+	// ValidateToken checks the validity of a token.
+	//
+	// It looks up the token by its identifier and returns it, or nil if it doesn't exist.
 	//
-	// Returns the user data, whether the user was found, and any error.
-	// If the user is not found, should return (zero value, false, nil).
-	Fetch(username string) (T, bool, error)
+	// Tokens should also be validated against a stored expiration date, and potentially
+	// invalidated in this function to support one-time use tokens.
+	ValidateToken(identifier string) (*MagicLinkToken[T], error)
 
 	// SendEmail sends a magic link email to the user.
 	//
@@ -46,6 +55,19 @@ type MagicLinkMailer[T any] interface {
 	SendEmail(userData T, token string) error
 }
 
+type MagicLinkToken[T any] struct {
+	// A unique identifier for this login token.
+	//
+	// This is the field that will be encrypted and embedded into the emailed login link.
+	Identifier string
+
+	// The user associated with the login token.
+	UserData T
+
+	// Path to which the user should be redirected after login.
+	RedirectPath string
+}
+
 // TokenLocation specifies where the magic link token should be retrieved from
 type TokenLocation struct{ location int }
 
@@ -63,9 +85,20 @@ func (mlc *MagicLinkConfig[T]) setDefaults() {
 	if mlc.TokenField == "" {
 		mlc.TokenField = "token"
 	}
-	// TokenLocation defaults to TokenLocationQuery (zero value)
-	if mlc.TokenExpiry == 0 {
-		mlc.TokenExpiry = 15 * time.Minute
+
+	if mlc.AfterSendHandler == nil {
+		mlc.AfterSendHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+			if _, err := w.Write([]byte("Magic link sent")); err != nil {
+				mlc.logError(err)
+			}
+		})
+	}
+
+	if mlc.RejectHandler == nil {
+		mlc.RejectHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			http.Error(w, "Authentication failed", http.StatusUnauthorized)
+		})
 	}
 }
 
@@ -75,38 +108,6 @@ func (mlc MagicLinkConfig[T]) logError(err error) {
 	}
 }
 
-type magicLinkToken struct {
-	Username  string
-	Redirect  string
-	ExpiresAt time.Time
-}
-
-func (t magicLinkToken) serialize() string {
-	return t.Redirect + "\x00" + t.ExpiresAt.Format(time.RFC3339) + "\x00" + t.Username
-}
-
-func parseMagicLinkToken(data string) (magicLinkToken, error) {
-	parts := strings.SplitN(data, "\x00", 3)
-	if len(parts) != 3 {
-		return magicLinkToken{}, errors.New("invalid token format")
-	}
-
-	redirect := parts[0]
-	expiresStr := parts[1]
-	username := parts[2]
-
-	expiresAt, err := time.Parse(time.RFC3339, expiresStr)
-	if err != nil {
-		return magicLinkToken{}, err
-	}
-
-	return magicLinkToken{
-		Redirect:  redirect,
-		ExpiresAt: expiresAt,
-		Username:  username,
-	}, nil
-}
-
 // MagicLinkLoginHandler returns an HTTP handler that processes magic link requests.
 //
 // It looks up the user, generates a token, and sends an email with the magic link.
@@ -130,39 +131,27 @@ func (a Auth[T]) MagicLinkLoginHandler(config MagicLinkConfig[T]) http.Handler {
 			return
 		}
 
-		userData, ok, err := config.Mailer.Fetch(username)
+		token, err := config.UserData.GenerateToken(username, config.Redirects.target(r))
 		if err != nil {
 			config.logError(err)
 			http.Error(w, "Error finding user", http.StatusInternalServerError)
 			return
 		}
-		if !ok {
+		if token == nil {
 			// Don't reveal whether user exists
-			w.WriteHeader(http.StatusOK)
-			if _, err := w.Write([]byte("Magic link sent")); err != nil {
-				config.logError(err)
-			}
+			config.AfterSendHandler.ServeHTTP(w, r)
 			return
 		}
 
-		tokenData := []byte(magicLinkToken{
-			Username:  username,
-			Redirect:  config.Redirects.target(r),
-			ExpiresAt: time.Now().Add(config.TokenExpiry),
-		}.serialize())
-
-		encryptedToken := a.enc.Encrypt(tokenData)
+		encryptedToken := a.enc.Encrypt([]byte(token.Identifier))
 
-		if err := config.Mailer.SendEmail(userData, encryptedToken); err != nil {
+		if err := config.UserData.SendEmail(token.UserData, encryptedToken); err != nil {
 			config.logError(err)
 			http.Error(w, "Failed to send email", http.StatusInternalServerError)
 			return
 		}
 
-		w.WriteHeader(http.StatusOK)
-		if _, err := w.Write([]byte("Magic link sent")); err != nil {
-			config.logError(err)
-		}
+		config.AfterSendHandler.ServeHTTP(w, r)
 	})
 }
 
@@ -192,42 +181,29 @@ func (a Auth[T]) MagicLinkVerifyHandler(config MagicLinkConfig[T]) http.Handler
 			return
 		}
 
-		tokenData, ok := a.enc.Decrypt(encryptedToken)
+		clearToken, ok := a.enc.Decrypt(encryptedToken)
 		if !ok {
 			http.Error(w, "Invalid token", http.StatusUnauthorized)
 			return
 		}
 
-		token, err := parseMagicLinkToken(string(tokenData))
-		if err != nil {
-			config.logError(err)
-			http.Error(w, "Invalid token", http.StatusUnauthorized)
-			return
-		}
-
-		if time.Now().After(token.ExpiresAt) {
-			http.Error(w, "Token expired", http.StatusUnauthorized)
-			return
-		}
-
-		userData, ok, err := config.Mailer.Fetch(token.Username)
+		token, err := config.UserData.ValidateToken(string(clearToken))
 		if err != nil {
 			config.logError(err)
-			http.Error(w, "Authentication failed", http.StatusUnauthorized)
+			config.RejectHandler.ServeHTTP(w, r)
 			return
 		}
-		if !ok {
-			http.Error(w, "Authentication failed", http.StatusUnauthorized)
+		if token == nil {
+			config.RejectHandler.ServeHTTP(w, r)
 			return
 		}
 
-		if err := a.Set(w, userData); err != nil {
+		if err := a.Set(w, token.UserData); err != nil {
 			config.logError(err)
 			http.Error(w, "Failed to set authentication", http.StatusInternalServerError)
 			return
 		}
 
-		http.Redirect(w, r, token.Redirect, http.StatusSeeOther)
+		http.Redirect(w, r, token.RedirectPath, http.StatusSeeOther)
 	})
 }
-