diff --git a/caskethttp/proxy/policy.go b/caskethttp/proxy/policy.go index 0f8fab44..9868a200 100644 --- a/caskethttp/proxy/policy.go +++ b/caskethttp/proxy/policy.go @@ -33,13 +33,13 @@ type Policy interface { } func init() { - RegisterPolicy("random", func(arg string) Policy { return &Random{} }) - RegisterPolicy("least_conn", func(arg string) Policy { return &LeastConn{} }) - RegisterPolicy("round_robin", func(arg string) Policy { return &RoundRobin{} }) - RegisterPolicy("ip_hash", func(arg string) Policy { return &IPHash{} }) - RegisterPolicy("first", func(arg string) Policy { return &First{} }) - RegisterPolicy("uri_hash", func(arg string) Policy { return &URIHash{} }) - RegisterPolicy("header", func(arg string) Policy { return &Header{arg} }) + RegisterPolicy("random", func(args []string) Policy { return &Random{} }) + RegisterPolicy("least_conn", func(args []string) Policy { return &LeastConn{} }) + RegisterPolicy("round_robin", func(args []string) Policy { return &RoundRobin{} }) + RegisterPolicy("ip_hash", func(args []string) Policy { return &IPHash{} }) + RegisterPolicy("first", func(args []string) Policy { return &First{} }) + RegisterPolicy("uri_hash", func(args []string) Policy { return &URIHash{} }) + RegisterPolicy("header", func(args []string) Policy { return &Header{args} }) } // Random is a policy that selects up hosts from a pool at random. @@ -183,17 +183,22 @@ func (r *First) Select(pool HostPool, request *http.Request) *UpstreamHost { type Header struct { // The name of the request header, the value of which will determine // how the request is routed - Name string + Names []string } var roundRobinPolicier RoundRobin // Select selects the host based on hashing the header value func (r *Header) Select(pool HostPool, request *http.Request) *UpstreamHost { - if r.Name == "" { + if r.Names == nil { return nil } - val := request.Header.Get(r.Name) + + val := "" + for _, name := range r.Names { + val += request.Header.Get(name) + } + if val == "" { // fallback to RoundRobin policy in case no Header in request return roundRobinPolicier.Select(pool, request) diff --git a/caskethttp/proxy/policy_test.go b/caskethttp/proxy/policy_test.go index be6ef947..6d333c6d 100644 --- a/caskethttp/proxy/policy_test.go +++ b/caskethttp/proxy/policy_test.go @@ -327,18 +327,18 @@ func TestHeaderPolicy(t *testing.T) { NilHost bool HostIndex int }{ - {"empty config", &Header{""}, "", "", true, 0}, - {"empty config+header+value", &Header{""}, "Affinity", "somevalue", true, 0}, - {"empty config+header", &Header{""}, "Affinity", "", true, 0}, - - {"no header(fallback to roundrobin)", &Header{"Affinity"}, "", "", false, 1}, - {"no header(fallback to roundrobin)", &Header{"Affinity"}, "", "", false, 2}, - {"no header(fallback to roundrobin)", &Header{"Affinity"}, "", "", false, 0}, - - {"hash route to host", &Header{"Affinity"}, "Affinity", "somevalue", false, 1}, - {"hash route to host", &Header{"Affinity"}, "Affinity", "somevalue2", false, 0}, - {"hash route to host", &Header{"Affinity"}, "Affinity", "somevalue3", false, 2}, - {"hash route with empty value", &Header{"Affinity"}, "Affinity", "", false, 1}, + {"empty config", &Header{nil}, "", "", true, 0}, + {"empty config+header+value", &Header{nil}, "Affinity", "somevalue", true, 0}, + {"empty config+header", &Header{nil}, "Affinity", "", true, 0}, + + {"no header(fallback to roundrobin)", &Header{[]string{"Affinity"}}, "", "", false, 1}, + {"no header(fallback to roundrobin)", &Header{[]string{"Affinity"}}, "", "", false, 2}, + {"no header(fallback to roundrobin)", &Header{[]string{"Affinity"}}, "", "", false, 0}, + + {"hash route to host", &Header{[]string{"Affinity"}}, "Affinity", "somevalue", false, 1}, + {"hash route to host", &Header{[]string{"Affinity"}}, "Affinity", "somevalue2", false, 0}, + {"hash route to host", &Header{[]string{"Affinity"}}, "Affinity", "somevalue3", false, 2}, + {"hash route with empty value", &Header{[]string{"Affinity"}}, "Affinity", "", false, 1}, } for idx, test := range tests { diff --git a/caskethttp/proxy/upstream.go b/caskethttp/proxy/upstream.go index 393055f0..10875ea8 100644 --- a/caskethttp/proxy/upstream.go +++ b/caskethttp/proxy/upstream.go @@ -36,12 +36,12 @@ import ( "crypto/tls" - "github.com/tmpim/casket/caskethttp/httpserver" "github.com/tmpim/casket/casketfile" + "github.com/tmpim/casket/caskethttp/httpserver" ) var ( - supportedPolicies = make(map[string]func(string) Policy) + supportedPolicies = make(map[string]func([]string) Policy) ) type staticUpstream struct { @@ -346,11 +346,11 @@ func parseBlock(c *casketfile.Dispenser, u *staticUpstream, hasSrv bool) error { if !ok { return c.ArgErr() } - arg := "" - if c.NextArg() { - arg = c.Val() + var args []string + for c.NextArg() { + args = append(args, c.Val()) } - u.Policy = policyCreateFunc(arg) + u.Policy = policyCreateFunc(args) case "fallback_delay": if !c.NextArg() { return c.ArgErr() @@ -780,7 +780,7 @@ func (u *staticUpstream) Stop() error { } // RegisterPolicy adds a custom policy to the proxy. -func RegisterPolicy(name string, policy func(string) Policy) { +func RegisterPolicy(name string, policy func([]string) Policy) { supportedPolicies[name] = policy } diff --git a/caskethttp/proxy/upstream_test.go b/caskethttp/proxy/upstream_test.go index 2736b621..4ca2bac3 100644 --- a/caskethttp/proxy/upstream_test.go +++ b/caskethttp/proxy/upstream_test.go @@ -127,7 +127,7 @@ func TestSelect(t *testing.T) { func TestRegisterPolicy(t *testing.T) { name := "custom" customPolicy := &customPolicy{} - RegisterPolicy(name, func(string) Policy { return customPolicy }) + RegisterPolicy(name, func([]string) Policy { return customPolicy }) if _, ok := supportedPolicies[name]; !ok { t.Error("Expected supportedPolicies to have a custom policy.") }