at-container-registry/pkg/labeler/auth.go

package labeler

import (
	"context"
	"crypto/rand"
	"crypto/subtle"
	"encoding/base64"
	"fmt"
	"html/template"
	"log/slog"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

const (
	sessionCookieName = "labeler_session"
	sessionTTL        = 24 * time.Hour
	csrfHeaderName    = "X-CSRF-Token"
	csrfFormField     = "csrf_token"
)

// Session represents an authenticated admin session. Restart wipes the in-memory
// map so any stolen cookie token becomes useless after a restart, by design.
//
// UserAgent and IPPrefix are captured at login and rechecked on every request —
// a stolen token replayed from a different browser or network prefix is rejected
// and the session is torn down. Empty bound values (Unix sockets, tests, unusual
// proxies) opt out rather than locking users out. Binding at /24 (IPv4) / /64
// (IPv6) tolerates DHCP renewals within a prefix without inviting cross-network
// replay.
type Session struct {
	DID       string
	Handle    string
	CSRFToken string
	CreatedAt time.Time
	UserAgent string
	IPPrefix  string
}

// Auth manages in-memory admin sessions for the labeler.
type Auth struct {
	ownerDID   string
	sessions   map[string]*Session
	sessionsMu sync.RWMutex
}

// NewAuth wires a fresh in-memory session store keyed to the configured owner DID.
func NewAuth(ownerDID string) *Auth {
	return &Auth{
		ownerDID: ownerDID,
		sessions: make(map[string]*Session),
	}
}

func randToken() (string, error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", fmt.Errorf("generate token: %w", err)
	}
	return base64.URLEncoding.EncodeToString(b), nil
}

// CreateSession installs a new in-memory session and returns its cookie token
// alongside the embedded CSRF token (for echoing into forms).
func (a *Auth) CreateSession(did, handle, userAgent, ipPrefix string) (string, *Session, error) {
	token, err := randToken()
	if err != nil {
		return "", nil, err
	}
	csrfToken, err := randToken()
	if err != nil {
		return "", nil, err
	}
	s := &Session{
		DID:       did,
		Handle:    handle,
		CSRFToken: csrfToken,
		CreatedAt: time.Now(),
		UserAgent: userAgent,
		IPPrefix:  ipPrefix,
	}
	a.sessionsMu.Lock()
	a.sessions[token] = s
	a.sessionsMu.Unlock()
	return token, s, nil
}

// GetSession returns the session for the cookie token, evicting expired entries on access.
func (a *Auth) GetSession(token string) *Session {
	a.sessionsMu.RLock()
	s := a.sessions[token]
	a.sessionsMu.RUnlock()
	if s == nil {
		return nil
	}
	if !s.CreatedAt.IsZero() && time.Since(s.CreatedAt) > sessionTTL {
		a.sessionsMu.Lock()
		delete(a.sessions, token)
		a.sessionsMu.Unlock()
		return nil
	}
	return s
}

// DeleteSession removes a session by cookie token (logout).
func (a *Auth) DeleteSession(token string) {
	a.sessionsMu.Lock()
	delete(a.sessions, token)
	a.sessionsMu.Unlock()
}

func setSessionCookie(w http.ResponseWriter, r *http.Request, token string) {
	secure := r.TLS != nil || r.Header.Get("X-Forwarded-Proto") == "https"
	http.SetCookie(w, &http.Cookie{
		Name:     sessionCookieName,
		Value:    token,
		Path:     "/",
		MaxAge:   int(sessionTTL.Seconds()),
		HttpOnly: true,
		Secure:   secure,
		SameSite: http.SameSiteLaxMode,
	})
}

func clearSessionCookie(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     sessionCookieName,
		Value:    "",
		Path:     "/",
		MaxAge:   -1,
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
	})
}

func getSessionCookie(r *http.Request) (string, bool) {
	cookie, err := r.Cookie(sessionCookieName)
	if err != nil {
		return "", false
	}
	return cookie.Value, true
}

type sessionContextKeyT struct{}

var sessionContextKey = sessionContextKeyT{}

// SessionFromContext returns the session attached to the request context, if any.
func SessionFromContext(ctx context.Context) *Session {
	s, _ := ctx.Value(sessionContextKey).(*Session)
	return s
}

// RequireOwner enforces a valid session bound to the owner DID, with UA / IP-prefix
// replay defense. State-mutating methods then go through the CSRF check below.
func (a *Auth) RequireOwner(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		token, ok := getSessionCookie(r)
		if !ok {
			http.Redirect(w, r, "/auth/login", http.StatusFound)
			return
		}
		session := a.GetSession(token)
		if session == nil {
			clearSessionCookie(w)
			http.Redirect(w, r, "/auth/login", http.StatusFound)
			return
		}
		if session.DID != a.ownerDID {
			a.DeleteSession(token)
			clearSessionCookie(w)
			http.Redirect(w, r, "/auth/login?error=access+denied", http.StatusFound)
			return
		}
		if session.UserAgent != "" && session.UserAgent != r.UserAgent() {
			slog.Warn("Admin session UA mismatch — suspected token replay", "did", session.DID)
			a.DeleteSession(token)
			clearSessionCookie(w)
			http.Redirect(w, r, "/auth/login?error=access+denied", http.StatusFound)
			return
		}
		if session.IPPrefix != "" {
			if now := clientIPPrefix(r); now != "" && now != session.IPPrefix {
				slog.Warn("Admin session IP-prefix mismatch — suspected token replay",
					"did", session.DID, "session", session.IPPrefix, "request", now)
				a.DeleteSession(token)
				clearSessionCookie(w)
				http.Redirect(w, r, "/auth/login?error=access+denied", http.StatusFound)
				return
			}
		}

		ctx := context.WithValue(r.Context(), sessionContextKey, session)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// RequireCSRF validates a per-session CSRF token on state-mutating requests. Safe
// methods pass through. Token comes from X-CSRF-Token header or the csrf_token form
// field for application/x-www-form-urlencoded bodies. Must run after RequireOwner.
func (a *Auth) RequireCSRF(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodGet, http.MethodHead, http.MethodOptions:
			next.ServeHTTP(w, r)
			return
		}
		session := SessionFromContext(r.Context())
		if session == nil || session.CSRFToken == "" {
			http.Error(w, "Forbidden: missing session", http.StatusForbidden)
			return
		}
		got := r.Header.Get(csrfHeaderName)
		if got == "" {
			contentType := r.Header.Get("Content-Type")
			if idx := strings.IndexByte(contentType, ';'); idx >= 0 {
				contentType = contentType[:idx]
			}
			contentType = strings.TrimSpace(strings.ToLower(contentType))
			if contentType == "application/x-www-form-urlencoded" {
				if err := r.ParseForm(); err == nil {
					got = r.PostFormValue(csrfFormField)
				}
			}
		}
		if subtle.ConstantTimeCompare([]byte(got), []byte(session.CSRFToken)) != 1 {
			slog.Warn("Labeler CSRF mismatch", "path", r.URL.Path, "did", session.DID)
			http.Error(w, "Forbidden: CSRF token mismatch — reload the page and try again.", http.StatusForbidden)
			return
		}
		next.ServeHTTP(w, r)
	})
}

// csrfInputHTML emits a hidden form input carrying the per-session CSRF token.
func csrfInputHTML(token string) template.HTML {
	escaped := template.HTMLEscapeString(token)
	return template.HTML(`<input type="hidden" name="` + csrfFormField + `" value="` + escaped + `">`)
}

// clientIPPrefix returns a stable prefix key for the request's client IP — /24 for
// IPv4, /64 for IPv6. Empty string means "don't bind" (avoids locking users behind
// unusual proxies / Unix sockets / tests).
func clientIPPrefix(r *http.Request) string {
	var host string
	if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
		if before, _, ok := strings.Cut(fwd, ","); ok {
			host = strings.TrimSpace(before)
		} else {
			host = strings.TrimSpace(fwd)
		}
	} else {
		h, _, err := net.SplitHostPort(r.RemoteAddr)
		if err == nil {
			host = h
		} else {
			host = r.RemoteAddr
		}
	}
	ip := net.ParseIP(host)
	if ip == nil {
		return ""
	}
	if v4 := ip.To4(); v4 != nil {
		return fmt.Sprintf("v4:%d.%d.%d", v4[0], v4[1], v4[2])
	}
	v6 := ip.To16()
	return fmt.Sprintf("v6:%02x%02x%02x%02x%02x%02x%02x%02x",
		v6[0], v6[1], v6[2], v6[3], v6[4], v6[5], v6[6], v6[7])
}