Skip to content

Commit

Permalink
reverseproxy: Add fallback for some policies, instead of always ran…
Browse files Browse the repository at this point in the history 
…dom (#5488)
  • Loading branch information
@francislavoie
francislavoie committed May 5, 2023
1 parent cdce452 commit 48598e1
Show file tree
Hide file tree
Showing 2 changed files with 237 additions and 39 deletions.
184 changes: 155 additions & 29 deletions modules/caddyhttp/reverseproxy/selectionpolicies.go
View file
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

0 comments on commit 48598e1

Please sign in to comment.