fiber/redirect.go

// ⚡️ Fiber is an Express inspired web framework written in Go with ☕️
// 📝 GitHub Repository: https://github.com/gofiber/fiber
// 📌 API Documentation: https://docs.gofiber.io

package fiber

import (
	"bytes"
	"encoding/hex"
	"sync"

	"github.com/gofiber/utils/v2"
	utilsbytes "github.com/gofiber/utils/v2/bytes"
	"github.com/valyala/bytebufferpool"
	"github.com/valyala/fasthttp"

	"github.com/zap-proto/fiber/v3/binder"
)

// Pool for redirection
var (
	redirectPool = sync.Pool{
		New: func() any {
			return &Redirect{
				status:   StatusSeeOther,
				messages: make(redirectionMsgs, 0),
			}
		},
	}
	oldInputPool = sync.Pool{
		New: func() any {
			return make(map[string]string)
		},
	}
)

const maxPoolableMapSize = 64

// FlashCookieName Cookie name to send flash messages when to use redirection.
const (
	FlashCookieName = "fiber_flash"
)

var flashCookieNeedle = []byte(FlashCookieName + "=")

// hasFlashCookie is on the request hot path and runs on every request/response cycle.
// Keep this cheap for users who don't use flash messages:
// 1) a fast raw-header prefilter to avoid unnecessary cookie parsing,
// 2) an exact cookie lookup to avoid prefix false positives (e.g. fiber_flashX).
func hasFlashCookie(header *fasthttp.RequestHeader) bool {
	rawHeaders := header.RawHeaders()
	if len(rawHeaders) == 0 {
		return false
	}

	if !bytes.Contains(rawHeaders, flashCookieNeedle) {
		return false
	}

	return header.Cookie(FlashCookieName) != nil
}

// redirectionMsgs is a struct that used to store flash messages and old input data in cookie using MSGP.
// msgp -file="redirect.go" -o="redirect_msgp.go" -unexported
// Cookie payloads are limited to ~4KB, so keep flash message counts bounded but usable.
//
//msgp:limit arrays:256 maps:32 marshal:true
//msgp:ignore Redirect RedirectConfig OldInputData FlashMessage
type redirectionMsg struct {
	key        string
	value      string
	level      uint8
	isOldInput bool
}

type redirectionMsgs []redirectionMsg

// OldInputData is a struct that holds the old input data.
type OldInputData struct {
	Key   string
	Value string
}

// FlashMessage is a struct that holds the flash message data.
type FlashMessage struct {
	Key   string
	Value string
	Level uint8
}

// Redirect is a struct that holds the redirect data.
type Redirect struct {
	c        *DefaultCtx     // Embed ctx
	messages redirectionMsgs // Flash messages and old input data
	status   int             // Status code of redirection. Default: 303 StatusSeeOther
}

// RedirectConfig A config to use with Redirect().Route()
// You can specify queries or route parameters.
// NOTE: We don't use net/url to parse parameters because of it has poor performance. You have to pass map.
type RedirectConfig struct {
	Params  Map               // Route parameters
	Queries map[string]string // Query map
}

// AcquireRedirect return default Redirect reference from the redirect pool
func AcquireRedirect() *Redirect {
	redirect, ok := redirectPool.Get().(*Redirect)
	if !ok {
		panic(errRedirectTypeAssertion)
	}

	return redirect
}

// ReleaseRedirect returns c acquired via Redirect to redirect pool.
//
// It is forbidden accessing req and/or its members after returning
// it to redirect pool.
func ReleaseRedirect(r *Redirect) {
	r.release()
	redirectPool.Put(r)
}

func (r *Redirect) release() {
	r.status = StatusSeeOther
	r.messages = r.messages[:0]
	r.c = nil
}

func acquireOldInput() map[string]string {
	oldInput, ok := oldInputPool.Get().(map[string]string)
	if !ok {
		return make(map[string]string)
	}

	return oldInput
}

func releaseOldInput(oldInput map[string]string) {
	if len(oldInput) > maxPoolableMapSize {
		return
	}

	clear(oldInput)
	oldInputPool.Put(oldInput)
}

// Status sets the status code of redirection.
// If status is not specified, status defaults to 303 See Other.
func (r *Redirect) Status(code int) *Redirect {
	r.status = code

	return r
}

// With You can send flash messages by using With().
// They will be sent as a cookie.
// You can get them by using: Redirect().Messages(), Redirect().Message()
// Note: You must use escape char before using ',' and ':' chars to avoid wrong parsing.
func (r *Redirect) With(key, value string, level ...uint8) *Redirect {
	// Get level
	var msgLevel uint8
	if len(level) > 0 {
		msgLevel = level[0]
	}

	// Override old message if exists
	for i, msg := range r.messages {
		if msg.key == key && !msg.isOldInput {
			r.messages[i].value = value
			r.messages[i].level = msgLevel

			return r
		}
	}

	r.messages = append(r.messages, redirectionMsg{
		key:   key,
		value: value,
		level: msgLevel,
	})

	return r
}

// WithInput You can send input data by using WithInput().
// They will be sent as a cookie.
// This method can send form, multipart form, query data to redirected route.
// You can get them by using: Redirect().OldInputs(), Redirect().OldInput()
func (r *Redirect) WithInput() *Redirect {
	// Get content-type
	ctype := utils.UnsafeString(utilsbytes.UnsafeToLower(r.c.RequestCtx().Request.Header.ContentType()))
	ctype = binder.FilterFlags(utils.ParseVendorSpecificContentType(ctype))

	oldInput := acquireOldInput()
	defer releaseOldInput(oldInput)

	switch ctype {
	case MIMEApplicationForm, MIMEMultipartForm:
		_ = r.c.Bind().Form(oldInput) //nolint:errcheck // not needed
	default:
		_ = r.c.Bind().Query(oldInput) //nolint:errcheck // not needed
	}

	// Add old input data
	for k, v := range oldInput {
		r.messages = append(r.messages, redirectionMsg{
			key:        k,
			value:      v,
			isOldInput: true,
		})
	}

	return r
}

// Messages Get flash messages.
func (r *Redirect) Messages() []FlashMessage {
	if len(r.c.flashMessages) == 0 {
		return nil
	}

	flashMessages := make([]FlashMessage, 0, len(r.c.flashMessages))
	writeIdx := 0

	for _, msg := range r.c.flashMessages {
		if msg.isOldInput {
			r.c.flashMessages[writeIdx] = msg
			writeIdx++
			continue
		}

		flashMessages = append(flashMessages, FlashMessage{
			Key:   msg.key,
			Value: msg.value,
			Level: msg.level,
		})
	}

	for i := writeIdx; i < len(r.c.flashMessages); i++ {
		r.c.flashMessages[i] = redirectionMsg{}
	}

	r.c.flashMessages = r.c.flashMessages[:writeIdx]

	return flashMessages
}

// Message Get flash message by key.
func (r *Redirect) Message(key string) FlashMessage {
	if len(r.c.flashMessages) == 0 {
		return FlashMessage{}
	}

	var flashMessage FlashMessage
	found := false
	writeIdx := 0

	for _, msg := range r.c.flashMessages {
		if msg.isOldInput || found || msg.key != key {
			r.c.flashMessages[writeIdx] = msg
			writeIdx++
			continue
		}

		flashMessage = FlashMessage{
			Key:   msg.key,
			Value: msg.value,
			Level: msg.level,
		}
		found = true
	}

	for i := writeIdx; i < len(r.c.flashMessages); i++ {
		r.c.flashMessages[i] = redirectionMsg{}
	}

	r.c.flashMessages = r.c.flashMessages[:writeIdx]

	return flashMessage
}

// OldInputs Get old input data.
func (r *Redirect) OldInputs() []OldInputData {
	// Count old inputs first to avoid allocation if none exist
	count := 0
	for _, msg := range r.c.flashMessages {
		if msg.isOldInput {
			count++
		}
	}

	if count == 0 {
		return nil
	}

	inputs := make([]OldInputData, 0, count)
	for _, msg := range r.c.flashMessages {
		if msg.isOldInput {
			inputs = append(inputs, OldInputData{
				Key:   msg.key,
				Value: msg.value,
			})
		}
	}

	return inputs
}

// OldInput Get old input data by key.
func (r *Redirect) OldInput(key string) OldInputData {
	msgs := r.c.flashMessages

	for _, msg := range msgs {
		if msg.key == key && msg.isOldInput {
			return OldInputData{
				Key:   msg.key,
				Value: msg.value,
			}
		}
	}

	return OldInputData{}
}

// To redirect to the URL derived from the specified path, with specified status.
func (r *Redirect) To(location string) error {
	r.c.setCanonical(HeaderLocation, location)
	r.c.Status(r.status)

	r.processFlashMessages()

	return nil
}

// Route redirects to the Route registered in the app with appropriate parameters.
// If you want to send queries or params to route, you should use config parameter.
func (r *Redirect) Route(name string, config ...RedirectConfig) error {
	// Check config
	cfg := RedirectConfig{}
	if len(config) > 0 {
		cfg = config[0]
	}

	// Get location from route name
	route := r.c.App().GetRoute(name)
	location, err := r.c.getLocationFromRoute(&route, cfg.Params)
	if err != nil {
		return err
	}

	// Check queries
	if len(cfg.Queries) > 0 {
		queryText := bytebufferpool.Get()
		defer bytebufferpool.Put(queryText)

		first := true
		for k, v := range cfg.Queries {
			if !first {
				queryText.WriteByte('&')
			}
			first = false
			queryText.WriteString(k)
			queryText.WriteByte('=')
			queryText.WriteString(v)
		}

		return r.To(location + "?" + r.c.app.toString(queryText.Bytes()))
	}

	return r.To(location)
}

// Back redirect to the URL to referer.
func (r *Redirect) Back(fallback ...string) error {
	location := r.c.Get(HeaderReferer)
	if location == "" {
		// Check fallback URL
		if len(fallback) == 0 {
			err := ErrRedirectBackNoFallback
			r.c.Status(err.Code)

			return err
		}
		location = fallback[0]
	}

	return r.To(location)
}

// parseAndClearFlashMessages is a method to get flash messages before they are getting removed
func (r *Redirect) parseAndClearFlashMessages() {
	// parse flash messages
	cookieValue, err := hex.DecodeString(r.c.Cookies(FlashCookieName))
	if err != nil {
		return
	}

	_, err = r.c.flashMessages.UnmarshalMsg(cookieValue)
	if err != nil {
		return
	}

	r.c.Cookie(&Cookie{
		Name:   FlashCookieName,
		Value:  "",
		Path:   "/",
		MaxAge: -1,
	})
}

// processFlashMessages is a helper function to process flash messages and old input data
// and set them as cookies
func (r *Redirect) processFlashMessages() {
	if len(r.messages) == 0 {
		return
	}

	val, err := r.messages.MarshalMsg(nil)
	if err != nil {
		return
	}

	dst := make([]byte, hex.EncodedLen(len(val)))
	hex.Encode(dst, val)

	r.c.Cookie(&Cookie{
		Name:        FlashCookieName,
		Value:       r.c.app.toString(dst),
		SessionOnly: true,
	})
}