diff options
Diffstat (limited to 'modules/caddyhttp/reverseproxy/selectionpolicies.go')
| -rw-r--r-- | modules/caddyhttp/reverseproxy/selectionpolicies.go | 360 |
1 files changed, 330 insertions, 30 deletions
diff --git a/modules/caddyhttp/reverseproxy/selectionpolicies.go b/modules/caddyhttp/reverseproxy/selectionpolicies.go index 0b7f50c..acb069a 100644 --- a/modules/caddyhttp/reverseproxy/selectionpolicies.go +++ b/modules/caddyhttp/reverseproxy/selectionpolicies.go @@ -18,17 +18,20 @@ import ( "crypto/hmac" "crypto/sha256" "encoding/hex" + "encoding/json" "fmt" "hash/fnv" weakrand "math/rand" "net" "net/http" "strconv" + "strings" "sync/atomic" - "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" ) func init() { @@ -36,13 +39,14 @@ func init() { caddy.RegisterModule(RandomChoiceSelection{}) caddy.RegisterModule(LeastConnSelection{}) caddy.RegisterModule(RoundRobinSelection{}) + caddy.RegisterModule(WeightedRoundRobinSelection{}) caddy.RegisterModule(FirstSelection{}) caddy.RegisterModule(IPHashSelection{}) + caddy.RegisterModule(ClientIPHashSelection{}) caddy.RegisterModule(URIHashSelection{}) + caddy.RegisterModule(QueryHashSelection{}) caddy.RegisterModule(HeaderHashSelection{}) caddy.RegisterModule(CookieHashSelection{}) - - weakrand.Seed(time.Now().UTC().UnixNano()) } // RandomSelection is a policy that selects @@ -72,6 +76,90 @@ func (r *RandomSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error { return nil } +// WeightedRoundRobinSelection is a policy that selects +// a host based on weighted round-robin ordering. +type WeightedRoundRobinSelection struct { + // The weight of each upstream in order, + // corresponding with the list of upstreams configured. + Weights []int `json:"weights,omitempty"` + index uint32 + totalWeight int +} + +// CaddyModule returns the Caddy module information. +func (WeightedRoundRobinSelection) CaddyModule() caddy.ModuleInfo { + return caddy.ModuleInfo{ + ID: "http.reverse_proxy.selection_policies.weighted_round_robin", + New: func() caddy.Module { + return new(WeightedRoundRobinSelection) + }, + } +} + +// UnmarshalCaddyfile sets up the module from Caddyfile tokens. +func (r *WeightedRoundRobinSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error { + for d.Next() { + args := d.RemainingArgs() + if len(args) == 0 { + return d.ArgErr() + } + + for _, weight := range args { + weightInt, err := strconv.Atoi(weight) + if err != nil { + return d.Errf("invalid weight value '%s': %v", weight, err) + } + if weightInt < 1 { + return d.Errf("invalid weight value '%s': weight should be non-zero and positive", weight) + } + r.Weights = append(r.Weights, weightInt) + } + } + return nil +} + +// Provision sets up r. +func (r *WeightedRoundRobinSelection) Provision(ctx caddy.Context) error { + for _, weight := range r.Weights { + r.totalWeight += weight + } + return nil +} + +// Select returns an available host, if any. +func (r *WeightedRoundRobinSelection) Select(pool UpstreamPool, _ *http.Request, _ http.ResponseWriter) *Upstream { + if len(pool) == 0 { + return nil + } + if len(r.Weights) < 2 { + return pool[0] + } + var index, totalWeight int + currentWeight := int(atomic.AddUint32(&r.index, 1)) % r.totalWeight + for i, weight := range r.Weights { + totalWeight += weight + if currentWeight < totalWeight { + index = i + break + } + } + + upstreams := make([]*Upstream, 0, len(r.Weights)) + for _, upstream := range pool { + if !upstream.Available() { + continue + } + upstreams = append(upstreams, upstream) + if len(upstreams) == cap(upstreams) { + break + } + } + if len(upstreams) == 0 { + return nil + } + return upstreams[index%len(upstreams)] +} + // RandomChoiceSelection is a policy that selects // two or more available hosts at random, then // chooses the one with the least load. @@ -181,7 +269,7 @@ func (LeastConnSelection) Select(pool UpstreamPool, _ *http.Request, _ http.Resp // sample: https://en.wikipedia.org/wiki/Reservoir_sampling if numReqs == leastReqs { count++ - if (weakrand.Int() % count) == 0 { //nolint:gosec + if count == 1 || (weakrand.Int()%count) == 0 { //nolint:gosec bestHost = host } } @@ -303,6 +391,39 @@ func (r *IPHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error { return nil } +// ClientIPHashSelection is a policy that selects a host +// based on hashing the client IP of the request, as determined +// by the HTTP app's trusted proxies settings. +type ClientIPHashSelection struct{} + +// CaddyModule returns the Caddy module information. +func (ClientIPHashSelection) CaddyModule() caddy.ModuleInfo { + return caddy.ModuleInfo{ + ID: "http.reverse_proxy.selection_policies.client_ip_hash", + New: func() caddy.Module { return new(ClientIPHashSelection) }, + } +} + +// Select returns an available host, if any. +func (ClientIPHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream { + address := caddyhttp.GetVar(req.Context(), caddyhttp.ClientIPVarKey).(string) + clientIP, _, err := net.SplitHostPort(address) + if err != nil { + clientIP = address // no port + } + return hostByHashing(pool, clientIP) +} + +// UnmarshalCaddyfile sets up the module from Caddyfile tokens. +func (r *ClientIPHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error { + for d.Next() { + if d.NextArg() { + return d.ArgErr() + } + } + return nil +} + // URIHashSelection is a policy that selects a // host by hashing the request URI. type URIHashSelection struct{} @@ -330,11 +451,95 @@ func (r *URIHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error { return nil } +// QueryHashSelection is a policy that selects +// a host based on a given request query parameter. +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. +func (QueryHashSelection) CaddyModule() caddy.ModuleInfo { + return caddy.ModuleInfo{ + ID: "http.reverse_proxy.selection_policies.query", + New: func() caddy.Module { return new(QueryHashSelection) }, + } +} + +// Provision sets up the module. +func (s *QueryHashSelection) Provision(ctx caddy.Context) error { + if s.Key == "" { + 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 + // for the same key, when the upstream only considers the first value. + // Keep in mind that a client changing the order of the values may + // affect which upstream is selected, but this is a semantically + // different request, because the order of the values is significant. + vals := strings.Join(req.URL.Query()[s.Key], ",") + if vals == "" { + return s.fallback.Select(pool, req, nil) + } + return hostByHashing(pool, vals) +} + +// UnmarshalCaddyfile sets up the module from Caddyfile tokens. +func (s *QueryHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error { + for d.Next() { + if !d.NextArg() { + return d.ArgErr() + } + 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 +} + // HeaderHashSelection is a policy that selects // a host based on a given request header. 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. @@ -345,12 +550,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 != "" { @@ -359,7 +576,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) } @@ -372,6 +589,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 } @@ -382,6 +617,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. @@ -392,15 +631,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 @@ -413,13 +685,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 { @@ -434,22 +708,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 @@ -512,6 +789,9 @@ func leastRequests(upstreams []*Upstream) *Upstream { if len(best) == 0 { return nil } + if len(best) == 1 { + return best[0] + } return best[weakrand.Intn(len(best))] //nolint:gosec } @@ -544,20 +824,40 @@ 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) _ Selector = (*RandomChoiceSelection)(nil) _ Selector = (*LeastConnSelection)(nil) _ Selector = (*RoundRobinSelection)(nil) + _ Selector = (*WeightedRoundRobinSelection)(nil) _ Selector = (*FirstSelection)(nil) _ Selector = (*IPHashSelection)(nil) + _ Selector = (*ClientIPHashSelection)(nil) _ Selector = (*URIHashSelection)(nil) + _ Selector = (*QueryHashSelection)(nil) _ Selector = (*HeaderHashSelection)(nil) _ Selector = (*CookieHashSelection)(nil) - _ caddy.Validator = (*RandomChoiceSelection)(nil) + _ caddy.Validator = (*RandomChoiceSelection)(nil) + _ caddy.Provisioner = (*RandomChoiceSelection)(nil) + _ caddy.Provisioner = (*WeightedRoundRobinSelection)(nil) _ caddyfile.Unmarshaler = (*RandomChoiceSelection)(nil) + _ caddyfile.Unmarshaler = (*WeightedRoundRobinSelection)(nil) ) |