at-container-registry/pkg/hold/admin/flash.go

package admin

import (
	"encoding/base64"
	"encoding/json"
	"net/http"
)

const flashCookieName = "hold_admin_flash"

// validFlashCategories bounds what can appear in alert class interpolation.
// Anything outside this set is coerced to "info" before the cookie is set.
var validFlashCategories = map[string]bool{
	"success": true,
	"error":   true,
	"warning": true,
	"info":    true,
}

// setFlash sets a flash message cookie. Category is coerced to "info" if
// not on the known allowlist — the value flows into an HTML class attribute
// in layout.html and must not carry arbitrary text.
func setFlash(w http.ResponseWriter, r *http.Request, category, message string) {
	if !validFlashCategories[category] {
		category = "info"
	}
	flash := Flash{
		Category: category,
		Message:  message,
	}

	data, err := json.Marshal(flash)
	if err != nil {
		return
	}

	encoded := base64.URLEncoding.EncodeToString(data)
	secure := r.TLS != nil || r.Header.Get("X-Forwarded-Proto") == "https"

	http.SetCookie(w, &http.Cookie{
		Name:     flashCookieName,
		Value:    encoded,
		Path:     "/admin",
		MaxAge:   60, // 1 minute - should be consumed on next page load
		HttpOnly: true,
		Secure:   secure,
		SameSite: http.SameSiteStrictMode,
	})
}

// getFlash retrieves the flash message. Callers that render the flash into
// the page layout should arrange for the cookie to be cleared after display
// via clearFlash — handlers typically defer clearFlash(w) at entry.
// getFlash also rejects flashes whose category has drifted off the
// allowlist (e.g. a forged or pre-upgrade cookie).
func getFlash(r *http.Request, ui *AdminUI) *Flash {
	cookie, err := r.Cookie(flashCookieName)
	if err != nil {
		return nil
	}

	data, err := base64.URLEncoding.DecodeString(cookie.Value)
	if err != nil {
		return nil
	}

	var flash Flash
	if err := json.Unmarshal(data, &flash); err != nil {
		return nil
	}

	if !validFlashCategories[flash.Category] {
		flash.Category = "info"
	}

	return &flash
}

// clearFlash clears the flash cookie (called after displaying)
func clearFlash(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     flashCookieName,
		Value:    "",
		Path:     "/admin",
		MaxAge:   -1,
		HttpOnly: true,
		SameSite: http.SameSiteStrictMode,
	})
}