Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

reverseproxy: Add fallback for some policies, instead of always random #5488

Merged
merged 1 commit into from
May 5, 2023
Merged

reverseproxy: Add fallback for some policies, instead of always random #5488

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
184 changes: 155 additions & 29 deletions modules/caddyhttp/reverseproxy/selectionpolicies.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (
"crypto/hmac"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"hash/fnv"
weakrand "math/rand"
Expand All @@ -29,6 +30,7 @@ import (
"time"

"github.com/caddyserver/caddy/v2"
"github.com/caddyserver/caddy/v2/caddyconfig"
"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
"github.com/caddyserver/caddy/v2/modules/caddyhttp"
)
Expand Down Expand Up @@ -372,6 +374,10 @@ func (r *URIHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
type QueryHashSelection struct {
// The query key whose value is to be hashed and used for upstream selection.
Key string `json:"key,omitempty"`

// The fallback policy to use if the query key is not present. Defaults to `random`.
FallbackRaw json.RawMessage `json:"fallback,omitempty" caddy:"namespace=http.reverse_proxy.selection_policies inline_key=policy"`
fallback Selector
}

// CaddyModule returns the Caddy module information.
Expand All @@ -382,12 +388,24 @@ func (QueryHashSelection) CaddyModule() caddy.ModuleInfo {
}
}

// Select returns an available host, if any.
func (s QueryHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream {
// Provision sets up the module.
func (s *QueryHashSelection) Provision(ctx caddy.Context) error {
if s.Key == "" {
return nil
return fmt.Errorf("query key is required")
}
if s.FallbackRaw == nil {
s.FallbackRaw = caddyconfig.JSONModuleObject(RandomSelection{}, "policy", "random", nil)
}
mod, err := ctx.LoadModule(s, "FallbackRaw")
if err != nil {
return fmt.Errorf("loading fallback selection policy: %s", err)
}
s.fallback = mod.(Selector)
return nil
}

// Select returns an available host, if any.
func (s QueryHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream {
// Since the query may have multiple values for the same key,
// we'll join them to avoid a problem where the user can control
// the upstream that the request goes to by sending multiple values
Expand All @@ -397,7 +415,7 @@ func (s QueryHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.
// different request, because the order of the values is significant.
vals := strings.Join(req.URL.Query()[s.Key], ",")
if vals == "" {
return RandomSelection{}.Select(pool, req, nil)
return s.fallback.Select(pool, req, nil)
}
return hostByHashing(pool, vals)
}
Expand All @@ -410,6 +428,24 @@ func (s *QueryHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
}
s.Key = d.Val()
}
for nesting := d.Nesting(); d.NextBlock(nesting); {
switch d.Val() {
case "fallback":
if !d.NextArg() {
return d.ArgErr()
}
if s.FallbackRaw != nil {
return d.Err("fallback selection policy already specified")
}
mod, err := loadFallbackPolicy(d)
if err != nil {
return err
}
s.FallbackRaw = mod
default:
return d.Errf("unrecognized option '%s'", d.Val())
}
}
return nil
}

Expand All @@ -418,6 +454,10 @@ func (s *QueryHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
type HeaderHashSelection struct {
// The HTTP header field whose value is to be hashed and used for upstream selection.
Field string `json:"field,omitempty"`

// The fallback policy to use if the header is not present. Defaults to `random`.
FallbackRaw json.RawMessage `json:"fallback,omitempty" caddy:"namespace=http.reverse_proxy.selection_policies inline_key=policy"`
fallback Selector
}

// CaddyModule returns the Caddy module information.
Expand All @@ -428,12 +468,24 @@ func (HeaderHashSelection) CaddyModule() caddy.ModuleInfo {
}
}

// Select returns an available host, if any.
func (s HeaderHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream {
// Provision sets up the module.
func (s *HeaderHashSelection) Provision(ctx caddy.Context) error {
if s.Field == "" {
return nil
return fmt.Errorf("header field is required")
}
if s.FallbackRaw == nil {
s.FallbackRaw = caddyconfig.JSONModuleObject(RandomSelection{}, "policy", "random", nil)
}
mod, err := ctx.LoadModule(s, "FallbackRaw")
if err != nil {
return fmt.Errorf("loading fallback selection policy: %s", err)
}
s.fallback = mod.(Selector)
return nil
}

// Select returns an available host, if any.
func (s HeaderHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream {
// The Host header should be obtained from the req.Host field
// since net/http removes it from the header map.
if s.Field == "Host" && req.Host != "" {
Expand All @@ -442,7 +494,7 @@ func (s HeaderHashSelection) Select(pool UpstreamPool, req *http.Request, _ http

val := req.Header.Get(s.Field)
if val == "" {
return RandomSelection{}.Select(pool, req, nil)
return s.fallback.Select(pool, req, nil)
}
return hostByHashing(pool, val)
}
Expand All @@ -455,6 +507,24 @@ func (s *HeaderHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
}
s.Field = d.Val()
}
for nesting := d.Nesting(); d.NextBlock(nesting); {
switch d.Val() {
case "fallback":
if !d.NextArg() {
return d.ArgErr()
}
if s.FallbackRaw != nil {
return d.Err("fallback selection policy already specified")
}
mod, err := loadFallbackPolicy(d)
if err != nil {
return err
}
s.FallbackRaw = mod
default:
return d.Errf("unrecognized option '%s'", d.Val())
}
}
return nil
}

Expand All @@ -465,6 +535,10 @@ type CookieHashSelection struct {
Name string `json:"name,omitempty"`
// Secret to hash (Hmac256) chosen upstream in cookie
Secret string `json:"secret,omitempty"`

// The fallback policy to use if the cookie is not present. Defaults to `random`.
FallbackRaw json.RawMessage `json:"fallback,omitempty" caddy:"namespace=http.reverse_proxy.selection_policies inline_key=policy"`
fallback Selector
}

// CaddyModule returns the Caddy module information.
Expand All @@ -475,15 +549,48 @@ func (CookieHashSelection) CaddyModule() caddy.ModuleInfo {
}
}

// Select returns an available host, if any.
func (s CookieHashSelection) Select(pool UpstreamPool, req *http.Request, w http.ResponseWriter) *Upstream {
// Provision sets up the module.
func (s *CookieHashSelection) Provision(ctx caddy.Context) error {
if s.Name == "" {
s.Name = "lb"
}
if s.FallbackRaw == nil {
s.FallbackRaw = caddyconfig.JSONModuleObject(RandomSelection{}, "policy", "random", nil)
}
mod, err := ctx.LoadModule(s, "FallbackRaw")
if err != nil {
return fmt.Errorf("loading fallback selection policy: %s", err)
}
s.fallback = mod.(Selector)
return nil
}

// Select returns an available host, if any.
func (s CookieHashSelection) Select(pool UpstreamPool, req *http.Request, w http.ResponseWriter) *Upstream {
// selects a new Host using the fallback policy (typically random)
// and write a sticky session cookie to the response.
selectNewHost := func() *Upstream {
upstream := s.fallback.Select(pool, req, w)
if upstream == nil {
return nil
}
sha, err := hashCookie(s.Secret, upstream.Dial)
if err != nil {
return upstream
}
http.SetCookie(w, &http.Cookie{
Name: s.Name,
Value: sha,
Path: "/",
Secure: false,
})
return upstream
}

cookie, err := req.Cookie(s.Name)
// If there's no cookie, select new random host
// If there's no cookie, select a host using the fallback policy
if err != nil || cookie == nil {
return selectNewHostWithCookieHashSelection(pool, w, s.Secret, s.Name)
return selectNewHost()
}
// If the cookie is present, loop over the available upstreams until we find a match
cookieValue := cookie.Value
Expand All @@ -496,13 +603,15 @@ func (s CookieHashSelection) Select(pool UpstreamPool, req *http.Request, w http
return upstream
}
}
// If there is no matching host, select new random host
return selectNewHostWithCookieHashSelection(pool, w, s.Secret, s.Name)
// If there is no matching host, select a host using the fallback policy
return selectNewHost()
}

// UnmarshalCaddyfile sets up the module from Caddyfile tokens. Syntax:
//
// lb_policy cookie [<name> [<secret>]]
// lb_policy cookie [<name> [<secret>]] {
// fallback <policy>
// }
//
// By default name is `lb`
func (s *CookieHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
Expand All @@ -517,22 +626,25 @@ func (s *CookieHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
default:
return d.ArgErr()
}
return nil
}

// Select a new Host randomly and add a sticky session cookie
func selectNewHostWithCookieHashSelection(pool []*Upstream, w http.ResponseWriter, cookieSecret string, cookieName string) *Upstream {
randomHost := selectRandomHost(pool)

if randomHost != nil {
// Hash (HMAC with some key for privacy) the upstream.Dial string as the cookie value
sha, err := hashCookie(cookieSecret, randomHost.Dial)
if err == nil {
// write the cookie.
http.SetCookie(w, &http.Cookie{Name: cookieName, Value: sha, Path: "/", Secure: false})
for nesting := d.Nesting(); d.NextBlock(nesting); {
switch d.Val() {
case "fallback":
if !d.NextArg() {
return d.ArgErr()
}
if s.FallbackRaw != nil {
return d.Err("fallback selection policy already specified")
}
mod, err := loadFallbackPolicy(d)
if err != nil {
return err
}
s.FallbackRaw = mod
default:
return d.Errf("unrecognized option '%s'", d.Val())
}
}
return randomHost
return nil
}

// hashCookie hashes (HMAC 256) some data with the secret
Expand Down Expand Up @@ -627,6 +739,20 @@ func hash(s string) uint32 {
return h.Sum32()
}

func loadFallbackPolicy(d *caddyfile.Dispenser) (json.RawMessage, error) {
name := d.Val()
modID := "http.reverse_proxy.selection_policies." + name
unm, err := caddyfile.UnmarshalModule(d, modID)
if err != nil {
return nil, err
}
sel, ok := unm.(Selector)
if !ok {
return nil, d.Errf("module %s (%T) is not a reverseproxy.Selector", modID, unm)
}
return caddyconfig.JSONModuleObject(sel, "policy", name, nil), nil
}

// Interface guards
var (
_ Selector = (*RandomSelection)(nil)
Expand Down