Yggdrasil/internal/middleware/handlers.go

package middleware

import (
	"net/http"
	"strings"
	"time"

	"github.com/Klagarge/yggdrasil/internal/oidc"
)

// loginState remembers where a pending OIDC flow must return.
type loginState struct {
	target  string
	expires time.Time
}

// session reads the session cookie and returns the session if valid.
func (ac *AccessControl) session(r *http.Request) (*Session, bool) {
	c, err := r.Cookie(sessionCookie)
	if err != nil || c.Value == "" {
		return nil, false
	}
	return ac.Sessions.Get(c.Value)
}

func (ac *AccessControl) setSessionCookie(w http.ResponseWriter, id string) {
	http.SetCookie(w, &http.Cookie{
		Name:     sessionCookie,
		Value:    id,
		Path:     "/",
		HttpOnly: true,
		Secure:   ac.CookieSecure,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   int(ac.Sessions.TTL() / time.Second),
	})
}

func (ac *AccessControl) clearSessionCookie(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     sessionCookie,
		Value:    "",
		Path:     "/",
		HttpOnly: true,
		Secure:   ac.CookieSecure,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   -1,
	})
}

// Login redirects to the OIDC provider, remembering where to return.
func (ac *AccessControl) Login(w http.ResponseWriter, r *http.Request) {
	if ac.Auth == nil {
		http.Error(w, "OIDC is not configured on this server", http.StatusServiceUnavailable)
		return
	}
	target := sanitizeTarget(r.URL.Query().Get("target"))
	state := randomID()
	ac.saveState(state, target)
	http.Redirect(w, r, ac.Auth.AuthCodeURL(state), http.StatusFound)
}

// Callback finishes the OIDC flow: it verifies the provider's response and
// starts a session.
func (ac *AccessControl) Callback(w http.ResponseWriter, r *http.Request) {
	if ac.Auth == nil {
		http.Error(w, "OIDC is not configured on this server", http.StatusServiceUnavailable)
		return
	}
	target, ok := ac.takeState(r.URL.Query().Get("state"))
	if !ok {
		http.Error(w, "unknown or expired login state; restart the login", http.StatusUnauthorized)
		return
	}
	code := r.URL.Query().Get("code")
	if code == "" {
		http.Error(w, "missing authorization code", http.StatusBadRequest)
		return
	}
	token, err := ac.Auth.Exchange(r.Context(), code)
	if err != nil {
		http.Error(w, "token exchange failed", http.StatusBadGateway)
		return
	}
	idToken, err := ac.Auth.VerifyIDToken(r.Context(), token)
	if err != nil {
		http.Error(w, "invalid identity token", http.StatusUnauthorized)
		return
	}
	profile, err := oidc.ProfileFromToken(idToken)
	if err != nil {
		http.Error(w, "unreadable identity token claims", http.StatusUnauthorized)
		return
	}
	ac.setSessionCookie(w, ac.Sessions.Create(profile))
	http.Redirect(w, r, target, http.StatusFound)
}

// Logout ends the session.
func (ac *AccessControl) Logout(w http.ResponseWriter, r *http.Request) {
	if c, err := r.Cookie(sessionCookie); err == nil && c.Value != "" {
		ac.Sessions.Delete(c.Value)
	}
	ac.clearSessionCookie(w)
	http.Redirect(w, r, "/", http.StatusFound)
}

func (ac *AccessControl) saveState(state, target string) {
	ac.stateMu.Lock()
	now := time.Now()
	for k, v := range ac.states { // opportunistic purge
		if now.After(v.expires) {
			delete(ac.states, k)
		}
	}
	ac.states[state] = loginState{target: target, expires: now.Add(stateTTL)}
	ac.stateMu.Unlock()
}

func (ac *AccessControl) takeState(state string) (string, bool) {
	if state == "" {
		return "", false
	}
	ac.stateMu.Lock()
	defer ac.stateMu.Unlock()
	s, ok := ac.states[state]
	if ok {
		delete(ac.states, state) // states are single use
	}
	if !ok || time.Now().After(s.expires) {
		return "", false
	}
	return s.target, true
}

// sanitizeTarget only accepts local redirect targets (open-redirect guard).
func sanitizeTarget(t string) string {
	if t == "" || !strings.HasPrefix(t, "/") || strings.HasPrefix(t, "//") || strings.HasPrefix(t, "/\\") {
		return "/"
	}
	return t
}