Vellum/internal/auth/middleware.go

package auth

import (
	"context"
	"database/sql"
	"log/slog"
	"net"
	"net/http"
	"time"

	"github.com/mithral/vellum/internal/httpx"
)

// contextKey is unexported so no other package can collide with our context values.
type contextKey string

const (
	userContextKey    contextKey = "auth.user"
	sessionContextKey contextKey = "auth.session"
)

// Authenticator loads sessions from cookies and enforces access.
type Authenticator struct {
	DB         *sql.DB
	CookieName string
	CookieTTL  time.Duration
	Secure     bool
	Log        *slog.Logger
}

// NewAuthenticator builds an Authenticator with sane cookie defaults.
func NewAuthenticator(db *sql.DB, cookieName string, ttl time.Duration, secure bool, log *slog.Logger) *Authenticator {
	if cookieName == "" {
		cookieName = "vellum_session"
	}
	return &Authenticator{
		DB:         db,
		CookieName: cookieName,
		CookieTTL:  ttl,
		Secure:     secure,
		Log:        log,
	}
}

// Middleware attaches the current user to the request context when a valid
// session cookie is present. It never rejects a request: enforcement is the
// job of RequireAuth, so public routes can also see the caller.
func (a *Authenticator) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()

		if token, err := a.TokenFromRequest(r); err == nil {
			if sess, err := SessionByToken(ctx, a.DB, token); err == nil {
				if u, err := GetUserByID(ctx, a.DB, sess.UserID); err == nil {
					ctx = context.WithValue(ctx, userContextKey, u)
					ctx = context.WithValue(ctx, sessionContextKey, sess)
					// Best-effort; must not fail the request.
					if err := TouchSession(ctx, a.DB, sess.ID); err != nil {
						a.Log.Warn("touch session", "err", err)
					}
					r = r.WithContext(ctx)
				} else if a.Log != nil {
					a.Log.Warn("session references missing user", "user_id", sess.UserID)
				}
			} else if err != ErrNoSession && a.Log != nil {
				a.Log.Warn("load session", "err", err)
			}
		}
		next.ServeHTTP(w, r)
	})
}

// RequireAuth rejects unauthenticated requests with 401.
func (a *Authenticator) RequireAuth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if _, ok := CurrentUser(r.Context()); !ok {
			httpx.WriteError(w, http.StatusUnauthorized, "authentication required")
			return
		}
		next.ServeHTTP(w, r)
	})
}

// RequireAdmin rejects authenticated non-admins with 403.
func (a *Authenticator) RequireAdmin(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		u, ok := CurrentUser(r.Context())
		if !ok {
			httpx.WriteError(w, http.StatusUnauthorized, "authentication required")
			return
		}
		if !u.IsAdmin {
			httpx.WriteError(w, http.StatusForbidden, "administrator access required")
			return
		}
		next.ServeHTTP(w, r)
	})
}

// Require calls fn with the current user, or 401s if there is none.
func (a *Authenticator) Require(fn func(http.ResponseWriter, *http.Request, User)) http.Handler {
	return a.RequireAuth(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		u, _ := CurrentUser(r.Context())
		fn(w, r, u)
	}))
}

// CurrentUser returns the authenticated user, if any.
func CurrentUser(ctx context.Context) (User, bool) {
	u, ok := ctx.Value(userContextKey).(User)
	return u, ok
}

// CurrentSession returns the session backing the authenticated user, if any.
func CurrentSession(ctx context.Context) (Session, bool) {
	s, ok := ctx.Value(sessionContextKey).(Session)
	return s, ok
}

// ── cookies ────────────────────────────────────────────────────────────

// SetSessionCookie issues the session cookie. HttpOnly and SameSite=Lax are
// always on; Secure follows configuration (it must be off for plain-HTTP dev).
func (a *Authenticator) SetSessionCookie(w http.ResponseWriter, token string, expires time.Time) {
	http.SetCookie(w, &http.Cookie{
		Name:     a.CookieName,
		Value:    token,
		Path:     "/",
		Expires:  expires,
		MaxAge:   int(time.Until(expires).Seconds()),
		HttpOnly: true,
		Secure:   a.Secure,
		SameSite: http.SameSiteLaxMode,
	})
}

// ClearSessionCookie removes the session cookie.
func (a *Authenticator) ClearSessionCookie(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     a.CookieName,
		Value:    "",
		Path:     "/",
		MaxAge:   -1,
		HttpOnly: true,
		Secure:   a.Secure,
		SameSite: http.SameSiteLaxMode,
	})
}

// TokenFromRequest reads the session token from the cookie, or the
// Authorization: Bearer header for non-browser clients (native apps, CLIs).
func (a *Authenticator) TokenFromRequest(r *http.Request) (string, error) {
	if c, err := r.Cookie(a.CookieName); err == nil && c.Value != "" {
		return c.Value, nil
	}
	if h := r.Header.Get("Authorization"); len(h) > 7 && h[:7] == "Bearer " {
		return h[7:], nil
	}
	return "", ErrNoSession
}

// ClientIP extracts a best-effort client address for logging and rate limiting.
// It deliberately does NOT trust X-Forwarded-For: behind an untrusted network
// that header is attacker-controlled, and this server is LAN-facing.
func ClientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}