oxy/roundrobin/stickysessions.go

package roundrobin

import (
	"errors"
	"net/http"
	"net/url"
	"time"

	"github.com/vulcand/oxy/v2/roundrobin/stickycookie"
)

// CookieOptions has all the options one would like to set on the affinity cookie.
type CookieOptions struct {
	HTTPOnly bool
	Secure   bool

	Path    string
	Domain  string
	Expires time.Time

	MaxAge   int
	SameSite http.SameSite
}

// StickySession is a mixin for load balancers that implements layer 7 (http cookie) session affinity.
type StickySession struct {
	cookieName  string
	cookieValue stickycookie.CookieValue
	options     CookieOptions
}

// NewStickySession creates a new StickySession.
func NewStickySession(cookieName string) *StickySession {
	return &StickySession{cookieName: cookieName, cookieValue: &stickycookie.RawValue{}}
}

// NewStickySessionWithOptions creates a new StickySession whilst allowing for options to
// shape its affinity cookie such as "httpOnly" or "secure".
func NewStickySessionWithOptions(cookieName string, options CookieOptions) *StickySession {
	return &StickySession{cookieName: cookieName, options: options, cookieValue: &stickycookie.RawValue{}}
}

// SetCookieValue set the CookieValue for the StickySession.
func (s *StickySession) SetCookieValue(value stickycookie.CookieValue) *StickySession {
	s.cookieValue = value
	return s
}

// GetBackend returns the backend URL stored in the sticky cookie, iff the backend is still in the valid list of servers.
func (s *StickySession) GetBackend(req *http.Request, servers []*url.URL) (*url.URL, bool, error) {
	cookie, err := req.Cookie(s.cookieName)
	if err != nil {
		if errors.Is(err, http.ErrNoCookie) {
			return nil, false, nil
		}

		return nil, false, err
	}

	server, err := s.cookieValue.FindURL(cookie.Value, servers)

	return server, server != nil, err
}

// StickBackend creates and sets the cookie.
func (s *StickySession) StickBackend(backend *url.URL, w http.ResponseWriter) {
	opt := s.options

	cp := "/"
	if opt.Path != "" {
		cp = opt.Path
	}

	cookie := &http.Cookie{
		Name:     s.cookieName,
		Value:    s.cookieValue.Get(backend),
		Path:     cp,
		Domain:   opt.Domain,
		Expires:  opt.Expires,
		MaxAge:   opt.MaxAge,
		Secure:   opt.Secure,
		HttpOnly: opt.HTTPOnly,
		SameSite: opt.SameSite,
	}
	http.SetCookie(w, cookie)
}